207 lines
6.8 KiB
Python
207 lines
6.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""Bit-exact workspace-reusing SwiGLU for the LTX diffusion VAE decoder."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from vllm.triton_utils import tl, triton
|
|
|
|
from . import is_ltx2_fusion_eligible
|
|
from ._numerics import mul_rn_f32, round_bf16_to_fp32
|
|
|
|
_POINTWISE_BLOCK = 1024
|
|
_BASE_TILE_ROWS = 16384
|
|
_MAX_TILE_HIDDEN_ELEMENTS = _BASE_TILE_ROWS * 4 * 2048
|
|
_FAILED_KEYS: set[tuple[int | None, int, int, int, int]] = set()
|
|
_VERIFIED_KEYS: set[tuple[int | None, int, int, int, int]] = set()
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _tile_rows(rows: int, hidden: int) -> int:
|
|
"""Bound projection storage while giving narrow MLPs larger GEMMs."""
|
|
|
|
budget_rows = _MAX_TILE_HIDDEN_ELEMENTS // hidden
|
|
aligned_rows = max(_BASE_TILE_ROWS, budget_rows // _BASE_TILE_ROWS * _BASE_TILE_ROWS)
|
|
return min(rows, aligned_rows)
|
|
|
|
|
|
@triton.jit
|
|
def _silu_mul_kernel(
|
|
gate_ptr,
|
|
up_ptr,
|
|
elements,
|
|
block: tl.constexpr,
|
|
):
|
|
offsets = tl.program_id(0).to(tl.int64) * block + tl.arange(0, block)
|
|
valid = offsets < elements
|
|
gate = tl.load(gate_ptr + offsets, mask=valid).to(tl.float32)
|
|
up = tl.load(up_ptr + offsets, mask=valid).to(tl.float32)
|
|
# Eager rounds SiLU to BF16 before the separate BF16 multiply.
|
|
silu = round_bf16_to_fp32(mul_rn_f32(gate, tl.sigmoid(gate)))
|
|
tl.store(gate_ptr + offsets, mul_rn_f32(silu, up), mask=valid)
|
|
|
|
|
|
def _weight_matches(
|
|
hidden_states: torch.Tensor,
|
|
weight: torch.Tensor,
|
|
shape: tuple[int, int],
|
|
) -> bool:
|
|
return (
|
|
weight.is_cuda
|
|
and weight.device == hidden_states.device
|
|
and weight.dtype is torch.bfloat16
|
|
and weight.shape == shape
|
|
and weight.is_contiguous()
|
|
)
|
|
|
|
|
|
def _supported_inputs(
|
|
hidden_states: torch.Tensor,
|
|
gate_weight: torch.Tensor,
|
|
up_weight: torch.Tensor,
|
|
down_weight: torch.Tensor,
|
|
) -> bool:
|
|
if hidden_states.ndim < 3:
|
|
return False
|
|
dim = hidden_states.shape[-1]
|
|
hidden = gate_weight.shape[0] if gate_weight.ndim == 2 else 0
|
|
return (
|
|
is_ltx2_fusion_eligible(hidden_states)
|
|
and hidden_states.dtype is torch.bfloat16
|
|
and hidden_states.shape[0] == 1
|
|
and hidden_states.numel() > 0
|
|
and hidden_states.is_contiguous()
|
|
and 256 <= dim <= 2048
|
|
and hidden == 4 * dim
|
|
and _weight_matches(hidden_states, gate_weight, (hidden, dim))
|
|
and _weight_matches(hidden_states, up_weight, (hidden, dim))
|
|
and _weight_matches(hidden_states, down_weight, (dim, hidden))
|
|
)
|
|
|
|
|
|
def _silu_mul_inplace(gate: torch.Tensor, up: torch.Tensor) -> None:
|
|
with torch.accelerator.device_index(gate.device.index):
|
|
_silu_mul_kernel[(triton.cdiv(gate.numel(), _POINTWISE_BLOCK),)](
|
|
gate,
|
|
up,
|
|
gate.numel(),
|
|
block=_POINTWISE_BLOCK,
|
|
num_warps=4,
|
|
)
|
|
|
|
|
|
def _reference_tile(
|
|
hidden_states: torch.Tensor,
|
|
gate_weight: torch.Tensor,
|
|
up_weight: torch.Tensor,
|
|
down_weight: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
return F.linear(
|
|
F.silu(F.linear(hidden_states, gate_weight)) * F.linear(hidden_states, up_weight),
|
|
down_weight,
|
|
)
|
|
|
|
|
|
def _launch(
|
|
hidden_states: torch.Tensor,
|
|
gate_weight: torch.Tensor,
|
|
up_weight: torch.Tensor,
|
|
down_weight: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
dim = hidden_states.shape[-1]
|
|
hidden = gate_weight.shape[0]
|
|
flat = hidden_states.reshape(-1, dim)
|
|
tile_rows = _tile_rows(flat.shape[0], hidden)
|
|
output = torch.empty_like(flat)
|
|
workspace_gate = torch.empty((tile_rows, hidden), device=flat.device, dtype=flat.dtype)
|
|
for start in range(0, flat.shape[0], tile_rows):
|
|
end = min(start + tile_rows, flat.shape[0])
|
|
tokens = end - start
|
|
tile = flat[start:end]
|
|
gate = workspace_gate[:tokens]
|
|
torch.mm(tile, gate_weight.t(), out=gate)
|
|
up = torch.mm(tile, up_weight.t())
|
|
_silu_mul_inplace(gate, up)
|
|
torch.mm(gate, down_weight.t(), out=output[start:end])
|
|
return output.reshape(hidden_states.shape)
|
|
|
|
|
|
def _verify_edge_tiles(
|
|
output: torch.Tensor,
|
|
hidden_states: torch.Tensor,
|
|
gate_weight: torch.Tensor,
|
|
up_weight: torch.Tensor,
|
|
down_weight: torch.Tensor,
|
|
) -> bool:
|
|
from diffusers.models.autoencoders.ltx2_diffusion_decoder import _SWIGLU_TILE_SIZE
|
|
|
|
dim = hidden_states.shape[-1]
|
|
flat = hidden_states.reshape(-1, dim)
|
|
actual = output.reshape(-1, dim)
|
|
tile_rows = _tile_rows(flat.shape[0], gate_weight.shape[0])
|
|
starts = {0, ((flat.shape[0] - 1) // tile_rows) * tile_rows}
|
|
for start in starts:
|
|
end = min(start + tile_rows, flat.shape[0])
|
|
# Match upstream's token tiling, not the optimized GEMM shape: changing
|
|
# the matmul row count can change its reduction order and BF16 output.
|
|
for reference_start in range(start, end, _SWIGLU_TILE_SIZE):
|
|
reference_end = min(reference_start + _SWIGLU_TILE_SIZE, end)
|
|
if not torch.equal(
|
|
actual[reference_start:reference_end],
|
|
_reference_tile(flat[reference_start:reference_end], gate_weight, up_weight, down_weight),
|
|
):
|
|
return False
|
|
return True
|
|
|
|
|
|
def try_swiglu_tiled_exact(
|
|
hidden_states: torch.Tensor,
|
|
gate_weight: torch.Tensor,
|
|
up_weight: torch.Tensor,
|
|
down_weight: torch.Tensor,
|
|
) -> torch.Tensor | None:
|
|
"""Run workspace-reusing tiled SwiGLU, or expose the original module."""
|
|
|
|
if not _supported_inputs(hidden_states, gate_weight, up_weight, down_weight):
|
|
return None
|
|
rows = hidden_states.numel() // hidden_states.shape[-1]
|
|
tile_rows = _tile_rows(rows, gate_weight.shape[0])
|
|
runtime_key = (
|
|
hidden_states.device.index,
|
|
hidden_states.shape[-1],
|
|
gate_weight.shape[0],
|
|
tile_rows,
|
|
(rows - 1) % tile_rows + 1,
|
|
)
|
|
if runtime_key in _FAILED_KEYS:
|
|
return None
|
|
try:
|
|
output = _launch(hidden_states, gate_weight, up_weight, down_weight)
|
|
if runtime_key not in _VERIFIED_KEYS:
|
|
if not _verify_edge_tiles(output, hidden_states, gate_weight, up_weight, down_weight):
|
|
_FAILED_KEYS.add(runtime_key)
|
|
logger.warning(
|
|
"Disabling LTX DiffVAE tiled SwiGLU fusion on %s after a bit-exactness mismatch",
|
|
hidden_states.device,
|
|
)
|
|
return None
|
|
_VERIFIED_KEYS.add(runtime_key)
|
|
except Exception as exc: # noqa: BLE001 - fail closed after optimized-path failure
|
|
_FAILED_KEYS.add(runtime_key)
|
|
logger.warning(
|
|
"Disabling LTX DiffVAE tiled SwiGLU fusion on %s after failure: %s",
|
|
hidden_states.device,
|
|
exc,
|
|
)
|
|
return None
|
|
return output
|
|
|
|
|
|
__all__ = ["try_swiglu_tiled_exact"]
|