diff --git a/python/sglang/kernels/ops/diffusion/triton/wan_temb_table_slices.py b/python/sglang/kernels/ops/diffusion/triton/wan_temb_table_slices.py new file mode 100644 index 000000000..b031b894c --- /dev/null +++ b/python/sglang/kernels/ops/diffusion/triton/wan_temb_table_slices.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fused, contiguous adaLN slices for Wan2.2-TI2V per-token modulation. + +The eager chain per block is + + ``(scale_shift_table.unsqueeze(0) + temb.float()).chunk(6, dim=2)`` + +which materializes the full ``(B, S, 6, D)`` tensor in fp32 (a widening copy +plus an add over ~8 GB at 704p/121f) and hands six **strided** slices to the +downstream fused-norm wrappers, whose ``.contiguous()`` calls then copy each +full ``(B, S, D)`` slice again. This kernel produces the six slices in one +pass over ``temb``, each naturally contiguous, so the downstream +``.contiguous()`` calls become no-ops. + +The math is a float32 add of the (exactly representable) widened ``temb`` +values — no rounding is involved at any step, so the outputs are bit-identical +to the eager chain by construction; callers still verify the first call and +fall back on any mismatch. +""" + +from __future__ import annotations + +import torch +import triton # type: ignore +import triton.language as tl # type: ignore + +from sglang.srt.utils.custom_op import register_custom_op + + +@triton.jit +def _temb_table_slices_kernel( + out_ptr, + temb_ptr, + table_ptr, + rows, + D: tl.constexpr, + BLOCK: tl.constexpr, + NCHUNK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) # over B * S + j = tl.program_id(1) # modulation slice index [0, 6) + for i in tl.static_range(NCHUNK): + cols = i * BLOCK + tl.arange(0, BLOCK) + mask = cols < D + t = tl.load(temb_ptr + (row * 6 + j) * D + cols, mask=mask, other=0.0).to( + tl.float32 + ) + w = tl.load(table_ptr + j * D + cols, mask=mask, other=0.0).to(tl.float32) + tl.store(out_ptr + (j * rows + row) * D + cols, w + t, mask=mask) + + +def can_use_fused_temb_table_slices(table: torch.Tensor, temb: torch.Tensor) -> bool: + return ( + temb.is_cuda + and temb.dtype in (torch.bfloat16, torch.float16, torch.float32) + and temb.dim() == 4 + and temb.shape[2] == 6 + and temb.is_contiguous() + and table.is_cuda + and table.device == temb.device + and table.dtype in (torch.bfloat16, torch.float16, torch.float32) + and table.shape == (1, 6, temb.shape[-1]) + and table.is_contiguous() + and temb.numel() > 0 + ) + + +def _fake_temb_table_slices(table: torch.Tensor, temb: torch.Tensor) -> torch.Tensor: + batch, seq_len, six, hidden = temb.shape + return temb.new_empty((six, batch, seq_len, hidden), dtype=torch.float32) + + +@register_custom_op( + op_name="triton_wan_temb_table_slices", + mutates_args=[], + fake_impl=_fake_temb_table_slices, +) +def fused_temb_table_slices(table: torch.Tensor, temb: torch.Tensor) -> torch.Tensor: + """``table.unsqueeze(0) + temb.float()`` as a ``(6, B, S, D)`` buffer. + + ``temb`` is ``(B, S, 6, D)``; ``table`` is the block's ``(1, 6, D)`` fp32 + adaLN table. ``out[j]`` is the ``j``-th modulation slice, contiguous. + """ + batch, seq_len, _, hidden = temb.shape + out = temb.new_empty((6, batch, seq_len, hidden), dtype=torch.float32) + rows = batch * seq_len + block = min(1024, triton.next_power_of_2(hidden)) + nchunk = (hidden + block - 1) // block + with torch.cuda.device(temb.device): + _temb_table_slices_kernel[(rows, 6)]( + out, + temb, + table, + rows, + D=hidden, + BLOCK=block, + NCHUNK=nchunk, + ) + return out diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index 3e671888a..d50ab44a7 100755 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py @@ -9,6 +9,14 @@ from typing import Any import torch import torch.nn as nn +from sglang.kernels.ops.diffusion.bitexact_gate import ( + BitExactFusionGate, + tensors_equal, +) +from sglang.kernels.ops.diffusion.triton.wan_temb_table_slices import ( + can_use_fused_temb_table_slices, + fused_temb_table_slices, +) from sglang.multimodal_gen.configs.models.dits import WanVideoConfig from sglang.multimodal_gen.configs.models.fsdp import is_block from sglang.multimodal_gen.runtime.distributed import ( @@ -331,6 +339,54 @@ class WanI2VCrossAttention(WanSelfAttention): return x +_WAN_TEMB_SLICES = BitExactFusionGate("Wan fused temb-table slices") + + +def _eager_temb_table_slices( + table: torch.Tensor, temb: torch.Tensor +) -> tuple[torch.Tensor, ...]: + parts = (table.unsqueeze(0) + temb.float()).chunk(6, dim=2) + return tuple(part.squeeze(2) for part in parts) + + +def _wan_temb_table_slices( + table: torch.Tensor, temb: torch.Tensor +) -> tuple[torch.Tensor, ...]: + """Per-token adaLN slices ``(table + temb.float()).chunk(6)`` in one pass. + + The fused kernel writes each ``(B, S, D)`` slice contiguously, so the + downstream fused-norm wrappers' ``.contiguous()`` calls stop copying the + full activation. A float32 add of widened values involves no rounding, + so the result is bit-identical to the eager chain; the first call still + verifies ``torch.equal`` and falls back permanently on mismatch. + """ + verified = _WAN_TEMB_SLICES.verified + if ( + not _WAN_TEMB_SLICES.disabled + and can_use_fused_temb_table_slices(table, temb) + and (verified or _WAN_TEMB_SLICES.can_attempt_once()) + ): + try: + buf = fused_temb_table_slices(table, temb) + except Exception as exc: + _WAN_TEMB_SLICES.on_exception(exc, logger=logger) + else: + out = tuple(buf[j] for j in range(6)) + if verified: + return out + return _WAN_TEMB_SLICES.accept_or_fallback( + out, + _eager_temb_table_slices(table, temb), + equal=tensors_equal, + logger=logger, + mismatch_msg=( + "Wan fused temb-table slices are not bit-exact on this " + "platform; falling back to eager" + ), + ) + return _eager_temb_table_slices(table, temb) + + class WanTransformerBlock(nn.Module): def __init__( self, @@ -490,6 +546,7 @@ class WanTransformerBlock(nn.Module): encoder_hidden_states: torch.Tensor, temb: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], + rope_cos_sin_cache: torch.Tensor | None = None, ) -> torch.Tensor: if hidden_states.dim() == 4: hidden_states = hidden_states.squeeze(1) @@ -497,16 +554,14 @@ class WanTransformerBlock(nn.Module): orig_dtype = hidden_states.dtype if temb.dim() == 4: # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v) - shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( - self.scale_shift_table.unsqueeze(0) + temb.float() - ).chunk(6, dim=2) - # batch_size, seq_len, 1, inner_dim - shift_msa = shift_msa.squeeze(2) - scale_msa = scale_msa.squeeze(2) - gate_msa = gate_msa.squeeze(2) - c_shift_msa = c_shift_msa.squeeze(2) - c_scale_msa = c_scale_msa.squeeze(2) - c_gate_msa = c_gate_msa.squeeze(2) + ( + shift_msa, + scale_msa, + gate_msa, + c_shift_msa, + c_scale_msa, + c_gate_msa, + ) = _wan_temb_table_slices(self.scale_shift_table, temb) else: # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B) e = self.scale_shift_table + temb.float() @@ -544,13 +599,17 @@ class WanTransformerBlock(nn.Module): # Apply rotary embeddings cos, sin = freqs_cis if _is_cuda and query.shape == key.shape: - cos_sin_cache = torch.cat( - [ - cos.to(dtype=torch.float32).contiguous(), - sin.to(dtype=torch.float32).contiguous(), - ], - dim=-1, - ) + # The concatenated cache only depends on freqs_cis, which is fixed + # for the whole forward; the transformer builds it once per call. + cos_sin_cache = rope_cos_sin_cache + if cos_sin_cache is None: + cos_sin_cache = torch.cat( + [ + cos.to(dtype=torch.float32).contiguous(), + sin.to(dtype=torch.float32).contiguous(), + ], + dim=-1, + ) query, key = apply_flashinfer_rope_qk_inplace( query, key, cos_sin_cache, is_neox=False ) @@ -757,6 +816,7 @@ class WanTransformerBlock_VSA(nn.Module): encoder_hidden_states: torch.Tensor, temb: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], + rope_cos_sin_cache: torch.Tensor | None = None, ) -> torch.Tensor: if hidden_states.dim() == 4: hidden_states = hidden_states.squeeze(1) @@ -791,13 +851,17 @@ class WanTransformerBlock_VSA(nn.Module): # Apply rotary embeddings cos, sin = freqs_cis if _is_cuda and query.shape == key.shape: - cos_sin_cache = torch.cat( - [ - cos.to(dtype=torch.float32).contiguous(), - sin.to(dtype=torch.float32).contiguous(), - ], - dim=-1, - ) + # The concatenated cache only depends on freqs_cis, which is fixed + # for the whole forward; the transformer builds it once per call. + cos_sin_cache = rope_cos_sin_cache + if cos_sin_cache is None: + cos_sin_cache = torch.cat( + [ + cos.to(dtype=torch.float32).contiguous(), + sin.to(dtype=torch.float32).contiguous(), + ], + dim=-1, + ) query, key = apply_flashinfer_rope_qk_inplace( query, key, cos_sin_cache, is_neox=False ) @@ -1158,9 +1222,23 @@ class WanTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin): if self.enable_teacache: original_hidden_states = hidden_states.clone() + rope_cos_sin_cache = None + if _is_cuda and freqs_cis is not None: + cos, sin = freqs_cis + rope_cos_sin_cache = torch.cat( + [ + cos.to(dtype=torch.float32).contiguous(), + sin.to(dtype=torch.float32).contiguous(), + ], + dim=-1, + ) for block in self.blocks: hidden_states = block( - hidden_states, encoder_hidden_states, timestep_proj, freqs_cis + hidden_states, + encoder_hidden_states, + timestep_proj, + freqs_cis, + rope_cos_sin_cache=rope_cos_sin_cache, ) # if teacache is enabled, we need to cache the original hidden states if self.enable_teacache: diff --git a/python/sglang/multimodal_gen/test/unit/test_wan_temb_table_slices.py b/python/sglang/multimodal_gen/test/unit/test_wan_temb_table_slices.py new file mode 100644 index 000000000..de54a5463 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_wan_temb_table_slices.py @@ -0,0 +1,48 @@ +import unittest + +import torch + +from sglang.multimodal_gen.runtime.models.dits.wanvideo import ( + _eager_temb_table_slices, + _wan_temb_table_slices, +) + + +class TestWanTembTableSlices(unittest.TestCase): + @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") + def test_fused_matches_eager_and_is_contiguous(self): + torch.manual_seed(0) + for batch, seq, hidden in [(1, 517, 3072), (2, 64, 1536)]: + temb = torch.randn( + batch, seq, 6, hidden, device="cuda", dtype=torch.bfloat16 + ) + table = torch.randn(1, 6, hidden, device="cuda", dtype=torch.float32) + reference = _eager_temb_table_slices(table, temb) + fused = _wan_temb_table_slices(table, temb) + self.assertEqual(len(fused), 6) + for ref, out in zip(reference, fused): + self.assertTrue(torch.equal(ref, out)) + self.assertTrue(out.is_contiguous()) + + @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") + def test_fp32_temb_also_supported(self): + torch.manual_seed(1) + temb = torch.randn(1, 33, 6, 512, device="cuda", dtype=torch.float32) + table = torch.randn(1, 6, 512, device="cuda", dtype=torch.float32) + reference = _eager_temb_table_slices(table, temb) + fused = _wan_temb_table_slices(table, temb) + for ref, out in zip(reference, fused): + self.assertTrue(torch.equal(ref, out)) + + def test_cpu_falls_back_to_eager(self): + torch.manual_seed(2) + temb = torch.randn(1, 9, 6, 64, dtype=torch.bfloat16) + table = torch.randn(1, 6, 64, dtype=torch.float32) + reference = _eager_temb_table_slices(table, temb) + out = _wan_temb_table_slices(table, temb) + for ref, got in zip(reference, out): + self.assertTrue(torch.equal(ref, got)) + + +if __name__ == "__main__": + unittest.main()