[diffusion] Wan2.2-TI2V: fuse per-token adaLN table add into contiguous slices + hoist rope cache (denoise -13.1% H100 / -12.6% H200, bit-exact; eager beats compile) (#34584)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
2026-08-13 17:26:46 +08:00
parent ba23846ccf
commit a23670ddbf
3 changed files with 250 additions and 25 deletions
@@ -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
@@ -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:
@@ -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()