[Diffusion] Fuse LTX2 Ada values (#29390)

python/sglang/jit_kernel/diffusion/triton/ltx2_ada_values.py

@@ -0,0 +1,184 @@
+# Adapted from NVlabs/Sana sol-engine LTX2 Ada-value fusion.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _ltx2_ada_values9_kernel(
+    temb_ptr,
+    table_ptr,
+    out0_ptr,
+    out1_ptr,
+    out2_ptr,
+    out3_ptr,
+    out4_ptr,
+    out5_ptr,
+    out6_ptr,
+    out7_ptr,
+    out8_ptr,
+    rows: tl.constexpr,
+    hidden: tl.constexpr,
+    total_params: tl.constexpr,
+    table_stride_p: tl.constexpr,
+    table_stride_d: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    cols = tl.arange(0, BLOCK_N)
+    mask = cols < hidden
+    temb_row = temb_ptr + row * total_params * hidden
+    base = row * hidden + cols
+
+    table0 = tl.load(
+        table_ptr + 0 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb0 = tl.load(
+        temb_row + (0 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table1 = tl.load(
+        table_ptr + 1 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb1 = tl.load(
+        temb_row + (1 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table2 = tl.load(
+        table_ptr + 2 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb2 = tl.load(
+        temb_row + (2 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table3 = tl.load(
+        table_ptr + 3 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb3 = tl.load(
+        temb_row + (3 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table4 = tl.load(
+        table_ptr + 4 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb4 = tl.load(
+        temb_row + (4 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table5 = tl.load(
+        table_ptr + 5 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb5 = tl.load(
+        temb_row + (5 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table6 = tl.load(
+        table_ptr + 6 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb6 = tl.load(
+        temb_row + (6 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table7 = tl.load(
+        table_ptr + 7 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb7 = tl.load(
+        temb_row + (7 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table8 = tl.load(
+        table_ptr + 8 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb8 = tl.load(
+        temb_row + (8 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+
+    tl.store(out0_ptr + base, (table0 + temb0).to(tl.bfloat16), mask=mask)
+    tl.store(out1_ptr + base, (table1 + temb1).to(tl.bfloat16), mask=mask)
+    tl.store(out2_ptr + base, (table2 + temb2).to(tl.bfloat16), mask=mask)
+    tl.store(out3_ptr + base, (table3 + temb3).to(tl.bfloat16), mask=mask)
+    tl.store(out4_ptr + base, (table4 + temb4).to(tl.bfloat16), mask=mask)
+    tl.store(out5_ptr + base, (table5 + temb5).to(tl.bfloat16), mask=mask)
+    tl.store(out6_ptr + base, (table6 + temb6).to(tl.bfloat16), mask=mask)
+    tl.store(out7_ptr + base, (table7 + temb7).to(tl.bfloat16), mask=mask)
+    tl.store(out8_ptr + base, (table8 + temb8).to(tl.bfloat16), mask=mask)
+
+
+def ltx2_ada_values9(
+    scale_shift_table: torch.Tensor,
+    timestep: torch.Tensor,
+) -> tuple[torch.Tensor, ...]:
+    if timestep.ndim != 3:
+        raise ValueError("timestep must have shape [B, S, 9 * D]")
+    if not timestep.is_cuda or timestep.dtype != torch.bfloat16:
+        raise ValueError("timestep must be a CUDA bfloat16 tensor")
+    if not timestep.is_contiguous():
+        raise ValueError("timestep must be contiguous")
+    if scale_shift_table.ndim != 2 or scale_shift_table.shape[0] != 9:
+        raise ValueError("scale_shift_table must have shape [9, D]")
+    if (
+        not scale_shift_table.is_cuda
+        or scale_shift_table.dtype not in (torch.bfloat16, torch.float32)
+        or scale_shift_table.stride(-1) != 1
+    ):
+        raise ValueError(
+            "scale_shift_table must be CUDA, bf16/fp32, last-dim contiguous"
+        )
+
+    total_params = int(scale_shift_table.shape[0])
+    hidden = int(scale_shift_table.shape[1])
+    if hidden <= 0 or timestep.shape[-1] != total_params * hidden:
+        raise ValueError("timestep last dim must equal 9 * hidden")
+    if hidden % 256 != 0 or hidden > 8192:
+        raise ValueError("hidden size is outside the supported LTX2 fast-path range")
+
+    batch, seq, _ = timestep.shape
+    rows = int(batch * seq)
+    outs = tuple(
+        torch.empty((batch, seq, hidden), device=timestep.device, dtype=timestep.dtype)
+        for _ in range(9)
+    )
+    _ltx2_ada_values9_kernel[(rows,)](
+        timestep,
+        scale_shift_table,
+        *outs,
+        rows,
+        hidden,
+        total_params,
+        scale_shift_table.stride(0),
+        scale_shift_table.stride(1),
+        BLOCK_N=triton.next_power_of_2(hidden),
+        num_warps=4 if hidden >= 4096 else 8,
+    )
+    return outs

python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py

@@ -49,6 +49,8 @@ logger = init_logger(__name__)
 ADALN_NUM_BASE_PARAMS = 6
 ADALN_NUM_CROSS_ATTN_PARAMS = 3
 
+_LTX2_FUSED_ADA_VALUES_RUNTIME_DISABLED = False
+
 
 def adaln_embedding_coefficient(cross_attention_adaln: bool) -> int:
     return ADALN_NUM_BASE_PARAMS + (
@@ -56,6 +58,48 @@ def adaln_embedding_coefficient(cross_attention_adaln: bool) -> int:
     )
 
 
+def _ltx2_disable_fused_ada_values(exc: Exception) -> None:
+    global _LTX2_FUSED_ADA_VALUES_RUNTIME_DISABLED
+    _LTX2_FUSED_ADA_VALUES_RUNTIME_DISABLED = True
+    logger.warning_once(f"Disabling LTX2 fused Ada values fast path: {exc}")
+
+
+def _ltx2_try_fused_ada_values9(
+    scale_shift_table: torch.Tensor,
+    batch_size: int,
+    timestep: torch.Tensor,
+) -> tuple[torch.Tensor, ...] | None:
+    if (
+        _LTX2_FUSED_ADA_VALUES_RUNTIME_DISABLED
+        or get_tp_world_size() != 1
+        or not timestep.is_cuda
+        or timestep.dtype != torch.bfloat16
+        or timestep.ndim != 3
+        or int(timestep.shape[0]) != int(batch_size)
+        or not timestep.is_contiguous()
+        or not scale_shift_table.is_cuda
+        or scale_shift_table.dtype not in (torch.bfloat16, torch.float32)
+        or scale_shift_table.ndim != 2
+        or int(scale_shift_table.shape[0]) != 9
+        or scale_shift_table.stride(-1) != 1
+    ):
+        return None
+
+    hidden = int(scale_shift_table.shape[1])
+    if hidden % 256 != 0 or hidden > 8192 or timestep.shape[-1] != 9 * hidden:
+        return None
+
+    try:
+        from sglang.jit_kernel.diffusion.triton.ltx2_ada_values import (
+            ltx2_ada_values9,
+        )
+
+        return ltx2_ada_values9(scale_shift_table, timestep)
+    except Exception as exc:
+        _ltx2_disable_fused_ada_values(exc)
+        return None
+
+
 def _ltx2_is_perturbed(
     perturbation_config: dict[str, object],
     key: str,
@@ -1037,13 +1081,21 @@ class LTX2TransformerBlock(nn.Module):
         audio_replicated_for_sp: bool = False,
         video_memory_prefix_len: int = 0,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-
         batch_size = hidden_states.size(0)
+        video_ada_values = _ltx2_try_fused_ada_values9(
+            self.scale_shift_table, batch_size, temb
+        )
+        audio_ada_values = _ltx2_try_fused_ada_values9(
+            self.audio_scale_shift_table, batch_size, temb_audio
+        )
 
         # 1. Video and Audio Self-Attention
-        vshift_msa, vscale_msa, vgate_msa = self.get_ada_values(
+        if video_ada_values is None:
-            self.scale_shift_table, batch_size, temb, slice(0, 3)
+            vshift_msa, vscale_msa, vgate_msa = self.get_ada_values(
-        )
+                self.scale_shift_table, batch_size, temb, slice(0, 3)
+            )
+        else:
+            vshift_msa, vscale_msa, vgate_msa = video_ada_values[0:3]
         norm_hidden_states = (
             self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_msa) + vshift_msa
         )
@@ -1058,9 +1110,12 @@ class LTX2TransformerBlock(nn.Module):
         )
         hidden_states = hidden_states + attn_hidden_states * vgate_msa
 
-        ashift_msa, ascale_msa, agate_msa = self.get_ada_values(
+        if audio_ada_values is None:
-            self.audio_scale_shift_table, batch_size, temb_audio, slice(0, 3)
+            ashift_msa, ascale_msa, agate_msa = self.get_ada_values(
-        )
+                self.audio_scale_shift_table, batch_size, temb_audio, slice(0, 3)
+            )
+        else:
+            ashift_msa, ascale_msa, agate_msa = audio_ada_values[0:3]
         norm_audio_hidden_states = (
             self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_msa)
             + ashift_msa
@@ -1081,9 +1136,12 @@ class LTX2TransformerBlock(nn.Module):
                 raise ValueError(
                     "cross_attention_adaln requires prompt modulation tensors."
                 )
-            vshift_q, vscale_q, vgate_q = self.get_ada_values(
+            if video_ada_values is None:
-                self.scale_shift_table, batch_size, temb, slice(6, 9)
+                vshift_q, vscale_q, vgate_q = self.get_ada_values(
-            )
+                    self.scale_shift_table, batch_size, temb, slice(6, 9)
+                )
+            else:
+                vshift_q, vscale_q, vgate_q = video_ada_values[6:9]
             v_prompt_shift, v_prompt_scale = self.get_ada_values(
                 self.prompt_scale_shift_table, batch_size, temb_prompt, slice(None)
             )
@@ -1100,9 +1158,12 @@ class LTX2TransformerBlock(nn.Module):
             )
             hidden_states = hidden_states + attn_hidden_states * vgate_q
 
-            ashift_q, ascale_q, agate_q = self.get_ada_values(
+            if audio_ada_values is None:
-                self.audio_scale_shift_table, batch_size, temb_audio, slice(6, 9)
+                ashift_q, ascale_q, agate_q = self.get_ada_values(
-            )
+                    self.audio_scale_shift_table, batch_size, temb_audio, slice(6, 9)
+                )
+            else:
+                ashift_q, ascale_q, agate_q = audio_ada_values[6:9]
             a_prompt_shift, a_prompt_scale = self.get_ada_values(
                 self.audio_prompt_scale_shift_table,
                 batch_size,
@@ -1251,18 +1312,24 @@ class LTX2TransformerBlock(nn.Module):
                 audio_hidden_states + v2a_gate * v2a_attn_hidden_states
             )
         # 4. Feedforward
-        vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values(
+        if video_ada_values is None:
-            self.scale_shift_table, batch_size, temb, slice(3, 6)
+            vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values(
-        )
+                self.scale_shift_table, batch_size, temb, slice(3, 6)
+            )
+        else:
+            vshift_mlp, vscale_mlp, vgate_mlp = video_ada_values[3:6]
         norm_hidden_states = (
             self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_mlp) + vshift_mlp
         )
         ff_output = self.ff(norm_hidden_states)
         hidden_states = hidden_states + ff_output * vgate_mlp
 
-        ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values(
+        if audio_ada_values is None:
-            self.audio_scale_shift_table, batch_size, temb_audio, slice(3, 6)
+            ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values(
-        )
+                self.audio_scale_shift_table, batch_size, temb_audio, slice(3, 6)
+            )
+        else:
+            ashift_mlp, ascale_mlp, agate_mlp = audio_ada_values[3:6]
         norm_audio_hidden_states = (
             self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_mlp)
             + ashift_mlp
@@ -1673,7 +1740,6 @@ class LTX2VideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
         late_audio_self_attention_mask: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
-
         batch_size = hidden_states.size(0)
         audio_timestep = audio_timestep if audio_timestep is not None else timestep
 

test/registered/jit/diffusion/test_ltx2_ada_values.py

@@ -0,0 +1,69 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.jit_kernel.diffusion.triton.ltx2_ada_values import ltx2_ada_values9
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=8, suite="base-b-kernel-unit-1-gpu-large")
+
+DEVICE = "cuda"
+
+
+@pytest.fixture(autouse=True)
+def cuda_setup():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    torch.cuda.manual_seed(0)
+
+
+def _reference(
+    scale_shift_table: torch.Tensor,
+    timestep: torch.Tensor,
+) -> tuple[torch.Tensor, ...]:
+    batch, seq, _ = timestep.shape
+    hidden = scale_shift_table.shape[1]
+    return (
+        scale_shift_table.to(device=timestep.device, dtype=timestep.dtype)
+        .view(1, 1, 9, hidden)
+        .add(timestep.reshape(batch, seq, 9, hidden))
+        .unbind(dim=2)
+    )
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("batch,seq,hidden", [(1, 1, 4096), (2, 3, 2048)])
+@pytest.mark.parametrize("table_dtype", [torch.bfloat16, torch.float32])
+def test_ltx2_ada_values9(
+    batch: int,
+    seq: int,
+    hidden: int,
+    table_dtype: torch.dtype,
+) -> None:
+    scale_shift_table = torch.randn(
+        9, hidden, device=DEVICE, dtype=table_dtype
+    ).contiguous()
+    timestep = torch.randn(
+        batch, seq, 9 * hidden, device=DEVICE, dtype=torch.bfloat16
+    ).contiguous()
+
+    actual = ltx2_ada_values9(scale_shift_table, timestep)
+    expected = _reference(scale_shift_table, timestep)
+
+    assert len(actual) == 9
+    for actual_value, expected_value in zip(actual, expected):
+        torch.testing.assert_close(actual_value, expected_value, atol=0, rtol=0)
+
+
+@torch.no_grad()
+def test_ltx2_ada_values9_rejects_unsupported_shape() -> None:
+    scale_shift_table = torch.randn(8, 4096, device=DEVICE, dtype=torch.bfloat16)
+    timestep = torch.randn(1, 1, 9 * 4096, device=DEVICE, dtype=torch.bfloat16)
+
+    with pytest.raises(ValueError, match="scale_shift_table"):
+        ltx2_ada_values9(scale_shift_table, timestep)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))