490 lines
15 KiB
Python
490 lines
15 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""Exact LTX DiffVAE residual, RMSNorm, and AdaLN Triton fusions."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from vllm.triton_utils import tl, triton
|
|
|
|
from . import is_ltx2_fusion_eligible
|
|
from ._numerics import (
|
|
add_rn_f32,
|
|
fma_rn_f32,
|
|
mul_rn_f32,
|
|
round_bf16_to_fp32,
|
|
rsqrt_approx_f32,
|
|
shfl_down_f32,
|
|
)
|
|
|
|
_HIDDEN_SIZE = 256
|
|
_POINTWISE_BLOCK = 1024
|
|
_VERIFY_ROWS = 256
|
|
_VERIFY_ELEMENTS = 1 << 20
|
|
_FAILED_ADALN_KEYS: set[tuple[int | None, int, float]] = set()
|
|
_VERIFIED_ADALN_KEYS: set[tuple[int | None, int, float]] = set()
|
|
_FAILED_ADD_DEVICES: set[int | None] = set()
|
|
_VERIFIED_ADD_DEVICES: set[int | None] = set()
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@triton.jit
|
|
def _two_warp_sum(values):
|
|
values = add_rn_f32(values, shfl_down_f32(values, 16))
|
|
values = add_rn_f32(values, shfl_down_f32(values, 8))
|
|
values = add_rn_f32(values, shfl_down_f32(values, 4))
|
|
values = add_rn_f32(values, shfl_down_f32(values, 2))
|
|
values = add_rn_f32(values, shfl_down_f32(values, 1))
|
|
threads = tl.arange(0, 64)
|
|
warp_0 = tl.sum(tl.where(threads == 0, values, 0.0))
|
|
warp_1 = tl.sum(tl.where(threads == 32, values, 0.0))
|
|
return add_rn_f32(warp_0, warp_1)
|
|
|
|
|
|
@triton.jit
|
|
def _load_ordered_residual(
|
|
x_ptr,
|
|
residual_a_ptr,
|
|
residual_b_ptr,
|
|
offset,
|
|
valid,
|
|
residual_count: tl.constexpr,
|
|
):
|
|
value = add_rn_f32(
|
|
tl.load(x_ptr + offset, mask=valid, other=0.0).to(tl.float32),
|
|
tl.load(residual_a_ptr + offset, mask=valid, other=0.0).to(tl.float32),
|
|
)
|
|
value = round_bf16_to_fp32(value)
|
|
if residual_count == 2:
|
|
value = add_rn_f32(
|
|
value,
|
|
tl.load(residual_b_ptr + offset, mask=valid, other=0.0).to(tl.float32),
|
|
)
|
|
value = round_bf16_to_fp32(value)
|
|
return value
|
|
|
|
|
|
@triton.jit
|
|
def _store_modulated(
|
|
output_ptr,
|
|
norm_weight_ptr,
|
|
scale_ptr,
|
|
shift_ptr,
|
|
value,
|
|
offset,
|
|
column,
|
|
modulation_base,
|
|
reciprocal_rms,
|
|
valid,
|
|
):
|
|
normalized = mul_rn_f32(value, reciprocal_rms)
|
|
norm_weight = tl.load(norm_weight_ptr + column, mask=valid, other=0.0).to(tl.float32)
|
|
normalized = round_bf16_to_fp32(mul_rn_f32(normalized, norm_weight))
|
|
scale = tl.load(scale_ptr + modulation_base + column, mask=valid, other=0.0).to(tl.float32)
|
|
shift = tl.load(shift_ptr + modulation_base + column, mask=valid, other=0.0).to(tl.float32)
|
|
one_plus_scale = round_bf16_to_fp32(add_rn_f32(1.0, scale))
|
|
product = round_bf16_to_fp32(mul_rn_f32(normalized, one_plus_scale))
|
|
tl.store(output_ptr + offset, add_rn_f32(product, shift), mask=valid)
|
|
|
|
|
|
@triton.jit
|
|
def _residual_rms_norm_modulate_kernel(
|
|
output_ptr,
|
|
x_ptr,
|
|
residual_a_ptr,
|
|
residual_b_ptr,
|
|
norm_weight_ptr,
|
|
scale_ptr,
|
|
shift_ptr,
|
|
rows_per_batch,
|
|
eps,
|
|
residual_count: tl.constexpr,
|
|
hidden_size: tl.constexpr,
|
|
):
|
|
row = tl.program_id(0).to(tl.int64)
|
|
batch = row // rows_per_batch
|
|
threads = tl.arange(0, 64)
|
|
row_base = row * hidden_size
|
|
|
|
# PyTorch's BF16 RMSNorm kernel uses four contiguous values per thread for
|
|
# hidden=256. Keep the four ordered BF16 residual values in registers so
|
|
# the output pass does not read the large activations a second time.
|
|
column_0 = threads * 4
|
|
column_1 = column_0 + 1
|
|
column_2 = column_0 + 2
|
|
column_3 = column_0 + 3
|
|
valid_0 = column_0 < hidden_size
|
|
valid_1 = column_1 < hidden_size
|
|
valid_2 = column_2 < hidden_size
|
|
valid_3 = column_3 < hidden_size
|
|
offset_0 = row_base + column_0
|
|
offset_1 = row_base + column_1
|
|
offset_2 = row_base + column_2
|
|
offset_3 = row_base + column_3
|
|
value_0 = _load_ordered_residual(x_ptr, residual_a_ptr, residual_b_ptr, offset_0, valid_0, residual_count)
|
|
value_1 = _load_ordered_residual(x_ptr, residual_a_ptr, residual_b_ptr, offset_1, valid_1, residual_count)
|
|
value_2 = _load_ordered_residual(x_ptr, residual_a_ptr, residual_b_ptr, offset_2, valid_2, residual_count)
|
|
value_3 = _load_ordered_residual(x_ptr, residual_a_ptr, residual_b_ptr, offset_3, valid_3, residual_count)
|
|
accumulator = fma_rn_f32(value_0, value_0, tl.zeros((64,), dtype=tl.float32))
|
|
accumulator = fma_rn_f32(value_1, value_1, accumulator)
|
|
accumulator = fma_rn_f32(value_2, value_2, accumulator)
|
|
accumulator = fma_rn_f32(value_3, value_3, accumulator)
|
|
|
|
total = _two_warp_sum(accumulator)
|
|
reciprocal_rms = rsqrt_approx_f32(total / hidden_size + eps)
|
|
modulation_base = batch * hidden_size
|
|
_store_modulated(
|
|
output_ptr,
|
|
norm_weight_ptr,
|
|
scale_ptr,
|
|
shift_ptr,
|
|
value_0,
|
|
offset_0,
|
|
column_0,
|
|
modulation_base,
|
|
reciprocal_rms,
|
|
valid_0,
|
|
)
|
|
_store_modulated(
|
|
output_ptr,
|
|
norm_weight_ptr,
|
|
scale_ptr,
|
|
shift_ptr,
|
|
value_1,
|
|
offset_1,
|
|
column_1,
|
|
modulation_base,
|
|
reciprocal_rms,
|
|
valid_1,
|
|
)
|
|
_store_modulated(
|
|
output_ptr,
|
|
norm_weight_ptr,
|
|
scale_ptr,
|
|
shift_ptr,
|
|
value_2,
|
|
offset_2,
|
|
column_2,
|
|
modulation_base,
|
|
reciprocal_rms,
|
|
valid_2,
|
|
)
|
|
_store_modulated(
|
|
output_ptr,
|
|
norm_weight_ptr,
|
|
scale_ptr,
|
|
shift_ptr,
|
|
value_3,
|
|
offset_3,
|
|
column_3,
|
|
modulation_base,
|
|
reciprocal_rms,
|
|
valid_3,
|
|
)
|
|
|
|
|
|
@triton.jit
|
|
def _residual_add3_kernel(
|
|
output_ptr,
|
|
x_ptr,
|
|
residual_a_ptr,
|
|
residual_b_ptr,
|
|
residual_c_ptr,
|
|
elements,
|
|
block: tl.constexpr,
|
|
):
|
|
offsets = tl.program_id(0).to(tl.int64) * block + tl.arange(0, block)
|
|
valid = offsets < elements
|
|
value = add_rn_f32(
|
|
tl.load(x_ptr + offsets, mask=valid, other=0.0).to(tl.float32),
|
|
tl.load(residual_a_ptr + offsets, mask=valid, other=0.0).to(tl.float32),
|
|
)
|
|
value = round_bf16_to_fp32(value)
|
|
value = add_rn_f32(
|
|
value,
|
|
tl.load(residual_b_ptr + offsets, mask=valid, other=0.0).to(tl.float32),
|
|
)
|
|
value = round_bf16_to_fp32(value)
|
|
value = add_rn_f32(
|
|
value,
|
|
tl.load(residual_c_ptr + offsets, mask=valid, other=0.0).to(tl.float32),
|
|
)
|
|
tl.store(output_ptr + offsets, value, mask=valid)
|
|
|
|
|
|
def _same_activation(x: torch.Tensor, residual: torch.Tensor) -> bool:
|
|
return (
|
|
residual.dtype is torch.bfloat16
|
|
and residual.is_cuda
|
|
and residual.device == x.device
|
|
and residual.shape == x.shape
|
|
and residual.is_contiguous()
|
|
)
|
|
|
|
|
|
def _modulation_matches(x: torch.Tensor, value: torch.Tensor) -> bool:
|
|
return (
|
|
value.dtype is torch.bfloat16
|
|
and value.is_cuda
|
|
and value.device == x.device
|
|
and value.shape[-1] == _HIDDEN_SIZE
|
|
and value.numel() == x.shape[0] * _HIDDEN_SIZE
|
|
and value.is_contiguous()
|
|
)
|
|
|
|
|
|
def _norm_weight_matches(x: torch.Tensor, weight: torch.Tensor) -> bool:
|
|
return (
|
|
weight.dtype is torch.bfloat16
|
|
and weight.is_cuda
|
|
and weight.device == x.device
|
|
and weight.shape == (_HIDDEN_SIZE,)
|
|
and weight.is_contiguous()
|
|
)
|
|
|
|
|
|
def _adaln_inputs_supported(
|
|
x: torch.Tensor,
|
|
residual_a: torch.Tensor,
|
|
residual_b: torch.Tensor | None,
|
|
norm_weight: torch.Tensor,
|
|
scale: torch.Tensor,
|
|
shift: torch.Tensor,
|
|
eps: float,
|
|
) -> bool:
|
|
return (
|
|
is_ltx2_fusion_eligible(x)
|
|
and x.dtype is torch.bfloat16
|
|
and x.ndim >= 2
|
|
and x.shape[-1] == _HIDDEN_SIZE
|
|
and x.numel() > 0
|
|
and x.is_contiguous()
|
|
and _same_activation(x, residual_a)
|
|
and (residual_b is None or _same_activation(x, residual_b))
|
|
and _norm_weight_matches(x, norm_weight)
|
|
and _modulation_matches(x, scale)
|
|
and _modulation_matches(x, shift)
|
|
and math.isfinite(eps)
|
|
and eps > 0
|
|
)
|
|
|
|
|
|
def _residual_reference(
|
|
x: torch.Tensor,
|
|
residual_a: torch.Tensor,
|
|
residual_b: torch.Tensor | None,
|
|
) -> torch.Tensor:
|
|
hidden_states = x + residual_a
|
|
if residual_b is not None:
|
|
hidden_states = hidden_states + residual_b
|
|
return hidden_states
|
|
|
|
|
|
def _adaln_reference(
|
|
x: torch.Tensor,
|
|
residual_a: torch.Tensor,
|
|
residual_b: torch.Tensor | None,
|
|
norm_weight: torch.Tensor,
|
|
scale: torch.Tensor,
|
|
shift: torch.Tensor,
|
|
eps: float,
|
|
) -> torch.Tensor:
|
|
hidden_states = _residual_reference(x, residual_a, residual_b)
|
|
normalized = F.rms_norm(hidden_states, (_HIDDEN_SIZE,), weight=norm_weight, eps=eps)
|
|
return normalized * (1 + scale) + shift
|
|
|
|
|
|
def _launch_adaln(
|
|
x: torch.Tensor,
|
|
residual_a: torch.Tensor,
|
|
residual_b: torch.Tensor | None,
|
|
norm_weight: torch.Tensor,
|
|
scale: torch.Tensor,
|
|
shift: torch.Tensor,
|
|
eps: float,
|
|
) -> torch.Tensor:
|
|
output = torch.empty_like(x)
|
|
rows = x.numel() // _HIDDEN_SIZE
|
|
rows_per_batch = rows // x.shape[0]
|
|
with torch.accelerator.device_index(x.device.index):
|
|
_residual_rms_norm_modulate_kernel[(rows,)](
|
|
output,
|
|
x,
|
|
residual_a,
|
|
x if residual_b is None else residual_b,
|
|
norm_weight,
|
|
scale,
|
|
shift,
|
|
rows_per_batch,
|
|
eps,
|
|
residual_count=1 if residual_b is None else 2,
|
|
hidden_size=_HIDDEN_SIZE,
|
|
num_warps=2,
|
|
)
|
|
return output
|
|
|
|
|
|
def _verify_adaln_prefix(
|
|
output: torch.Tensor,
|
|
x: torch.Tensor,
|
|
residual_a: torch.Tensor,
|
|
residual_b: torch.Tensor | None,
|
|
norm_weight: torch.Tensor,
|
|
scale: torch.Tensor,
|
|
shift: torch.Tensor,
|
|
eps: float,
|
|
) -> bool:
|
|
rows_per_batch = x.numel() // (x.shape[0] * _HIDDEN_SIZE)
|
|
rows = min(rows_per_batch, _VERIFY_ROWS)
|
|
x_rows = x.reshape(x.shape[0], rows_per_batch, _HIDDEN_SIZE)[0, :rows]
|
|
residual_a_rows = residual_a.reshape(x.shape[0], rows_per_batch, _HIDDEN_SIZE)[0, :rows]
|
|
residual_b_rows = (
|
|
None if residual_b is None else residual_b.reshape(x.shape[0], rows_per_batch, _HIDDEN_SIZE)[0, :rows]
|
|
)
|
|
scale_row = scale.reshape(x.shape[0], 1, _HIDDEN_SIZE)[0]
|
|
shift_row = shift.reshape(x.shape[0], 1, _HIDDEN_SIZE)[0]
|
|
reference = _adaln_reference(
|
|
x_rows,
|
|
residual_a_rows,
|
|
residual_b_rows,
|
|
norm_weight,
|
|
scale_row,
|
|
shift_row,
|
|
eps,
|
|
)
|
|
actual = output.reshape(x.shape[0], rows_per_batch, _HIDDEN_SIZE)[0, :rows]
|
|
return torch.equal(actual, reference)
|
|
|
|
|
|
def try_residual_rms_norm_modulate_exact(
|
|
x: torch.Tensor,
|
|
residual_a: torch.Tensor,
|
|
residual_b: torch.Tensor | None,
|
|
norm_weight: torch.Tensor,
|
|
scale: torch.Tensor,
|
|
shift: torch.Tensor,
|
|
eps: float,
|
|
) -> torch.Tensor | None:
|
|
"""Fuse one or two ordered BF16 residuals with RMSNorm and AdaLN."""
|
|
|
|
if not _adaln_inputs_supported(x, residual_a, residual_b, norm_weight, scale, shift, eps):
|
|
return None
|
|
residual_count = 1 if residual_b is None else 2
|
|
runtime_key = (x.device.index, residual_count, float(eps))
|
|
if runtime_key in _FAILED_ADALN_KEYS:
|
|
return None
|
|
try:
|
|
output = _launch_adaln(x, residual_a, residual_b, norm_weight, scale, shift, eps)
|
|
if runtime_key not in _VERIFIED_ADALN_KEYS:
|
|
if not _verify_adaln_prefix(
|
|
output,
|
|
x,
|
|
residual_a,
|
|
residual_b,
|
|
norm_weight,
|
|
scale,
|
|
shift,
|
|
eps,
|
|
):
|
|
_FAILED_ADALN_KEYS.add(runtime_key)
|
|
logger.warning(
|
|
"Disabling LTX DiffVAE residual-AdaLN fusion on %s after a bit-exactness mismatch",
|
|
x.device,
|
|
)
|
|
return None
|
|
_VERIFIED_ADALN_KEYS.add(runtime_key)
|
|
except Exception as exc: # noqa: BLE001 - fail closed after optimized-path failure
|
|
_FAILED_ADALN_KEYS.add(runtime_key)
|
|
logger.warning(
|
|
"Disabling LTX DiffVAE residual-AdaLN fusion on %s after failure: %s",
|
|
x.device,
|
|
exc,
|
|
)
|
|
return None
|
|
return output
|
|
|
|
|
|
def _add_inputs_supported(x: torch.Tensor, residuals: tuple[torch.Tensor, ...]) -> bool:
|
|
return (
|
|
is_ltx2_fusion_eligible(x)
|
|
and x.dtype is torch.bfloat16
|
|
and x.numel() > 0
|
|
and x.is_contiguous()
|
|
and all(_same_activation(x, residual) for residual in residuals)
|
|
)
|
|
|
|
|
|
def _launch_add3(
|
|
x: torch.Tensor,
|
|
residual_a: torch.Tensor,
|
|
residual_b: torch.Tensor,
|
|
residual_c: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
output = torch.empty_like(x)
|
|
with torch.accelerator.device_index(x.device.index):
|
|
_residual_add3_kernel[(triton.cdiv(x.numel(), _POINTWISE_BLOCK),)](
|
|
output,
|
|
x,
|
|
residual_a,
|
|
residual_b,
|
|
residual_c,
|
|
x.numel(),
|
|
block=_POINTWISE_BLOCK,
|
|
num_warps=4,
|
|
)
|
|
return output
|
|
|
|
|
|
def try_residual_add3_exact(
|
|
x: torch.Tensor,
|
|
residual_a: torch.Tensor,
|
|
residual_b: torch.Tensor,
|
|
residual_c: torch.Tensor,
|
|
) -> torch.Tensor | None:
|
|
"""Fuse three ordered BF16 residual additions into one pointwise pass."""
|
|
|
|
residuals = (residual_a, residual_b, residual_c)
|
|
if not _add_inputs_supported(x, residuals) or x.device.index in _FAILED_ADD_DEVICES:
|
|
return None
|
|
try:
|
|
output = _launch_add3(x, *residuals)
|
|
if x.device.index not in _VERIFIED_ADD_DEVICES:
|
|
elements = min(x.numel(), _VERIFY_ELEMENTS)
|
|
reference = (
|
|
_residual_reference(
|
|
x.reshape(-1)[:elements],
|
|
residual_a.reshape(-1)[:elements],
|
|
residual_b.reshape(-1)[:elements],
|
|
)
|
|
+ residual_c.reshape(-1)[:elements]
|
|
)
|
|
if not torch.equal(output.reshape(-1)[:elements], reference):
|
|
_FAILED_ADD_DEVICES.add(x.device.index)
|
|
logger.warning(
|
|
"Disabling LTX DiffVAE residual-add fusion on %s after a bit-exactness mismatch",
|
|
x.device,
|
|
)
|
|
return None
|
|
_VERIFIED_ADD_DEVICES.add(x.device.index)
|
|
except Exception as exc: # noqa: BLE001 - fail closed after optimized-path failure
|
|
_FAILED_ADD_DEVICES.add(x.device.index)
|
|
logger.warning(
|
|
"Disabling LTX DiffVAE residual-add fusion on %s after failure: %s",
|
|
x.device,
|
|
exc,
|
|
)
|
|
return None
|
|
return output
|
|
|
|
|
|
__all__ = [
|
|
"try_residual_add3_exact",
|
|
"try_residual_rms_norm_modulate_exact",
|
|
]
|