[kernels] Reorganize ops/diffusion by operator domain behind a lazy facade (#35114)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

test/registered/kernels/benchmark/diffusion/bench_causal_conv3d_cat_pad.py

@@ -3,12 +3,12 @@ from dataclasses import dataclass
 import torch
 
 from sglang.kernels.jit.benchmark import marker
-from sglang.kernels.ops.diffusion.causal_conv3d_cat_pad import (
-    fused_causal_conv3d_cat_pad_cuda,
-)
-from sglang.kernels.ops.diffusion.triton.causal_conv3d_pad import (
+from sglang.kernels.ops.diffusion import (
     fused_causal_conv3d_cat_pad as fused_causal_conv3d_cat_pad_triton,
 )
+from sglang.kernels.ops.diffusion import (
+    fused_causal_conv3d_cat_pad_cuda,
+)
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(

test/registered/kernels/benchmark/diffusion/bench_group_norm_silu.py

@@ -10,7 +10,7 @@ import torch
 import torch.nn.functional as F
 import triton.testing
 
-from sglang.kernels.ops.diffusion.triton.group_norm_silu import triton_group_norm_silu
+from sglang.kernels.ops.diffusion import triton_group_norm_silu
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.utils import is_in_ci
 

test/registered/kernels/benchmark/diffusion/bench_ltx2_qknorm_split_rope.py

@@ -4,9 +4,7 @@ from dataclasses import dataclass
 
 import torch
 
-from sglang.kernels.ops.diffusion.ltx2_qknorm_split_rope import (
-    ltx2_qknorm_split_rope_cuda,
-)
+from sglang.kernels.ops.diffusion import ltx2_qknorm_split_rope_cuda
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci
 

test/registered/kernels/benchmark/diffusion/bench_norm_impls.py

@@ -15,8 +15,11 @@ import torch.nn.functional as F
 
 from sglang.kernels.jit.benchmark.utils import DEFAULT_DEVICE
 from sglang.kernels.jit.utils import KERNEL_PATH
-from sglang.kernels.ops.diffusion.triton.norm import norm_infer, rms_norm_fn
-from sglang.kernels.ops.diffusion.triton.rmsnorm_onepass import triton_one_pass_rms_norm
+from sglang.kernels.ops.diffusion import (
+    norm_infer,
+    rms_norm_fn,
+    triton_one_pass_rms_norm,
+)
 from sglang.kernels.ops.layernorm.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm
 from sglang.kernels.ops.layernorm.norm import rmsnorm as jit_rmsnorm
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

test/registered/kernels/benchmark/diffusion/bench_qknorm_rope.py

@@ -153,7 +153,7 @@ def split_qknorm_rope(inputs: dict[str, torch.Tensor | bool]) -> None:
 
 
 def fused_qknorm_rope(inputs: dict[str, torch.Tensor | bool]) -> None:
-    from sglang.kernels.ops.diffusion.qknorm_rope import fused_inplace_qknorm_rope
+    from sglang.kernels.ops.diffusion import fused_inplace_qknorm_rope
 
     fused_inplace_qknorm_rope(
         inputs["q"],

test/registered/kernels/benchmark/diffusion/bench_qwen_image_modulation.py

@@ -4,10 +4,10 @@ import torch
 import triton.testing
 
 from sglang.kernels.jit.benchmark.utils import run_benchmark_no_cudagraph
-from sglang.kernels.ops.diffusion.triton.norm import norm_infer
-from sglang.kernels.ops.diffusion.triton.scale_shift import (
+from sglang.kernels.ops.diffusion import (
     fuse_layernorm_scale_shift_gate_select01_kernel,
     fuse_residual_layernorm_scale_shift_gate_select01_kernel,
+    norm_infer,
 )
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.utils import is_in_ci

test/registered/kernels/benchmark/diffusion/bench_residual_gate_add.py

@@ -4,8 +4,7 @@ from dataclasses import dataclass
 
 import torch
 
-from sglang.kernels.ops.diffusion.residual_gate_add import residual_gate_add_cuda
-from sglang.kernels.ops.diffusion.triton.scale_shift import fuse_scale_shift_kernel
+from sglang.kernels.ops.diffusion import fuse_scale_shift_kernel, residual_gate_add_cuda
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci
 

test/registered/kernels/benchmark/diffusion/bench_timestep_embedding.py

@@ -0,0 +1,70 @@
+"""CUDA vs eager-diffusers timing for the sinusoidal timestep embedding.
+
+Moved out of ``ops/diffusion/test_timestep_embedding.py``: it asserted nothing
+and was skipped unless ``SGLANG_RUN_JIT_KERNEL_PERF_TESTS=1``, so it belonged
+with the other benchmarks rather than in the correctness suite.
+"""
+
+import sys
+
+import torch
+
+from sglang.kernels.ops.diffusion import timestep_embedding
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(
+    est_time=20, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)
+
+BATCHES = [1, 2, 8, 63, 256, 512, 613, 1024, 1536]
+DIMS = [32, 64, 128, 256, 512, 1024, 2048, 4096]
+
+
+def _reference(timesteps, dim, max_period=10000):
+    half_dim = dim // 2
+    exponent = -torch.log(
+        torch.tensor(max_period, dtype=torch.float32, device=timesteps.device)
+    ) * torch.arange(0, half_dim, dtype=torch.float32, device=timesteps.device)
+    emb = torch.exp(exponent / (half_dim - 1))
+    emb = timesteps[:, None].float() * emb[None, :]
+    return torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+
+
+def _time_ms(fn, *args, warmup=4, repeat=20):
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    for _ in range(warmup):
+        fn(*args)
+    torch.cuda.synchronize()
+    start.record()
+    for _ in range(repeat):
+        fn(*args)
+    end.record()
+    end.synchronize()
+    return start.elapsed_time(end) / repeat
+
+
+def main() -> int:
+    if not torch.cuda.is_available():
+        print("CUDA required")
+        return 0
+    print(f"{'batch':>7} {'dim':>6} {'torch(ms)':>11} {'cuda(ms)':>10} {'speedup':>8}")
+    speedups = []
+    for batch in BATCHES:
+        for dim in DIMS:
+            t = torch.linspace(
+                0, max(100000, batch), steps=batch, device="cuda", dtype=torch.float32
+            )
+            torch_ms = _time_ms(_reference, t, dim)
+            cuda_ms = _time_ms(timestep_embedding, t, dim)
+            speedups.append(torch_ms / cuda_ms)
+            print(
+                f"{batch:>7} {dim:>6} {torch_ms:>11.6f} {cuda_ms:>10.6f} "
+                f"{speedups[-1]:>8.3f}"
+            )
+    print(f"average speedup: {sum(speedups) / len(speedups):.4f}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

test/registered/kernels/ops/diffusion/test_activation.py

@@ -0,0 +1,112 @@
+"""``diffusion.activation``: activation-function fusions.
+
+All of these are bit-exact by construction -- they are elementwise chains with
+no reduction, so reproducing aten's per-op fp32-opmath / round-to-bf16
+boundaries is enough and ``torch.equal`` is the assertion.
+
+The cublasLt linear+tanh-GELU epilogue is *not* here: it is not bit-exact and
+is therefore quality-gated, so it is tested through its mount protocol in
+``test_sites.py``.
+"""
+
+import sys
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from sglang.kernels.ops.diffusion import (
+    can_use_fused_bias_glu,
+    can_use_fused_bias_silu,
+    can_use_fused_silu_mul,
+    fused_bias_glu,
+    fused_bias_silu,
+    fused_packed_silu_mul_bitexact,
+    fused_silu_mul_bitexact,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=6, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+@pytest.mark.parametrize("channels", [2240, 11200])
+def test_sana_bias_silu_is_bit_exact(channels):
+    torch.manual_seed(0)
+    x = torch.randn(
+        (1, channels, 7, 5),
+        device="cuda",
+        dtype=torch.bfloat16,
+    ).to(memory_format=torch.channels_last)
+    bias = torch.randn(channels, device="cuda", dtype=torch.bfloat16)
+
+    assert can_use_fused_bias_silu(x, bias)
+    actual = fused_bias_silu(x, bias)
+    expected = F.silu(x + bias[None, :, None, None])
+
+    assert actual.is_contiguous(memory_format=torch.channels_last)
+    assert torch.equal(actual, expected)
+
+
+@pytest.mark.parametrize("channels", [2240, 5600])
+def test_sana_bias_glu_is_bit_exact(channels):
+    torch.manual_seed(1)
+    x = torch.randn(
+        (1, 2 * channels, 7, 5),
+        device="cuda",
+        dtype=torch.bfloat16,
+    ).to(memory_format=torch.channels_last)
+    bias = torch.randn(2 * channels, device="cuda", dtype=torch.bfloat16)
+
+    assert can_use_fused_bias_glu(x, bias)
+    actual = fused_bias_glu(x, bias)
+    biased = x + bias[None, :, None, None]
+    hidden, gate = torch.chunk(biased, 2, dim=1)
+    expected = hidden * F.silu(gate)
+
+    assert actual.is_contiguous(memory_format=torch.channels_last)
+    assert torch.equal(actual, expected)
+
+
+# ---------------------------------------------------------------------------
+# silu(a) * b for split-projection SwiGLU
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("shape", [(1, 4096, 3072), (2, 17, 512)])
+def test_silu_mul_is_bit_exact(shape):
+    # Separate gate/up GEMMs, so the concatenated ``silu_and_mul`` kernels do
+    # not apply without an extra full-width cat -- this kernel replaces the
+    # eager ``F.silu(a) * b`` pair instead.
+    torch.manual_seed(0)
+    a = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
+    b = torch.randn_like(a)
+
+    assert can_use_fused_silu_mul(a, b)
+    assert torch.equal(fused_silu_mul_bitexact(a, b), F.silu(a) * b)
+
+
+@pytest.mark.parametrize("hidden", [384, 3072])
+@pytest.mark.parametrize("strided", [False, True])
+def test_packed_silu_mul_is_bit_exact(hidden, strided):
+    # The packed form splits one [.., 2 * hidden] projection in-kernel; it must
+    # accept the strided view a wider projection slice produces.
+    torch.manual_seed(1)
+    if strided:
+        x = torch.randn(1, 19, 3 * hidden, device="cuda", dtype=torch.bfloat16)
+        x = x[..., : 2 * hidden]
+    else:
+        x = torch.randn(1, 19, 2 * hidden, device="cuda", dtype=torch.bfloat16)
+
+    expected = F.silu(x[..., :hidden]) * x[..., hidden:]
+    assert torch.equal(fused_packed_silu_mul_bitexact(x), expected)
+
+
+def test_silu_mul_rejects_mismatched_operands():
+    a = torch.randn(1, 8, 64, device="cuda", dtype=torch.bfloat16)
+    assert not can_use_fused_silu_mul(a, a.float())  # mixed dtypes
+    assert not can_use_fused_silu_mul(a, a[:, :-1])  # mismatched shapes
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/diffusion/test_autoencoder_kl_fastpath.py

@@ -1,62 +0,0 @@
-"""Install-path checks for the generic AutoencoderKL CUDA fast path."""
-
-import sys
-
-import pytest
-import torch
-
-from sglang.multimodal_gen.configs.models.vaes.stablediffusion3 import (
-    StableDiffusion3VAEConfig,
-)
-from sglang.multimodal_gen.runtime.models.vaes import flux2_vae_cuda_opt as vae_opt
-from sglang.multimodal_gen.runtime.models.vaes.autoencoder import AutoencoderKL
-from sglang.multimodal_gen.runtime.models.vaes.fast_path_gate import (
-    use_vae_fast_path,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-
-
-def _small_config():
-    config = StableDiffusion3VAEConfig()
-    config.arch_config.latent_channels = 2
-    config.arch_config.block_out_channels = (4, 4)
-    config.arch_config.down_block_types = ("DownEncoderBlock2D",) * 2
-    config.arch_config.up_block_types = ("UpDecoderBlock2D",) * 2
-    config.arch_config.layers_per_block = 1
-    config.arch_config.norm_num_groups = 1
-    config.arch_config.sample_size = 8
-    return config
-
-
-@torch.no_grad()
-def test_autoencoder_kl_fastpath_install():
-    torch.manual_seed(0)
-    vae = AutoencoderKL(_small_config()).to("cuda", torch.bfloat16).eval()
-    ref_names = {n for n, _ in vae.named_parameters()}
-    ref_sd = {k: v.clone() for k, v in vae.state_dict().items()}
-    z = torch.randn(1, 2, 8, 8, device="cuda", dtype=torch.bfloat16)
-    ref = vae.decode(z)
-
-    opt = vae_opt.maybe_optimize_autoencoder_kl(vae)
-    # Wrappers must not change parameter FQNs; strict load must round-trip.
-    assert {n for n, _ in opt.named_parameters()} == ref_names
-    opt.load_state_dict(ref_sd, strict=True)
-    # Gate off: bit-for-bit the original path.
-    assert torch.equal(opt.decode(z), ref)
-    # use_vae_fast_path() is a no-op when nothing registered a gate, so check
-    # the wrappers went in before relying on it to switch paths.
-    assert any(
-        isinstance(m, (vae_opt.FusedGroupNormSiLU, vae_opt.FusedUpsample2xConv2d))
-        for m in opt.modules()
-    )
-    # Gate on: fast path runs and stays close; leaving the scope restores exact.
-    with use_vae_fast_path(opt, True):
-        torch.testing.assert_close(opt.decode(z).float(), ref.float(), atol=0.1, rtol=0)
-    assert torch.equal(opt.decode(z), ref)
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/diffusion/test_bitexact_gate.py

@@ -1,159 +0,0 @@
-import sys
-from types import ModuleType
-from unittest.mock import MagicMock, patch
-
-import pytest
-import torch
-
-from sglang.kernels.ops.diffusion.bitexact_gate import (
-    BitExactFusionGate,
-    flashinfer_rmsnorm_diagnostic_hint,
-    tensors_equal,
-)
-from sglang.test.ci.ci_register import register_cpu_ci
-from sglang.test.test_utils import CustomTestCase
-
-register_cpu_ci(est_time=2, suite="base-a-test-cpu")
-
-
-def test_bitexact_gate_once_mode_verifies_then_reuses():
-    gate = BitExactFusionGate("once")
-    calls = {"fused": 0, "ref": 0}
-
-    def fused():
-        calls["fused"] += 1
-        return torch.tensor([1.0])
-
-    def ref():
-        calls["ref"] += 1
-        return torch.tensor([1.0])
-
-    assert torch.equal(gate.accept_or_fallback(fused(), ref()), torch.tensor([1.0]))
-    assert gate.verified and not gate.disabled and calls == {"fused": 1, "ref": 1}
-    assert torch.equal(fused(), torch.tensor([1.0]))
-    assert calls == {"fused": 2, "ref": 1}
-
-
-def test_bitexact_gate_mismatch_disables_permanently():
-    gate = BitExactFusionGate("mismatch")
-
-    out = gate.accept_or_fallback(
-        torch.tensor([1.0]),
-        torch.tensor([2.0]),
-        mismatch_msg="mismatch",
-    )
-    assert torch.equal(out, torch.tensor([2.0]))
-    assert gate.disabled and not gate.verified
-
-
-def test_bitexact_gate_per_signature_tracks_each_sig():
-    gate = BitExactFusionGate("sig", per_signature=True)
-    a = torch.tensor([1.0])
-    assert torch.equal(gate.accept_or_fallback(a, a, sig=("a",)), a)
-    assert gate.is_verified(("a",))
-    assert not gate.is_verified(("b",))
-    assert torch.equal(gate.accept_or_fallback(a, a, sig=("b",)), a)
-    assert gate.verified_sigs == {("a",), ("b",)}
-
-
-def test_bitexact_gate_skips_first_sight_during_graph_capture(monkeypatch):
-    # Negative-branch contract: an unverified gate must not attempt first-sight
-    # verification inside CUDA graph capture — the eager-reference host sync
-    # would abort the capture (and BCG would permanently block the signature).
-    gate = BitExactFusionGate("capture")
-    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
-    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
-    assert not gate.can_attempt_once()
-    # A verified gate replays the fused kernel alone, which is capture-safe.
-    gate.mark_verified()
-    assert gate.can_attempt_once()
-
-
-def test_tensors_equal_supports_sequences():
-    assert tensors_equal(
-        (torch.tensor([1.0]), torch.tensor([2.0])),
-        (torch.tensor([1.0]), torch.tensor([2.0])),
-    )
-    assert not tensors_equal(
-        (torch.tensor([1.0]), torch.tensor([2.0])),
-        (torch.tensor([1.0]), torch.tensor([3.0])),
-    )
-
-
-class TestBitExactFallbackDiagnostics(CustomTestCase):
-    def test_mismatch_warning_is_actionable_and_diagnostic_is_lazy(self):
-        logger = MagicMock()
-        diagnostic = MagicMock(return_value="backend=CuTe DSL")
-        gate = BitExactFusionGate("diagnostic")
-
-        matched = gate.accept_or_fallback(
-            torch.tensor([1.0]),
-            torch.tensor([1.0]),
-            logger=logger,
-            diagnostic_hint=diagnostic,
-        )
-        self.assertTrue(torch.equal(matched, torch.tensor([1.0])))
-        diagnostic.assert_not_called()
-        logger.warning_once.assert_not_called()
-
-        gate = BitExactFusionGate("diagnostic")
-        fallback = gate.accept_or_fallback(
-            torch.tensor([1.0]),
-            torch.tensor([2.0]),
-            logger=logger,
-            diagnostic_hint=diagnostic,
-        )
-
-        self.assertTrue(torch.equal(fallback, torch.tensor([2.0])))
-        diagnostic.assert_called_once_with()
-        warning = logger.warning_once.call_args.args[0]
-        self.assertIn("Correctness is preserved", warning)
-        self.assertIn("reference kernel or reduction-order change", warning)
-        self.assertIn("backend=CuTe DSL", warning)
-
-    def test_diagnostic_failure_cannot_break_the_eager_fallback(self):
-        logger = MagicMock()
-
-        def broken_diagnostic():
-            raise RuntimeError("diagnostics unavailable")
-
-        gate = BitExactFusionGate("diagnostic")
-        fallback = gate.accept_or_fallback(
-            torch.tensor([1.0]),
-            torch.tensor([2.0]),
-            logger=logger,
-            diagnostic_hint=broken_diagnostic,
-        )
-
-        self.assertTrue(torch.equal(fallback, torch.tensor([2.0])))
-        self.assertTrue(gate.disabled)
-        self.assertIn("Correctness is preserved", logger.warning_once.call_args.args[0])
-
-    def test_flashinfer_rmsnorm_hint_reports_backend_and_versions(self):
-        flashinfer = ModuleType("flashinfer")
-        flashinfer_norm = ModuleType("flashinfer.norm")
-        flashinfer_norm._USE_CUDA_NORM = False
-        versions = {
-            "flashinfer-python": "0.6.12",
-            "flashinfer-cubin": "0.6.12",
-            "flashinfer-jit-cache": "0.6.12+cu130",
-        }
-
-        with (
-            patch.dict(
-                sys.modules,
-                {"flashinfer": flashinfer, "flashinfer.norm": flashinfer_norm},
-            ),
-            patch("importlib.metadata.version", side_effect=versions.__getitem__),
-            patch.dict("os.environ", {"FLASHINFER_USE_CUDA_NORM": "0"}),
-        ):
-            hint = flashinfer_rmsnorm_diagnostic_hint()
-
-        self.assertIn("backend=CuTe DSL", hint)
-        self.assertIn("FLASHINFER_USE_CUDA_NORM=0", hint)
-        for package, version in versions.items():
-            self.assertIn(f"{package}={version}", hint)
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/diffusion/test_causal_conv3d_cat_pad.py

@@ -1,89 +0,0 @@
-import sys
-
-import pytest
-import torch
-
-from sglang.kernels.jit.utils import get_ci_test_range
-from sglang.kernels.ops.diffusion.causal_conv3d_cat_pad import (
-    fused_causal_conv3d_cat_pad_cuda,
-)
-from sglang.kernels.ops.diffusion.triton.causal_conv3d_pad import (
-    fused_causal_conv3d_cat_pad as fused_causal_conv3d_cat_pad_triton,
-)
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-
-register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
-register_amd_ci(est_time=10, stage="jit-kernel-unit", runner_config="amd")
-
-DEVICE = "cuda"
-DTYPE = torch.bfloat16
-
-COSMOS3_CASES = get_ci_test_range(
-    [
-        (1024, 1, 30, 52, 1),
-        (1024, 1, 30, 52, 2),
-        (1024, 2, 60, 104, 1),
-        (1024, 2, 60, 104, 2),
-        (512, 4, 120, 208, 1),
-        (512, 4, 120, 208, 2),
-        (256, 4, 240, 416, 1),
-        (256, 4, 240, 416, 2),
-    ],
-    [(1024, 1, 30, 52, 1), (512, 4, 120, 208, 2)],
-)
-
-
-def _make_inputs(
-    channels: int,
-    t_size: int,
-    h_size: int,
-    w_size: int,
-    cache_t: int,
-) -> tuple[torch.Tensor, torch.Tensor, tuple[int, ...]]:
-    generator = torch.Generator(device=DEVICE)
-    generator.manual_seed(channels * 1009 + t_size * 251 + h_size + cache_t)
-    x = torch.randn(
-        (1, channels, t_size, h_size, w_size),
-        device=DEVICE,
-        dtype=DTYPE,
-        generator=generator,
-    )
-    cache_x = torch.randn(
-        (1, channels, cache_t, h_size, w_size),
-        device=DEVICE,
-        dtype=DTYPE,
-        generator=generator,
-    )
-    padding = (1, 1, 1, 1, cache_t, 0)
-    return x, cache_x, padding
-
-
-@pytest.mark.parametrize("channels,t_size,h_size,w_size,cache_t", COSMOS3_CASES)
-def test_causal_conv3d_cat_pad(
-    channels: int,
-    t_size: int,
-    h_size: int,
-    w_size: int,
-    cache_t: int,
-) -> None:
-    x, cache_x, padding = _make_inputs(channels, t_size, h_size, w_size, cache_t)
-    actual = fused_causal_conv3d_cat_pad_cuda(x, cache_x, padding)
-    expected = fused_causal_conv3d_cat_pad_triton(x, cache_x, padding)
-    torch.testing.assert_close(actual, expected, atol=0, rtol=0)
-
-
-def test_causal_conv3d_cat_pad_torch_compile() -> None:
-    x, cache_x, padding = _make_inputs(1024, 1, 30, 52, 1)
-
-    @torch.compile(fullgraph=True)
-    def fn(x: torch.Tensor, cache_x: torch.Tensor) -> torch.Tensor:
-        return fused_causal_conv3d_cat_pad_cuda(x, cache_x, padding)
-
-    actual = fn(x, cache_x)
-    expected = fused_causal_conv3d_cat_pad_triton(x, cache_x, padding)
-    torch.testing.assert_close(actual, expected, atol=0, rtol=0)
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__]))

test/registered/kernels/ops/diffusion/test_ernie_norm_scale_shift.py

@@ -1,135 +0,0 @@
-"""ERNIE fused norm/scale/shift fast paths must stay bit-exact vs eager."""
-
-import sys
-from unittest.mock import patch
-
-import pytest
-import torch
-
-import sglang.multimodal_gen.runtime.models.dits.ernie_image as ernie_image
-from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
-from sglang.multimodal_gen.runtime.models.dits.ernie_image import (
-    _ernie_gated_norm_scale_shift,
-    _ernie_norm_scale_shift,
-    _ernie_qknorm_rope,
-    _ernie_qknorm_rope_reference,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=4, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-
-
-@pytest.mark.parametrize("shape", [(1, 4216, 4096), (2, 1140, 4096), (1, 128, 2048)])
-def test_fused_norm_scale_shift_is_bit_exact(shape):
-    # (1, 4216, 4096) is the real ERNIE-Image shape (1024^2 image + text
-    # tokens, hidden 4096); 2048 covers the threads_per_row=32 regime.
-    torch.manual_seed(0)
-    batch, seq, hidden = shape
-    norm = RMSNorm(hidden, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
-    with torch.no_grad():
-        norm.weight.copy_(torch.randn(hidden))
-    x = torch.randn(batch, seq, hidden, device="cuda", dtype=torch.bfloat16)
-    residual = torch.randn_like(x)
-    update = torch.randn_like(x)
-    scale = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16) * 0.1
-    shift = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16) * 0.1
-    gate = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16)
-
-    with torch.no_grad():
-        out = _ernie_norm_scale_shift(norm, x, scale, shift)
-        ref = norm(x) * (1 + scale) + shift
-        assert torch.equal(out, ref)
-
-        out2, res = _ernie_gated_norm_scale_shift(
-            norm, residual, update, gate, scale, shift
-        )
-        res_ref = residual + gate * update
-        ref2 = norm(res_ref) * (1 + scale) + shift
-        assert torch.equal(res, res_ref)
-        assert torch.equal(out2, ref2)
-
-    # the fast paths must actually be in use (not silently disabled)
-    assert ernie_image._ERNIE_NORM.verified
-    assert ernie_image._ERNIE_GATED_NORM.verified
-    assert not ernie_image._ERNIE_NORM.disabled
-    assert not ernie_image._ERNIE_GATED_NORM.disabled
-
-
-def test_fused_qknorm_rope_is_bit_exact():
-    torch.manual_seed(1)
-    ernie_image._ERNIE_QKNORM_ROPE.disabled = False
-    ernie_image._ERNIE_QKNORM_ROPE.verified = False
-    batch, seq, heads, head_dim = 1, 257, 32, 128
-    q = torch.randn(batch, seq, heads, head_dim, device="cuda", dtype=torch.bfloat16)
-    k = torch.randn_like(q)
-    q_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
-    k_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
-    cos = torch.randn(seq, head_dim, device="cuda", dtype=torch.bfloat16)
-    sin = torch.randn_like(cos)
-    cache = torch.cat((cos, sin), dim=-1).contiguous()
-    positions = torch.arange(seq, device="cuda", dtype=torch.long)
-
-    q_ref, k_ref = _ernie_qknorm_rope_reference(
-        q.clone(), k.clone(), q_norm, k_norm, head_dim, cos, sin
-    )
-    q_out, k_out = _ernie_qknorm_rope(
-        q,
-        k,
-        q_norm,
-        k_norm,
-        head_dim,
-        cos,
-        sin,
-        cache,
-        positions,
-    )
-
-    assert torch.equal(q_out, q_ref)
-    assert torch.equal(k_out, k_ref)
-    assert ernie_image._ERNIE_QKNORM_ROPE.verified
-    assert not ernie_image._ERNIE_QKNORM_ROPE.disabled
-
-
-def test_qknorm_rope_first_attempt_exception_uses_pristine_inputs():
-    torch.manual_seed(2)
-    ernie_image._ERNIE_QKNORM_ROPE.disabled = False
-    ernie_image._ERNIE_QKNORM_ROPE.verified = False
-    batch, seq, heads, head_dim = 1, 17, 4, 128
-    q = torch.randn(batch, seq, heads, head_dim, device="cuda", dtype=torch.bfloat16)
-    k = torch.randn_like(q)
-    q_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
-    k_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
-    cos = torch.randn(seq, head_dim, device="cuda", dtype=torch.bfloat16)
-    sin = torch.randn_like(cos)
-    cache = torch.cat((cos, sin), dim=-1).contiguous()
-    positions = torch.arange(seq, device="cuda", dtype=torch.long)
-    q_ref, k_ref = _ernie_qknorm_rope_reference(
-        q.clone(), k.clone(), q_norm, k_norm, head_dim, cos, sin
-    )
-
-    def mutate_then_raise(**kwargs):
-        kwargs["q"].zero_()
-        kwargs["k"].zero_()
-        raise RuntimeError("synthetic kernel failure")
-
-    with patch.object(ernie_image, "apply_qk_norm_rope", mutate_then_raise):
-        q_out, k_out = _ernie_qknorm_rope(
-            q,
-            k,
-            q_norm,
-            k_norm,
-            head_dim,
-            cos,
-            sin,
-            cache,
-            positions,
-        )
-
-    assert torch.equal(q_out, q_ref)
-    assert torch.equal(k_out, k_ref)
-    assert ernie_image._ERNIE_QKNORM_ROPE.disabled
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__]))

test/registered/kernels/ops/diffusion/test_flux2_eager_fusions.py

@@ -1,101 +0,0 @@
-"""FLUX.2 eager fusions must be bit-exact for real packed/view layouts."""
-
-import unittest
-from unittest.mock import patch
-
-import torch
-import torch.nn.functional as F
-
-import sglang.multimodal_gen.runtime.models.dits.flux_2 as flux2
-from sglang.multimodal_gen.runtime.models.dits.flux_2 import (
-    _flux2_norm_modulate,
-    _flux2_swiglu,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-from sglang.test.test_utils import CustomTestCase
-
-register_cuda_ci(est_time=12, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-
-
-@unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
-class TestFlux2EagerFusions(CustomTestCase):
-    def setUp(self):
-        flux2._FLUX2_LN_MOD.disabled = False
-        flux2._FLUX2_LN_MOD.verified = False
-        flux2._FLUX2_LN_MOD_SIGS.clear()
-        flux2._FLUX2_SWIGLU.disabled = False
-        flux2._FLUX2_SWIGLU.verified = False
-        flux2._FLUX2_SWIGLU_SIGS.clear()
-
-    def test_norm_modulate_is_bit_exact_across_sequence_lengths(self):
-        torch.manual_seed(0)
-        hidden = 256
-        norm = torch.nn.LayerNorm(
-            hidden, eps=1e-6, elementwise_affine=False, device="cuda"
-        )
-        # FLUX.2 modulation values are views of one packed projection.
-        params = torch.randn(1, 1, 6 * hidden, device="cuda").bfloat16()
-        shift, scale = params.chunk(6, dim=-1)[:2]
-
-        for seq in (17, 65):
-            x = torch.randn(1, seq, hidden, device="cuda").bfloat16()
-            expected = norm(x) * (1 + scale) + shift
-            actual = _flux2_norm_modulate(norm, x, scale, shift)
-            self.assertTrue(torch.equal(actual, expected))
-
-        self.assertFalse(flux2._FLUX2_LN_MOD.disabled)
-        self.assertEqual(len(flux2._FLUX2_LN_MOD_SIGS), 1)
-
-    def test_packed_swiglu_is_bit_exact_for_contiguous_and_strided_views(self):
-        torch.manual_seed(1)
-        hidden = 384
-        inputs = [
-            torch.randn(1, 19, 2 * hidden, device="cuda").bfloat16(),
-            torch.randn(1, 19, 3 * hidden, device="cuda").bfloat16()[..., : 2 * hidden],
-        ]
-        for x in inputs:
-            expected = F.silu(x[..., :hidden]) * x[..., hidden:]
-            actual = _flux2_swiglu(x)
-            self.assertTrue(torch.equal(actual, expected))
-
-        self.assertFalse(flux2._FLUX2_SWIGLU.disabled)
-        self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 2)
-
-    def test_fp16_preserves_reference_path(self):
-        x = torch.randn(1, 17, 512, device="cuda", dtype=torch.float16)
-        expected = F.silu(x[..., :256]) * x[..., 256:]
-        actual = _flux2_swiglu(x)
-        self.assertTrue(torch.equal(actual, expected))
-        self.assertFalse(flux2._FLUX2_SWIGLU.disabled)
-
-    def test_packed_swiglu_rejects_non_dense_outer_stride(self):
-        base = torch.randn(2, 23, 512, device="cuda", dtype=torch.bfloat16)
-        x = base[:, :19]
-        self.assertNotEqual(x.stride(0), x.shape[1] * x.stride(1))
-
-        expected = F.silu(x[..., :256]) * x[..., 256:]
-        actual = _flux2_swiglu(x)
-        self.assertTrue(torch.equal(actual, expected))
-        self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 0)
-
-    def test_new_swiglu_signature_is_not_verified_during_graph_capture(self):
-        first = torch.randn(1, 17, 512, device="cuda", dtype=torch.bfloat16)
-        self.assertTrue(
-            torch.equal(
-                _flux2_swiglu(first),
-                F.silu(first[..., :256]) * first[..., 256:],
-            )
-        )
-        self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 1)
-
-        second = torch.randn(1, 19, 768, device="cuda", dtype=torch.bfloat16)
-        with patch("torch.cuda.is_current_stream_capturing", return_value=True):
-            actual = _flux2_swiglu(second)
-
-        expected = F.silu(second[..., :384]) * second[..., 384:]
-        self.assertTrue(torch.equal(actual, expected))
-        self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 1)
-
-
-if __name__ == "__main__":
-    unittest.main()

test/registered/kernels/ops/diffusion/test_flux2_vae_fastpath.py

@@ -1,64 +0,0 @@
-"""Focused correctness checks for the FLUX.2 VAE CUDA fast path."""
-
-import sys
-
-import pytest
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from diffusers.models.upsampling import Upsample2D
-
-from sglang.kernels.ops.diffusion.triton import group_norm_silu_twopass as gn_kernel
-from sglang.multimodal_gen.runtime.models.vaes import flux2_vae_cuda_opt as vae_opt
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-
-
-@torch.no_grad()
-def test_flux2_vae_fastpath():
-    torch.manual_seed(0)
-    gate = vae_opt.VaeFastPathGate()
-    gn = nn.GroupNorm(32, 128, eps=1e-6).to("cuda", torch.bfloat16)
-    x = torch.randn(1, 128, 64, 64, device="cuda", dtype=torch.bfloat16).to(
-        memory_format=torch.channels_last
-    )
-    ref = F.silu(gn(x))
-    fused_gn = vae_opt.FusedGroupNormSiLU(gn, gate)
-    assert set(fused_gn.state_dict()) == {"weight", "bias"}
-    assert torch.equal(fused_gn(x), ref)
-    assert (
-        gn_kernel.group_norm_silu_4d(x.contiguous(), gn.weight, gn.bias, 32, 1e-6)
-        is None
-    )
-    assert gn_kernel.group_norm_silu_4d(x, gn.weight.cpu(), gn.bias, 32, 1e-6) is None
-    assert (
-        gn_kernel.group_norm_silu_4d(x[..., :0, :], gn.weight, gn.bias, 32, 1e-6)
-        is None
-    )
-
-    gate.enabled = True
-    fast = fused_gn(x)
-    assert fast.is_contiguous(memory_format=torch.channels_last)
-    torch.testing.assert_close(fast.float(), ref.float(), atol=0.06, rtol=0)
-
-    gate.enabled = False
-    up = Upsample2D(channels=32, use_conv=True).to("cuda", torch.bfloat16)
-    fused_up = vae_opt.FusedUpsample2xConv2d(up, gate)
-    assert set(fused_up.state_dict()) == {"conv.weight", "conv.bias"}
-    x = torch.randn(2, 32, 33, 29, device="cuda", dtype=torch.bfloat16)
-    ref = up(x)
-    assert torch.equal(fused_up(x), ref)
-    assert fused_up._fused_weight is None
-
-    gate.enabled = True
-    fast = fused_up(x)
-    assert fused_up._fused_weight is not None
-    ref_range = ref.float().max() - ref.float().min()
-    relative_mse = F.mse_loss(fast.float(), ref.float()) / ref_range.square()
-    assert relative_mse < 3.2e-5
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/diffusion/test_flux_ln_modulate.py

@@ -1,75 +0,0 @@
-"""FLUX.1 fused LN+modulate fast path must stay bit-exact vs eager."""
-
-import pytest
-import torch
-
-import sglang.multimodal_gen.runtime.models.dits.flux as flux
-from sglang.kernels.ops.diffusion.fused_ln_modulate import (
-    mark_fused_ln_modulate_site,
-    mount_fused_ln_modulate,
-)
-from sglang.multimodal_gen.runtime.models.dits.flux import (
-    _flux_fused_ln_modulate,
-    _flux_norm_modulate,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=4, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-
-
-def _eager(norm, x, scale, shift):
-    return norm(x) * (1 + scale[:, None]) + shift[:, None]
-
-
-def _make_site_inputs(shape, chunks, seed):
-    torch.manual_seed(seed)
-    batch, seq, hidden = shape
-    norm = torch.nn.LayerNorm(hidden, eps=1e-6, elementwise_affine=False).cuda()
-    x = (torch.randn(batch, seq, hidden, device="cuda") * 8).bfloat16()
-    emb = torch.randn(batch, chunks * hidden, device="cuda").bfloat16()
-    parts = emb.chunk(chunks, dim=1)  # strided adaLN projection views
-    return norm, x, parts[0], parts[1]
-
-
-@pytest.mark.parametrize(
-    "shape,chunks",
-    [
-        ((1, 4096, 3072), 6),  # dual-stream image tokens (1024^2), chunk(6)
-        ((1, 512, 3072), 6),  # dual-stream text tokens
-        ((1, 4608, 3072), 3),  # single-stream concat, chunk(3)
-        ((2, 300, 3072), 6),  # CFG batch, odd seq
-    ],
-)
-def test_flux_fused_ln_modulate_is_bit_exact(shape, chunks):
-    # Every distinct (shape, stride, eps) signature the FLUX.1 sites emit
-    # must verify torch.equal on first sight and stay enabled.
-    norm, x, shift, scale = _make_site_inputs(shape, chunks, seed=0)
-    out = _flux_fused_ln_modulate(norm, x, scale, shift)
-    assert out is not None
-    assert torch.equal(out, _eager(norm, x, scale, shift))
-    assert not flux._FLUX_LN_MOD.disabled
-    assert flux._FLUX_LN_MOD.verified
-
-
-def test_flux_norm_modulate_bitexact_supersedes_high_fold():
-    # With the quality="high" affine fold mounted, the bit-exact kernel
-    # still takes priority, so the site output stays lossless.
-    site = torch.nn.Module()
-    mark_fused_ln_modulate_site(site)
-    assert mount_fused_ln_modulate(site)
-    norm, x, shift, scale = _make_site_inputs((1, 128, 3072), 6, seed=1)
-    out = _flux_norm_modulate(site, norm, x, scale, shift)
-    assert torch.equal(out, _eager(norm, x, scale, shift))
-
-
-def test_flux_fused_ln_modulate_rejects_unsupported_hidden():
-    # hidden % 4 != 0 is outside the kernel contract and must bail out.
-    norm, x, shift, scale = _make_site_inputs((1, 64, 3070), 6, seed=2)
-    assert _flux_fused_ln_modulate(norm, x, scale, shift) is None
-
-
-if __name__ == "__main__":
-    import sys
-
-    sys.exit(pytest.main([__file__]))

test/registered/kernels/ops/diffusion/test_flydsl_fused_norm.py

@@ -1,133 +0,0 @@
-import sys
-
-import pytest
-import torch
-import torch.nn.functional as F
-
-from sglang.test.ci.ci_register import register_amd_ci
-
-register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd")
-
-DEVICE = "cuda"
-D = 5120
-EPS = 1e-6
-
-
-def _ref_rms_norm(x_f32, weight, eps):
-    var = x_f32.pow(2).mean(-1, keepdim=True)
-    return x_f32 * torch.rsqrt(var + eps)
-
-
-def _ref_fused_residual_norm_ss(
-    residual, x, gate, weight, bias, scale, shift, norm_type, eps
-):
-    ref_res = residual.float() + x.float() * (gate.float() if gate is not None else 1)
-    ref_res_bf16 = ref_res.to(torch.bfloat16)
-    if norm_type == "layer":
-        normed = F.layer_norm(ref_res_bf16.float(), (D,), weight, bias, eps)
-    else:
-        normed = _ref_rms_norm(ref_res_bf16.float(), weight, eps) * weight.float()
-    y = (normed * (1.0 + scale.float()) + shift.float()).to(torch.bfloat16)
-    return y, ref_res_bf16
-
-
-def _ref_norm_ss(x, weight, bias, scale, shift, norm_type, eps):
-    if norm_type == "layer":
-        normed = F.layer_norm(x.float(), (D,), weight, bias, eps)
-    else:
-        normed = _ref_rms_norm(x.float(), weight, eps) * weight.float()
-    return (normed * (1.0 + scale.float()) + shift.float()).to(torch.bfloat16)
-
-
-@pytest.fixture(autouse=True)
-def cuda_setup():
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA required")
-    if not hasattr(torch.version, "hip") or not torch.version.hip:
-        pytest.skip("ROCm/HIP required for FlyDSL kernels")
-    torch.manual_seed(42)
-
-
-FUSED_CASES = [
-    ("rms", 1, 16),
-    ("rms", 2, 16),
-    ("layer", 2, 16),
-    ("rms", 1, 90000),
-]
-
-
-@pytest.mark.parametrize("norm_type,B,L", FUSED_CASES)
-def test_fused_residual_norm_scale_shift(norm_type, B, L):
-    from sglang.kernels.ops.diffusion.flydsl.fused_residual_norm import (
-        flydsl_fused_residual_norm_scale_shift,
-    )
-
-    residual = torch.randn(B, L, D, device=DEVICE, dtype=torch.bfloat16)
-    x = torch.randn(B, L, D, device=DEVICE, dtype=torch.bfloat16)
-    gate = torch.randn(B, 1, D, device=DEVICE, dtype=torch.bfloat16)
-    weight = torch.randn(D, device=DEVICE, dtype=torch.float32)
-    bias = (
-        torch.randn(D, device=DEVICE, dtype=torch.float32)
-        if norm_type == "layer"
-        else None
-    )
-    scale = torch.randn(B, 1, D, device=DEVICE, dtype=torch.bfloat16)
-    shift = torch.randn(B, 1, D, device=DEVICE, dtype=torch.bfloat16)
-
-    y, res_out = flydsl_fused_residual_norm_scale_shift(
-        residual,
-        x,
-        gate,
-        weight,
-        bias,
-        scale,
-        shift,
-        norm_type,
-        EPS,
-    )
-    y_ref, res_ref = _ref_fused_residual_norm_ss(
-        residual,
-        x,
-        gate,
-        weight,
-        bias,
-        scale,
-        shift,
-        norm_type,
-        EPS,
-    )
-    torch.testing.assert_close(res_out, res_ref, atol=5e-2, rtol=5e-2)
-    torch.testing.assert_close(y, y_ref, atol=1.0, rtol=5e-2)
-
-
-NSS_CASES = [
-    ("rms", 2, 16),
-    ("layer", 2, 16),
-    ("rms", 1, 90000),
-    ("layer", 1, 90000),
-]
-
-
-@pytest.mark.parametrize("norm_type,B,L", NSS_CASES)
-def test_norm_scale_shift(norm_type, B, L):
-    from sglang.kernels.ops.diffusion.flydsl.fused_residual_norm import (
-        flydsl_norm_scale_shift,
-    )
-
-    x = torch.randn(B, L, D, device=DEVICE, dtype=torch.bfloat16)
-    weight = torch.randn(D, device=DEVICE, dtype=torch.float32)
-    bias = (
-        torch.randn(D, device=DEVICE, dtype=torch.float32)
-        if norm_type == "layer"
-        else None
-    )
-    scale = torch.randn(B, 1, D, device=DEVICE, dtype=torch.bfloat16)
-    shift = torch.randn(B, 1, D, device=DEVICE, dtype=torch.bfloat16)
-
-    y = flydsl_norm_scale_shift(x, weight, bias, scale, shift, norm_type, EPS)
-    y_ref = _ref_norm_ss(x, weight, bias, scale, shift, norm_type, EPS)
-    torch.testing.assert_close(y, y_ref, atol=1.0, rtol=5e-2)
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/diffusion/test_fused_gate_rmsnorm.py

@@ -1,54 +0,0 @@
-"""Core checks for the quality-gated fused gate-RMSNorm path."""
-
-import sys
-
-import pytest
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from sglang.kernels.ops.diffusion import fused_gate_rmsnorm as fgn
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=4, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-DIM, EPS = 4608, 1e-5  # Ideogram 4 hidden size / norm_eps
-
-
-class _Site(nn.Module):
-    def __init__(self, dtype=torch.bfloat16):
-        super().__init__()
-        self.norm = nn.RMSNorm(DIM, eps=EPS, device="cuda", dtype=dtype)
-        fgn.mark_fused_gate_rmsnorm_site(self, ("norm",))
-
-
-def test_fused_matches_ideogram_reference():
-    torch.manual_seed(0)
-    site = _Site()
-    w = site.norm.weight.data
-    x = torch.randn(1, 64, DIM, device="cuda", dtype=torch.bfloat16)
-    residual = torch.randn_like(x)
-    # adaln-style strided chunks, as produced by Ideogram's modulation .chunk()
-    mods = torch.randn(1, 1, 2 * DIM, device="cuda", dtype=torch.bfloat16)
-    scale, gate = mods.chunk(2, dim=-1)
-    assert fgn.mount_fused_gate_rmsnorm(site)
-    got_scale = fgn.fused_rmsnorm_scale(x, w, 1.0 + scale, EPS)
-    got_gate = fgn.fused_rmsnorm_tanh_residual(x, gate, residual, w, EPS)
-    ref_scale = F.rms_norm(x, (DIM,), w, EPS) * (1.0 + scale)
-    ref_gate = residual + torch.tanh(gate) * F.rms_norm(x, (DIM,), w, EPS)
-    # fused path uses bf16-native norm statistics: close, not bit-exact
-    torch.testing.assert_close(got_scale, ref_scale, atol=8e-2, rtol=4e-2)
-    torch.testing.assert_close(got_gate, ref_gate, atol=8e-2, rtol=4e-2)
-
-
-def test_mount_guards_all_or_nothing():
-    good, bad = _Site(), _Site(torch.float32)
-    assert not fgn.mount_fused_gate_rmsnorm(nn.ModuleList([good, bad]))
-    assert not fgn.fused_gate_rmsnorm_active(good)
-    assert fgn.mount_fused_gate_rmsnorm(good)
-    fgn.unmount_fused_gate_rmsnorm(good)
-    assert not fgn.fused_gate_rmsnorm_active(good)
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__]))

test/registered/kernels/ops/diffusion/test_fused_linear_gelu.py

@@ -1,86 +0,0 @@
-"""Core checks for the quality-gated linear + tanh-GELU fusion."""
-
-import sys
-
-import pytest
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from sglang.kernels.ops.diffusion import fused_linear_gelu as gelu
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=4, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-
-
-class _Site(nn.Module):
-    def __init__(self, dtype=torch.bfloat16, bias=True):
-        super().__init__()
-        self.proj = nn.Linear(64, 256, bias=bias, device="cuda", dtype=dtype)
-        gelu.mark_fused_gelu_site(self, "proj")
-
-    def forward(self, x):
-        if gelu.fused_gelu_active(self) and gelu.can_fuse_linear_gelu(self.proj, x):
-            return gelu.fused_linear_gelu_tanh(x, self.proj.weight, self.proj.bias)
-        return F.gelu(self.proj(x), approximate="tanh")
-
-
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
-def test_fused_matches_reference(dtype):
-    torch.manual_seed(0)
-    site = _Site(dtype)
-    x = torch.randn(512, 64, device="cuda", dtype=dtype)
-    ref = site(x)
-    assert gelu.mount_fused_linear_gelu(site)
-    atol = 2e-2 if dtype == torch.bfloat16 else 4e-3
-    torch.testing.assert_close(site(x), ref, atol=atol, rtol=2e-2)
-
-
-def test_flux_gelu_proj_site():
-    """FLUX.1 shared-FF site: gate off is bit-exact, gate on is close."""
-    from sglang.multimodal_gen.runtime.models.dits.flux import FluxFusedGELUProj
-
-    torch.manual_seed(0)
-    proj = nn.Linear(3072, 12288, device="cuda", dtype=torch.bfloat16)
-    site = FluxFusedGELUProj(proj)
-    x = torch.randn(1, 512, 3072, device="cuda", dtype=torch.bfloat16)
-    ref = F.gelu(proj(x), approximate="tanh")
-
-    assert torch.equal(site(x), ref)  # unmounted default: bit-exact reference
-    assert gelu.mount_fused_linear_gelu(site)
-    torch.testing.assert_close(site(x), ref, atol=2e-2, rtol=2e-2)
-    gelu.unmount_fused_linear_gelu(site)
-    assert torch.equal(site(x), ref)
-
-
-def test_mount_guards_and_lossless_path():
-    torch.manual_seed(0)
-    good, bad = _Site(), _Site(torch.float32)
-    model = nn.ModuleList([good, bad])
-    assert not gelu.mount_fused_linear_gelu(model)
-    assert not gelu.fused_gelu_active(good)
-
-    x = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16)
-    ref = good(x)
-    assert gelu.mount_fused_linear_gelu(good)
-    gelu.unmount_fused_linear_gelu(good)
-    assert torch.equal(good(x), ref)
-
-    no_bias = nn.Linear(8, 8, bias=False, device="cuda", dtype=torch.bfloat16)
-    assert not gelu.can_fuse_linear_gelu_static(no_bias)
-    assert not gelu.can_fuse_linear_gelu(good.proj, x.float())
-
-
-@torch.no_grad()
-def test_mounted_site_torch_compile_fullgraph():
-    site = _Site()
-    x = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16)
-    assert gelu.mount_fused_linear_gelu(site)
-    expected = site(x)
-    actual = torch.compile(site, fullgraph=True)(x)
-    torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__]))

test/registered/kernels/ops/diffusion/test_fused_ln_modulate.py

@@ -1,82 +0,0 @@
-import pytest
-import torch
-import torch.nn as nn
-
-from sglang.kernels.ops.diffusion.fused_ln_modulate import (
-    can_fuse_ln_modulate,
-    fused_ln_modulate,
-    fused_ln_modulate_active,
-    mark_fused_ln_modulate_site,
-    mount_fused_ln_modulate,
-    unmount_fused_ln_modulate,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-
-
-@pytest.fixture(autouse=True)
-def cuda_setup():
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA required")
-    torch.cuda.manual_seed(0)
-
-
-@pytest.mark.parametrize("seq_len", [4096, 512])
-def test_fused_ln_modulate_matches_reference(seq_len):
-    x = torch.randn((1, seq_len, 3072), device="cuda", dtype=torch.bfloat16)
-    scale = torch.randn((1, 3072), device="cuda", dtype=torch.bfloat16)
-    shift = torch.randn_like(scale)
-    assert can_fuse_ln_modulate(x, scale, shift)
-    out = fused_ln_modulate(x, scale, shift, eps=1e-6)
-    norm = nn.LayerNorm(3072, eps=1e-6, elementwise_affine=False).cuda()
-    ref = norm(x) * (1 + scale[:, None]) + shift[:, None]
-    # Contract: bf16 rounding-order-level difference only, not bit-exact.
-    torch.testing.assert_close(out, ref, atol=0.0625, rtol=0.05)
-
-
-def test_fused_ln_modulate_guards_and_mount_protocol():
-    x = torch.randn((2, 64, 3072), device="cuda", dtype=torch.bfloat16)
-    row = torch.randn((2, 3072), device="cuda", dtype=torch.bfloat16)
-    assert not can_fuse_ln_modulate(x, row, row)  # folded affine needs B == 1
-    root = nn.Module()
-    root.child = nn.Module()
-    mark_fused_ln_modulate_site(root.child)
-    assert not fused_ln_modulate_active(root.child)
-    assert mount_fused_ln_modulate(root)
-    assert fused_ln_modulate_active(root.child)
-    unmount_fused_ln_modulate(root)
-    assert not fused_ln_modulate_active(root.child)
-    assert not mount_fused_ln_modulate(nn.Module())  # no marked sites
-
-
-@torch.no_grad()
-def test_mounted_ln_modulate_site_torch_compile_fullgraph():
-    class Site(nn.Module):
-        def __init__(self):
-            super().__init__()
-            mark_fused_ln_modulate_site(self)
-
-        def forward(self, x, scale, shift):
-            if fused_ln_modulate_active(self) and can_fuse_ln_modulate(x, scale, shift):
-                return fused_ln_modulate(x, scale, shift, eps=1e-6)
-            return (
-                nn.functional.layer_norm(x, (x.shape[-1],), eps=1e-6)
-                * (1 + scale[:, None])
-                + shift[:, None]
-            )
-
-    site = Site()
-    assert mount_fused_ln_modulate(site)
-    x = torch.randn(1, 64, 128, device="cuda", dtype=torch.bfloat16)
-    scale = torch.randn(1, 128, device="cuda", dtype=torch.bfloat16)
-    shift = torch.randn_like(scale)
-    expected = site(x, scale, shift)
-    actual = torch.compile(site, fullgraph=True)(x, scale, shift)
-    torch.testing.assert_close(actual, expected, atol=0.0625, rtol=0.05)
-
-
-if __name__ == "__main__":
-    import sys
-
-    sys.exit(pytest.main([__file__]))

test/registered/kernels/ops/diffusion/test_fused_norm_scale_shift.py

@@ -1,251 +0,0 @@
-import sys
-from typing import Optional, Tuple
-
-import pytest
-import torch
-from einops import rearrange
-from torch import Tensor
-
-from sglang.kernels.ops.diffusion.cutedsl.scale_residual_norm_scale_shift import (
-    fused_norm_scale_shift,
-    fused_scale_residual_norm_scale_shift,
-    validate_scale_shift,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=28, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-
-DEVICE = "cuda"
-SHAPE_MAP = {
-    "1": lambda B, S, F, D: (1,),
-    "D": lambda B, S, F, D: (D,),
-    "1D": lambda B, S, F, D: (1, D),
-    "BD": lambda B, S, F, D: (B, D),
-    "11D": lambda B, S, F, D: (1, 1, D),
-    "B1D": lambda B, S, F, D: (B, 1, D),
-    "1SD": lambda B, S, F, D: (1, S, D),
-    "BSD": lambda B, S, F, D: (B, S, D),
-    "BF1D": lambda B, S, F, D: (B, F, 1, D),
-}
-SHAPES = [
-    # (B, S, F, D)
-    (1, 115200, 1, 3072),  # Hunyuan
-    (1, 32760, 1, 1536),  # Wan
-    (1, 6, 1, 3072),  # Qwen
-    (1, 1024, 8, 3072),
-    (4, 512, 16, 3072),
-]
-DTYPES = [torch.float16, torch.bfloat16, torch.float32]
-NORM_TYPES = ["layer", "rms"]
-AFFINE_MODES = ["D", "NAT"]
-INDEX_MODES = ["BSD", "1", "1SD", "BD", "B1D", "D", "1D", "11D", "BF1D"]
-
-
-def _tol(dtype: torch.dtype):
-    return 1e-5 if dtype == torch.float32 else 5e-2
-
-
-@pytest.fixture(autouse=True)
-def cuda_setup():
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA required")
-    torch.cuda.manual_seed(0)
-
-
-def _apply_scale_shift(y: Tensor, scale: Tensor, shift: Tensor) -> Tensor:
-    if scale.ndim == 4:
-        num_frame = scale.shape[1]
-        return rearrange(
-            rearrange(y, "b (f l) d -> b f l d", f=num_frame) * (1 + scale) + shift,
-            "b f l d -> b (f l) d",
-        )
-    else:
-        scale = rearrange(scale, "b d -> b 1 d") if scale.ndim == 2 else scale
-        shift = rearrange(shift, "b d -> b 1 d") if shift.ndim == 2 else shift
-        return y * (1 + scale) + shift
-
-
-def fused_norm_scale_shift_ref(
-    x: Tensor,
-    weight: Optional[Tensor],
-    bias: Optional[Tensor],
-    scale: Tensor,
-    shift: Tensor,
-    norm_type: str,
-    eps: float,
-) -> Tensor:
-    original_dtype = x.dtype
-    x, weight, bias, scale, shift = (
-        v.float() if v is not None else v for v in [x, weight, bias, scale, shift]
-    )
-    if norm_type == "layer":
-        norm = torch.layer_norm(x, x.shape[-1:], eps=eps, weight=weight, bias=bias)
-    else:
-        norm = torch.rms_norm(x, x.shape[-1:], eps=eps, weight=weight)
-    return _apply_scale_shift(norm, scale, shift).to(original_dtype)
-
-
-def fused_scale_residual_norm_scale_shift_ref(
-    residual: Tensor,
-    x: Tensor,
-    gate: Optional[Tensor] | int,
-    weight: Optional[Tensor],
-    bias: Optional[Tensor],
-    scale: Tensor,
-    shift: Tensor,
-    norm_type: str,
-    eps: float,
-):
-    original_dtype = x.dtype
-    residual, x, gate, weight, bias, scale, shift = (
-        v.float() if isinstance(v, Tensor) else v
-        for v in [residual, x, gate, weight, bias, scale, shift]
-    )
-    if isinstance(gate, int):
-        x = residual + gate * x
-    else:
-        if gate.ndim == 4:
-            num_frame = gate.shape[1]
-            x_fld = rearrange(x, "b (f l) d -> b f l d", f=num_frame)
-            x = residual + rearrange(x_fld * gate, "b f l d -> b (f l) d")
-        else:
-            gate = rearrange(gate, "b d -> b 1 d") if gate.ndim == 2 else gate
-            x = residual + gate * x
-    if norm_type == "layer":
-        norm = torch.layer_norm(x, x.shape[-1:], eps=eps, weight=weight, bias=bias)
-    else:
-        norm = torch.rms_norm(x, x.shape[-1:], eps=eps, weight=weight)
-    y_ref = _apply_scale_shift(norm, scale, shift)
-    return y_ref.to(original_dtype), x.to(original_dtype)
-
-
-def _make_tensor(index_mode: str, shape: Tuple, dtype: torch.dtype):
-    if index_mode == "NAT":
-        return None
-    return torch.randn(*SHAPE_MAP[index_mode](*shape), device=DEVICE, dtype=dtype)
-
-
-def test_validate_scale_shift_rejects_non_divisible_frames():
-    with pytest.raises(ValueError, match=r"S\(10\) must be divisible by F\(4\)"):
-        validate_scale_shift(
-            torch.empty((1, 4, 1, 256), device=DEVICE, dtype=torch.float16),
-            1,
-            10,
-            256,
-        )
-
-
-@torch.no_grad()
-def run_norm_scale_shift(
-    shape=SHAPES[0],
-    dtype=DTYPES[0],
-    affine_dtype=DTYPES[0],
-    scale_dtype=DTYPES[0],
-    shift_dtype=DTYPES[0],
-    norm_type=NORM_TYPES[0],
-    affine_mode=AFFINE_MODES[0],
-    scale_mode="BSD",
-    shift_mode="BSD",
-    eps=1e-5,
-):
-    x = _make_tensor("BSD", shape, dtype)
-    weight = _make_tensor(affine_mode, shape, affine_dtype)
-    bias = _make_tensor(affine_mode, shape, affine_dtype)
-    scale = _make_tensor(scale_mode, shape, scale_dtype)
-    shift = _make_tensor(shift_mode, shape, shift_dtype)
-    y_dev = fused_norm_scale_shift(x, weight, bias, scale, shift, norm_type, eps)
-    y_ref = fused_norm_scale_shift_ref(x, weight, bias, scale, shift, norm_type, eps)
-    torch.testing.assert_close(y_dev, y_ref, atol=_tol(dtype), rtol=_tol(dtype))
-
-
-@torch.no_grad()
-def run_scale_resi_norm_scale_shift(
-    shape=SHAPES[0],
-    dtype=DTYPES[0],
-    affine_dtype=DTYPES[0],
-    scale_dtype=DTYPES[0],
-    shift_dtype=DTYPES[0],
-    norm_type=NORM_TYPES[0],
-    affine_mode=AFFINE_MODES[0],
-    gate_mode="B1D",
-    scale_mode="BSD",
-    shift_mode="BSD",
-    eps=1e-5,
-):
-    residual = _make_tensor("BSD", shape, dtype)
-    x = _make_tensor("BSD", shape, dtype)
-    gate = _make_tensor(gate_mode, shape, dtype)
-    weight = _make_tensor(affine_mode, shape, affine_dtype)
-    bias = _make_tensor(affine_mode, shape, affine_dtype)
-    scale = _make_tensor(scale_mode, shape, scale_dtype)
-    shift = _make_tensor(shift_mode, shape, shift_dtype)
-    y_dev, res_dev = fused_scale_residual_norm_scale_shift(
-        residual, x, gate, weight, bias, scale, shift, norm_type, eps
-    )
-    y_ref, res_ref = fused_scale_residual_norm_scale_shift_ref(
-        residual, x, gate, weight, bias, scale, shift, norm_type, eps
-    )
-    torch.testing.assert_close(y_dev, y_ref, atol=_tol(dtype), rtol=_tol(dtype))
-    torch.testing.assert_close(res_dev, res_ref, atol=_tol(dtype), rtol=_tol(dtype))
-
-
-@pytest.mark.parametrize("norm_type", NORM_TYPES)
-class TestFusedNormScaleShift:
-    @pytest.mark.parametrize("shape", SHAPES)
-    @pytest.mark.parametrize("dtype", DTYPES)
-    def test_shape_dtype(self, shape, dtype, norm_type):
-        run_norm_scale_shift(shape=shape, dtype=dtype, norm_type=norm_type)
-
-    @pytest.mark.parametrize("dtype", DTYPES)
-    def test_dtype_0(self, dtype, norm_type):
-        run_norm_scale_shift(affine_dtype=dtype, norm_type=norm_type)
-
-    @pytest.mark.parametrize("dtype", DTYPES)
-    def test_dtype_1(self, dtype, norm_type):
-        run_norm_scale_shift(scale_dtype=dtype, shift_dtype=dtype, norm_type=norm_type)
-
-    @pytest.mark.parametrize("affine_mode", AFFINE_MODES)
-    def test_normtype_affine(self, affine_mode, norm_type):
-        run_norm_scale_shift(affine_mode=affine_mode, norm_type=norm_type)
-
-    @pytest.mark.parametrize("index_mode", INDEX_MODES)
-    def test_index_mode(self, index_mode, norm_type):
-        run_norm_scale_shift(
-            scale_mode=index_mode, shift_mode=index_mode, norm_type=norm_type
-        )
-
-
-@pytest.mark.parametrize("norm_type", NORM_TYPES)
-class TestFusedScaleResidualNormScaleShift:
-    @pytest.mark.parametrize("shape", SHAPES)
-    @pytest.mark.parametrize("dtype", DTYPES)
-    def test_shape_dtype(self, shape, dtype, norm_type):
-        run_scale_resi_norm_scale_shift(shape=shape, dtype=dtype, norm_type=norm_type)
-
-    @pytest.mark.parametrize("dtype", DTYPES)
-    def test_dtype_0(self, dtype, norm_type):
-        run_scale_resi_norm_scale_shift(affine_dtype=dtype, norm_type=norm_type)
-
-    @pytest.mark.parametrize("dtype", DTYPES)
-    def test_dtype_1(self, dtype, norm_type):
-        run_scale_resi_norm_scale_shift(
-            scale_dtype=dtype, shift_dtype=dtype, norm_type=norm_type
-        )
-
-    @pytest.mark.parametrize("affine_mode", AFFINE_MODES)
-    def test_normtype_affine(self, affine_mode, norm_type):
-        run_scale_resi_norm_scale_shift(affine_mode=affine_mode, norm_type=norm_type)
-
-    @pytest.mark.parametrize("index_mode", INDEX_MODES)
-    def test_scale_shift_index_mode(self, index_mode, norm_type):
-        run_scale_resi_norm_scale_shift(
-            scale_mode=index_mode, shift_mode=index_mode, norm_type=norm_type
-        )
-
-    @pytest.mark.parametrize("index_mode", INDEX_MODES)
-    def test_gate_index_mode(self, index_mode, norm_type):
-        run_scale_resi_norm_scale_shift(gate_mode=index_mode, norm_type=norm_type)
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/diffusion/test_glm_image_ln_modulate.py

@@ -1,55 +0,0 @@
-"""GLM-Image fused LN+modulate / qk-LN fast paths must stay bit-exact vs eager."""
-
-import pytest
-import torch
-
-import sglang.multimodal_gen.runtime.models.dits.glm_image as glm_image
-from sglang.multimodal_gen.runtime.models.dits.glm_image import (
-    _eager_ln_modulate,
-    _glm_ln_modulate,
-    _glm_qk_layernorm,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=4, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-
-
-@pytest.mark.parametrize("shape", [(1, 4096, 4096), (2, 301, 4096), (1, 1, 2560)])
-def test_fused_ln_modulate_is_bit_exact(shape):
-    # (1, 4096, 4096) is the real GLM-Image image-stream shape (1024^2,
-    # hidden 4096); the others cover the text stream and another hidden.
-    torch.manual_seed(0)
-    batch, seq, hidden = shape
-    norm = torch.nn.LayerNorm(hidden, eps=1e-5, elementwise_affine=False).cuda()
-    x = (torch.randn(batch, seq, hidden, device="cuda") * 8).bfloat16()
-    emb = torch.randn(batch, 12 * hidden, device="cuda").bfloat16()
-    chunks = emb.chunk(12, dim=1)  # strided adaLN projection views
-    shift, scale = chunks[0], chunks[2]
-    out = _glm_ln_modulate(norm, x, scale, shift, x.dtype)
-    assert torch.equal(out, _eager_ln_modulate(norm, x, scale, shift, x.dtype))
-    assert glm_image._GLM_LN_MOD.verified
-    assert not glm_image._GLM_LN_MOD.disabled
-
-
-@pytest.mark.parametrize("shape", [(1, 4360, 32, 128), (2, 37, 3, 40), (1, 129, 5, 64)])
-def test_fused_qk_head_layernorm_is_bit_exact(shape):
-    # (1, 4360, 32, 128) is the real GLM-Image q/k shape (text + image
-    # tokens, 32 heads of dim 128); the others cover partially-filled warps.
-    torch.manual_seed(1)
-    batch, seq, heads, head_dim = shape
-    norm_q = torch.nn.LayerNorm(head_dim, eps=1e-5, elementwise_affine=False).cuda()
-    norm_k = torch.nn.LayerNorm(head_dim, eps=1e-5, elementwise_affine=False).cuda()
-    q = (torch.randn(batch, seq, heads, head_dim, device="cuda") * 5).bfloat16()
-    k = (torch.randn(batch, seq, heads, head_dim, device="cuda") * 5).bfloat16()
-    q_out, k_out = _glm_qk_layernorm(norm_q, norm_k, q, k, q.dtype)
-    assert torch.equal(q_out, norm_q(q).to(q.dtype))
-    assert torch.equal(k_out, norm_k(k).to(k.dtype))
-    assert glm_image._GLM_QK_LN.verified
-    assert not glm_image._GLM_QK_LN.disabled
-
-
-if __name__ == "__main__":
-    import sys
-
-    sys.exit(pytest.main([__file__]))

test/registered/kernels/ops/diffusion/test_group_norm_silu.py

@@ -1,104 +0,0 @@
-import sys
-
-import pytest
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from sglang.kernels.ops.diffusion.group_norm_silu import apply_group_norm_silu
-from sglang.kernels.ops.diffusion.triton.group_norm_silu import triton_group_norm_silu
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-
-register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-register_amd_ci(est_time=15, suite="nightly-amd-kernel-1-gpu", nightly=True)
-
-DEVICE = "cuda"
-DTYPES = [torch.float16, torch.bfloat16, torch.float32]
-TEST_CASES = [
-    pytest.param((2, 64, 32, 32), 32, id="image_2d"),
-    pytest.param((1, 64, 4, 16, 16), 32, id="video_3d"),
-    pytest.param((4, 128), 32, id="token_2d"),
-]
-LARGE_TILE_CASE = ((1, 128, 20, 256, 256), 32)
-
-
-def _tol(dtype: torch.dtype) -> tuple[float, float]:
-    if dtype == torch.float32:
-        return 1e-5, 1e-5
-    if dtype == torch.bfloat16:
-        return 7e-2, 2e-2
-    return 3e-3, 3e-3
-
-
-@pytest.fixture(autouse=True)
-def cuda_setup():
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA required")
-    torch.cuda.manual_seed(0)
-
-
-def _reference(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    bias: torch.Tensor,
-    num_groups: int,
-    eps: float = 1e-5,
-) -> torch.Tensor:
-    return F.silu(F.group_norm(x, num_groups, weight=weight, bias=bias, eps=eps))
-
-
-@torch.no_grad()
-@pytest.mark.parametrize("shape,num_groups", TEST_CASES)
-@pytest.mark.parametrize("dtype", DTYPES)
-def test_triton_group_norm_silu(
-    shape: tuple[int, ...], num_groups: int, dtype: torch.dtype
-) -> None:
-    channels = shape[1]
-    x = torch.randn(shape, device=DEVICE, dtype=dtype)
-    weight = torch.randn(channels, device=DEVICE, dtype=dtype)
-    bias = torch.randn(channels, device=DEVICE, dtype=dtype)
-
-    actual = triton_group_norm_silu(x, weight, bias, num_groups=num_groups)
-    expected = _reference(x, weight, bias, num_groups)
-
-    atol, rtol = _tol(dtype)
-    torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
-
-
-@torch.no_grad()
-@pytest.mark.parametrize("shape,num_groups", TEST_CASES[:2])
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-def test_apply_group_norm_silu(
-    shape: tuple[int, ...],
-    num_groups: int,
-    dtype: torch.dtype,
-) -> None:
-    norm = nn.GroupNorm(num_groups, shape[1], eps=1e-5, affine=True).to(
-        device=DEVICE, dtype=dtype
-    )
-    activation = nn.SiLU()
-    hidden_states = torch.randn(shape, device=DEVICE, dtype=dtype)
-
-    actual = apply_group_norm_silu(hidden_states, norm, activation)
-    expected = activation(norm(hidden_states))
-
-    atol, rtol = _tol(dtype)
-    torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
-
-
-@torch.no_grad()
-def test_triton_group_norm_silu_large_tile_bf16() -> None:
-    shape, num_groups = LARGE_TILE_CASE
-    x = torch.randn(shape, device=DEVICE, dtype=torch.bfloat16)
-    weight = torch.randn(shape[1], device=DEVICE, dtype=torch.bfloat16)
-    bias = torch.randn(shape[1], device=DEVICE, dtype=torch.bfloat16)
-
-    actual = triton_group_norm_silu(x, weight, bias, num_groups=num_groups)
-    expected = _reference(x, weight, bias, num_groups)
-
-    atol, rtol = _tol(torch.bfloat16)
-    torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/diffusion/test_hunyuanvideo_eager_fusions.py

@@ -1,94 +0,0 @@
-"""HunyuanVideo eager QKV/RoPE and quality-gated QKNorm tests."""
-
-import sys
-from unittest.mock import patch
-
-import pytest
-import torch
-
-import sglang.kernels.ops.diffusion.hunyuan_qknorm as hunyuan_qknorm
-from sglang.kernels.ops.diffusion.hunyuan_qknorm import (
-    mark_hunyuan_qknorm_site,
-    mount_hunyuan_qknorm,
-    unmount_hunyuan_qknorm,
-)
-from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
-from sglang.multimodal_gen.runtime.layers.rotary_embedding.utils import (
-    _apply_rotary_emb,
-)
-from sglang.multimodal_gen.runtime.models.dits.hunyuanvideo import (
-    _hunyuan_pack_qkv,
-    _hunyuan_qknorm,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-
-
-@pytest.mark.parametrize("img_tokens,txt_tokens", [(257, 31), (4096, 256)])
-def test_hunyuan_qkv_rope_pack_is_bit_exact(img_tokens, txt_tokens):
-    torch.manual_seed(0)
-    shape_img = (1, img_tokens, 24, 128)
-    shape_txt = (1, txt_tokens, 24, 128)
-    img_q, img_k, img_v = (
-        torch.randn(shape_img, device="cuda", dtype=torch.bfloat16) for _ in range(3)
-    )
-    txt_q, txt_k, txt_v = (
-        torch.randn(shape_txt, device="cuda", dtype=torch.bfloat16) for _ in range(3)
-    )
-    cos = torch.randn(img_tokens, 64, device="cuda")
-    sin = torch.randn_like(cos)
-
-    q, k, v = _hunyuan_pack_qkv(img_q, img_k, img_v, txt_q, txt_k, txt_v, cos, sin)
-    q_ref = torch.cat(
-        (_apply_rotary_emb(img_q, cos, sin, is_neox_style=False), txt_q), dim=1
-    )
-    k_ref = torch.cat(
-        (_apply_rotary_emb(img_k, cos, sin, is_neox_style=False), txt_k), dim=1
-    )
-    v_ref = torch.cat((img_v, txt_v), dim=1)
-
-    assert torch.equal(q, q_ref)
-    assert torch.equal(k, k_ref)
-    assert torch.equal(v, v_ref)
-
-
-def test_hunyuan_quality_qknorm_matches_rmsnorm():
-    torch.manual_seed(1)
-    site = torch.nn.Module()
-    mark_hunyuan_qknorm_site(site)
-    q_norm = RMSNorm(128, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
-    k_norm = RMSNorm(128, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
-    packed = torch.randn(1, 257, 3, 24, 128, device="cuda", dtype=torch.bfloat16)
-    q, k = packed[:, :, 0], packed[:, :, 1]
-    q_ref = q_norm(q.contiguous()).to(q)
-    k_ref = k_norm(k.contiguous()).to(k)
-
-    q_unmounted, k_unmounted = _hunyuan_qknorm(site, q, k, q_norm, k_norm)
-    assert torch.equal(q_unmounted, q_ref)
-    assert torch.equal(k_unmounted, k_ref)
-
-    assert mount_hunyuan_qknorm(site)
-    q_out, k_out = _hunyuan_qknorm(site, q, k, q_norm, k_norm)
-    torch.testing.assert_close(q_out, q_ref, atol=2e-2, rtol=2e-2)
-    torch.testing.assert_close(k_out, k_ref, atol=2e-2, rtol=2e-2)
-
-    unmount_hunyuan_qknorm(site)
-    q_unmounted, k_unmounted = _hunyuan_qknorm(site, q, k, q_norm, k_norm)
-    assert torch.equal(q_unmounted, q_ref)
-    assert torch.equal(k_unmounted, k_ref)
-
-
-def test_hunyuan_quality_qknorm_stays_unmounted_without_cute_kernel():
-    site = torch.nn.Module()
-    mark_hunyuan_qknorm_site(site)
-
-    with patch.object(hunyuan_qknorm, "_get_qk_rmsnorm_cute", return_value=None):
-        assert not mount_hunyuan_qknorm(site)
-
-    assert not hunyuan_qknorm._FUSION.is_enabled(site)
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/diffusion/test_import_surface.py

@@ -0,0 +1,210 @@
+"""Guards that keep the ``diffusion`` package's import surface from eroding.
+
+The reorganization only stays useful if two invariants hold:
+
+1. runtime code imports from ``sglang.kernels.ops.diffusion`` and not from a
+   submodule, so the internal layout can move without touching call sites;
+2. the facade's ``_EXPORTS`` table and the registry's ``_SPECS`` table both
+   point at symbols that actually exist.
+
+Neither is checkable by the type system, and both fail silently -- a stale
+``_EXPORTS`` entry only raises when some model happens to call that kernel, on
+a GPU, at serving time.  These are pure-CPU tests: they read the tables and
+resolve them with ``importlib``/``ast`` without importing torch backends.
+"""
+
+import ast
+import importlib
+import pathlib
+import subprocess
+import sys
+
+import pytest
+
+from sglang.kernels.ops.diffusion import _EXPORTS, _SPECS
+from sglang.kernels.registry import registry
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=4, suite="base-a-test-cpu")
+
+PACKAGE = "sglang.kernels.ops.diffusion"
+_PACKAGE_DIR = pathlib.Path(importlib.import_module(PACKAGE).__file__ or "").parent
+_REPO_ROOT = _PACKAGE_DIR.parents[4]  # <repo>/python/sglang/kernels/ops/diffusion
+
+# Backend-specific test files may name a leaf module on purpose; everything
+# else -- all runtime code -- must go through the facade.
+_DEEP_IMPORT_ALLOWLIST = {
+    "python/sglang/multimodal_gen/test/unit/test_latent_upsampler_group_norm_silu.py",
+    "test/registered/kernels/ops/diffusion/test_model_fast_paths.py",
+    "test/registered/kernels/ops/diffusion/test_sites.py",
+}
+
+
+def _module_defines(module_path: str) -> set[str]:
+    """Top-level names bound by a submodule, without importing it.
+
+    Importing would pull in Triton / CuTe-DSL / FlyDSL, none of which are
+    installed on the CPU CI lane -- so this reads the source instead.
+    """
+    path = _PACKAGE_DIR / (module_path.replace(".", "/") + ".py")
+    if not path.exists():
+        path = _PACKAGE_DIR / module_path.replace(".", "/") / "__init__.py"
+    assert path.exists(), f"{PACKAGE}.{module_path} does not exist"
+
+    names: set[str] = set()
+    for node in ast.parse(path.read_text(encoding="utf-8")).body:
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
+            names.add(node.name)
+        elif isinstance(node, ast.Assign):
+            names.update(t.id for t in node.targets if isinstance(t, ast.Name))
+        elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name):
+            names.add(node.target.id)
+        elif isinstance(node, (ast.Import, ast.ImportFrom)):
+            names.update((a.asname or a.name).split(".")[0] for a in node.names)
+        elif isinstance(node, (ast.If, ast.Try)):
+            # Platform-conditional rebinds (``x = select_impl(...)``) and
+            # guarded defs still bind a public name.
+            for inner in ast.walk(node):
+                if isinstance(inner, (ast.FunctionDef, ast.ClassDef)):
+                    names.add(inner.name)
+                elif isinstance(inner, ast.Assign):
+                    names.update(t.id for t in inner.targets if isinstance(t, ast.Name))
+    return names
+
+
+def test_every_export_resolves_to_a_real_symbol():
+    missing = [
+        f"{symbol} -> {module}"
+        for symbol, module in sorted(_EXPORTS.items())
+        if symbol not in _module_defines(module)
+    ]
+    assert not missing, f"stale _EXPORTS entries: {missing}"
+
+
+def test_every_symbol_imported_from_the_facade_is_exported():
+    """The reverse of the check above, and the one that actually bites.
+
+    A missing ``_EXPORTS`` entry raises ``ImportError`` at module import, so a
+    module-level ``from ...diffusion import x`` fails loudly.  A *function-local*
+    one -- the pattern used for optional backends -- fails only when that test
+    or code path runs, on the platform that has the backend.  Enumerating the
+    call sites catches it here instead.
+    """
+    unexported = set()
+    for root in ("python/sglang", "test", "benchmark"):
+        root_dir = _REPO_ROOT / root
+        if not root_dir.exists():
+            continue
+        for path in root_dir.rglob("*.py"):
+            rel = path.relative_to(_REPO_ROOT).as_posix()
+            if rel.startswith("python/sglang/kernels/ops/diffusion/"):
+                continue
+            try:
+                tree = ast.parse(path.read_text(encoding="utf-8"))
+            except (SyntaxError, UnicodeDecodeError):
+                continue
+            for node in ast.walk(tree):
+                if isinstance(node, ast.ImportFrom) and node.module == PACKAGE:
+                    unexported.update(
+                        a.name
+                        for a in node.names
+                        if a.name not in _EXPORTS and not a.name.startswith("_")
+                    )
+    assert not unexported, f"imported but not in _EXPORTS: {sorted(unexported)}"
+
+
+def test_every_registered_spec_target_resolves():
+    missing = []
+    for _op, _backend, target, _caps, _description in _SPECS:
+        module, _, attr = target.partition(":")
+        if attr not in _module_defines(module):
+            missing.append(target)
+    assert not missing, f"stale _SPECS targets: {missing}"
+
+
+def test_registry_holds_the_diffusion_ops():
+    # Registration happens at package import, is metadata-only, and is what
+    # ``select_kernel`` / the tracing tools read.
+    registered = {op for op in registry.ops() if op.startswith("diffusion.")}
+    assert {op for op, *_ in _SPECS} <= registered
+
+
+def test_facade_rejects_unknown_attributes():
+    module = sys.modules[PACKAGE]
+    with pytest.raises(AttributeError):
+        module.definitely_not_a_kernel
+    assert set(module.__all__) == set(_EXPORTS)
+    assert set(_EXPORTS) <= set(dir(module))
+
+
+def test_importing_the_package_does_not_import_any_leaf_module():
+    """The reason ``__getattr__`` is lazy rather than a block of re-exports.
+
+    The backends have disjoint, heavy, mutually-exclusive dependencies --
+    Triton (CUDA/ROCm), CUTLASS/CuTe-DSL, FlyDSL (gfx950), MLX (Apple).  If
+    ``_EXPORTS`` ever degrades into eager ``from .norm.x import y`` lines, all
+    of them become import-time requirements on every platform, which is how a
+    CPU-only or Apple install starts failing at ``import sglang``.
+
+    Asserted on this package's own leaf modules rather than on ``triton`` in
+    ``sys.modules``: sibling operator groups import Triton for their own
+    reasons, so a global check would not isolate this package's behavior.
+    Run in a fresh interpreter because this process has already resolved
+    exports through the facade.
+    """
+    code = (
+        "import importlib, sys\n"
+        f"importlib.import_module('{PACKAGE}')\n"
+        f"prefix = '{PACKAGE}.'\n"
+        "leaves = [m for m in sys.modules if m.startswith(prefix)"
+        " and not m.endswith('__init__')]\n"
+        "print(','.join(sorted(m for m in leaves if '.' in m[len(prefix):]"
+        " or sys.modules[m].__file__ and not sys.modules[m].__file__"
+        ".endswith('__init__.py'))))\n"
+    )
+    result = subprocess.run(
+        [sys.executable, "-c", code], capture_output=True, text=True, timeout=600
+    )
+    assert result.returncode == 0, result.stderr
+    leaked = [m for m in result.stdout.strip().split(",") if m]
+    assert not leaked, f"importing {PACKAGE} eagerly imported: {leaked}"
+
+
+@pytest.mark.parametrize("root", ["python/sglang", "test", "benchmark"])
+def test_runtime_code_imports_only_through_the_facade(root):
+    root_dir = _REPO_ROOT / root
+    if not root_dir.exists():  # source checkouts only
+        pytest.skip(f"{root} not present in this install")
+
+    offenders = []
+    for path in root_dir.rglob("*.py"):
+        rel = path.relative_to(_REPO_ROOT).as_posix()
+        if rel.startswith("python/sglang/kernels/ops/diffusion/"):
+            continue  # intra-package imports are the point of the subpackages
+        if rel in _DEEP_IMPORT_ALLOWLIST:
+            continue
+        try:
+            tree = ast.parse(path.read_text(encoding="utf-8"))
+        except (SyntaxError, UnicodeDecodeError):
+            continue
+        for node in ast.walk(tree):
+            if (
+                isinstance(node, ast.ImportFrom)
+                and node.module
+                and node.module.startswith(f"{PACKAGE}.")
+            ):
+                offenders.append(f"{rel}:{node.lineno} imports {node.module}")
+            elif isinstance(node, ast.Import):
+                offenders.extend(
+                    f"{rel}:{node.lineno} imports {a.name}"
+                    for a in node.names
+                    if a.name.startswith(f"{PACKAGE}.")
+                )
+    assert not offenders, (
+        "import from sglang.kernels.ops.diffusion instead of a submodule:\n  "
+        + "\n  ".join(offenders)
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/diffusion/test_layout.py

@@ -0,0 +1,597 @@
+"""``diffusion.layout``: data-movement kernels.
+
+Every kernel in this domain only moves values (plus zero fill, plus at most
+one same-order add), so each one is *bitwise* identical to the aten chain it
+replaces.  That makes ``torch.equal`` -- not ``assert_close`` -- the right
+assertion throughout this file; a tolerance here would hide a real bug.
+
+Covered: USP output head merge, Ulysses destination-major QKV pack, varlen
+pack/scatter, causal Conv3d cat+pad (CUDA and Triton), and the Wan causal-VAE
+cache kernels.
+"""
+
+import sys
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.attention.flash_attention import flash_attn_varlen_func
+from sglang.kernels.ops.diffusion import (
+    build_inv_indices,
+    can_use_usp_merge_heads,
+    cat_pad_channels_last_3d,
+    dup_up3d_add,
+)
+from sglang.kernels.ops.diffusion import (
+    fused_causal_conv3d_cat_pad as fused_causal_conv3d_cat_pad_triton,
+)
+from sglang.kernels.ops.diffusion import (
+    fused_causal_conv3d_cat_pad_cuda,
+    fused_pack_qkv,
+    fused_scatter_to_padded,
+    pack_qkv_destination_major,
+    usp_merge_heads,
+)
+from sglang.multimodal_gen.runtime.layers.attention.backends import (
+    flash_attn as _fa_backend,
+)
+from sglang.multimodal_gen.runtime.layers.attention.layer import build_varlen_mask_meta
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+
+register_cuda_ci(est_time=110, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+# Nightly is not redundant: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1, which
+# expands the get_ci_test_range sweeps below.
+register_cuda_ci(est_time=20, stage="nightly", runner_config="1-gpu-large")
+register_amd_ci(est_time=10, stage="jit-kernel-unit", runner_config="amd")
+register_amd_ci(est_time=15, suite="nightly-amd-kernel-1-gpu", nightly=True)
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+DEVICE = "cuda"
+
+
+def _cl3d(shape, dtype):
+    return torch.randn(shape, device=DEVICE, dtype=dtype).contiguous(
+        memory_format=torch.channels_last_3d
+    )
+
+
+# ---------------------------------------------------------------------------
+# USP output head merge
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "world,seq,batch,h_local,head_dim",
+    [
+        (4, 7936, 1, 14, 128),  # H3 768p production shape (Ulysses 4)
+        (2, 64, 3, 4, 64),  # batched
+        (4, 33, 2, 4, 100),  # scalar fallback inside the CUDA kernel
+    ],
+)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+@pytest.mark.skipif(
+    bool(torch.version.hip),
+    reason="the USP merge-heads JIT fast path is CUDA-only by design -- "
+    "can_use_usp_merge_heads() returns False under HIP, and the aten fallback "
+    "it degrades to is covered by the unsupported-inputs test below",
+)
+def test_usp_merge_heads_bitwise(dtype, world, seq, batch, h_local, head_dim):
+    generator = torch.Generator(device=DEVICE).manual_seed(4321)
+    x = torch.randn(
+        world,
+        seq,
+        batch,
+        h_local,
+        head_dim,
+        dtype=dtype,
+        device=DEVICE,
+        generator=generator,
+    )
+    assert can_use_usp_merge_heads(x)
+    out = usp_merge_heads(x)
+    ref = x.permute(2, 1, 0, 3, 4).contiguous()
+    assert out.shape == ref.shape
+    assert torch.equal(out, ref)
+
+
+def test_usp_merge_heads_unsupported_inputs_use_exact_fallback():
+    # The wrapper degrades to the aten permute for anything the fast path
+    # rejects -- a wrong rank, a transposed view, an empty leading dim, or a
+    # ROCm build -- so callers never need their own guard.
+    x = torch.randn(2, 4, 1, 4, 64, dtype=torch.bfloat16, device=DEVICE)
+    for value in (x.transpose(0, 1), x[:0], x[0]):
+        assert not can_use_usp_merge_heads(value)
+        if value.dim() == 5:
+            assert torch.equal(
+                usp_merge_heads(value), value.permute(2, 1, 0, 3, 4).contiguous()
+            )
+
+    with patch.object(torch.version, "hip", "6.3"):
+        assert not can_use_usp_merge_heads(x)
+        assert torch.equal(usp_merge_heads(x), x.permute(2, 1, 0, 3, 4).contiguous())
+
+
+# ---------------------------------------------------------------------------
+# Ulysses destination-major QKV pack
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_pack_qkv_destination_major_is_bit_exact(dtype):
+    torch.manual_seed(0)
+    rows, world_size, global_heads, head_size = 17, 4, 12, 64
+    q, k, v = (
+        torch.randn(rows, global_heads, head_size, device=DEVICE, dtype=dtype)
+        for _ in range(3)
+    )
+
+    local_heads = global_heads // world_size
+    expected = torch.empty(
+        world_size, rows, local_heads, 3 * head_size, device=DEVICE, dtype=dtype
+    )
+    for index, tensor in enumerate((q, k, v)):
+        shards = tensor.view(rows, world_size, local_heads, head_size).permute(
+            1, 0, 2, 3
+        )
+        expected[..., index * head_size : (index + 1) * head_size].copy_(shards)
+
+    assert torch.equal(pack_qkv_destination_major(q, k, v, world_size), expected)
+
+
+def test_pack_qkv_destination_major_validates_inputs():
+    q = torch.empty(2, 4, 8, device=DEVICE, dtype=torch.bfloat16)
+    with pytest.raises(ValueError, match="same 3D shape"):
+        pack_qkv_destination_major(q, q[:, :-1], q, 2)
+    with pytest.raises(ValueError, match="divide global_heads"):
+        pack_qkv_destination_major(q, q, q, 3)
+    with pytest.raises(ValueError, match="expected shape"):
+        pack_qkv_destination_major(q, q, q, 2, out=torch.empty_like(q))
+
+
+# ---------------------------------------------------------------------------
+# Varlen pack / scatter
+# ---------------------------------------------------------------------------
+
+VARLEN_DTYPES = get_ci_test_range([torch.bfloat16, torch.float16], [torch.bfloat16])
+# (name, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens)
+VARLEN_SHAPES = get_ci_test_range(
+    [
+        ("small_c2", 2, 64, 128, 4, 64, [32, 48]),
+        ("prod_c2", 2, 256, 1024, 24, 128, [128, 200]),
+        ("all_valid_b1", 1, 64, 128, 4, 64, [64]),
+        ("all_valid_b4", 4, 64, 128, 4, 64, [64, 64, 64, 64]),
+        ("c8_prod", 8, 256, 4096, 24, 128, [128, 200, 256, 100, 50, 256, 256, 50]),
+        # one batch with zero valid text tokens (image side still valid)
+        ("zero_txt_one_batch", 2, 64, 128, 4, 64, [0, 32]),
+        # bs=1 with no text validity (only image rows packed)
+        ("bs1_zero_txt", 1, 64, 128, 4, 64, [0]),
+    ],
+    [
+        ("small_c2", 2, 64, 128, 4, 64, [32, 48]),
+        ("prod_c2", 2, 256, 1024, 24, 128, [128, 200]),
+        ("all_valid_b4", 4, 64, 128, 4, 64, [64, 64, 64, 64]),
+    ],
+)
+
+
+def _build_mask(bs, s_txt, s_img, valid_txt_lens):
+    mask = torch.zeros(bs, s_txt + s_img, dtype=torch.bool, device=DEVICE)
+    for b, vt in enumerate(valid_txt_lens):
+        mask[b, :vt] = True
+        mask[b, s_txt:] = True
+    return mask
+
+
+def _build_meta(mask):
+    bs, seq = mask.shape
+    indices = mask.reshape(-1).nonzero(as_tuple=False).flatten()
+    return indices, build_inv_indices(indices, bs * seq)
+
+
+@pytest.mark.parametrize("dtype", VARLEN_DTYPES)
+@pytest.mark.parametrize("shape", VARLEN_SHAPES, ids=lambda s: s[0])
+def test_varlen_pack_matches_index_select(dtype, shape):
+    _, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens = shape
+    torch.manual_seed(0)
+    s = s_txt + s_img
+    indices, _ = _build_meta(_build_mask(bs, s_txt, s_img, valid_txt_lens))
+
+    q, k, v = (
+        torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
+        for _ in range(3)
+    )
+    fused = fused_pack_qkv(q, k, v, indices)
+    for got, src in zip(fused, (q, k, v), strict=True):
+        want = src.reshape(bs * s, num_heads, head_dim).index_select(0, indices)
+        assert torch.equal(got, want)
+
+
+@pytest.mark.parametrize("dtype", VARLEN_DTYPES)
+@pytest.mark.parametrize("shape", VARLEN_SHAPES, ids=lambda s: s[0])
+def test_varlen_scatter_matches_index_copy(dtype, shape):
+    _, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens = shape
+    torch.manual_seed(1)
+    s = s_txt + s_img
+    mask = _build_mask(bs, s_txt, s_img, valid_txt_lens)
+    indices, inv_indices = _build_meta(mask)
+
+    out_unpad = torch.randn(
+        indices.shape[0], num_heads, head_dim, dtype=dtype, device=DEVICE
+    )
+    flat = torch.zeros(bs * s, num_heads, head_dim, dtype=dtype, device=DEVICE)
+    flat.index_copy_(0, indices, out_unpad)
+    out_ref = flat.view(bs, s, num_heads, head_dim)
+
+    out_fused = fused_scatter_to_padded(out_unpad, inv_indices, bs, s)
+    assert torch.equal(out_ref, out_fused)
+    invalid = ~mask
+    if invalid.any():
+        # Padding rows must be exactly zero, not merely small.
+        assert out_fused[invalid].abs().max().item() == 0.0
+
+
+def test_varlen_pack_handles_non_contiguous_input():
+    # Q/K/V arrive as (B, H, S, D) permutes from attention; the helper must
+    # make them contiguous itself rather than reading the wrong strides.
+    torch.manual_seed(2)
+    bs, s_txt, s_img, num_heads, head_dim = 2, 64, 128, 4, 64
+    indices, _ = _build_meta(_build_mask(bs, s_txt, s_img, [32, 48]))
+
+    pre = torch.randn(
+        bs, num_heads, s_txt + s_img, head_dim, dtype=torch.bfloat16, device=DEVICE
+    )
+    q, k, v = (torch.randn_like(pre).permute(0, 2, 1, 3) for _ in range(3))
+    assert not q.is_contiguous()
+
+    fused = fused_pack_qkv(q, k, v, indices)
+    for got, src in zip(fused, (q, k, v), strict=True):
+        want = src.contiguous().flatten(0, 1).index_select(0, indices)
+        assert torch.equal(got, want)
+
+
+def test_build_inv_indices_matches_manual():
+    torch.manual_seed(3)
+    bs, s = 2, 32
+    mask = torch.bernoulli(torch.full((bs, s), 0.6, device=DEVICE)).to(torch.bool)
+    indices = mask.reshape(-1).nonzero(as_tuple=False).flatten()
+
+    manual = torch.full((bs * s,), -1, dtype=torch.int32, device=DEVICE)
+    if indices.numel():
+        manual[indices.long()] = torch.arange(
+            indices.numel(), dtype=torch.int32, device=DEVICE
+        )
+    assert torch.equal(build_inv_indices(indices, bs * s), manual)
+
+
+def test_varlen_empty_valid_set_handled():
+    # An all-False mask is reachable (a request whose text side is fully
+    # masked): pack must return empty tensors and scatter an all-zero dense
+    # output rather than launching a degenerate grid.
+    bs, s, num_heads, head_dim = 2, 16, 4, 64
+    indices = torch.zeros(bs, s, dtype=torch.bool, device=DEVICE).reshape(-1).nonzero()
+    indices = indices.flatten()
+    inv_indices = build_inv_indices(indices, bs * s)
+    assert indices.numel() == 0
+
+    q = torch.randn(bs, s, num_heads, head_dim, dtype=torch.bfloat16, device=DEVICE)
+    unpad = fused_pack_qkv(q, q.clone(), q.clone(), indices)
+    assert all(t.shape == (0, num_heads, head_dim) for t in unpad)
+
+    out_padded = fused_scatter_to_padded(unpad[0], inv_indices, bs, s)
+    assert out_padded.shape == (bs, s, num_heads, head_dim)
+    assert out_padded.abs().max().item() == 0.0
+
+
+# The kernels above are unit-tested against index_select/index_copy_; this
+# section drives them through the production USPAttention masked branch, where
+# a wrong index layout would produce plausible-looking attention output rather
+# than an obvious mismatch.
+
+
+def _sdpa_with_key_mask(q, k, v, key_mask, softmax_scale):
+    """Reference: SDPA with a ``[B, S]`` key mask broadcast to ``[B, 1, 1, S]``."""
+    q_ = q.transpose(1, 2)
+    k_ = k.transpose(1, 2)
+    v_ = v.transpose(1, 2)
+    mask = key_mask.to(dtype=q.dtype)[:, None, None, :]
+    mask = (mask - 1.0) * torch.finfo(q.dtype).max
+    out = F.scaled_dot_product_attention(
+        q_,
+        k_,
+        v_,
+        attn_mask=mask,
+        dropout_p=0.0,
+        is_causal=False,
+        scale=softmax_scale,
+    )
+    return out.transpose(1, 2)
+
+
+def _varlen_path(q, k, v, key_mask, softmax_scale):
+    """Production varlen path matching USPAttention.forward masked branch."""
+    bs, seq = q.shape[0], q.shape[1]
+    meta = build_varlen_mask_meta(key_mask)
+    indices = meta["indices"]
+    if indices.shape[0] == 0:
+        return torch.zeros_like(q)
+    q_unpad, k_unpad, v_unpad = fused_pack_qkv(q, k, v, indices)
+    try:
+        out_unpad = flash_attn_varlen_func(
+            q=q_unpad,
+            k=k_unpad,
+            v=v_unpad,
+            cu_seqlens_q=meta["cu_seqlens"],
+            cu_seqlens_k=meta["cu_seqlens"],
+            max_seqlen_q=meta["max_seqlen"],
+            max_seqlen_k=meta["max_seqlen"],
+            softmax_scale=softmax_scale,
+            causal=False,
+            ver=_fa_backend.fa_ver,
+        )
+    except ImportError as exc:  # pragma: no cover - image-dependent
+        # ``flash_attn_varlen_func`` resolves its backend lazily, so an image
+        # without the selected FlashAttention build raises here rather than at
+        # import.  This file also runs on the B200 lane (for the causal-Conv3d
+        # section), which ships no ``flash_attn`` -- skip only this end-to-end
+        # comparison there; the pack/scatter kernels themselves are covered
+        # unit-wise above on every lane.
+        pytest.skip(f"FlashAttention varlen v{_fa_backend.fa_ver} unavailable: {exc}")
+    return fused_scatter_to_padded(out_unpad, meta["inv_indices"], bs, seq)
+
+
+@pytest.mark.parametrize("dtype", VARLEN_DTYPES)
+@pytest.mark.parametrize("shape", VARLEN_SHAPES, ids=lambda s: s[0])
+def test_varlen_path_matches_sdpa_on_valid_rows(dtype, shape):
+    """Valid rows: varlen output ≈ SDPA output within FA tolerance."""
+    _, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens = shape
+    torch.manual_seed(0)
+    s = s_txt + s_img
+    softmax_scale = head_dim**-0.5
+    mask = _build_mask(bs, s_txt, s_img, valid_txt_lens)
+    q = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
+    k = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
+    v = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
+
+    out_sdpa = _sdpa_with_key_mask(q, k, v, mask, softmax_scale)
+    out_varlen = _varlen_path(q, k, v, mask, softmax_scale)
+
+    valid = mask[..., None, None].expand_as(out_sdpa)
+    rtol = 1e-2 if dtype == torch.bfloat16 else 5e-3
+    atol = 5e-2 if dtype == torch.bfloat16 else 1e-2
+    torch.testing.assert_close(
+        out_sdpa[valid],
+        out_varlen[valid],
+        rtol=rtol,
+        atol=atol,
+    )
+
+
+@pytest.mark.parametrize("dtype", VARLEN_DTYPES)
+@pytest.mark.parametrize("shape", VARLEN_SHAPES, ids=lambda s: s[0])
+def test_varlen_path_zeros_masked_rows(dtype, shape):
+    """Masked rows: varlen path produces exact zeros (documented contract)."""
+    _, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens = shape
+    torch.manual_seed(1)
+    s = s_txt + s_img
+    softmax_scale = head_dim**-0.5
+    mask = _build_mask(bs, s_txt, s_img, valid_txt_lens)
+    q = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
+    k = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
+    v = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
+
+    out_varlen = _varlen_path(q, k, v, mask, softmax_scale)
+
+    invalid = ~mask
+    if invalid.any():
+        assert (out_varlen[invalid] == 0).all(), "masked rows must be zero-filled"
+
+
+# ---------------------------------------------------------------------------
+# Causal Conv3d cat + pad (CUDA JIT vs Triton)
+# ---------------------------------------------------------------------------
+
+CONV3D_CASES = get_ci_test_range(
+    [
+        (1024, 1, 30, 52, 1),
+        (1024, 1, 30, 52, 2),
+        (1024, 2, 60, 104, 1),
+        (1024, 2, 60, 104, 2),
+        (512, 4, 120, 208, 1),
+        (512, 4, 120, 208, 2),
+        (256, 4, 240, 416, 1),
+        (256, 4, 240, 416, 2),
+    ],
+    [(1024, 1, 30, 52, 1), (512, 4, 120, 208, 2)],
+)
+
+
+def _conv3d_inputs(channels, t_size, h_size, w_size, cache_t):
+    generator = torch.Generator(device=DEVICE)
+    generator.manual_seed(channels * 1009 + t_size * 251 + h_size + cache_t)
+    x = torch.randn(
+        (1, channels, t_size, h_size, w_size),
+        device=DEVICE,
+        dtype=torch.bfloat16,
+        generator=generator,
+    )
+    cache_x = torch.randn(
+        (1, channels, cache_t, h_size, w_size),
+        device=DEVICE,
+        dtype=torch.bfloat16,
+        generator=generator,
+    )
+    return x, cache_x, (1, 1, 1, 1, cache_t, 0)
+
+
+@pytest.mark.parametrize("channels,t_size,h_size,w_size,cache_t", CONV3D_CASES)
+def test_causal_conv3d_cat_pad_cuda_matches_triton(
+    channels, t_size, h_size, w_size, cache_t
+):
+    x, cache_x, padding = _conv3d_inputs(channels, t_size, h_size, w_size, cache_t)
+    actual = fused_causal_conv3d_cat_pad_cuda(x, cache_x, padding)
+    expected = fused_causal_conv3d_cat_pad_triton(x, cache_x, padding)
+    assert torch.equal(actual, expected)
+
+
+def test_causal_conv3d_cat_pad_torch_compile():
+    # The CUDA path is a registered custom op, so a fullgraph compile must not
+    # graph-break on it.
+    x, cache_x, padding = _conv3d_inputs(1024, 1, 30, 52, 1)
+
+    @torch.compile(fullgraph=True)
+    def fn(x, cache_x):
+        return fused_causal_conv3d_cat_pad_cuda(x, cache_x, padding)
+
+    assert torch.equal(
+        fn(x, cache_x), fused_causal_conv3d_cat_pad_triton(x, cache_x, padding)
+    )
+
+
+# ---------------------------------------------------------------------------
+# Wan causal VAE cache kernels
+# ---------------------------------------------------------------------------
+
+
+def _ref_cat_pad(x, cache, padding):
+    p = list(padding)
+    if cache is not None:
+        x = torch.cat([cache, x], dim=2)
+        p[4] -= cache.shape[2]
+    if any(p):
+        x = F.pad(x, p)
+    return x.contiguous(memory_format=torch.channels_last_3d)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize(
+    "c,t,h,w,cache_t,pads",
+    [
+        (96, 1, 10, 14, 0, (1, 1, 1, 1, 2, 0)),  # first chunk, zero-fill front
+        (96, 1, 10, 14, 1, (1, 1, 1, 1, 2, 0)),  # legacy 1-frame cache
+        (96, 1, 10, 14, 2, (1, 1, 1, 1, 2, 0)),  # steady state k3 conv
+        (64, 1, 10, 14, 2, (0, 0, 0, 0, 2, 0)),  # time_conv (temporal only)
+        (48, 4, 10, 14, 2, (1, 1, 1, 1, 2, 0)),  # encoder-style T=4 chunk
+    ],
+)
+def test_cat_pad_channels_last_3d_bitwise(dtype, c, t, h, w, cache_t, pads):
+    torch.cuda.manual_seed(0)
+    x = _cl3d((1, c, t, h, w), dtype)
+    cache = None
+    if cache_t:
+        # Strided interior view: caches may arrive as non-contiguous slices.
+        ph, pw = pads[2], pads[0]
+        buf = _cl3d((1, c, cache_t, h + 2 * ph, w + 2 * pw), dtype)
+        cache = buf[:, :, :, ph : ph + h, pw : pw + w]
+    ref = _ref_cat_pad(x, cache, pads)
+
+    out = cat_pad_channels_last_3d(x, cache, pads)
+    assert out is not None and out.shape == ref.shape
+    assert out.is_contiguous(memory_format=torch.channels_last_3d)
+    assert torch.equal(out, ref)
+
+    # Dual-output mode: the same pass also emits the compact feature cache
+    # (unpadded interior of the last frames), bitwise equal to the slice.
+    pair = cat_pad_channels_last_3d(x, cache, pads, keep_cache_t=2)
+    assert pair is not None
+    out2, keep = pair
+    assert torch.equal(out2, ref)
+    ph, pw = pads[2], pads[0]
+    keep_t = min(2, ref.shape[2])
+    want = ref[:, :, ref.shape[2] - keep_t :, ph : ph + h, pw : pw + w]
+    assert keep.shape == want.shape
+    assert keep.is_contiguous(memory_format=torch.channels_last_3d)
+    assert torch.equal(keep, want)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize(
+    "c_in,c_out,t,h,w,ft,fs,drop",
+    [
+        (128, 64, 1, 10, 14, 2, 2, False),
+        (128, 64, 1, 10, 14, 2, 2, True),  # first_chunk slicing
+        (64, 32, 2, 10, 14, 1, 2, False),
+    ],
+)
+def test_dup_up3d_add_bitwise(dtype, c_in, c_out, t, h, w, ft, fs, drop):
+    torch.cuda.manual_seed(0)
+    repeats = c_out * ft * fs * fs // c_in
+    src = _cl3d((1, c_in, t, h, w), dtype)
+    t_out = t * ft - (ft - 1 if drop else 0)
+    # Main arm as a permuted view, like the WanResample 2D output.
+    main = torch.randn(
+        (1, t_out, c_out, h * fs, w * fs), device=DEVICE, dtype=dtype
+    ).permute(0, 2, 1, 3, 4)
+
+    dup = src.repeat_interleave(repeats, dim=1)
+    dup = dup.view(1, c_out, ft, fs, fs, t, h, w)
+    dup = dup.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
+    dup = dup.view(1, c_out, t * ft, h * fs, w * fs)
+    if drop:
+        dup = dup[:, :, ft - 1 :, :, :]
+    ref = main + dup
+
+    out = dup_up3d_add(main, src, ft, fs, repeats, drop)
+    assert out is not None and out.shape == ref.shape
+    # Layout must match the aten add output exactly (downstream reductions
+    # are layout-sensitive), and every value must be bitwise identical.
+    assert out.stride() == ref.stride()
+    assert torch.equal(out, ref)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("pads_temporal_only", [False, True])
+def test_wan_cached_conv_chunk_loop_bitwise(pads_temporal_only):
+    """The fused conv-input/compact-cache scheme must reproduce the original
+    clone/cat bookkeeping bitwise across a chunked decode, including the
+    first-chunk zero fill and the "Rep" marker start used by WanResample."""
+    from sglang.multimodal_gen.runtime.models.vaes import wanvae
+    from sglang.multimodal_gen.runtime.models.vaes.wanvae import (
+        CACHE_T,
+        WanCausalConv3d,
+        _cache_payload,
+        _run_cached_causal_conv,
+    )
+
+    torch.cuda.manual_seed(0)
+    c = 64
+    if pads_temporal_only:
+        conv = WanCausalConv3d(c, 2 * c, (3, 1, 1), padding=(1, 0, 0))
+    else:
+        conv = WanCausalConv3d(c, c, 3, padding=1)
+    conv = conv.to(device=DEVICE, dtype=torch.float32)
+    conv.weight.data = conv.weight.data.contiguous(memory_format=torch.channels_last_3d)
+    chunks = [_cl3d((1, c, 1, 10, 14), torch.float32) for _ in range(4)]
+
+    def run(force_fallback, start):
+        cache = [start]
+        outs = []
+        orig = wanvae.cat_pad_channels_last_3d
+        if force_fallback:
+            wanvae.cat_pad_channels_last_3d = None
+        try:
+            for x in chunks:
+                outs.append(_run_cached_causal_conv(conv, x, cache, 0))
+        finally:
+            wanvae.cat_pad_channels_last_3d = orig
+        return outs, cache[0]
+
+    for start in (None, "Rep"):
+        fused_outs, fused_cache = run(False, start)
+        ref_outs, ref_cache = run(True, start)
+        for got, want in zip(fused_outs, ref_outs, strict=True):
+            assert torch.equal(got, want)
+        got_payload = _cache_payload(fused_cache)
+        assert got_payload is not None and got_payload.shape[2] == CACHE_T
+        # Reference cache holds the last CACHE_T unpadded frames.
+        assert torch.equal(got_payload, ref_cache[:, :, -CACHE_T:])
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/diffusion/test_ltx2_ada_values.py

@@ -1,92 +0,0 @@
-import sys
-
-import pytest
-import torch
-
-from sglang.kernels.ops.diffusion.triton.ltx2_ada_values import ltx2_ada_values9
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-
-register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-register_amd_ci(est_time=8, suite="nightly-amd-kernel-1-gpu", nightly=True)
-
-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):
-        assert actual_value.is_contiguous()
-        torch.testing.assert_close(actual_value, expected_value, atol=0, rtol=0)
-
-
-@torch.no_grad()
-def test_ltx2_ada_values9_torch_compile_fullgraph() -> None:
-    hidden = 4096
-    scale_shift_table = torch.randn(
-        9, hidden, device=DEVICE, dtype=torch.bfloat16
-    ).contiguous()
-    timestep = torch.randn(
-        1, 1, 9 * hidden, device=DEVICE, dtype=torch.bfloat16
-    ).contiguous()
-
-    actual = torch.compile(ltx2_ada_values9, fullgraph=True)(
-        scale_shift_table, timestep
-    )
-    expected = _reference(scale_shift_table, timestep)
-
-    assert len(actual) == 9
-    for actual_value, expected_value in zip(actual, expected):
-        assert actual_value.is_contiguous()
-        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"]))

test/registered/kernels/ops/diffusion/test_ltx2_rms_norm_modulate.py

@@ -1,88 +0,0 @@
-"""LTX-2 quality=high RMSNorm+modulate fusion: gated, close to eager."""
-
-import sys
-
-import pytest
-import torch
-from torch import nn
-
-import sglang.multimodal_gen.runtime.models.dits.ltx_2 as ltx2_module
-from sglang.kernels.ops.diffusion.ltx2_rmsnorm_modulate import (
-    fused_ltx2_rms_norm_modulate,
-    mark_ltx2_rms_norm_modulate_site,
-    mount_ltx2_rms_norm_modulate,
-    unmount_ltx2_rms_norm_modulate,
-)
-from sglang.multimodal_gen.runtime.layers.layernorm import RMSNormNoWeight
-from sglang.multimodal_gen.runtime.models.dits.ltx_2 import _ltx2_rms_norm_modulate
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-
-register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-register_amd_ci(est_time=8, suite="nightly-amd-kernel-1-gpu", nightly=True)
-
-
-@pytest.fixture(autouse=True)
-def _setup():
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA required")
-    torch.cuda.manual_seed(0)
-
-
-def _eager(rms, x, scale, shift, eps):
-    return rms(x, eps) * (1 + scale) + shift
-
-
-def _inputs(hidden, batch=1, seq=4096):
-    rms = RMSNormNoWeight()
-    x = torch.randn(batch, seq, hidden, device="cuda", dtype=torch.bfloat16)
-    scale = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16) * 0.1
-    shift = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16) * 0.1
-    return rms, x, scale, shift
-
-
-# hidden 4096 = LTX-2 video stream, 2048 = audio stream.
-@pytest.mark.parametrize("hidden", [4096, 2048])
-def test_lossless_default_is_bitexact(hidden):
-    # A marked-but-unmounted site uses only the self-verified bit-exact
-    # modulate fast path after the reference aten RMSNorm.
-    block = nn.Module()
-    mark_ltx2_rms_norm_modulate_site(block)
-    rms, x, scale, shift = _inputs(hidden)
-    out = _ltx2_rms_norm_modulate(block, rms, x, scale, shift, 1e-6)
-    assert torch.equal(out, _eager(rms, x, scale, shift, 1e-6))
-
-
-def test_lossless_compile_keeps_expression_visible_to_inductor(monkeypatch):
-    block = nn.Module()
-    mark_ltx2_rms_norm_modulate_site(block)
-    rms, x, scale, shift = _inputs(2048, seq=126)
-    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
-    monkeypatch.setattr(
-        ltx2_module,
-        "_ltx2_modulate",
-        lambda *_args: pytest.fail("compiled path must not call the opaque custom op"),
-    )
-
-    out = _ltx2_rms_norm_modulate(block, rms, x, scale, shift, 1e-6)
-    assert torch.equal(out, _eager(rms, x, scale, shift, 1e-6))
-
-
-@pytest.mark.parametrize("hidden", [4096, 2048])
-def test_mounted_high_uses_fused_kernel(hidden):
-    block = nn.Module()
-    mark_ltx2_rms_norm_modulate_site(block)
-    assert mount_ltx2_rms_norm_modulate(block)
-    try:
-        rms, x, scale, shift = _inputs(hidden)
-        out = _ltx2_rms_norm_modulate(block, rms, x, scale, shift, 1e-6)
-        # The mounted path routes through the fused kernel exactly.
-        assert torch.equal(out, fused_ltx2_rms_norm_modulate(x, scale, shift, 1e-6))
-        # And stays within half-precision rounding of the eager reference.
-        ref = _eager(rms, x, scale, shift, 1e-6)
-        assert torch.allclose(out.float(), ref.float(), atol=3e-2, rtol=1e-2)
-    finally:
-        unmount_ltx2_rms_norm_modulate(block)
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/diffusion/test_model_fast_paths.py

@@ -0,0 +1,712 @@
+"""Per-model fast paths: each model wrapper must reproduce its own reference.
+
+A diffusion kernel is only as good as the wrapper that decides when to use it,
+and that decision is model-specific: FLUX.1 and GLM-Image feed different adaLN
+layouts to the same LayerNorm+modulate kernel, Sana only engages on non-default
+streams, ERNIE runs a bit-exact chain unconditionally.  Kernel-level numerics
+live in ``test_norm.py`` / ``test_modulate.py`` / ``test_rope.py``; this file
+covers the wiring: right kernel, right reference, gate ends in the right state.
+
+Two assertion styles appear, and the difference is the point:
+
+- ``torch.equal`` for the bit-exact default-on paths.  These self-verify at
+  runtime, so a tolerance here would let a real regression through -- the gate
+  would silently fall back to eager and the fusion would simply stop running.
+- a tolerance for the quality-gated paths, which are *documented* as differing
+  from eager at half-precision rounding-order level.
+
+Each section also asserts the gate ended up ``verified`` / not ``disabled``:
+without it a test still passes when the fast path never engaged at all.
+"""
+
+import sys
+import unittest
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.upsampling import Upsample2D
+
+import sglang.kernels.ops.diffusion.sites.hunyuan_qknorm_site as hunyuan_qknorm
+import sglang.multimodal_gen.runtime.models.dits.ernie_image as ernie_image
+import sglang.multimodal_gen.runtime.models.dits.flux as flux
+import sglang.multimodal_gen.runtime.models.dits.flux_2 as flux2
+import sglang.multimodal_gen.runtime.models.dits.glm_image as glm_image
+import sglang.multimodal_gen.runtime.models.dits.ltx_2 as ltx2_module
+import sglang.multimodal_gen.runtime.models.dits.sana as sana
+from sglang.kernels.ops.diffusion import (
+    can_use_wan_rmsnorm_silu,
+    fused_ltx2_rms_norm_modulate,
+    mark_fused_ln_modulate_site,
+    mark_hunyuan_qknorm_site,
+    mark_ltx2_rms_norm_modulate_site,
+    mount_fused_ln_modulate,
+    mount_hunyuan_qknorm,
+    mount_ltx2_rms_norm_modulate,
+    unmount_hunyuan_qknorm,
+    unmount_ltx2_rms_norm_modulate,
+    wan_rmsnorm_silu,
+)
+from sglang.multimodal_gen.configs.models.vaes.stablediffusion3 import (
+    StableDiffusion3VAEConfig,
+)
+from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm, RMSNormNoWeight
+from sglang.multimodal_gen.runtime.layers.rotary_embedding.utils import (
+    _apply_rotary_emb,
+)
+from sglang.multimodal_gen.runtime.models.dits.ernie_image import (
+    _ernie_gated_norm_scale_shift,
+    _ernie_norm_scale_shift,
+    _ernie_qknorm_rope,
+    _ernie_qknorm_rope_reference,
+)
+from sglang.multimodal_gen.runtime.models.dits.flux import (
+    _flux_fused_ln_modulate,
+    _flux_norm_modulate,
+)
+from sglang.multimodal_gen.runtime.models.dits.flux_2 import (
+    _flux2_norm_modulate,
+    _flux2_swiglu,
+)
+from sglang.multimodal_gen.runtime.models.dits.glm_image import (
+    _eager_ln_modulate as _glm_eager_ln_modulate,
+)
+from sglang.multimodal_gen.runtime.models.dits.glm_image import (
+    _glm_ln_modulate,
+    _glm_qk_layernorm,
+)
+from sglang.multimodal_gen.runtime.models.dits.hunyuanvideo import (
+    _hunyuan_pack_qkv,
+    _hunyuan_qknorm,
+)
+from sglang.multimodal_gen.runtime.models.dits.ltx_2 import _ltx2_rms_norm_modulate
+from sglang.multimodal_gen.runtime.models.dits.sana import (
+    _eager_ln_modulate as _sana_eager_ln_modulate,
+)
+from sglang.multimodal_gen.runtime.models.dits.sana import (
+    _sana_ln_modulate,
+)
+from sglang.multimodal_gen.runtime.models.vaes import flux2_vae_cuda_opt as vae_opt
+from sglang.multimodal_gen.runtime.models.vaes.autoencoder import AutoencoderKL
+from sglang.multimodal_gen.runtime.models.vaes.fast_path_gate import use_vae_fast_path
+from sglang.multimodal_gen.runtime.models.vaes.wan_vae_cuda_opt import (
+    FusedWanRMSNormSiLU,
+    VaeFastPathGate,
+)
+from sglang.multimodal_gen.runtime.models.vaes.wanvae import WanRMS_norm
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=95, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_amd_ci(est_time=8, suite="nightly-amd-kernel-1-gpu", nightly=True)
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+@pytest.fixture(autouse=True)
+def _seed_cuda():
+    """Every wrapper below asserts against a reference computed from the same
+    random draw, so the seed must be fixed per test, not per module."""
+    torch.cuda.manual_seed(0)
+
+
+# -------------------------------------------------------------------------
+# FLUX.1 -- LayerNorm + adaLN modulate, and the shared-FF GELU site
+# -------------------------------------------------------------------------
+
+
+def _flux_eager(norm, x, scale, shift):
+    return norm(x) * (1 + scale[:, None]) + shift[:, None]
+
+
+def _flux_site_inputs(shape, chunks, seed):
+    torch.manual_seed(seed)
+    batch, seq, hidden = shape
+    norm = torch.nn.LayerNorm(hidden, eps=1e-6, elementwise_affine=False).cuda()
+    x = (torch.randn(batch, seq, hidden, device="cuda") * 8).bfloat16()
+    emb = torch.randn(batch, chunks * hidden, device="cuda").bfloat16()
+    parts = emb.chunk(chunks, dim=1)  # strided adaLN projection views
+    return norm, x, parts[0], parts[1]
+
+
+@pytest.mark.parametrize(
+    "shape,chunks",
+    [
+        ((1, 4096, 3072), 6),  # dual-stream image tokens (1024^2), chunk(6)
+        ((1, 512, 3072), 6),  # dual-stream text tokens
+        ((1, 4608, 3072), 3),  # single-stream concat, chunk(3)
+        ((2, 300, 3072), 6),  # CFG batch, odd seq
+    ],
+)
+def test_flux_fused_ln_modulate_is_bit_exact(shape, chunks):
+    # Every distinct (shape, stride, eps) signature the FLUX.1 sites emit
+    # must verify torch.equal on first sight and stay enabled.
+    norm, x, shift, scale = _flux_site_inputs(shape, chunks, seed=0)
+    out = _flux_fused_ln_modulate(norm, x, scale, shift)
+    assert out is not None
+    assert torch.equal(out, _flux_eager(norm, x, scale, shift))
+    assert not flux._FLUX_LN_MOD.disabled
+    assert flux._FLUX_LN_MOD.verified
+
+
+def test_flux_norm_modulate_bitexact_supersedes_high_fold():
+    # With the quality="high" affine fold mounted, the bit-exact kernel
+    # still takes priority, so the site output stays lossless.
+    site = torch.nn.Module()
+    mark_fused_ln_modulate_site(site)
+    assert mount_fused_ln_modulate(site)
+    norm, x, shift, scale = _flux_site_inputs((1, 128, 3072), 6, seed=1)
+    out = _flux_norm_modulate(site, norm, x, scale, shift)
+    assert torch.equal(out, _flux_eager(norm, x, scale, shift))
+
+
+def test_flux_fused_ln_modulate_rejects_unsupported_hidden():
+    # hidden % 4 != 0 is outside the kernel contract and must bail out.
+    norm, x, shift, scale = _flux_site_inputs((1, 64, 3070), 6, seed=2)
+    assert _flux_fused_ln_modulate(norm, x, scale, shift) is None
+
+
+# -------------------------------------------------------------------------
+# FLUX.2 -- packed norm+modulate and packed SwiGLU views
+# -------------------------------------------------------------------------
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
+class TestFlux2EagerFusions(CustomTestCase):
+    def setUp(self):
+        flux2._FLUX2_LN_MOD.disabled = False
+        flux2._FLUX2_LN_MOD.verified = False
+        flux2._FLUX2_LN_MOD_SIGS.clear()
+        flux2._FLUX2_SWIGLU.disabled = False
+        flux2._FLUX2_SWIGLU.verified = False
+        flux2._FLUX2_SWIGLU_SIGS.clear()
+
+    def test_norm_modulate_is_bit_exact_across_sequence_lengths(self):
+        torch.manual_seed(0)
+        hidden = 256
+        norm = torch.nn.LayerNorm(
+            hidden, eps=1e-6, elementwise_affine=False, device="cuda"
+        )
+        # FLUX.2 modulation values are views of one packed projection.
+        params = torch.randn(1, 1, 6 * hidden, device="cuda").bfloat16()
+        shift, scale = params.chunk(6, dim=-1)[:2]
+
+        for seq in (17, 65):
+            x = torch.randn(1, seq, hidden, device="cuda").bfloat16()
+            expected = norm(x) * (1 + scale) + shift
+            actual = _flux2_norm_modulate(norm, x, scale, shift)
+            self.assertTrue(torch.equal(actual, expected))
+
+        self.assertFalse(flux2._FLUX2_LN_MOD.disabled)
+        self.assertEqual(len(flux2._FLUX2_LN_MOD_SIGS), 1)
+
+    def test_packed_swiglu_is_bit_exact_for_contiguous_and_strided_views(self):
+        torch.manual_seed(1)
+        hidden = 384
+        inputs = [
+            torch.randn(1, 19, 2 * hidden, device="cuda").bfloat16(),
+            torch.randn(1, 19, 3 * hidden, device="cuda").bfloat16()[..., : 2 * hidden],
+        ]
+        for x in inputs:
+            expected = F.silu(x[..., :hidden]) * x[..., hidden:]
+            actual = _flux2_swiglu(x)
+            self.assertTrue(torch.equal(actual, expected))
+
+        self.assertFalse(flux2._FLUX2_SWIGLU.disabled)
+        self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 2)
+
+    def test_fp16_preserves_reference_path(self):
+        x = torch.randn(1, 17, 512, device="cuda", dtype=torch.float16)
+        expected = F.silu(x[..., :256]) * x[..., 256:]
+        actual = _flux2_swiglu(x)
+        self.assertTrue(torch.equal(actual, expected))
+        self.assertFalse(flux2._FLUX2_SWIGLU.disabled)
+
+    def test_packed_swiglu_rejects_non_dense_outer_stride(self):
+        base = torch.randn(2, 23, 512, device="cuda", dtype=torch.bfloat16)
+        x = base[:, :19]
+        self.assertNotEqual(x.stride(0), x.shape[1] * x.stride(1))
+
+        expected = F.silu(x[..., :256]) * x[..., 256:]
+        actual = _flux2_swiglu(x)
+        self.assertTrue(torch.equal(actual, expected))
+        self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 0)
+
+    def test_new_swiglu_signature_is_not_verified_during_graph_capture(self):
+        first = torch.randn(1, 17, 512, device="cuda", dtype=torch.bfloat16)
+        self.assertTrue(
+            torch.equal(
+                _flux2_swiglu(first),
+                F.silu(first[..., :256]) * first[..., 256:],
+            )
+        )
+        self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 1)
+
+        second = torch.randn(1, 19, 768, device="cuda", dtype=torch.bfloat16)
+        with patch("torch.cuda.is_current_stream_capturing", return_value=True):
+            actual = _flux2_swiglu(second)
+
+        expected = F.silu(second[..., :384]) * second[..., 384:]
+        self.assertTrue(torch.equal(actual, expected))
+        self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 1)
+
+
+# -------------------------------------------------------------------------
+# GLM-Image -- LayerNorm + modulate and per-head qk LayerNorm
+# -------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("shape", [(1, 4096, 4096), (2, 301, 4096), (1, 1, 2560)])
+def test_glm_ln_modulate_is_bit_exact(shape):
+    # (1, 4096, 4096) is the real GLM-Image image-stream shape (1024^2,
+    # hidden 4096); the others cover the text stream and another hidden.
+    torch.manual_seed(0)
+    batch, seq, hidden = shape
+    norm = torch.nn.LayerNorm(hidden, eps=1e-5, elementwise_affine=False).cuda()
+    x = (torch.randn(batch, seq, hidden, device="cuda") * 8).bfloat16()
+    emb = torch.randn(batch, 12 * hidden, device="cuda").bfloat16()
+    chunks = emb.chunk(12, dim=1)  # strided adaLN projection views
+    shift, scale = chunks[0], chunks[2]
+    out = _glm_ln_modulate(norm, x, scale, shift, x.dtype)
+    assert torch.equal(out, _glm_eager_ln_modulate(norm, x, scale, shift, x.dtype))
+    assert glm_image._GLM_LN_MOD.verified
+    assert not glm_image._GLM_LN_MOD.disabled
+
+
+@pytest.mark.parametrize("shape", [(1, 4360, 32, 128), (2, 37, 3, 40), (1, 129, 5, 64)])
+def test_glm_qk_head_layernorm_is_bit_exact(shape):
+    # (1, 4360, 32, 128) is the real GLM-Image q/k shape (text + image
+    # tokens, 32 heads of dim 128); the others cover partially-filled warps.
+    torch.manual_seed(1)
+    batch, seq, heads, head_dim = shape
+    norm_q = torch.nn.LayerNorm(head_dim, eps=1e-5, elementwise_affine=False).cuda()
+    norm_k = torch.nn.LayerNorm(head_dim, eps=1e-5, elementwise_affine=False).cuda()
+    q = (torch.randn(batch, seq, heads, head_dim, device="cuda") * 5).bfloat16()
+    k = (torch.randn(batch, seq, heads, head_dim, device="cuda") * 5).bfloat16()
+    q_out, k_out = _glm_qk_layernorm(norm_q, norm_k, q, k, q.dtype)
+    assert torch.equal(q_out, norm_q(q).to(q.dtype))
+    assert torch.equal(k_out, norm_k(k).to(k.dtype))
+    assert glm_image._GLM_QK_LN.verified
+    assert not glm_image._GLM_QK_LN.disabled
+
+
+# -------------------------------------------------------------------------
+# Sana -- stream-conditional LayerNorm + modulate
+# -------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "shape,nmod,transposed",
+    [
+        ((2, 1024, 2240), 6, False),
+        ((2, 1024, 2240), 2, False),
+        ((1, 1024, 2240), 6, True),
+        ((1, 37, 2240), 6, False),
+    ],
+)
+def test_sana_fused_ln_modulate_is_bit_exact(shape, nmod, transposed):
+    # (., 1024, 2240) is the real Sana 1024px shape; hidden 2240 % 512 != 0
+    # exercises the kernel's partial tail chunk.  nmod mirrors the two adaLN
+    # chunk layouts, transposed the permuted layout the Sana DiT serves.
+    torch.manual_seed(0)
+    batch, seq, hidden = shape
+    norm = torch.nn.LayerNorm(hidden, eps=1e-6, elementwise_affine=False).cuda()
+    x = (torch.randn(batch, seq, hidden, device="cuda") * 4).bfloat16()
+    if transposed:
+        x = x.permute(0, 2, 1).contiguous().permute(0, 2, 1)
+    emb = torch.randn(batch, nmod, hidden, device="cuda").bfloat16()
+    shift, scale = emb.chunk(nmod, dim=1)[0], emb.chunk(nmod, dim=1)[-1]
+    # default-stream eager serving must stay on the untouched eager chain
+    n_sigs = len(sana._SANA_LN_MOD.verified_sigs)
+    _sana_ln_modulate(norm, x, scale, shift)
+    assert len(sana._SANA_LN_MOD.verified_sigs) == n_sigs
+    # The fusion engages on non-default streams (the BCG warmup/capture path).
+    # x/scale/shift were filled on the default stream, so the side stream must
+    # wait for that work before reading them -- without this the fused kernel
+    # can read a half-written tensor, the first-sight torch.equal check fails,
+    # and the gate disables itself *permanently*, which then breaks every later
+    # parametrization too. It only loses the race when the GPU is contended,
+    # which is why it shows up on shared CI runners and not on an idle box.
+    side = torch.cuda.Stream()
+    side.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(side):
+        out = _sana_ln_modulate(norm, x, scale, shift)
+        assert len(sana._SANA_LN_MOD.verified_sigs) == n_sigs + 1  # verified
+        out2 = _sana_ln_modulate(norm, x, scale, shift)  # verified-sig lane
+    torch.cuda.current_stream().wait_stream(side)
+    torch.cuda.synchronize()
+    assert torch.equal(out, _sana_eager_ln_modulate(norm, x, scale, shift))
+    assert torch.equal(out2, out) and not sana._SANA_LN_MOD.disabled
+
+
+# -------------------------------------------------------------------------
+# ERNIE-Image -- bit-exact RMSNorm scale/shift and rotate-half RoPE
+# -------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("shape", [(1, 4216, 4096), (2, 1140, 4096), (1, 128, 2048)])
+def test_ernie_norm_scale_shift_is_bit_exact(shape):
+    # (1, 4216, 4096) is the real ERNIE-Image shape (1024^2 image + text
+    # tokens, hidden 4096); 2048 covers the threads_per_row=32 regime.
+    torch.manual_seed(0)
+    batch, seq, hidden = shape
+    norm = RMSNorm(hidden, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
+    with torch.no_grad():
+        norm.weight.copy_(torch.randn(hidden))
+    x = torch.randn(batch, seq, hidden, device="cuda", dtype=torch.bfloat16)
+    residual = torch.randn_like(x)
+    update = torch.randn_like(x)
+    scale = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16) * 0.1
+    shift = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16) * 0.1
+    gate = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16)
+
+    with torch.no_grad():
+        out = _ernie_norm_scale_shift(norm, x, scale, shift)
+        ref = norm(x) * (1 + scale) + shift
+        assert torch.equal(out, ref)
+
+        out2, res = _ernie_gated_norm_scale_shift(
+            norm, residual, update, gate, scale, shift
+        )
+        res_ref = residual + gate * update
+        ref2 = norm(res_ref) * (1 + scale) + shift
+        assert torch.equal(res, res_ref)
+        assert torch.equal(out2, ref2)
+
+    # the fast paths must actually be in use (not silently disabled)
+    assert ernie_image._ERNIE_NORM.verified
+    assert ernie_image._ERNIE_GATED_NORM.verified
+    assert not ernie_image._ERNIE_NORM.disabled
+    assert not ernie_image._ERNIE_GATED_NORM.disabled
+
+
+def test_ernie_qknorm_rope_is_bit_exact():
+    torch.manual_seed(1)
+    ernie_image._ERNIE_QKNORM_ROPE.disabled = False
+    ernie_image._ERNIE_QKNORM_ROPE.verified = False
+    batch, seq, heads, head_dim = 1, 257, 32, 128
+    q = torch.randn(batch, seq, heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    q_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
+    k_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
+    cos = torch.randn(seq, head_dim, device="cuda", dtype=torch.bfloat16)
+    sin = torch.randn_like(cos)
+    cache = torch.cat((cos, sin), dim=-1).contiguous()
+    positions = torch.arange(seq, device="cuda", dtype=torch.long)
+
+    q_ref, k_ref = _ernie_qknorm_rope_reference(
+        q.clone(), k.clone(), q_norm, k_norm, head_dim, cos, sin
+    )
+    q_out, k_out = _ernie_qknorm_rope(
+        q,
+        k,
+        q_norm,
+        k_norm,
+        head_dim,
+        cos,
+        sin,
+        cache,
+        positions,
+    )
+
+    assert torch.equal(q_out, q_ref)
+    assert torch.equal(k_out, k_ref)
+    assert ernie_image._ERNIE_QKNORM_ROPE.verified
+    assert not ernie_image._ERNIE_QKNORM_ROPE.disabled
+
+
+def test_ernie_qknorm_rope_first_attempt_exception_uses_pristine_inputs():
+    torch.manual_seed(2)
+    ernie_image._ERNIE_QKNORM_ROPE.disabled = False
+    ernie_image._ERNIE_QKNORM_ROPE.verified = False
+    batch, seq, heads, head_dim = 1, 17, 4, 128
+    q = torch.randn(batch, seq, heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    q_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
+    k_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
+    cos = torch.randn(seq, head_dim, device="cuda", dtype=torch.bfloat16)
+    sin = torch.randn_like(cos)
+    cache = torch.cat((cos, sin), dim=-1).contiguous()
+    positions = torch.arange(seq, device="cuda", dtype=torch.long)
+    q_ref, k_ref = _ernie_qknorm_rope_reference(
+        q.clone(), k.clone(), q_norm, k_norm, head_dim, cos, sin
+    )
+
+    def mutate_then_raise(**kwargs):
+        kwargs["q"].zero_()
+        kwargs["k"].zero_()
+        raise RuntimeError("synthetic kernel failure")
+
+    with patch.object(ernie_image, "apply_qk_norm_rope", mutate_then_raise):
+        q_out, k_out = _ernie_qknorm_rope(
+            q,
+            k,
+            q_norm,
+            k_norm,
+            head_dim,
+            cos,
+            sin,
+            cache,
+            positions,
+        )
+
+    assert torch.equal(q_out, q_ref)
+    assert torch.equal(k_out, k_ref)
+    assert ernie_image._ERNIE_QKNORM_ROPE.disabled
+
+
+# -------------------------------------------------------------------------
+# LTX-2 -- weightless RMSNorm + modulate (quality-gated)
+# -------------------------------------------------------------------------
+
+
+def _ltx2_eager(rms, x, scale, shift, eps):
+    return rms(x, eps) * (1 + scale) + shift
+
+
+def _ltx2_inputs(hidden, batch=1, seq=4096):
+    rms = RMSNormNoWeight()
+    x = torch.randn(batch, seq, hidden, device="cuda", dtype=torch.bfloat16)
+    scale = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16) * 0.1
+    shift = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16) * 0.1
+    return rms, x, scale, shift
+
+
+# hidden 4096 = LTX-2 video stream, 2048 = audio stream.
+@pytest.mark.parametrize("hidden", [4096, 2048])
+def test_ltx2_lossless_default_is_bitexact(hidden):
+    # A marked-but-unmounted site uses only the self-verified bit-exact
+    # modulate fast path after the reference aten RMSNorm.
+    block = nn.Module()
+    mark_ltx2_rms_norm_modulate_site(block)
+    rms, x, scale, shift = _ltx2_inputs(hidden)
+    out = _ltx2_rms_norm_modulate(block, rms, x, scale, shift, 1e-6)
+    assert torch.equal(out, _ltx2_eager(rms, x, scale, shift, 1e-6))
+
+
+def test_ltx2_lossless_compile_keeps_expression_visible_to_inductor(monkeypatch):
+    block = nn.Module()
+    mark_ltx2_rms_norm_modulate_site(block)
+    rms, x, scale, shift = _ltx2_inputs(2048, seq=126)
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
+    monkeypatch.setattr(
+        ltx2_module,
+        "_ltx2_modulate",
+        lambda *_args: pytest.fail("compiled path must not call the opaque custom op"),
+    )
+
+    out = _ltx2_rms_norm_modulate(block, rms, x, scale, shift, 1e-6)
+    assert torch.equal(out, _ltx2_eager(rms, x, scale, shift, 1e-6))
+
+
+@pytest.mark.parametrize("hidden", [4096, 2048])
+def test_ltx2_mounted_high_uses_fused_kernel(hidden):
+    block = nn.Module()
+    mark_ltx2_rms_norm_modulate_site(block)
+    assert mount_ltx2_rms_norm_modulate(block)
+    try:
+        rms, x, scale, shift = _ltx2_inputs(hidden)
+        out = _ltx2_rms_norm_modulate(block, rms, x, scale, shift, 1e-6)
+        # The mounted path routes through the fused kernel exactly.
+        assert torch.equal(out, fused_ltx2_rms_norm_modulate(x, scale, shift, 1e-6))
+        # And stays within half-precision rounding of the eager reference.
+        ref = _ltx2_eager(rms, x, scale, shift, 1e-6)
+        assert torch.allclose(out.float(), ref.float(), atol=3e-2, rtol=1e-2)
+    finally:
+        unmount_ltx2_rms_norm_modulate(block)
+
+
+# -------------------------------------------------------------------------
+# HunyuanVideo -- QKV/RoPE pack and quality-gated QK RMSNorm
+# -------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("img_tokens,txt_tokens", [(257, 31), (4096, 256)])
+def test_hunyuan_qkv_rope_pack_is_bit_exact(img_tokens, txt_tokens):
+    torch.manual_seed(0)
+    shape_img = (1, img_tokens, 24, 128)
+    shape_txt = (1, txt_tokens, 24, 128)
+    img_q, img_k, img_v = (
+        torch.randn(shape_img, device="cuda", dtype=torch.bfloat16) for _ in range(3)
+    )
+    txt_q, txt_k, txt_v = (
+        torch.randn(shape_txt, device="cuda", dtype=torch.bfloat16) for _ in range(3)
+    )
+    cos = torch.randn(img_tokens, 64, device="cuda")
+    sin = torch.randn_like(cos)
+
+    q, k, v = _hunyuan_pack_qkv(img_q, img_k, img_v, txt_q, txt_k, txt_v, cos, sin)
+    q_ref = torch.cat(
+        (_apply_rotary_emb(img_q, cos, sin, is_neox_style=False), txt_q), dim=1
+    )
+    k_ref = torch.cat(
+        (_apply_rotary_emb(img_k, cos, sin, is_neox_style=False), txt_k), dim=1
+    )
+    v_ref = torch.cat((img_v, txt_v), dim=1)
+
+    assert torch.equal(q, q_ref)
+    assert torch.equal(k, k_ref)
+    assert torch.equal(v, v_ref)
+
+
+def test_hunyuan_quality_qknorm_matches_rmsnorm():
+    torch.manual_seed(1)
+    site = torch.nn.Module()
+    mark_hunyuan_qknorm_site(site)
+    q_norm = RMSNorm(128, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
+    k_norm = RMSNorm(128, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
+    packed = torch.randn(1, 257, 3, 24, 128, device="cuda", dtype=torch.bfloat16)
+    q, k = packed[:, :, 0], packed[:, :, 1]
+    q_ref = q_norm(q.contiguous()).to(q)
+    k_ref = k_norm(k.contiguous()).to(k)
+
+    q_unmounted, k_unmounted = _hunyuan_qknorm(site, q, k, q_norm, k_norm)
+    assert torch.equal(q_unmounted, q_ref)
+    assert torch.equal(k_unmounted, k_ref)
+
+    assert mount_hunyuan_qknorm(site)
+    q_out, k_out = _hunyuan_qknorm(site, q, k, q_norm, k_norm)
+    torch.testing.assert_close(q_out, q_ref, atol=2e-2, rtol=2e-2)
+    torch.testing.assert_close(k_out, k_ref, atol=2e-2, rtol=2e-2)
+
+    unmount_hunyuan_qknorm(site)
+    q_unmounted, k_unmounted = _hunyuan_qknorm(site, q, k, q_norm, k_norm)
+    assert torch.equal(q_unmounted, q_ref)
+    assert torch.equal(k_unmounted, k_ref)
+
+
+def test_hunyuan_quality_qknorm_stays_unmounted_without_cute_kernel():
+    site = torch.nn.Module()
+    mark_hunyuan_qknorm_site(site)
+
+    with patch.object(hunyuan_qknorm, "_get_qk_rmsnorm_cute", return_value=None):
+        assert not mount_hunyuan_qknorm(site)
+
+    assert not hunyuan_qknorm._FUSION.is_enabled(site)
+
+
+# -------------------------------------------------------------------------
+# Wan VAE -- fused RMSNorm+SiLU module gate
+# -------------------------------------------------------------------------
+
+
+def _wan_cl3d(shape, dtype):
+    return torch.randn(shape, device="cuda", dtype=dtype).contiguous(
+        memory_format=torch.channels_last_3d
+    )
+
+
+@torch.no_grad()
+def test_wan_vae_gate_dispatch() -> None:
+    # Gate off must stay bit-exact; gate on must route to the fused kernel.
+    torch.cuda.manual_seed(0)
+    norm = WanRMS_norm(96, images=False).to(device="cuda", dtype=torch.bfloat16)
+    norm.gamma.add_(torch.randn_like(norm.gamma))
+    gate = VaeFastPathGate()
+    fused = FusedWanRMSNormSiLU(norm, gate)
+    # Parameter names must not change (weight transfer matches by name).
+    assert [n for n, _ in fused.named_parameters()] == ["gamma"]
+    x = _wan_cl3d((1, 96, 3, 10, 14), torch.bfloat16)
+    assert torch.equal(fused(x), nn.SiLU()(norm(x)))
+    gate.enabled = True
+    expected = wan_rmsnorm_silu(x, norm.gamma, rms_scale=float(norm.scale))
+    assert torch.equal(fused(x), expected)
+
+
+@torch.no_grad()
+def test_wan_vae_rejects_empty_input() -> None:
+    x = torch.empty(1, 96, 0, 2, 2, device="cuda", dtype=torch.bfloat16).to(
+        memory_format=torch.channels_last_3d
+    )
+    gamma = torch.ones(96, 1, 1, 1, device="cuda", dtype=torch.bfloat16)
+    assert not can_use_wan_rmsnorm_silu(x, gamma, None)
+
+
+# -------------------------------------------------------------------------
+# FLUX.2 VAE -- fused GroupNorm+SiLU and folded 2x upsample conv
+# -------------------------------------------------------------------------
+
+
+@torch.no_grad()
+def test_flux2_vae_fast_path():
+    torch.manual_seed(0)
+    gate = vae_opt.VaeFastPathGate()
+    gn = nn.GroupNorm(32, 128, eps=1e-6).to("cuda", torch.bfloat16)
+    x = torch.randn(1, 128, 64, 64, device="cuda", dtype=torch.bfloat16).to(
+        memory_format=torch.channels_last
+    )
+    ref = F.silu(gn(x))
+    fused_gn = vae_opt.FusedGroupNormSiLU(gn, gate)
+    assert set(fused_gn.state_dict()) == {"weight", "bias"}
+    assert torch.equal(fused_gn(x), ref)  # gate off: bit-exact reference
+
+    gate.enabled = True
+    fast = fused_gn(x)
+    assert fast.is_contiguous(memory_format=torch.channels_last)
+    torch.testing.assert_close(fast.float(), ref.float(), atol=0.06, rtol=0)
+
+    gate.enabled = False
+    up = Upsample2D(channels=32, use_conv=True).to("cuda", torch.bfloat16)
+    fused_up = vae_opt.FusedUpsample2xConv2d(up, gate)
+    assert set(fused_up.state_dict()) == {"conv.weight", "conv.bias"}
+    x = torch.randn(2, 32, 33, 29, device="cuda", dtype=torch.bfloat16)
+    ref = up(x)
+    assert torch.equal(fused_up(x), ref)
+    assert fused_up._fused_weight is None
+
+    gate.enabled = True
+    fast = fused_up(x)
+    assert fused_up._fused_weight is not None
+    ref_range = ref.float().max() - ref.float().min()
+    relative_mse = F.mse_loss(fast.float(), ref.float()) / ref_range.square()
+    assert relative_mse < 3.2e-5
+
+
+# ---------------------------------------------------------------------------
+# AutoencoderKL (generic) -- fast-path install must not disturb the checkpoint
+# ---------------------------------------------------------------------------
+
+
+def _small_config():
+    config = StableDiffusion3VAEConfig()
+    config.arch_config.latent_channels = 2
+    config.arch_config.block_out_channels = (4, 4)
+    config.arch_config.down_block_types = ("DownEncoderBlock2D",) * 2
+    config.arch_config.up_block_types = ("UpDecoderBlock2D",) * 2
+    config.arch_config.layers_per_block = 1
+    config.arch_config.norm_num_groups = 1
+    config.arch_config.sample_size = 8
+    return config
+
+
+@torch.no_grad()
+def test_autoencoder_kl_fastpath_install():
+    torch.manual_seed(0)
+    vae = AutoencoderKL(_small_config()).to("cuda", torch.bfloat16).eval()
+    ref_names = {n for n, _ in vae.named_parameters()}
+    ref_sd = {k: v.clone() for k, v in vae.state_dict().items()}
+    z = torch.randn(1, 2, 8, 8, device="cuda", dtype=torch.bfloat16)
+    ref = vae.decode(z)
+
+    opt = vae_opt.maybe_optimize_autoencoder_kl(vae)
+    # Wrappers must not change parameter FQNs; strict load must round-trip.
+    assert {n for n, _ in opt.named_parameters()} == ref_names
+    opt.load_state_dict(ref_sd, strict=True)
+    # Gate off: bit-for-bit the original path.
+    assert torch.equal(opt.decode(z), ref)
+    # use_vae_fast_path() is a no-op when nothing registered a gate, so check
+    # the wrappers went in before relying on it to switch paths.
+    assert any(
+        isinstance(m, (vae_opt.FusedGroupNormSiLU, vae_opt.FusedUpsample2xConv2d))
+        for m in opt.modules()
+    )
+    # Gate on: fast path runs and stays close; leaving the scope restores exact.
+    with use_vae_fast_path(opt, True):
+        torch.testing.assert_close(opt.decode(z).float(), ref.float(), atol=0.1, rtol=0)
+    assert torch.equal(opt.decode(z), ref)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/diffusion/test_modulate.py

@@ -0,0 +1,424 @@
+"""``diffusion.modulate``: adaLN modulation, gating and timestep conditioning.
+
+The bit-exact kernels here (``modulate_scale_shift``, ``residual_gate_add``,
+``ltx2_ada_values9``, ``try_fused_scaled_residual_add_exact``) reproduce every
+aten rounding boundary, so they are asserted with ``torch.equal``.  The
+select-0/1 LayerNorm fusions compute their statistics differently from the
+reference chain and are asserted with a tolerance.
+"""
+
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.diffusion import (
+    can_use_modulate_scale_shift_cuda,
+    can_use_residual_gate_add_cuda,
+    fuse_layernorm_scale_shift_gate_select01_kernel,
+    fuse_residual_layernorm_scale_shift_gate_select01_kernel,
+    ltx2_ada_values9,
+    modulate_scale_shift,
+    modulate_scale_shift_cuda,
+    norm_infer,
+    residual_gate_add,
+    residual_gate_add_cuda,
+    timestep_embedding,
+    try_fused_scaled_residual_add_exact,
+)
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+
+register_cuda_ci(est_time=75, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+# Nightly is not redundant: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1, which
+# expands the get_ci_test_range sweeps below.
+register_cuda_ci(est_time=50, stage="nightly", runner_config="1-gpu-large")
+register_amd_ci(est_time=38, suite="nightly-amd-kernel-1-gpu", nightly=True)
+
+DEVICE = "cuda"
+
+
+@pytest.fixture(autouse=True)
+def cuda_setup():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    torch.cuda.manual_seed(0)
+
+
+# ---------------------------------------------------------------------------
+# modulate: x * (1 + scale) + shift
+# ---------------------------------------------------------------------------
+
+# FLUX.1 1024^2 adaLN shapes (D=3072) plus batched and odd-length coverage.
+MODULATE_CASES = [
+    (1, 4096, 3072),
+    (1, 512, 3072),
+    (1, 4608, 3072),
+    (2, 1024, 3072),
+    (1, 17, 64),
+]
+
+
+def _eager_modulate(x, scale, shift):
+    return x * (1 + scale[:, None]) + shift[:, None]
+
+
+@pytest.mark.parametrize("shape", MODULATE_CASES)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_modulate_scale_shift_matches_eager(shape, dtype):
+    x = torch.randn(shape, device=DEVICE, dtype=dtype)
+    scale = torch.randn((shape[0], shape[-1]), device=DEVICE, dtype=dtype)
+    shift = torch.randn_like(scale)
+    assert torch.equal(
+        modulate_scale_shift_cuda(x, scale, shift), _eager_modulate(x, scale, shift)
+    )
+
+
+def test_modulate_scale_shift_accepts_adaln_chunk_views():
+    # Production feeds strided ``emb.chunk(6)`` views, not fresh tensors.
+    x = torch.randn((1, 4096, 3072), device=DEVICE, dtype=torch.bfloat16)
+    emb = torch.randn((1, 6 * 3072), device=DEVICE, dtype=torch.bfloat16)
+    shift, scale = emb.chunk(6, dim=1)[:2]
+    assert can_use_modulate_scale_shift_cuda(x, scale, shift)
+    assert torch.equal(
+        modulate_scale_shift_cuda(x, scale, shift), _eager_modulate(x, scale, shift)
+    )
+
+
+def test_modulate_scale_shift_guards_reject_fp32():
+    x = torch.randn((1, 64, 64), device=DEVICE, dtype=torch.float32)
+    row = torch.randn((1, 64), device=DEVICE, dtype=torch.float32)
+    assert not can_use_modulate_scale_shift_cuda(x, row, row)
+    # The public wrapper still returns the eager result on a rejected input.
+    assert torch.equal(modulate_scale_shift(x, row, row), _eager_modulate(x, row, row))
+
+
+# ---------------------------------------------------------------------------
+# residual + gate * update
+# ---------------------------------------------------------------------------
+
+GATE_CASES = [
+    ((1, 1024, 4096), (1, 1, 4096)),
+    ((1, 512, 4096), (1, 512, 4096)),
+    ((1, 17, 65), (1, 1, 65)),
+    ((1, 17, 65), (1, 17, 65)),
+    # FLUX.1 / FLUX.2-klein 1024^2 shapes (D=3072): dual-stream image/text
+    # and single-stream/joint concat; gates are [1, 1, D] modulation rows.
+    ((1, 4096, 3072), (1, 1, 3072)),
+    ((1, 512, 3072), (1, 1, 3072)),
+    ((1, 4608, 3072), (1, 1, 3072)),
+    # FLUX.2-dev (D=6144) joint sequence.
+    ((1, 4608, 6144), (1, 1, 6144)),
+    # ERNIE-4.5-VL 1024^2 image tokens plus text tokens.
+    ((1, 4216, 4096), (1, 1, 4096)),
+]
+
+
+def _assert_gate_add(out, ref):
+    if ref.dtype == torch.float32:
+        # fp32 has no rounding boundary to reproduce; the kernel keeps the
+        # accumulation in fp32 and only order may differ.
+        torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
+    else:
+        assert torch.equal(out, ref)
+
+
+@pytest.mark.parametrize("residual_shape,gate_shape", GATE_CASES)
+def test_residual_gate_add_matches_torch(residual_shape, gate_shape):
+    residual = torch.randn(residual_shape, device=DEVICE, dtype=torch.bfloat16)
+    update = torch.randn_like(residual)
+    gate = torch.randn(gate_shape, device=DEVICE, dtype=torch.bfloat16)
+
+    ref = residual + update * gate
+    _assert_gate_add(residual_gate_add_cuda(residual, update, gate), ref)
+    assert torch.equal(residual_gate_add(residual, update, gate), ref)
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("gate_shape", [(1, 1, 64), (1, 9, 64)])
+def test_residual_gate_add_dtypes(dtype, gate_shape):
+    residual = torch.randn((1, 9, 64), device=DEVICE, dtype=dtype)
+    update = torch.randn_like(residual)
+    gate = torch.randn(gate_shape, device=DEVICE, dtype=dtype)
+    _assert_gate_add(
+        residual_gate_add_cuda(residual, update, gate), residual + update * gate
+    )
+
+
+def test_residual_gate_add_guards_and_eager_fallback():
+    residual = torch.randn((1, 8, 64), device=DEVICE, dtype=torch.bfloat16)
+    update = torch.randn_like(residual)
+    gate = torch.randn((1, 1, 64), device=DEVICE, dtype=torch.bfloat16)
+    assert can_use_residual_gate_add_cuda(residual, update, gate)
+
+    rejected = [
+        (residual.cpu(), update, gate),  # not on device
+        (residual, update.float(), gate),  # mixed dtypes
+        (residual, update[:, ::2], gate),  # strided update
+        (residual, update, gate[:, :, ::2]),  # strided gate
+        (residual[:, :0], update[:, :0], gate),  # empty token dim
+    ]
+    for args in rejected:
+        assert not can_use_residual_gate_add_cuda(*args)
+
+    # Only [1, ..., 1, D] row-broadcast gates are supported; a batched
+    # [B>1, 1, D] gate is not row-broadcast here and must fall back.
+    batched = torch.randn((2, 8, 64), device=DEVICE, dtype=torch.bfloat16)
+    batched_update = torch.randn_like(batched)
+    batched_gate = torch.randn((2, 1, 64), device=DEVICE, dtype=torch.bfloat16)
+    assert not can_use_residual_gate_add_cuda(batched, batched_update, batched_gate)
+    assert torch.equal(
+        residual_gate_add(batched, batched_update, batched_gate),
+        batched + batched_update * batched_gate,
+    )
+
+
+def test_residual_gate_add_torch_compile_fullgraph():
+    residual = torch.randn((1, 32, 128), device=DEVICE, dtype=torch.bfloat16)
+    update = torch.randn_like(residual)
+    gate = torch.randn((1, 1, 128), device=DEVICE, dtype=torch.bfloat16)
+    compiled = torch.compile(residual_gate_add, fullgraph=True)
+    assert torch.equal(compiled(residual, update, gate), residual + update * gate)
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_scaled_residual_add_is_bit_exact(dtype):
+    # fp32 residual accumulator + half-precision update, as the DiT blocks
+    # that keep their residual stream in fp32 emit it.
+    residual = torch.randn(2, 17, 64, device=DEVICE, dtype=torch.float32)
+    x = torch.randn(2, 17, 64, device=DEVICE, dtype=dtype)
+    scale = torch.randn(64, device=DEVICE, dtype=torch.float32)
+
+    actual = try_fused_scaled_residual_add_exact(residual, x, scale)
+    assert actual is not None
+    assert torch.equal(actual, residual + x * scale)
+
+
+@torch.no_grad()
+def test_scaled_residual_add_rejects_unsupported_inputs():
+    residual = torch.empty(2, 3, 8, device=DEVICE, dtype=torch.float32)
+    x = torch.empty_like(residual)
+    scale = torch.empty(8, device=DEVICE, dtype=torch.float32)
+    # A too-small hidden dim and a mismatched scale length both bail out;
+    # ``try_`` returning None is this helper's documented contract.
+    assert try_fused_scaled_residual_add_exact(residual, x, scale) is None
+    assert try_fused_scaled_residual_add_exact(residual, x.half(), scale[:-1]) is None
+
+
+# ---------------------------------------------------------------------------
+# LTX-2 nine-way adaLN value split
+# ---------------------------------------------------------------------------
+
+
+def _ltx2_reference(scale_shift_table, timestep):
+    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])
+@pytest.mark.parametrize("compiled", [False, True])
+def test_ltx2_ada_values9(batch, seq, hidden, table_dtype, compiled):
+    scale_shift_table = torch.randn(9, hidden, device=DEVICE, dtype=table_dtype)
+    timestep = torch.randn(batch, seq, 9 * hidden, device=DEVICE, dtype=torch.bfloat16)
+
+    fn = (
+        torch.compile(ltx2_ada_values9, fullgraph=True)
+        if compiled
+        else ltx2_ada_values9
+    )
+    actual = fn(scale_shift_table, timestep)
+    expected = _ltx2_reference(scale_shift_table, timestep)
+
+    assert len(actual) == 9
+    for got, want in zip(actual, expected, strict=True):
+        # Each slice must come out naturally contiguous -- that is the point
+        # of the kernel; a strided slice would re-add the downstream copy.
+        assert got.is_contiguous()
+        assert torch.equal(got, want)
+
+
+@torch.no_grad()
+def test_ltx2_ada_values9_rejects_unsupported_shape():
+    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)
+
+
+# ---------------------------------------------------------------------------
+# select-0/1 LayerNorm modulation (Qwen-Image)
+# ---------------------------------------------------------------------------
+
+SELECT01_DTYPES = get_ci_test_range(
+    [torch.float16, torch.bfloat16, torch.float32], [torch.float16, torch.bfloat16]
+)
+SELECT01_SHAPES = get_ci_test_range(
+    [(b, s, h) for b in (1, 2, 4) for s in (6, 33, 128, 257) for h in (512, 3072)],
+    [(1, 6, 512), (2, 128, 3072)],
+)
+EPS = 1e-6
+
+
+def _select01_reference(normalized, mods, index):
+    scale0, shift0, gate0, scale1, shift1, gate1 = mods
+    idx = index.bool().unsqueeze(-1)
+    scale = torch.where(idx, scale1.unsqueeze(1), scale0.unsqueeze(1))
+    shift = torch.where(idx, shift1.unsqueeze(1), shift0.unsqueeze(1))
+    gate = torch.where(idx, gate1.unsqueeze(1), gate0.unsqueeze(1))
+    return normalized * (1 + scale) + shift, gate
+
+
+@pytest.mark.parametrize("dtype", SELECT01_DTYPES)
+@pytest.mark.parametrize("shape", SELECT01_SHAPES)
+@pytest.mark.parametrize("with_residual", [False, True])
+def test_layernorm_scale_shift_gate_select01(dtype, shape, with_residual):
+    batch_size, seq_len, hidden_size = shape
+    x = torch.randn(batch_size, seq_len, hidden_size, device=DEVICE, dtype=dtype)
+    weight = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
+    bias = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
+    index = torch.randint(0, 2, (batch_size, seq_len), device=DEVICE, dtype=torch.int32)
+    mods = tuple(
+        torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
+        for _ in range(6)
+    )
+    scale0, shift0, gate0, scale1, shift1, gate1 = mods
+
+    if with_residual:
+        residual = torch.randn_like(x)
+        residual_gate = torch.randn_like(x)
+        residual_ref = residual + residual_gate * x
+        normalized = norm_infer(
+            residual_ref.flatten(0, 1), weight, bias, eps=EPS, is_rms_norm=False
+        ).view_as(residual_ref)
+        out_ref, gate_ref = _select01_reference(normalized, mods, index)
+        out, residual_out, gate = (
+            fuse_residual_layernorm_scale_shift_gate_select01_kernel(
+                x.contiguous(),
+                residual=residual.contiguous(),
+                residual_gate=residual_gate.contiguous(),
+                weight=weight,
+                bias=bias,
+                scale0=scale0,
+                shift0=shift0,
+                gate0=gate0,
+                scale1=scale1,
+                shift1=shift1,
+                gate1=gate1,
+                index=index,
+                eps=EPS,
+            )
+        )
+    else:
+        normalized = norm_infer(
+            x.flatten(0, 1), weight, bias, eps=EPS, is_rms_norm=False
+        ).view_as(x)
+        out_ref, gate_ref = _select01_reference(normalized, mods, index)
+        residual_ref = residual_out = None
+        out, gate = fuse_layernorm_scale_shift_gate_select01_kernel(
+            x.contiguous(),
+            weight=weight,
+            bias=bias,
+            scale0=scale0,
+            shift0=shift0,
+            gate0=gate0,
+            scale1=scale1,
+            shift1=shift1,
+            gate1=gate1,
+            index=index,
+            eps=EPS,
+        )
+
+    tol = 1e-5 if dtype == torch.float32 else 5e-2
+    torch.testing.assert_close(out, out_ref, atol=tol, rtol=tol)
+    torch.testing.assert_close(gate, gate_ref, atol=tol, rtol=tol)
+    if with_residual:
+        torch.testing.assert_close(residual_out, residual_ref, atol=tol, rtol=tol)
+
+
+# ---------------------------------------------------------------------------
+# Sinusoidal timestep embedding
+# ---------------------------------------------------------------------------
+
+TIMESTEP_BATCHES = get_ci_test_range(
+    [1, 2, 8, 128, 256, 512, 1536, 2048, 4096, 11008, 16384], [1, 128, 2048, 16384]
+)
+TIMESTEP_DIMS = get_ci_test_range(
+    [32, 128, 256, 512, 1536, 2048, 4096, 8192], [32, 512, 8192]
+)
+TIMESTEP_DTYPES = get_ci_test_range(
+    [torch.float16, torch.bfloat16, torch.float32], [torch.float16, torch.bfloat16]
+)
+
+
+def timestep_embedding_reference(
+    timesteps,
+    dim,
+    *,
+    flip_sin_to_cos=False,
+    downscale_freq_shift=1,
+    scale=1,
+    max_period=10000,
+):
+    """diffusers' ``get_timestep_embedding``, kept verbatim as the oracle."""
+    assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
+    timesteps = timesteps.to(torch.float32)
+    half_dim = dim // 2
+    exponent = -torch.log(
+        torch.tensor(max_period, dtype=torch.float32, device=timesteps.device)
+    ) * torch.arange(
+        start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
+    )
+    exponent = exponent / (half_dim - downscale_freq_shift)
+
+    emb = torch.exp(exponent)
+    emb = timesteps[:, None].float() * emb[None, :]
+    emb = scale * emb
+    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+    if flip_sin_to_cos:
+        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+    if dim % 2 == 1:
+        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+    return emb
+
+
+@pytest.mark.parametrize("batch_size", TIMESTEP_BATCHES)
+@pytest.mark.parametrize("dim", TIMESTEP_DIMS)
+@pytest.mark.parametrize("dtype", TIMESTEP_DTYPES)
+@pytest.mark.parametrize(
+    "flip_sin_to_cos,downscale_freq_shift,scale",
+    [
+        (True, 0, 1),  # the sgl-diffusion default
+        (False, 1, 1),  # the diffusers default
+        (True, 1, 0.01),  # scaled variant used by the SD-style embedders
+    ],
+)
+def test_timestep_embedding_matches_diffusers(
+    batch_size, dim, dtype, flip_sin_to_cos, downscale_freq_shift, scale
+):
+    t = torch.randint(low=0, high=1000, size=(batch_size,), device=DEVICE).to(dtype)
+    kwargs = dict(
+        flip_sin_to_cos=flip_sin_to_cos,
+        downscale_freq_shift=downscale_freq_shift,
+        scale=scale,
+        max_period=10000,
+    )
+    torch.testing.assert_close(
+        timestep_embedding(t, dim, **kwargs),
+        timestep_embedding_reference(t, dim, **kwargs),
+        atol=1e-3,
+        rtol=1e-3,
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/diffusion/test_modulate_scale_shift.py

@@ -1,58 +0,0 @@
-import pytest
-import torch
-
-from sglang.kernels.ops.diffusion.modulate_scale_shift import (
-    can_use_modulate_scale_shift_cuda,
-    modulate_scale_shift,
-    modulate_scale_shift_cuda,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
-
-# FLUX.1 1024^2 adaLN shapes (D=3072) plus batched and odd-length coverage.
-CASES = [(1, 4096, 3072), (1, 512, 3072), (1, 4608, 3072), (2, 1024, 3072), (1, 17, 64)]
-
-
-@pytest.fixture(autouse=True)
-def cuda_setup():
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA required")
-    torch.cuda.manual_seed(0)
-
-
-def _eager(x, scale, shift):
-    return x * (1 + scale[:, None]) + shift[:, None]
-
-
-@pytest.mark.parametrize("shape", CASES)
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
-def test_modulate_scale_shift_matches_eager(shape, dtype):
-    x = torch.randn(shape, device="cuda", dtype=dtype)
-    scale = torch.randn((shape[0], shape[-1]), device="cuda", dtype=dtype)
-    shift = torch.randn_like(scale)
-    out = modulate_scale_shift_cuda(x, scale, shift)
-    assert torch.equal(out, _eager(x, scale, shift))  # bitwise contract
-
-
-def test_modulate_scale_shift_adaln_chunk_views():
-    x = torch.randn((1, 4096, 3072), device="cuda", dtype=torch.bfloat16)
-    emb = torch.randn((1, 6 * 3072), device="cuda", dtype=torch.bfloat16)
-    shift, scale = emb.chunk(6, dim=1)[:2]
-    assert can_use_modulate_scale_shift_cuda(x, scale, shift)
-    out = modulate_scale_shift_cuda(x, scale, shift)
-    assert torch.equal(out, _eager(x, scale, shift))
-
-
-def test_modulate_scale_shift_guards_reject_fp32():
-    x = torch.randn((1, 64, 64), device="cuda", dtype=torch.float32)
-    row = torch.randn((1, 64), device="cuda", dtype=torch.float32)
-    assert not can_use_modulate_scale_shift_cuda(x, row, row)
-    assert torch.equal(modulate_scale_shift(x, row, row), _eager(x, row, row))
-
-
-if __name__ == "__main__":
-    import sys
-
-    sys.exit(pytest.main([__file__]))

test/registered/kernels/ops/diffusion/test_native_bf16_rmsnorm.py

@@ -1,87 +0,0 @@
-import pytest
-import torch
-
-from sglang.kernels.ops.diffusion.triton.native_bf16_rmsnorm import (
-    rmsnorm_scale,
-    rmsnorm_tanh_residual,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-
-EPS = 1e-5
-
-
-def _native_bf16_rmsnorm(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
-    square = (x * x).to(torch.bfloat16)
-    mean_square = square.mean(dim=-1, keepdim=True).to(torch.bfloat16)
-    rstd = torch.rsqrt((mean_square + EPS).to(torch.bfloat16).float()).to(
-        torch.bfloat16
-    )
-    return ((x * rstd).to(torch.bfloat16) * weight).to(torch.bfloat16)
-
-
-def test_native_bf16_rmsnorm_rejects_unsupported_inputs():
-    x = torch.randn(2, 3, 16, dtype=torch.bfloat16)
-    weight = torch.randn(16, dtype=torch.bfloat16)
-    modulation = torch.randn(2, 1, 16, dtype=torch.bfloat16)
-    residual = torch.randn_like(x)
-
-    assert rmsnorm_scale(x, weight, modulation, EPS) is None
-    assert rmsnorm_tanh_residual(x, modulation, residual, weight, EPS) is None
-    assert rmsnorm_scale(x, weight[:-1], modulation, EPS) is None
-    assert rmsnorm_tanh_residual(x, modulation, residual[..., :-1], weight, EPS) is None
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-@pytest.mark.parametrize("shape", [(1, 32, 2560), (2, 17, 256)])
-def test_rmsnorm_scale_matches_native_bf16(shape):
-    torch.manual_seed(0)
-    batch, _, dim = shape
-    x = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(dim, device="cuda", dtype=torch.bfloat16)
-    scale = torch.randn(batch, 1, dim, device="cuda", dtype=torch.bfloat16)
-
-    actual = rmsnorm_scale(x, weight, scale, EPS)
-    expected = (_native_bf16_rmsnorm(x, weight) * scale).to(torch.bfloat16)
-
-    assert actual is not None
-    torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-@pytest.mark.parametrize("shape", [(1, 32, 2560), (2, 17, 256)])
-def test_rmsnorm_tanh_residual_matches_native_bf16(shape):
-    torch.manual_seed(0)
-    batch, _, dim = shape
-    x = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
-    gate = torch.randn(batch, 1, dim, device="cuda", dtype=torch.bfloat16)
-    residual = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
-    weight = torch.randn(dim, device="cuda", dtype=torch.bfloat16)
-
-    actual = rmsnorm_tanh_residual(x, gate, residual, weight, EPS)
-    norm = _native_bf16_rmsnorm(x, weight)
-    gated = (torch.tanh(gate.float()).to(torch.bfloat16) * norm).to(torch.bfloat16)
-    expected = (residual + gated).to(torch.bfloat16)
-
-    assert actual is not None
-    # Triton's exp-based tanh can differ slightly from torch.tanh in BF16.
-    torch.testing.assert_close(actual, expected, atol=4e-2, rtol=2e-2)
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-def test_native_bf16_rmsnorm_rejects_hidden_size_above_limit():
-    dim = 8448
-    x = torch.empty(1, 1, dim, device="cuda", dtype=torch.bfloat16)
-    weight = torch.empty(dim, device="cuda", dtype=torch.bfloat16)
-    modulation = torch.empty(1, 1, dim, device="cuda", dtype=torch.bfloat16)
-    residual = torch.empty_like(x)
-
-    assert rmsnorm_scale(x, weight, modulation, EPS) is None
-    assert rmsnorm_tanh_residual(x, modulation, residual, weight, EPS) is None
-
-
-if __name__ == "__main__":
-    import sys
-
-    sys.exit(pytest.main([__file__]))

test/registered/kernels/ops/diffusion/test_norm.py

@@ -0,0 +1,466 @@
+"""``diffusion.norm``: GroupNorm / RMSNorm / LayerNorm and their fused epilogues.
+
+This domain has the most implementations of any in the package (see the
+selection matrix in ``sglang/kernels/ops/diffusion/README.md``), so the suite
+is organized by *kernel*, and each section states which oracle it is held to:
+
+- ``triton_group_norm_silu`` / ``apply_group_norm_silu`` -> ``F.group_norm`` +
+  ``F.silu`` with a per-dtype tolerance (fp32 statistics, different reduction).
+- the two-pass channels-last GroupNorm -> same oracle, plus its support
+  predicates (the kernels raise on an unsupported input rather than returning
+  ``None``).
+- ``rmsnorm_scale`` / ``rmsnorm_tanh_residual`` -> a bf16-native reference that
+  reproduces Z-Image's own norm, with a tolerance for Triton's exp-based tanh.
+- the CuTe-DSL fused norm+scale/shift -> an fp32 reference chain.
+
+The FlyDSL norms live in ``test_norm_flydsl.py``: they are ROCm gfx950-only,
+so they run on a CI lane this file does not, and keeping them here dragged the
+CUDA-only CuTe-DSL cases onto the AMD runner.
+
+The *bit-exact* norms (``fused_rmsnorm_scale_shift_bitexact``,
+``fused_layernorm_modulate``, ``zimage_qk_rmsnorm_native``) are exercised
+through their model wrappers in ``test_model_fast_paths.py``, where the live
+eager chain they must reproduce is available.
+"""
+
+import sys
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from sglang.kernels.ops.diffusion import (
+    apply_group_norm_silu,
+    can_use_group_norm_silu_4d,
+    can_use_wan_rmsnorm_silu,
+    group_norm_silu_4d,
+    rmsnorm_scale,
+    rmsnorm_tanh_residual,
+    triton_group_norm_silu,
+    wan_rmsnorm_silu,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=70, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+DEVICE = "cuda"
+DTYPES = [torch.float16, torch.bfloat16, torch.float32]
+EPS = 1e-5
+
+
+def _tol(dtype: torch.dtype) -> tuple[float, float]:
+    if dtype == torch.float32:
+        return 1e-5, 1e-5
+    if dtype == torch.bfloat16:
+        return 7e-2, 2e-2
+    return 3e-3, 3e-3
+
+
+@pytest.fixture(autouse=True)
+def cuda_setup():
+    torch.cuda.manual_seed(0)
+
+
+def _cl3d(shape, dtype):
+    return torch.randn(shape, device=DEVICE, dtype=dtype).contiguous(
+        memory_format=torch.channels_last_3d
+    )
+
+
+# ---------------------------------------------------------------------------
+# GroupNorm + SiLU
+# ---------------------------------------------------------------------------
+
+GN_CASES = [
+    pytest.param((2, 64, 32, 32), 32, id="image_2d"),
+    pytest.param((1, 64, 4, 16, 16), 32, id="video_3d"),
+    pytest.param((4, 128), 32, id="token_2d"),
+]
+
+
+def _gn_silu_reference(x, weight, bias, num_groups, eps=EPS):
+    return F.silu(F.group_norm(x, num_groups, weight=weight, bias=bias, eps=eps))
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("shape,num_groups", GN_CASES)
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_triton_group_norm_silu(shape, num_groups, dtype):
+    channels = shape[1]
+    x = torch.randn(shape, device=DEVICE, dtype=dtype)
+    weight = torch.randn(channels, device=DEVICE, dtype=dtype)
+    bias = torch.randn(channels, device=DEVICE, dtype=dtype)
+
+    atol, rtol = _tol(dtype)
+    torch.testing.assert_close(
+        triton_group_norm_silu(x, weight, bias, num_groups=num_groups),
+        _gn_silu_reference(x, weight, bias, num_groups),
+        atol=atol,
+        rtol=rtol,
+    )
+
+
+@torch.no_grad()
+def test_triton_group_norm_silu_large_tile_bf16():
+    # A tile large enough to force the chunked launch path (128 channels over
+    # 20x256x256), which the small cases above never reach.
+    shape, num_groups = (1, 128, 20, 256, 256), 32
+    x = torch.randn(shape, device=DEVICE, dtype=torch.bfloat16)
+    weight = torch.randn(shape[1], device=DEVICE, dtype=torch.bfloat16)
+    bias = torch.randn(shape[1], device=DEVICE, dtype=torch.bfloat16)
+
+    atol, rtol = _tol(torch.bfloat16)
+    torch.testing.assert_close(
+        triton_group_norm_silu(x, weight, bias, num_groups=num_groups),
+        _gn_silu_reference(x, weight, bias, num_groups),
+        atol=atol,
+        rtol=rtol,
+    )
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("shape,num_groups", GN_CASES[:2])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_apply_group_norm_silu_module_wrapper(shape, num_groups, dtype):
+    # The nn.Module-taking wrapper must match the eager module pair it stands
+    # in for, including its own guard set (affine=True, non-inplace SiLU).
+    norm = nn.GroupNorm(num_groups, shape[1], eps=EPS, affine=True).to(
+        device=DEVICE, dtype=dtype
+    )
+    activation = nn.SiLU()
+    x = torch.randn(shape, device=DEVICE, dtype=dtype)
+
+    atol, rtol = _tol(dtype)
+    torch.testing.assert_close(
+        apply_group_norm_silu(x, norm, activation),
+        activation(norm(x)),
+        atol=atol,
+        rtol=rtol,
+    )
+
+
+@torch.no_grad()
+def test_group_norm_silu_4d_channels_last_and_guards():
+    gn = nn.GroupNorm(32, 128, eps=1e-6).to(DEVICE, torch.bfloat16)
+    x = torch.randn(1, 128, 64, 64, device=DEVICE, dtype=torch.bfloat16).to(
+        memory_format=torch.channels_last
+    )
+    assert can_use_group_norm_silu_4d(x, gn.weight, gn.bias, 32)
+    out = group_norm_silu_4d(x, gn.weight, gn.bias, 32, 1e-6)
+    assert out.is_contiguous(memory_format=torch.channels_last)
+    torch.testing.assert_close(out.float(), F.silu(gn(x)).float(), atol=0.06, rtol=0)
+
+    # Guards: the kernel exists only for channels_last inputs with device-side
+    # affine params and a non-empty spatial extent.  Each rejected case must
+    # fail the predicate *and* raise if called anyway -- silently returning
+    # ``None`` is what this protocol replaced.
+    rejected = [
+        (x.contiguous(), gn.weight, gn.bias),  # contiguous (NCHW) layout
+        (x, gn.weight.cpu(), gn.bias),  # host-side affine
+        (x[..., :0, :], gn.weight, gn.bias),  # empty spatial extent
+    ]
+    for args in rejected:
+        assert not can_use_group_norm_silu_4d(*args, 32)
+        with pytest.raises(ValueError):
+            group_norm_silu_4d(*args, 32, 1e-6)
+
+
+# ---------------------------------------------------------------------------
+# BF16-native RMSNorm fusions (Z-Image / Ideogram)
+# ---------------------------------------------------------------------------
+
+
+def _native_bf16_rmsnorm(x, weight):
+    """Z-Image's own norm: every step materialized in bf16, no fp32 carry."""
+    square = (x * x).to(torch.bfloat16)
+    mean_square = square.mean(dim=-1, keepdim=True).to(torch.bfloat16)
+    rstd = torch.rsqrt((mean_square + EPS).to(torch.bfloat16).float()).to(
+        torch.bfloat16
+    )
+    return ((x * rstd).to(torch.bfloat16) * weight).to(torch.bfloat16)
+
+
+@pytest.mark.parametrize("shape", [(1, 32, 2560), (2, 17, 256)])
+def test_rmsnorm_scale_matches_native_bf16(shape):
+    batch, _, dim = shape
+    x = torch.randn(shape, device=DEVICE, dtype=torch.bfloat16)
+    weight = torch.randn(dim, device=DEVICE, dtype=torch.bfloat16)
+    scale = torch.randn(batch, 1, dim, device=DEVICE, dtype=torch.bfloat16)
+
+    actual = rmsnorm_scale(x, weight, scale, EPS)
+    assert actual is not None
+    torch.testing.assert_close(
+        actual,
+        (_native_bf16_rmsnorm(x, weight) * scale).to(torch.bfloat16),
+        atol=2e-2,
+        rtol=2e-2,
+    )
+
+
+@pytest.mark.parametrize("shape", [(1, 32, 2560), (2, 17, 256)])
+def test_rmsnorm_tanh_residual_matches_native_bf16(shape):
+    batch, _, dim = shape
+    x = torch.randn(shape, device=DEVICE, dtype=torch.bfloat16)
+    gate = torch.randn(batch, 1, dim, device=DEVICE, dtype=torch.bfloat16)
+    residual = torch.randn(shape, device=DEVICE, dtype=torch.bfloat16)
+    weight = torch.randn(dim, device=DEVICE, dtype=torch.bfloat16)
+
+    actual = rmsnorm_tanh_residual(x, gate, residual, weight, EPS)
+    norm = _native_bf16_rmsnorm(x, weight)
+    gated = (torch.tanh(gate.float()).to(torch.bfloat16) * norm).to(torch.bfloat16)
+
+    assert actual is not None
+    # Triton's exp-based tanh can differ slightly from torch.tanh in bf16.
+    torch.testing.assert_close(
+        actual, (residual + gated).to(torch.bfloat16), atol=4e-2, rtol=2e-2
+    )
+
+
+@pytest.mark.parametrize("on_host", [True, False])
+def test_native_bf16_rmsnorm_rejects_unsupported_inputs(on_host):
+    # Host tensors and a hidden size past the kernel limit are both outside
+    # the contract; these entry points signal that by returning None (they
+    # are internal fast-path probes, not public predicate+kernel pairs).
+    device = "cpu" if on_host else DEVICE
+    dim = 16 if on_host else 8448
+    x = torch.randn(2, 3, dim, dtype=torch.bfloat16, device=device)
+    weight = torch.randn(dim, dtype=torch.bfloat16, device=device)
+    modulation = torch.randn(2, 1, dim, dtype=torch.bfloat16, device=device)
+    residual = torch.randn_like(x)
+
+    assert rmsnorm_scale(x, weight, modulation, EPS) is None
+    assert rmsnorm_tanh_residual(x, modulation, residual, weight, EPS) is None
+    if on_host:
+        # Mismatched trailing dims are rejected too.
+        assert rmsnorm_scale(x, weight[:-1], modulation, EPS) is None
+        assert (
+            rmsnorm_tanh_residual(x, modulation, residual[..., :-1], weight, EPS)
+            is None
+        )
+
+
+# ---------------------------------------------------------------------------
+# Wan VAE channels_last_3d RMSNorm + SiLU
+# ---------------------------------------------------------------------------
+
+
+@torch.no_grad()
+@pytest.mark.parametrize(
+    "x_dtype,affine_dtype,atol,rtol",
+    [
+        (torch.float32, torch.float32, 1e-5, 1e-5),  # FastWan2.2 fp32 decode
+        (torch.bfloat16, torch.float32, 1.5e-1, 3e-2),  # Wan2.1 bf16 autocast
+    ],
+)
+def test_wan_rmsnorm_silu_numerics(x_dtype, affine_dtype, atol, rtol):
+    x = _cl3d((1, 96, 3, 10, 14), x_dtype)
+    gamma = torch.randn((96, 1, 1, 1), device=DEVICE, dtype=affine_dtype)
+    for bias in (None, torch.randn_like(gamma)):
+        expected = F.silu(
+            F.normalize(x, dim=1) * 96**0.5 * gamma + (0 if bias is None else bias)
+        )
+        actual = wan_rmsnorm_silu(x, gamma, bias)
+        assert actual.dtype == expected.dtype
+        # The kernel must preserve the channels_last_3d layout; a relayout
+        # here would undo the reason the decoder runs in that format.
+        assert actual.stride() == x.stride()
+        torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
+
+
+@torch.no_grad()
+def test_wan_rmsnorm_silu_rejects_empty_input():
+    x = torch.empty(1, 96, 0, 2, 2, device=DEVICE, dtype=torch.bfloat16).to(
+        memory_format=torch.channels_last_3d
+    )
+    gamma = torch.ones(96, 1, 1, 1, device=DEVICE, dtype=torch.bfloat16)
+    assert not can_use_wan_rmsnorm_silu(x, gamma, None)
+    with pytest.raises(ValueError):
+        wan_rmsnorm_silu(x, gamma)
+
+
+# ---------------------------------------------------------------------------
+# CuTe-DSL fused (residual +) norm + scale/shift
+# ---------------------------------------------------------------------------
+
+SHAPE_MAP = {
+    "1": lambda B, S, F_, D: (1,),
+    "D": lambda B, S, F_, D: (D,),
+    "1D": lambda B, S, F_, D: (1, D),
+    "BD": lambda B, S, F_, D: (B, D),
+    "11D": lambda B, S, F_, D: (1, 1, D),
+    "B1D": lambda B, S, F_, D: (B, 1, D),
+    "1SD": lambda B, S, F_, D: (1, S, D),
+    "BSD": lambda B, S, F_, D: (B, S, D),
+    "BF1D": lambda B, S, F_, D: (B, F_, 1, D),
+}
+# (B, S, F, D)
+CUTE_SHAPES = [
+    (1, 115200, 1, 3072),  # HunyuanVideo
+    (1, 32760, 1, 1536),  # Wan
+    (1, 6, 1, 3072),  # Qwen-Image
+    (1, 1024, 8, 3072),
+    (4, 512, 16, 3072),
+]
+NORM_TYPES = ["layer", "rms"]
+AFFINE_MODES = ["D", "NAT"]
+INDEX_MODES = ["BSD", "1", "1SD", "BD", "B1D", "D", "1D", "11D", "BF1D"]
+
+
+def _import_cutedsl():
+    """Import the CuTe-DSL entry points, skipping when the backend is absent.
+
+    This file is registered on the AMD lane for its FlyDSL section, but the
+    CuTe-DSL norms need cuda-python and CUTLASS, which the ROCm image does not
+    ship.  Guarded per test rather than by dropping this file from the AMD
+    lane, so the Triton and FlyDSL sections keep running there.
+    """
+    try:
+        from sglang.kernels.ops.diffusion import (
+            fused_norm_scale_shift,
+            fused_scale_residual_norm_scale_shift,
+        )
+    except ImportError as exc:  # pragma: no cover - platform-dependent
+        pytest.skip(f"CuTe-DSL backend unavailable: {exc}")
+    return fused_norm_scale_shift, fused_scale_residual_norm_scale_shift
+
+
+def _make_tensor(index_mode, shape, dtype):
+    if index_mode == "NAT":
+        return None
+    return torch.randn(*SHAPE_MAP[index_mode](*shape), device=DEVICE, dtype=dtype)
+
+
+def _apply_scale_shift(y, scale, shift):
+    if scale.ndim == 4:
+        num_frame = scale.shape[1]
+        return rearrange(
+            rearrange(y, "b (f l) d -> b f l d", f=num_frame) * (1 + scale) + shift,
+            "b f l d -> b (f l) d",
+        )
+    scale = rearrange(scale, "b d -> b 1 d") if scale.ndim == 2 else scale
+    shift = rearrange(shift, "b d -> b 1 d") if shift.ndim == 2 else shift
+    return y * (1 + scale) + shift
+
+
+def _cute_reference(residual, x, gate, weight, bias, scale, shift, norm_type, eps):
+    """fp32 oracle for both variants; ``residual is None`` = no-residual form."""
+    original_dtype = x.dtype
+    residual, x, gate, weight, bias, scale, shift = (
+        v.float() if isinstance(v, torch.Tensor) else v
+        for v in (residual, x, gate, weight, bias, scale, shift)
+    )
+    residual_out = None
+    if residual is not None:
+        if isinstance(gate, int):
+            x = residual + gate * x
+        elif gate.ndim == 4:
+            folded = rearrange(x, "b (f l) d -> b f l d", f=gate.shape[1])
+            x = residual + rearrange(folded * gate, "b f l d -> b (f l) d")
+        else:
+            g = rearrange(gate, "b d -> b 1 d") if gate.ndim == 2 else gate
+            x = residual + g * x
+        residual_out = x.to(original_dtype)
+    if norm_type == "layer":
+        norm = torch.layer_norm(x, x.shape[-1:], eps=eps, weight=weight, bias=bias)
+    else:
+        norm = torch.rms_norm(x, x.shape[-1:], eps=eps, weight=weight)
+    return _apply_scale_shift(norm, scale, shift).to(original_dtype), residual_out
+
+
+@torch.no_grad()
+def _run_cute(
+    with_residual,
+    shape=CUTE_SHAPES[0],
+    dtype=DTYPES[0],
+    affine_dtype=DTYPES[0],
+    mod_dtype=DTYPES[0],
+    norm_type=NORM_TYPES[0],
+    affine_mode=AFFINE_MODES[0],
+    gate_mode="B1D",
+    index_mode="BSD",
+    eps=EPS,
+):
+    fused_norm_scale_shift, fused_scale_residual_norm_scale_shift = _import_cutedsl()
+
+    x = _make_tensor("BSD", shape, dtype)
+    weight = _make_tensor(affine_mode, shape, affine_dtype)
+    bias = _make_tensor(affine_mode, shape, affine_dtype)
+    scale = _make_tensor(index_mode, shape, mod_dtype)
+    shift = _make_tensor(index_mode, shape, mod_dtype)
+    tol = 1e-5 if dtype == torch.float32 else 5e-2
+
+    if with_residual:
+        residual = _make_tensor("BSD", shape, dtype)
+        gate = _make_tensor(gate_mode, shape, dtype)
+        y, res = fused_scale_residual_norm_scale_shift(
+            residual, x, gate, weight, bias, scale, shift, norm_type, eps
+        )
+        y_ref, res_ref = _cute_reference(
+            residual, x, gate, weight, bias, scale, shift, norm_type, eps
+        )
+        torch.testing.assert_close(res, res_ref, atol=tol, rtol=tol)
+    else:
+        y = fused_norm_scale_shift(x, weight, bias, scale, shift, norm_type, eps)
+        y_ref, _ = _cute_reference(
+            None, x, None, weight, bias, scale, shift, norm_type, eps
+        )
+    torch.testing.assert_close(y, y_ref, atol=tol, rtol=tol)
+
+
+@pytest.mark.parametrize("with_residual", [False, True])
+@pytest.mark.parametrize("norm_type", NORM_TYPES)
+@pytest.mark.parametrize("shape", CUTE_SHAPES)
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_cutedsl_norm_scale_shift_shapes(with_residual, norm_type, shape, dtype):
+    _run_cute(with_residual, shape=shape, dtype=dtype, norm_type=norm_type)
+
+
+@pytest.mark.parametrize("with_residual", [False, True])
+@pytest.mark.parametrize("norm_type", NORM_TYPES)
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("operand", ["affine", "modulation"])
+def test_cutedsl_norm_scale_shift_mixed_operand_dtypes(
+    with_residual, norm_type, dtype, operand
+):
+    # The affine params and the modulation rows may each arrive in a dtype
+    # other than the activation's; both combinations must dispatch.
+    kwargs = {"affine_dtype" if operand == "affine" else "mod_dtype": dtype}
+    _run_cute(with_residual, norm_type=norm_type, **kwargs)
+
+
+@pytest.mark.parametrize("with_residual", [False, True])
+@pytest.mark.parametrize("norm_type", NORM_TYPES)
+@pytest.mark.parametrize("affine_mode", AFFINE_MODES)
+def test_cutedsl_norm_scale_shift_affine_modes(with_residual, norm_type, affine_mode):
+    _run_cute(with_residual, norm_type=norm_type, affine_mode=affine_mode)
+
+
+@pytest.mark.parametrize("with_residual", [False, True])
+@pytest.mark.parametrize("norm_type", NORM_TYPES)
+@pytest.mark.parametrize("index_mode", INDEX_MODES)
+def test_cutedsl_norm_scale_shift_index_modes(with_residual, norm_type, index_mode):
+    _run_cute(with_residual, norm_type=norm_type, index_mode=index_mode)
+
+
+@pytest.mark.parametrize("norm_type", NORM_TYPES)
+@pytest.mark.parametrize("index_mode", INDEX_MODES)
+def test_cutedsl_scale_residual_gate_index_modes(norm_type, index_mode):
+    _run_cute(True, norm_type=norm_type, gate_mode=index_mode)
+
+
+def test_validate_scale_shift_rejects_non_divisible_frames():
+    _import_cutedsl()
+    from sglang.kernels.ops.diffusion import validate_scale_shift
+
+    with pytest.raises(ValueError, match=r"S\(10\) must be divisible by F\(4\)"):
+        validate_scale_shift(
+            torch.empty((1, 4, 1, 256), device=DEVICE, dtype=torch.float16), 1, 10, 256
+        )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/diffusion/test_norm_flydsl.py

@@ -0,0 +1,98 @@
+"""``diffusion.norm``: the FlyDSL fused norm + scale/shift kernels (ROCm).
+
+Split out of ``test_norm.py`` rather than merged with the other norm backends:
+FlyDSL is an AMD gfx950-only compiler, so these run on the AMD CI lane and
+nothing else in that file does.  Keeping them together forced the CUDA-only
+CuTe-DSL cases onto the ROCm runner, where cuda-python does not exist.
+
+Oracle: an fp32 reference chain, with a tolerance -- the kernel keeps fp32
+statistics but reorders the reduction.
+"""
+
+import sys
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from sglang.test.ci.ci_register import register_amd_ci
+
+register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd")
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="GPU required")
+
+DEVICE = "cuda"
+
+FLYDSL_D = 5120
+FLYDSL_EPS = 1e-6
+
+
+def _require_rocm():
+    if not torch.version.hip:
+        pytest.skip("ROCm/HIP required for FlyDSL kernels")
+
+
+def _flydsl_reference(residual, x, gate, weight, bias, scale, shift, norm_type, eps):
+    if residual is not None:
+        x = (residual.float() + x.float() * gate.float()).to(torch.bfloat16)
+        residual_out = x
+    else:
+        residual_out = None
+    if norm_type == "layer":
+        normed = F.layer_norm(x.float(), (FLYDSL_D,), weight, bias, eps)
+    else:
+        var = x.float().pow(2).mean(-1, keepdim=True)
+        normed = x.float() * torch.rsqrt(var + eps) * weight.float()
+    y = (normed * (1.0 + scale.float()) + shift.float()).to(torch.bfloat16)
+    return y, residual_out
+
+
+@pytest.mark.parametrize("with_residual", [False, True])
+@pytest.mark.parametrize(
+    "norm_type,B,L",
+    [("rms", 1, 16), ("rms", 2, 16), ("layer", 2, 16), ("rms", 1, 90000)],
+)
+def test_flydsl_norm_scale_shift(with_residual, norm_type, B, L):
+    _require_rocm()
+    # Imported inside the test: the FlyDSL compiler only exists on ROCm, and
+    # the facade resolves an export the moment it is named -- a module-level
+    # import here would fail collection of this whole file on CUDA.
+    from sglang.kernels.ops.diffusion import (
+        flydsl_fused_residual_norm_scale_shift,
+        flydsl_norm_scale_shift,
+    )
+
+    torch.manual_seed(42)
+    shape = (B, L, FLYDSL_D)
+    x = torch.randn(shape, device=DEVICE, dtype=torch.bfloat16)
+    weight = torch.randn(FLYDSL_D, device=DEVICE, dtype=torch.float32)
+    bias = (
+        torch.randn(FLYDSL_D, device=DEVICE, dtype=torch.float32)
+        if norm_type == "layer"
+        else None
+    )
+    scale = torch.randn(B, 1, FLYDSL_D, device=DEVICE, dtype=torch.bfloat16)
+    shift = torch.randn(B, 1, FLYDSL_D, device=DEVICE, dtype=torch.bfloat16)
+
+    if with_residual:
+        residual = torch.randn(shape, device=DEVICE, dtype=torch.bfloat16)
+        gate = torch.randn(B, 1, FLYDSL_D, device=DEVICE, dtype=torch.bfloat16)
+        y, res = flydsl_fused_residual_norm_scale_shift(
+            residual, x, gate, weight, bias, scale, shift, norm_type, FLYDSL_EPS
+        )
+        y_ref, res_ref = _flydsl_reference(
+            residual, x, gate, weight, bias, scale, shift, norm_type, FLYDSL_EPS
+        )
+        torch.testing.assert_close(res, res_ref, atol=5e-2, rtol=5e-2)
+    else:
+        y = flydsl_norm_scale_shift(
+            x, weight, bias, scale, shift, norm_type, FLYDSL_EPS
+        )
+        y_ref, _ = _flydsl_reference(
+            None, x, None, weight, bias, scale, shift, norm_type, FLYDSL_EPS
+        )
+    torch.testing.assert_close(y, y_ref, atol=1.0, rtol=5e-2)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/diffusion/test_quality_gate.py

@@ -1,47 +0,0 @@
-import sys
-
-import pytest
-import torch.nn as nn
-
-from sglang.kernels.ops.diffusion.quality_gate import QualityGatedFusion
-from sglang.test.ci.ci_register import register_cpu_ci
-
-register_cpu_ci(est_time=2, suite="base-a-test-cpu")
-
-
-def test_quality_gate_mounts_and_unmounts_all_sites():
-    fusion = QualityGatedFusion(
-        name="test fusion",
-        marker_attr="_test_fusion_site",
-        enabled_attr="_test_fusion_enabled",
-    )
-    root = nn.ModuleList([nn.Module(), nn.Module()])
-    for index, site in enumerate(root):
-        fusion.mark(site, index)
-
-    assert [fusion.metadata(site) for site in root] == [0, 1]
-    assert fusion.mount(root)
-    assert all(fusion.is_enabled(site) for site in root)
-    fusion.unmount(root)
-    assert not any(fusion.is_enabled(site) for site in root)
-
-
-def test_quality_gate_rejection_is_all_or_nothing():
-    fusion = QualityGatedFusion(
-        name="test fusion",
-        marker_attr="_test_fusion_site",
-        enabled_attr="_test_fusion_enabled",
-    )
-    root = nn.ModuleList([nn.Module(), nn.Module()])
-    for index, site in enumerate(root):
-        fusion.mark(site, index)
-
-    assert not fusion.mount(
-        root, reject_reason=lambda site: "rejected" if fusion.metadata(site) else None
-    )
-    assert not any(fusion.is_enabled(site) for site in root)
-    assert not fusion.mount(nn.Module())
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/diffusion/test_qwen_image_modulation.py

@@ -1,228 +0,0 @@
-import sys
-
-import pytest
-import torch
-import triton
-
-from sglang.kernels.jit.utils import get_ci_test_range
-from sglang.kernels.ops.diffusion.triton.norm import norm_infer
-from sglang.kernels.ops.diffusion.triton.scale_shift import (
-    fuse_layernorm_scale_shift_gate_select01_kernel,
-    fuse_residual_layernorm_scale_shift_gate_select01_kernel,
-)
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-
-register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-# Nightly is not redundant here: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1 to expand get_ci_test_range sweeps.
-register_cuda_ci(est_time=30, stage="nightly", runner_config="1-gpu-large")
-register_amd_ci(est_time=30, suite="nightly-amd-kernel-1-gpu", nightly=True)
-
-DEVICE = "cuda"
-DTYPES = get_ci_test_range(
-    [torch.float16, torch.bfloat16, torch.float32], [torch.float16, torch.bfloat16]
-)
-BATCH_SIZES = get_ci_test_range([1, 2, 4], [1, 2])
-SEQ_LENS = get_ci_test_range([6, 33, 128, 257], [6, 128])
-HIDDEN_SIZES = get_ci_test_range([512, 1024, 1536, 3072], [512, 3072])
-EPS = 1e-6
-
-
-def _tol(dtype: torch.dtype) -> tuple[float, float]:
-    if dtype == torch.float32:
-        return 1e-5, 1e-5
-    return 5e-2, 5e-2
-
-
-def _make_modulation_tensors(batch_size: int, hidden_size: int, dtype: torch.dtype):
-    scale0 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
-    shift0 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
-    gate0 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
-    scale1 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
-    shift1 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
-    gate1 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
-    return scale0, shift0, gate0, scale1, shift1, gate1
-
-
-def _baseline_select01_modulation(
-    x: torch.Tensor,
-    weight: torch.Tensor | None,
-    bias: torch.Tensor | None,
-    scale0: torch.Tensor,
-    shift0: torch.Tensor,
-    gate0: torch.Tensor,
-    scale1: torch.Tensor,
-    shift1: torch.Tensor,
-    gate1: torch.Tensor,
-    index: torch.Tensor,
-    eps: float,
-):
-    normalized = norm_infer(
-        x.view(-1, x.shape[-1]),
-        weight,
-        bias,
-        eps=eps,
-        is_rms_norm=False,
-    ).view_as(x)
-    return _apply_select01_modulation(
-        normalized, scale0, shift0, gate0, scale1, shift1, gate1, index
-    )
-
-
-def _baseline_residual_select01_modulation(
-    x: torch.Tensor,
-    residual: torch.Tensor,
-    residual_gate: torch.Tensor,
-    weight: torch.Tensor | None,
-    bias: torch.Tensor | None,
-    scale0: torch.Tensor,
-    shift0: torch.Tensor,
-    gate0: torch.Tensor,
-    scale1: torch.Tensor,
-    shift1: torch.Tensor,
-    gate1: torch.Tensor,
-    index: torch.Tensor,
-    eps: float,
-):
-    residual_out = residual + residual_gate * x
-    normalized = norm_infer(
-        residual_out.view(-1, residual_out.shape[-1]),
-        weight,
-        bias,
-        eps=eps,
-        is_rms_norm=False,
-    ).view_as(residual_out)
-    output, gate_out = _apply_select01_modulation(
-        normalized, scale0, shift0, gate0, scale1, shift1, gate1, index
-    )
-    return output, residual_out, gate_out
-
-
-def _apply_select01_modulation(
-    x: torch.Tensor,
-    scale0: torch.Tensor,
-    shift0: torch.Tensor,
-    gate0: torch.Tensor,
-    scale1: torch.Tensor,
-    shift1: torch.Tensor,
-    gate1: torch.Tensor,
-    index: torch.Tensor,
-):
-    idx = index.bool().unsqueeze(-1)
-    scale = torch.where(idx, scale1.unsqueeze(1), scale0.unsqueeze(1))
-    shift = torch.where(idx, shift1.unsqueeze(1), shift0.unsqueeze(1))
-    gate = torch.where(idx, gate1.unsqueeze(1), gate0.unsqueeze(1))
-    return x * (1 + scale) + shift, gate
-
-
-@pytest.fixture(autouse=True)
-def cuda_setup():
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA required")
-    torch.cuda.manual_seed(0)
-
-
-@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize("batch_size", BATCH_SIZES)
-@pytest.mark.parametrize("seq_len", SEQ_LENS)
-@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
-def test_fused_layernorm_scale_shift_gate_select01(
-    dtype, batch_size, seq_len, hidden_size
-):
-    x = torch.randn(batch_size, seq_len, hidden_size, device=DEVICE, dtype=dtype)
-    weight = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
-    bias = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
-    index = torch.randint(0, 2, (batch_size, seq_len), device=DEVICE, dtype=torch.int32)
-    scale0, shift0, gate0, scale1, shift1, gate1 = _make_modulation_tensors(
-        batch_size, hidden_size, dtype
-    )
-
-    out_ref, gate_ref = _baseline_select01_modulation(
-        x,
-        weight,
-        bias,
-        scale0,
-        shift0,
-        gate0,
-        scale1,
-        shift1,
-        gate1,
-        index,
-        EPS,
-    )
-    out_fused, gate_fused = fuse_layernorm_scale_shift_gate_select01_kernel(
-        x.contiguous(),
-        weight=weight,
-        bias=bias,
-        scale0=scale0,
-        shift0=shift0,
-        gate0=gate0,
-        scale1=scale1,
-        shift1=shift1,
-        gate1=gate1,
-        index=index,
-        eps=EPS,
-    )
-
-    atol, rtol = _tol(dtype)
-    triton.testing.assert_close(out_ref, out_fused, atol=atol, rtol=rtol)
-    triton.testing.assert_close(gate_ref, gate_fused, atol=atol, rtol=rtol)
-
-
-@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize("batch_size", BATCH_SIZES)
-@pytest.mark.parametrize("seq_len", SEQ_LENS)
-@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
-def test_fused_residual_layernorm_scale_shift_gate_select01(
-    dtype, batch_size, seq_len, hidden_size
-):
-    x = torch.randn(batch_size, seq_len, hidden_size, device=DEVICE, dtype=dtype)
-    residual = torch.randn_like(x)
-    residual_gate = torch.randn_like(x)
-    weight = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
-    bias = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
-    index = torch.randint(0, 2, (batch_size, seq_len), device=DEVICE, dtype=torch.int32)
-    scale0, shift0, gate0, scale1, shift1, gate1 = _make_modulation_tensors(
-        batch_size, hidden_size, dtype
-    )
-
-    out_ref, residual_ref, gate_ref = _baseline_residual_select01_modulation(
-        x,
-        residual,
-        residual_gate,
-        weight,
-        bias,
-        scale0,
-        shift0,
-        gate0,
-        scale1,
-        shift1,
-        gate1,
-        index,
-        EPS,
-    )
-    out_fused, residual_fused, gate_fused = (
-        fuse_residual_layernorm_scale_shift_gate_select01_kernel(
-            x.contiguous(),
-            residual=residual.contiguous(),
-            residual_gate=residual_gate.contiguous(),
-            weight=weight,
-            bias=bias,
-            scale0=scale0,
-            shift0=shift0,
-            gate0=gate0,
-            scale1=scale1,
-            shift1=shift1,
-            gate1=gate1,
-            index=index,
-            eps=EPS,
-        )
-    )
-
-    atol, rtol = _tol(dtype)
-    triton.testing.assert_close(out_ref, out_fused, atol=atol, rtol=rtol)
-    triton.testing.assert_close(residual_ref, residual_fused, atol=atol, rtol=rtol)
-    triton.testing.assert_close(gate_ref, gate_fused, atol=atol, rtol=rtol)
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/diffusion/test_residual_gate_add.py

@@ -1,123 +0,0 @@
-import sys
-
-import pytest
-import torch
-
-from sglang.kernels.ops.diffusion.residual_gate_add import (
-    can_use_residual_gate_add_cuda,
-    residual_gate_add,
-    residual_gate_add_cuda,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
-
-
-CASES = [
-    ((1, 1024, 4096), (1, 1, 4096)),
-    ((1, 512, 4096), (1, 512, 4096)),
-    ((1, 17, 65), (1, 1, 65)),
-    ((1, 17, 65), (1, 17, 65)),
-    # FLUX.1 / FLUX.2-klein 1024^2 shapes (D=3072): dual-stream image/text
-    # and single-stream/joint concat; gates are [1, 1, D] modulation rows.
-    ((1, 4096, 3072), (1, 1, 3072)),
-    ((1, 512, 3072), (1, 1, 3072)),
-    ((1, 4608, 3072), (1, 1, 3072)),
-    # FLUX.2-dev (D=6144) joint sequence.
-    ((1, 4608, 6144), (1, 1, 6144)),
-    # ERNIE-4.5-VL 1024^2 image tokens plus text tokens.
-    ((1, 4216, 4096), (1, 1, 4096)),
-]
-
-
-def _tol(dtype: torch.dtype) -> float:
-    return 1e-5 if dtype == torch.float32 else 5e-2
-
-
-def _assert_matches_torch(out: torch.Tensor, ref: torch.Tensor) -> None:
-    if ref.dtype == torch.float32:
-        torch.testing.assert_close(out, ref, atol=_tol(ref.dtype), rtol=_tol(ref.dtype))
-    else:
-        torch.testing.assert_close(out, ref, atol=0, rtol=0)
-
-
-@pytest.fixture(autouse=True)
-def cuda_setup():
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA required")
-    torch.cuda.manual_seed(0)
-
-
-@pytest.mark.parametrize("residual_shape,gate_shape", CASES)
-def test_residual_gate_add_matches_torch(residual_shape, gate_shape):
-    residual = torch.randn(residual_shape, device="cuda", dtype=torch.bfloat16)
-    update = torch.randn_like(residual)
-    gate = torch.randn(gate_shape, device="cuda", dtype=torch.bfloat16)
-
-    out = residual_gate_add_cuda(residual, update, gate)
-    ref = residual + update * gate
-    _assert_matches_torch(out, ref)
-    assert torch.equal(residual_gate_add(residual, update, gate), ref)
-
-
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
-@pytest.mark.parametrize("gate_shape", [(1, 1, 64), (1, 9, 64)])
-def test_residual_gate_add_dtypes(dtype, gate_shape):
-    residual = torch.randn((1, 9, 64), device="cuda", dtype=dtype)
-    update = torch.randn_like(residual)
-    gate = torch.randn(gate_shape, device="cuda", dtype=dtype)
-
-    out = residual_gate_add_cuda(residual, update, gate)
-    ref = residual + update * gate
-    _assert_matches_torch(out, ref)
-
-
-def test_can_use_residual_gate_add_cuda_rejects_unsupported_inputs():
-    residual = torch.randn((1, 8, 64), device="cuda", dtype=torch.bfloat16)
-    update = torch.randn_like(residual)
-    gate = torch.randn((1, 1, 64), device="cuda", dtype=torch.bfloat16)
-
-    assert can_use_residual_gate_add_cuda(residual, update, gate)
-    assert not can_use_residual_gate_add_cuda(residual.cpu(), update, gate)
-    assert not can_use_residual_gate_add_cuda(residual, update.float(), gate)
-    assert not can_use_residual_gate_add_cuda(residual, update[:, ::2], gate)
-    assert not can_use_residual_gate_add_cuda(residual, update, gate[:, :, ::2])
-    empty_residual = residual[:, :0]
-    empty_update = update[:, :0]
-    assert not can_use_residual_gate_add_cuda(empty_residual, empty_update, gate)
-    assert torch.equal(
-        residual_gate_add(empty_residual, empty_update, gate),
-        empty_residual + empty_update * gate,
-    )
-
-    # Only [1, ..., 1, D] row-broadcast gates are supported; a batched
-    # [B>1, 1, D] gate is not row-broadcast here and must fall back.
-    batched_residual = torch.randn((2, 8, 64), device="cuda", dtype=torch.bfloat16)
-    batched_update = torch.randn_like(batched_residual)
-    batched_gate = torch.randn((2, 1, 64), device="cuda", dtype=torch.bfloat16)
-    assert not can_use_residual_gate_add_cuda(
-        batched_residual, batched_update, batched_gate
-    )
-    assert torch.equal(
-        residual_gate_add(batched_residual, batched_update, batched_gate),
-        batched_residual + batched_update * batched_gate,
-    )
-
-
-def test_residual_gate_add_custom_op_torch_compile_fullgraph():
-    residual = torch.randn((1, 32, 128), device="cuda", dtype=torch.bfloat16)
-    update = torch.randn_like(residual)
-    gate = torch.randn((1, 1, 128), device="cuda", dtype=torch.bfloat16)
-
-    def fn(residual, update, gate):
-        return residual_gate_add(residual, update, gate)
-
-    compiled = torch.compile(fn, fullgraph=True)
-    out = compiled(residual, update, gate)
-    ref = residual + update * gate
-    _assert_matches_torch(out, ref)
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/diffusion/test_rope.py

@@ -1,3 +1,20 @@
+"""``diffusion.rope``: rotary embeddings and the QK-norm chains fused into them.
+
+Two families with different oracles:
+
+- ``fused_inplace_qknorm_rope`` / ``fused_qknorm_rope_pack_kv`` are compared
+  against the *split* baseline (a separate qknorm kernel plus FlashInfer or
+  sgl_kernel RoPE).  In the default mode the two differ by about one bf16
+  rounding step, so those cases use a tolerance; with
+  ``round_norm_before_rope=True`` the fused kernel reproduces the split
+  rounding exactly and ``torch.equal`` applies.
+The LTX-2 split-RoPE kernel lives in ``test_rope_ltx2.py``: it is validated on
+B200 and registered on that lane alone, which the cases here cannot share --
+their oracle is the *split* baseline (a separate qknorm kernel plus sgl_kernel
+or FlashInfer RoPE), whose dispatch differs on Blackwell, so the bit-exact
+assertions below do not hold there.
+"""
+
 import itertools
 import sys
 
@@ -6,12 +23,20 @@ import torch
 import triton
 
 from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.diffusion import (
+    can_use_fused_inplace_qknorm_rope,
+    fused_inplace_qknorm_rope,
+    fused_qknorm_rope_pack_kv,
+)
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=44, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-# Nightly is not redundant here: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1 to expand get_ci_test_range sweeps.
+# Nightly is not redundant: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1, which
+# expands the get_ci_test_range sweeps below.
 register_cuda_ci(est_time=220, stage="nightly", runner_config="1-gpu-large")
 
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
 DEVICE = "cuda"
 DTYPE = torch.bfloat16
 MAX_SEQ_LEN = 131072
@@ -70,8 +95,6 @@ def fused_qknorm_rope(
     positions: torch.Tensor,
     is_neox: bool,
 ) -> None:
-    from sglang.kernels.ops.diffusion.qknorm_rope import fused_inplace_qknorm_rope
-
     fused_inplace_qknorm_rope(
         q,
         k,
@@ -85,10 +108,6 @@ def fused_qknorm_rope(
 
 
 def test_qknorm_rope_rejects_unsupported_dtypes() -> None:
-    from sglang.kernels.ops.diffusion.qknorm_rope import (
-        can_use_fused_inplace_qknorm_rope,
-    )
-
     assert not can_use_fused_inplace_qknorm_rope(128, 128, False, torch.float32)
     assert not can_use_fused_inplace_qknorm_rope(
         128, 128, False, torch.bfloat16, torch.float64
@@ -164,9 +183,6 @@ def test_qknorm_rope(
 def test_qknorm_rope_preserves_split_bf16_rounding() -> None:
     from sgl_kernel import rotary_embedding
 
-    from sglang.kernels.ops.diffusion.qknorm_rope import (
-        fused_inplace_qknorm_rope,
-    )
     from sglang.kernels.ops.layernorm.norm import fused_inplace_qknorm
 
     num_tokens, num_heads, head_dim, rope_dim = 257, 28, 128, 96
@@ -217,7 +233,6 @@ def test_qknorm_rope_preserves_split_bf16_rounding() -> None:
 
 
 def test_qknorm_rope_preserves_full_width_neox_cache() -> None:
-    from sglang.kernels.ops.diffusion.qknorm_rope import fused_inplace_qknorm_rope
     from sglang.kernels.ops.layernorm.norm import fused_inplace_qknorm
 
     num_tokens, num_heads, head_dim = 257, 32, 128
@@ -256,9 +271,6 @@ def test_qknorm_rope_preserves_full_width_neox_cache() -> None:
 
 
 def test_qknorm_rope_requires_opt_in_for_strided_packed_gqa() -> None:
-    from sglang.kernels.ops.diffusion.qknorm_rope import (
-        fused_inplace_qknorm_rope,
-    )
     from sglang.multimodal_gen.runtime.layers.layernorm import (
         RMSNorm,
         apply_qk_norm_rope,
@@ -329,10 +341,6 @@ def test_qknorm_rope_requires_opt_in_for_strided_packed_gqa() -> None:
 
 
 def test_qknorm_rope_pack_kv_matches_separate_ops() -> None:
-    from sglang.kernels.ops.diffusion.qknorm_rope import (
-        fused_inplace_qknorm_rope,
-        fused_qknorm_rope_pack_kv,
-    )
 
     batch_size = 2
     prefix_tokens, suffix_tokens = 17, 257
@@ -416,9 +424,6 @@ def test_qknorm_rope_pack_kv_matches_separate_ops() -> None:
 def test_qknorm_rope_pack_kv_preserves_split_bf16_rounding() -> None:
     from sgl_kernel import rotary_embedding
 
-    from sglang.kernels.ops.diffusion.qknorm_rope import (
-        fused_qknorm_rope_pack_kv,
-    )
     from sglang.kernels.ops.layernorm.norm import fused_inplace_qknorm
 
     batch_size = 1
@@ -507,7 +512,6 @@ def test_qknorm_rope_pack_kv_preserves_split_bf16_rounding() -> None:
 
 
 def test_qknorm_rope_accepts_empty_token_dimension() -> None:
-    from sglang.kernels.ops.diffusion.qknorm_rope import fused_inplace_qknorm_rope
 
     num_heads, head_dim = 8, 128
     q = torch.empty(0, num_heads, head_dim, device=DEVICE, dtype=DTYPE)
@@ -530,4 +534,4 @@ def test_qknorm_rope_accepts_empty_token_dimension() -> None:
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, "-v", "-s"]))
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/diffusion/test_rope_ltx2.py

@@ -1,10 +1,19 @@
+"""``diffusion.rope``: the LTX-2 QK-norm + split-RoPE CUDA kernel.
+
+Split out of ``test_rope.py`` rather than merged with the other RoPE kernels:
+this one is validated on B200 and registered on that lane alone, while the
+``fused_inplace_qknorm_rope`` cases there are held to the *split* baseline,
+whose sgl_kernel / FlashInfer dispatch differs on Blackwell -- their bit-exact
+assertions fail on B200.  One file cannot carry both lane sets.
+"""
+
 import sys
 
 import pytest
 import torch
 import torch.nn.functional as F
 
-from sglang.kernels.ops.diffusion.ltx2_qknorm_split_rope import (
+from sglang.kernels.ops.diffusion import (
     can_use_ltx2_qknorm_split_rope_cuda,
     ltx2_qknorm_split_rope_cuda,
 )
@@ -12,23 +21,18 @@ from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
 
+DEVICE = "cuda"
 BF16_FUSED_ATOL = 1.6e-1
 
 
-def _require_cuda_b200() -> None:
+def _require_b200() -> None:
     if not torch.cuda.is_available():
         pytest.skip("CUDA required")
     if torch.cuda.get_device_capability()[0] < 10:
         pytest.skip("LTX2 QKNorm split-RoPE CUDA path is validated on B200")
 
 
-@pytest.fixture(autouse=True)
-def cuda_setup():
-    _require_cuda_b200()
-    torch.cuda.manual_seed(20260630)
-
-
-def _make_cos_sin(
+def _ltx2_make_cos_sin(
     batch: int, seq_len: int, num_heads: int, head_dim: int
 ) -> tuple[torch.Tensor, torch.Tensor]:
     half_dim = head_dim // 2
@@ -64,7 +68,7 @@ def _apply_split_rotary_ref(
     return out.swapaxes(1, 2).reshape(batch, seq_len, -1).to(dtype=x_dtype)
 
 
-def _reference(
+def _ltx2_reference(
     q: torch.Tensor,
     k: torch.Tensor,
     q_cos: torch.Tensor,
@@ -95,12 +99,14 @@ def _reference(
 def test_ltx2_qknorm_split_rope_matches_torch_exactly(
     batch: int, q_seq: int, k_seq: int, num_heads: int, head_dim: int
 ) -> None:
+    _require_b200()
+    torch.cuda.manual_seed(20260630)
     hidden = num_heads * head_dim
     eps = 1e-6
     q = torch.randn(batch, q_seq, hidden, device="cuda", dtype=torch.bfloat16)
     k = torch.randn(batch, k_seq, hidden, device="cuda", dtype=torch.bfloat16)
-    q_cos, q_sin = _make_cos_sin(batch, q_seq, num_heads, head_dim)
-    k_cos, k_sin = _make_cos_sin(batch, k_seq, num_heads, head_dim)
+    q_cos, q_sin = _ltx2_make_cos_sin(batch, q_seq, num_heads, head_dim)
+    k_cos, k_sin = _ltx2_make_cos_sin(batch, k_seq, num_heads, head_dim)
     q_weight = torch.randn(hidden, device="cuda", dtype=torch.bfloat16)
     k_weight = torch.randn(hidden, device="cuda", dtype=torch.bfloat16)
 
@@ -117,7 +123,9 @@ def test_ltx2_qknorm_split_rope_matches_torch_exactly(
         head_dim=head_dim,
     )
 
-    q_ref, k_ref = _reference(q, k, q_cos, q_sin, k_cos, k_sin, q_weight, k_weight, eps)
+    q_ref, k_ref = _ltx2_reference(
+        q, k, q_cos, q_sin, k_cos, k_sin, q_weight, k_weight, eps
+    )
     q_out, k_out = ltx2_qknorm_split_rope_cuda(
         q,
         q_cos,
@@ -138,9 +146,11 @@ def test_ltx2_qknorm_split_rope_matches_torch_exactly(
 
 
 def test_ltx2_qknorm_split_rope_rejects_unsupported_inputs() -> None:
+    _require_b200()
+    torch.cuda.manual_seed(20260630)
     q = torch.randn((1, 3, 4096), device="cuda", dtype=torch.bfloat16)
     k = torch.randn_like(q)
-    q_cos, q_sin = _make_cos_sin(1, 3, 32, 128)
+    q_cos, q_sin = _ltx2_make_cos_sin(1, 3, 32, 128)
     q_weight = torch.randn(4096, device="cuda", dtype=torch.bfloat16)
     k_weight = torch.randn(4096, device="cuda", dtype=torch.bfloat16)
 
@@ -183,12 +193,14 @@ def test_ltx2_qknorm_split_rope_rejects_unsupported_inputs() -> None:
 
 
 def test_ltx2_qknorm_split_rope_custom_op_torch_compile_fullgraph() -> None:
+    _require_b200()
+    torch.cuda.manual_seed(20260630)
     batch, q_seq, k_seq, num_heads, head_dim = 1, 3, 2, 32, 64
     hidden = num_heads * head_dim
     q = torch.randn(batch, q_seq, hidden, device="cuda", dtype=torch.bfloat16)
     k = torch.randn(batch, k_seq, hidden, device="cuda", dtype=torch.bfloat16)
-    q_cos, q_sin = _make_cos_sin(batch, q_seq, num_heads, head_dim)
-    k_cos, k_sin = _make_cos_sin(batch, k_seq, num_heads, head_dim)
+    q_cos, q_sin = _ltx2_make_cos_sin(batch, q_seq, num_heads, head_dim)
+    k_cos, k_sin = _ltx2_make_cos_sin(batch, k_seq, num_heads, head_dim)
     q_weight = torch.randn(hidden, device="cuda", dtype=torch.bfloat16)
     k_weight = torch.randn(hidden, device="cuda", dtype=torch.bfloat16)
 
@@ -209,7 +221,7 @@ def test_ltx2_qknorm_split_rope_custom_op_torch_compile_fullgraph() -> None:
 
     compiled = torch.compile(fn, fullgraph=True)
     q_out, k_out = compiled(q, k, q_cos, q_sin, k_cos, k_sin, q_weight, k_weight)
-    q_ref, k_ref = _reference(
+    q_ref, k_ref = _ltx2_reference(
         q, k, q_cos, q_sin, k_cos, k_sin, q_weight, k_weight, 1e-6
     )
     torch.cuda.synchronize()
@@ -218,4 +230,4 @@ def test_ltx2_qknorm_split_rope_custom_op_torch_compile_fullgraph() -> None:
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, "-v", "-s"]))
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/diffusion/test_sana_conv_post.py

@@ -1,60 +0,0 @@
-"""Sana GLUMB post-processing fusions must match the eager bf16 chain."""
-
-import pytest
-import torch
-import torch.nn.functional as F
-
-from sglang.kernels.ops.diffusion.triton.sana_conv_post import (
-    can_use_fused_bias_glu,
-    can_use_fused_bias_silu,
-    fused_bias_glu,
-    fused_bias_silu,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=3, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-
-
-@pytest.mark.parametrize("channels", [2240, 11200])
-def test_sana_bias_silu_is_bit_exact(channels):
-    torch.manual_seed(0)
-    x = torch.randn(
-        (1, channels, 7, 5),
-        device="cuda",
-        dtype=torch.bfloat16,
-    ).to(memory_format=torch.channels_last)
-    bias = torch.randn(channels, device="cuda", dtype=torch.bfloat16)
-
-    assert can_use_fused_bias_silu(x, bias)
-    actual = fused_bias_silu(x, bias)
-    expected = F.silu(x + bias[None, :, None, None])
-
-    assert actual.is_contiguous(memory_format=torch.channels_last)
-    assert torch.equal(actual, expected)
-
-
-@pytest.mark.parametrize("channels", [2240, 5600])
-def test_sana_bias_glu_is_bit_exact(channels):
-    torch.manual_seed(1)
-    x = torch.randn(
-        (1, 2 * channels, 7, 5),
-        device="cuda",
-        dtype=torch.bfloat16,
-    ).to(memory_format=torch.channels_last)
-    bias = torch.randn(2 * channels, device="cuda", dtype=torch.bfloat16)
-
-    assert can_use_fused_bias_glu(x, bias)
-    actual = fused_bias_glu(x, bias)
-    biased = x + bias[None, :, None, None]
-    hidden, gate = torch.chunk(biased, 2, dim=1)
-    expected = hidden * F.silu(gate)
-
-    assert actual.is_contiguous(memory_format=torch.channels_last)
-    assert torch.equal(actual, expected)
-
-
-if __name__ == "__main__":
-    import sys
-
-    sys.exit(pytest.main([__file__]))

test/registered/kernels/ops/diffusion/test_sana_ln_modulate.py

@@ -1,55 +0,0 @@
-"""Sana fused LN+modulate fast path must stay bit-exact vs eager."""
-
-import pytest
-import torch
-
-import sglang.multimodal_gen.runtime.models.dits.sana as sana
-from sglang.multimodal_gen.runtime.models.dits.sana import (
-    _eager_ln_modulate,
-    _sana_ln_modulate,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=3, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-
-
-@pytest.mark.parametrize(
-    "shape,nmod,transposed",
-    [
-        ((2, 1024, 2240), 6, False),
-        ((2, 1024, 2240), 2, False),
-        ((1, 1024, 2240), 6, True),
-        ((1, 37, 2240), 6, False),
-    ],
-)
-def test_sana_fused_ln_modulate_is_bit_exact(shape, nmod, transposed):
-    # (., 1024, 2240) is the real Sana 1024px shape; hidden 2240 % 512 != 0
-    # exercises the kernel's partial tail chunk.  nmod mirrors the two adaLN
-    # chunk layouts, transposed the permuted layout the Sana DiT serves.
-    torch.manual_seed(0)
-    batch, seq, hidden = shape
-    norm = torch.nn.LayerNorm(hidden, eps=1e-6, elementwise_affine=False).cuda()
-    x = (torch.randn(batch, seq, hidden, device="cuda") * 4).bfloat16()
-    if transposed:
-        x = x.permute(0, 2, 1).contiguous().permute(0, 2, 1)
-    emb = torch.randn(batch, nmod, hidden, device="cuda").bfloat16()
-    shift, scale = emb.chunk(nmod, dim=1)[0], emb.chunk(nmod, dim=1)[-1]
-    # default-stream eager serving must stay on the untouched eager chain
-    n_sigs = len(sana._SANA_LN_MOD.verified_sigs)
-    _sana_ln_modulate(norm, x, scale, shift)
-    assert len(sana._SANA_LN_MOD.verified_sigs) == n_sigs
-    # the fusion engages on non-default streams (the BCG warmup/capture path)
-    with torch.cuda.stream(torch.cuda.Stream()):
-        out = _sana_ln_modulate(norm, x, scale, shift)
-        assert len(sana._SANA_LN_MOD.verified_sigs) == n_sigs + 1  # verified
-        out2 = _sana_ln_modulate(norm, x, scale, shift)  # verified-sig lane
-    torch.cuda.synchronize()
-    assert torch.equal(out, _eager_ln_modulate(norm, x, scale, shift))
-    assert torch.equal(out2, out) and not sana._SANA_LN_MOD.disabled
-
-
-if __name__ == "__main__":
-    import sys
-
-    sys.exit(pytest.main([__file__]))

test/registered/kernels/ops/diffusion/test_scale_shift.py

@@ -1,40 +0,0 @@
-import sys
-
-import pytest
-import torch
-
-from sglang.kernels.ops.diffusion.triton.scale_shift import (
-    try_fused_scaled_residual_add_exact,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=5, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-
-
-@torch.no_grad()
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-def test_scaled_residual_add_is_bit_exact(dtype):
-    torch.manual_seed(0)
-    residual = torch.randn(2, 17, 64, device="cuda", dtype=torch.float32)
-    x = torch.randn(2, 17, 64, device="cuda", dtype=dtype)
-    scale = torch.randn(64, device="cuda", dtype=torch.float32)
-
-    actual = try_fused_scaled_residual_add_exact(residual, x, scale)
-    expected = residual + x * scale
-    assert actual is not None
-    assert torch.equal(actual, expected)
-
-
-@torch.no_grad()
-def test_scaled_residual_add_rejects_unsupported_inputs():
-    residual = torch.empty(2, 3, 8, device="cuda", dtype=torch.float32)
-    x = torch.empty_like(residual)
-    scale = torch.empty(8, device="cuda", dtype=torch.float32)
-
-    assert try_fused_scaled_residual_add_exact(residual, x, scale) is None
-    assert try_fused_scaled_residual_add_exact(residual, x.half(), scale[:-1]) is None
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/diffusion/test_sites.py

@@ -0,0 +1,416 @@
+"""``diffusion.sites``: the two gate protocols and the mount lifecycle.
+
+Nothing here is a kernel.  ``sites`` decides *whether* a fusion is allowed to
+run for a given request and model:
+
+- :class:`QualityGatedFusion` -- for fusions that are **not** bit-exact.  They
+  mount onto marked ``nn.Module`` sites only for ``quality="high"``, at batch
+  boundaries, all-or-nothing per transformer.
+- :class:`BitExactFusionGate` -- for fusions that **are** bit-exact.  They run
+  by default but verify themselves against the live eager chain on first
+  sight and disable permanently on any mismatch.
+
+The protocol tests are pure-CPU.  The mount-lifecycle tests below use
+synthetic sites (a bare ``nn.Module`` with the marker attribute) so they test
+the protocol rather than any one model; the real model wrappers live in
+``test_model_fast_paths.py``.
+"""
+
+import sys
+from types import ModuleType
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import sglang.kernels.ops.diffusion.sites.fused_gate_rmsnorm_site as gate_rmsnorm
+import sglang.kernels.ops.diffusion.sites.fused_linear_gelu_site as linear_gelu
+from sglang.kernels.ops.diffusion import (
+    BitExactFusionGate,
+    QualityGatedFusion,
+    can_use_ln_modulate,
+    flashinfer_rmsnorm_diagnostic_hint,
+    fused_ln_modulate,
+    fused_ln_modulate_active,
+    mark_fused_ln_modulate_site,
+    mount_fused_ln_modulate,
+    tensors_equal,
+    unmount_fused_ln_modulate,
+)
+from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=3, suite="base-a-test-cpu")
+register_cuda_ci(est_time=38, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+requires_cuda = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA required"
+)
+
+
+# ---------------------------------------------------------------------------
+# QualityGatedFusion protocol (CPU)
+# ---------------------------------------------------------------------------
+
+
+def test_quality_gate_mounts_and_unmounts_all_sites():
+    fusion = QualityGatedFusion(
+        name="test fusion",
+        marker_attr="_test_fusion_site",
+        enabled_attr="_test_fusion_enabled",
+    )
+    root = nn.ModuleList([nn.Module(), nn.Module()])
+    for index, site in enumerate(root):
+        fusion.mark(site, index)
+
+    assert [fusion.metadata(site) for site in root] == [0, 1]
+    assert fusion.mount(root)
+    assert all(fusion.is_enabled(site) for site in root)
+    fusion.unmount(root)
+    assert not any(fusion.is_enabled(site) for site in root)
+
+
+def test_quality_gate_rejection_is_all_or_nothing():
+    fusion = QualityGatedFusion(
+        name="test fusion",
+        marker_attr="_test_fusion_site",
+        enabled_attr="_test_fusion_enabled",
+    )
+    root = nn.ModuleList([nn.Module(), nn.Module()])
+    for index, site in enumerate(root):
+        fusion.mark(site, index)
+
+    assert not fusion.mount(
+        root, reject_reason=lambda site: "rejected" if fusion.metadata(site) else None
+    )
+    assert not any(fusion.is_enabled(site) for site in root)
+    assert not fusion.mount(nn.Module())
+
+
+# -------------------------------------------------------------------------
+# BitExactFusionGate protocol (CPU)
+# -------------------------------------------------------------------------
+
+
+def test_bitexact_gate_once_mode_verifies_then_reuses():
+    gate = BitExactFusionGate("once")
+    calls = {"fused": 0, "ref": 0}
+
+    def fused():
+        calls["fused"] += 1
+        return torch.tensor([1.0])
+
+    def ref():
+        calls["ref"] += 1
+        return torch.tensor([1.0])
+
+    assert torch.equal(gate.accept_or_fallback(fused(), ref()), torch.tensor([1.0]))
+    assert gate.verified and not gate.disabled and calls == {"fused": 1, "ref": 1}
+    assert torch.equal(fused(), torch.tensor([1.0]))
+    assert calls == {"fused": 2, "ref": 1}
+
+
+def test_bitexact_gate_mismatch_disables_permanently():
+    gate = BitExactFusionGate("mismatch")
+
+    out = gate.accept_or_fallback(
+        torch.tensor([1.0]),
+        torch.tensor([2.0]),
+        mismatch_msg="mismatch",
+    )
+    assert torch.equal(out, torch.tensor([2.0]))
+    assert gate.disabled and not gate.verified
+
+
+def test_bitexact_gate_per_signature_tracks_each_sig():
+    gate = BitExactFusionGate("sig", per_signature=True)
+    a = torch.tensor([1.0])
+    assert torch.equal(gate.accept_or_fallback(a, a, sig=("a",)), a)
+    assert gate.is_verified(("a",))
+    assert not gate.is_verified(("b",))
+    assert torch.equal(gate.accept_or_fallback(a, a, sig=("b",)), a)
+    assert gate.verified_sigs == {("a",), ("b",)}
+
+
+def test_bitexact_gate_skips_first_sight_during_graph_capture(monkeypatch):
+    # Negative-branch contract: an unverified gate must not attempt first-sight
+    # verification inside CUDA graph capture — the eager-reference host sync
+    # would abort the capture (and BCG would permanently block the signature).
+    gate = BitExactFusionGate("capture")
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
+    assert not gate.can_attempt_once()
+    # A verified gate replays the fused kernel alone, which is capture-safe.
+    gate.mark_verified()
+    assert gate.can_attempt_once()
+
+
+def test_tensors_equal_supports_sequences():
+    assert tensors_equal(
+        (torch.tensor([1.0]), torch.tensor([2.0])),
+        (torch.tensor([1.0]), torch.tensor([2.0])),
+    )
+    assert not tensors_equal(
+        (torch.tensor([1.0]), torch.tensor([2.0])),
+        (torch.tensor([1.0]), torch.tensor([3.0])),
+    )
+
+
+class TestBitExactFallbackDiagnostics(CustomTestCase):
+    def test_mismatch_warning_is_actionable_and_diagnostic_is_lazy(self):
+        logger = MagicMock()
+        diagnostic = MagicMock(return_value="backend=CuTe DSL")
+        gate = BitExactFusionGate("diagnostic")
+
+        matched = gate.accept_or_fallback(
+            torch.tensor([1.0]),
+            torch.tensor([1.0]),
+            logger=logger,
+            diagnostic_hint=diagnostic,
+        )
+        self.assertTrue(torch.equal(matched, torch.tensor([1.0])))
+        diagnostic.assert_not_called()
+        logger.warning_once.assert_not_called()
+
+        gate = BitExactFusionGate("diagnostic")
+        fallback = gate.accept_or_fallback(
+            torch.tensor([1.0]),
+            torch.tensor([2.0]),
+            logger=logger,
+            diagnostic_hint=diagnostic,
+        )
+
+        self.assertTrue(torch.equal(fallback, torch.tensor([2.0])))
+        diagnostic.assert_called_once_with()
+        warning = logger.warning_once.call_args.args[0]
+        self.assertIn("Correctness is preserved", warning)
+        self.assertIn("reference kernel or reduction-order change", warning)
+        self.assertIn("backend=CuTe DSL", warning)
+
+    def test_diagnostic_failure_cannot_break_the_eager_fallback(self):
+        logger = MagicMock()
+
+        def broken_diagnostic():
+            raise RuntimeError("diagnostics unavailable")
+
+        gate = BitExactFusionGate("diagnostic")
+        fallback = gate.accept_or_fallback(
+            torch.tensor([1.0]),
+            torch.tensor([2.0]),
+            logger=logger,
+            diagnostic_hint=broken_diagnostic,
+        )
+
+        self.assertTrue(torch.equal(fallback, torch.tensor([2.0])))
+        self.assertTrue(gate.disabled)
+        self.assertIn("Correctness is preserved", logger.warning_once.call_args.args[0])
+
+    def test_flashinfer_rmsnorm_hint_reports_backend_and_versions(self):
+        flashinfer = ModuleType("flashinfer")
+        flashinfer_norm = ModuleType("flashinfer.norm")
+        flashinfer_norm._USE_CUDA_NORM = False
+        versions = {
+            "flashinfer-python": "0.6.12",
+            "flashinfer-cubin": "0.6.12",
+            "flashinfer-jit-cache": "0.6.12+cu130",
+        }
+
+        with (
+            patch.dict(
+                sys.modules,
+                {"flashinfer": flashinfer, "flashinfer.norm": flashinfer_norm},
+            ),
+            patch("importlib.metadata.version", side_effect=versions.__getitem__),
+            patch.dict("os.environ", {"FLASHINFER_USE_CUDA_NORM": "0"}),
+        ):
+            hint = flashinfer_rmsnorm_diagnostic_hint()
+
+        self.assertIn("backend=CuTe DSL", hint)
+        self.assertIn("FLASHINFER_USE_CUDA_NORM=0", hint)
+        for package, version in versions.items():
+            self.assertIn(f"{package}={version}", hint)
+
+
+# ---------------------------------------------------------------------------
+# Mount lifecycle on synthetic sites (CUDA)
+# ---------------------------------------------------------------------------
+
+LN_DIM = 3072  # FLUX.1 hidden size
+
+
+@requires_cuda
+@pytest.mark.parametrize("seq_len", [4096, 512])
+def test_fused_ln_modulate_matches_reference(seq_len):
+    torch.cuda.manual_seed(0)
+    x = torch.randn((1, seq_len, LN_DIM), device="cuda", dtype=torch.bfloat16)
+    scale = torch.randn((1, LN_DIM), device="cuda", dtype=torch.bfloat16)
+    shift = torch.randn_like(scale)
+    assert can_use_ln_modulate(x, scale, shift)
+
+    norm = nn.LayerNorm(LN_DIM, eps=1e-6, elementwise_affine=False).cuda()
+    ref = norm(x) * (1 + scale[:, None]) + shift[:, None]
+    # Contract: bf16 rounding-order-level difference only, not bit-exact --
+    # which is exactly why this fusion is quality-gated.
+    torch.testing.assert_close(
+        fused_ln_modulate(x, scale, shift, eps=1e-6), ref, atol=0.0625, rtol=0.05
+    )
+
+
+@requires_cuda
+def test_fused_ln_modulate_guards_and_mount_protocol():
+    x = torch.randn((2, 64, LN_DIM), device="cuda", dtype=torch.bfloat16)
+    row = torch.randn((2, LN_DIM), device="cuda", dtype=torch.bfloat16)
+    assert not can_use_ln_modulate(x, row, row)  # folded affine needs B == 1
+
+    root = nn.Module()
+    root.child = nn.Module()
+    mark_fused_ln_modulate_site(root.child)
+    assert not fused_ln_modulate_active(root.child)
+    assert mount_fused_ln_modulate(root)
+    assert fused_ln_modulate_active(root.child)
+    unmount_fused_ln_modulate(root)
+    assert not fused_ln_modulate_active(root.child)
+    assert not mount_fused_ln_modulate(nn.Module())  # no marked sites
+
+
+@requires_cuda
+@torch.no_grad()
+def test_mounted_ln_modulate_site_compiles_fullgraph():
+    class Site(nn.Module):
+        def __init__(self):
+            super().__init__()
+            mark_fused_ln_modulate_site(self)
+
+        def forward(self, x, scale, shift):
+            if fused_ln_modulate_active(self) and can_use_ln_modulate(x, scale, shift):
+                return fused_ln_modulate(x, scale, shift, eps=1e-6)
+            return (
+                F.layer_norm(x, (x.shape[-1],), eps=1e-6) * (1 + scale[:, None])
+                + shift[:, None]
+            )
+
+    site = Site()
+    assert mount_fused_ln_modulate(site)
+    x = torch.randn(1, 64, 128, device="cuda", dtype=torch.bfloat16)
+    scale = torch.randn(1, 128, device="cuda", dtype=torch.bfloat16)
+    shift = torch.randn_like(scale)
+    expected = site(x, scale, shift)
+    torch.testing.assert_close(
+        torch.compile(site, fullgraph=True)(x, scale, shift),
+        expected,
+        atol=0.0625,
+        rtol=0.05,
+    )
+
+
+GATE_RMSNORM_DIM, GATE_RMSNORM_EPS = 4608, 1e-5  # Ideogram 4 hidden / norm_eps
+
+
+class _GateRMSNormSite(nn.Module):
+    def __init__(self, dtype=torch.bfloat16):
+        super().__init__()
+        self.norm = nn.RMSNorm(
+            GATE_RMSNORM_DIM, eps=GATE_RMSNORM_EPS, device="cuda", dtype=dtype
+        )
+        gate_rmsnorm.mark_fused_gate_rmsnorm_site(self, ("norm",))
+
+
+@requires_cuda
+def test_fused_gate_rmsnorm_matches_ideogram_reference():
+    torch.manual_seed(0)
+    site = _GateRMSNormSite()
+    w = site.norm.weight.data
+    dim = GATE_RMSNORM_DIM
+    x = torch.randn(1, 64, dim, device="cuda", dtype=torch.bfloat16)
+    residual = torch.randn_like(x)
+    # adaln-style strided chunks, as produced by Ideogram's modulation .chunk()
+    mods = torch.randn(1, 1, 2 * dim, device="cuda", dtype=torch.bfloat16)
+    scale, gate = mods.chunk(2, dim=-1)
+
+    assert gate_rmsnorm.mount_fused_gate_rmsnorm(site)
+    got_scale = gate_rmsnorm.fused_rmsnorm_scale(x, w, 1.0 + scale, GATE_RMSNORM_EPS)
+    got_gate = gate_rmsnorm.fused_rmsnorm_tanh_residual(
+        x, gate, residual, w, GATE_RMSNORM_EPS
+    )
+    norm = F.rms_norm(x, (dim,), w, GATE_RMSNORM_EPS)
+    # The fused path uses bf16-native norm statistics: close, not bit-exact.
+    torch.testing.assert_close(got_scale, norm * (1.0 + scale), atol=8e-2, rtol=4e-2)
+    torch.testing.assert_close(
+        got_gate, residual + torch.tanh(gate) * norm, atol=8e-2, rtol=4e-2
+    )
+
+
+@requires_cuda
+def test_fused_gate_rmsnorm_mount_is_all_or_nothing():
+    good, bad = _GateRMSNormSite(), _GateRMSNormSite(torch.float32)
+    # One fp32 norm anywhere in the tree keeps *every* site on the reference.
+    assert not gate_rmsnorm.mount_fused_gate_rmsnorm(nn.ModuleList([good, bad]))
+    assert not gate_rmsnorm.fused_gate_rmsnorm_active(good)
+    assert gate_rmsnorm.mount_fused_gate_rmsnorm(good)
+    gate_rmsnorm.unmount_fused_gate_rmsnorm(good)
+    assert not gate_rmsnorm.fused_gate_rmsnorm_active(good)
+
+
+class _GeluSite(nn.Module):
+    def __init__(self, dtype=torch.bfloat16, bias=True):
+        super().__init__()
+        self.proj = nn.Linear(64, 256, bias=bias, device="cuda", dtype=dtype)
+        linear_gelu.mark_fused_gelu_site(self, "proj")
+
+    def forward(self, x):
+        if linear_gelu.fused_gelu_active(self) and linear_gelu.can_use_linear_gelu(
+            self.proj, x
+        ):
+            return linear_gelu.fused_linear_gelu_tanh(
+                x, self.proj.weight, self.proj.bias
+            )
+        return F.gelu(self.proj(x), approximate="tanh")
+
+
+@requires_cuda
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_fused_linear_gelu_matches_reference(dtype):
+    torch.manual_seed(0)
+    site = _GeluSite(dtype)
+    x = torch.randn(512, 64, device="cuda", dtype=dtype)
+    ref = site(x)  # unmounted: the reference chain
+    assert linear_gelu.mount_fused_linear_gelu(site)
+    atol = 2e-2 if dtype == torch.bfloat16 else 4e-3
+    torch.testing.assert_close(site(x), ref, atol=atol, rtol=2e-2)
+
+
+@requires_cuda
+def test_fused_linear_gelu_guards_and_lossless_path():
+    torch.manual_seed(0)
+    good, bad = _GeluSite(), _GeluSite(torch.float32)
+    assert not linear_gelu.mount_fused_linear_gelu(nn.ModuleList([good, bad]))
+    assert not linear_gelu.fused_gelu_active(good)
+
+    x = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16)
+    ref = good(x)
+    assert linear_gelu.mount_fused_linear_gelu(good)
+    linear_gelu.unmount_fused_linear_gelu(good)
+    # Unmounting must restore the reference path bit-for-bit.
+    assert torch.equal(good(x), ref)
+
+    no_bias = nn.Linear(8, 8, bias=False, device="cuda", dtype=torch.bfloat16)
+    assert not linear_gelu.can_use_linear_gelu_static(no_bias)
+    assert not linear_gelu.can_use_linear_gelu(good.proj, x.float())
+
+
+@requires_cuda
+@torch.no_grad()
+def test_mounted_gelu_site_compiles_fullgraph():
+    site = _GeluSite()
+    x = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16)
+    assert linear_gelu.mount_fused_linear_gelu(site)
+    expected = site(x)
+    torch.testing.assert_close(
+        torch.compile(site, fullgraph=True)(x), expected, atol=2e-2, rtol=2e-2
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/diffusion/test_timestep_embedding.py

@@ -1,184 +0,0 @@
-import os
-import sys
-
-import numpy as np
-import pytest
-import torch
-
-try:
-    import tabulate
-except Exception:
-    tabulate = None
-
-from sglang.kernels.jit.utils import get_ci_test_range
-from sglang.kernels.ops.diffusion.timestep_embedding import (
-    timestep_embedding as timestep_embedding_cuda,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=16, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-# Nightly is not redundant here: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1 to expand get_ci_test_range sweeps.
-register_cuda_ci(est_time=20, stage="nightly", runner_config="1-gpu-large")
-
-CORRECTNESS_BATCH_SIZES = get_ci_test_range(
-    [1, 2, 8, 128, 256, 512, 1536, 2048, 4096, 11008, 16384],
-    [1, 128, 2048, 16384],
-)
-CORRECTNESS_DIMS = get_ci_test_range(
-    [32, 128, 256, 512, 1536, 2048, 4096, 8192],
-    [32, 512, 8192],
-)
-DIFFUSERS_BATCH_SIZES = get_ci_test_range(
-    [1, 2, 8, 128, 256, 512, 1536, 2048, 16384],
-    [1, 512, 16384],
-)
-DIFFUSERS_DIMS = get_ci_test_range([32, 256, 512, 1536, 8192], [32, 512, 8192])
-DTYPES = get_ci_test_range(
-    [torch.float16, torch.bfloat16, torch.float32],
-    [torch.float16, torch.bfloat16],
-)
-SCALES = get_ci_test_range([1, 0.01], [1, 0.01])
-
-
-def get_timestep_embedding_reference(
-    timesteps: torch.Tensor,
-    dim: int,
-    *,
-    flip_sin_to_cos: bool = False,
-    downscale_freq_shift: float = 1,
-    scale: float = 1,
-    max_period: int = 10000,
-):
-    assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
-
-    timesteps = timesteps.to(torch.float32)
-    half_dim = dim // 2
-    exponent = -torch.log(
-        torch.tensor(max_period, dtype=torch.float32, device=timesteps.device)
-    ) * torch.arange(
-        start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
-    )
-    exponent = exponent / (half_dim - downscale_freq_shift)
-
-    emb = torch.exp(exponent)
-    emb = timesteps[:, None].float() * emb[None, :]
-
-    emb = scale * emb
-
-    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
-    if flip_sin_to_cos:
-        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
-    if dim % 2 == 1:
-        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
-    return emb
-
-
-@pytest.mark.parametrize("batch_size", CORRECTNESS_BATCH_SIZES)
-@pytest.mark.parametrize("dim", CORRECTNESS_DIMS)
-@pytest.mark.parametrize("dtype", DTYPES)
-def test_timestep_embedding_correctness_with_sgld(batch_size, dim, dtype):
-    device = "cuda"
-    t = torch.randint(low=0, high=1000, size=(batch_size,), device=device).to(dtype)
-    torch_output = get_timestep_embedding_reference(
-        t, dim, flip_sin_to_cos=True, downscale_freq_shift=0
-    )
-    cuda_output = timestep_embedding_cuda(
-        t, dim, flip_sin_to_cos=True, downscale_freq_shift=0
-    )
-    torch.testing.assert_close(torch_output, cuda_output, atol=1e-3, rtol=1e-3)
-
-
-@pytest.mark.parametrize("batch_size", DIFFUSERS_BATCH_SIZES)
-@pytest.mark.parametrize("dim", DIFFUSERS_DIMS)
-@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize("flip_sin_to_cos", [False, True])
-@pytest.mark.parametrize("downscale_freq_shift", [0, 1])
-@pytest.mark.parametrize("scale", SCALES)
-def test_timestep_embedding_correctness_with_diffusers(
-    batch_size, dim, flip_sin_to_cos, downscale_freq_shift, scale, dtype
-):
-    device = "cuda"
-    t = torch.randint(low=0, high=1000, size=(batch_size,), device=device).to(dtype)
-    torch_output = get_timestep_embedding_reference(
-        t,
-        dim,
-        flip_sin_to_cos=flip_sin_to_cos,
-        downscale_freq_shift=downscale_freq_shift,
-        scale=scale,
-        max_period=10000,
-    )
-    cuda_output = timestep_embedding_cuda(
-        t,
-        dim,
-        flip_sin_to_cos=flip_sin_to_cos,
-        downscale_freq_shift=downscale_freq_shift,
-        scale=scale,
-        max_period=10000,
-    )
-    torch.testing.assert_close(torch_output, cuda_output, atol=1e-3, rtol=1e-3)
-
-
-def test_timestep_embedding_perf():
-    if os.environ.get("SGLANG_RUN_JIT_KERNEL_PERF_TESTS") != "1":
-        pytest.skip("Perf test disabled by default")
-    if tabulate is None:
-        pytest.skip("Optional dependency 'tabulate' is not installed")
-
-    NUM_BATCH = [1, 2, 8, 63, 256, 512, 613, 1024, 1536]
-    NUM_DIM = [32, 64, 128, 256, 512, 1024, 2048, 4096]
-
-    def perf_kernel_fn(kernel_fn: callable, *args, **kwargs):
-        warmup_times = 4
-        repeat_times = 20
-        start = torch.cuda.Event(enable_timing=True)
-        end = torch.cuda.Event(enable_timing=True)
-
-        for _ in range(warmup_times):
-            kernel_fn(*args, **kwargs)
-        torch.cuda.synchronize()
-
-        start.record()
-        for _ in range(repeat_times):
-            kernel_fn(*args, **kwargs)
-        end.record()
-        end.synchronize()
-        return start.elapsed_time(end) / repeat_times
-
-    device = "cuda"
-    results = []
-
-    cuda_speedups = []
-    for B in NUM_BATCH:
-        for dim in NUM_DIM:
-            t = torch.linspace(0, max(100000, B), steps=B, device=device).to(
-                torch.float32
-            )
-            time_torch = perf_kernel_fn(get_timestep_embedding_reference, t, dim)
-            time_cuda = perf_kernel_fn(timestep_embedding_cuda, t, dim)
-            speedup_cuda = time_torch / time_cuda
-
-            results.append(
-                {
-                    "Batch Size": B,
-                    "Dimension": dim,
-                    "Torch Time (ms)": time_torch,
-                    "CUDA Time (ms)": time_cuda,
-                    "Speedup (CUDA)": speedup_cuda,
-                }
-            )
-            cuda_speedups.append(speedup_cuda)
-
-    print("=== Timestep Embedding Benchmark Results ===")
-    print(
-        tabulate.tabulate(
-            results,
-            headers="keys",
-            tablefmt="fancy_grid",
-            floatfmt=(".0f", ".0f", ".6f", ".6f", ".5f"),
-        )
-    )
-    print(f"Average Speedup(cuda): {np.mean(cuda_speedups):.4f}")
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/diffusion/test_ulysses_qkv.py

@@ -1,54 +0,0 @@
-import sys
-
-import pytest
-import torch
-
-from sglang.kernels.ops.diffusion.triton.ulysses_qkv import (
-    pack_qkv_destination_major,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-
-
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-def test_pack_qkv_destination_major_is_bit_exact(dtype):
-    torch.manual_seed(0)
-    rows, world_size, global_heads, head_size = 17, 4, 12, 64
-    q, k, v = (
-        torch.randn(rows, global_heads, head_size, device="cuda", dtype=dtype)
-        for _ in range(3)
-    )
-
-    local_heads = global_heads // world_size
-    expected = torch.empty(
-        world_size,
-        rows,
-        local_heads,
-        3 * head_size,
-        device="cuda",
-        dtype=dtype,
-    )
-    for index, tensor in enumerate((q, k, v)):
-        shards = tensor.view(rows, world_size, local_heads, head_size).permute(
-            1, 0, 2, 3
-        )
-        expected[..., index * head_size : (index + 1) * head_size].copy_(shards)
-
-    actual = pack_qkv_destination_major(q, k, v, world_size)
-    assert torch.equal(actual, expected)
-
-
-def test_pack_qkv_destination_major_validates_inputs():
-    q = torch.empty(2, 4, 8, device="cuda", dtype=torch.bfloat16)
-    with pytest.raises(ValueError, match="same 3D shape"):
-        pack_qkv_destination_major(q, q[:, :-1], q, 2)
-    with pytest.raises(ValueError, match="divide global_heads"):
-        pack_qkv_destination_major(q, q, q, 3)
-    with pytest.raises(ValueError, match="expected shape"):
-        pack_qkv_destination_major(q, q, q, 2, out=torch.empty_like(q))
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/diffusion/test_usp_relayout.py

@@ -1,69 +0,0 @@
-"""Bitwise tests for the generic Ulysses output head-merge fast path."""
-
-import sys
-from unittest.mock import patch
-
-import pytest
-import torch
-
-from sglang.kernels.ops.diffusion.usp_relayout import (
-    can_use_usp_merge_heads,
-    usp_merge_heads,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=4, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-
-DEVICE = "cuda"
-
-
-@pytest.mark.parametrize(
-    "world,seq,batch,h_local,head_dim",
-    [
-        (4, 7936, 1, 14, 128),  # H3 768p production shape (Ulysses 4)
-        (2, 64, 3, 4, 64),  # batched
-        (4, 33, 2, 4, 100),  # scalar fallback inside the CUDA kernel
-    ],
-)
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
-def test_usp_merge_heads_bitwise(dtype, world, seq, batch, h_local, head_dim):
-    generator = torch.Generator(device=DEVICE).manual_seed(4321)
-    x = torch.randn(
-        world,
-        seq,
-        batch,
-        h_local,
-        head_dim,
-        dtype=dtype,
-        device=DEVICE,
-        generator=generator,
-    )
-    assert can_use_usp_merge_heads(x)
-    out = usp_merge_heads(x)
-    ref = x.permute(2, 1, 0, 3, 4).contiguous()
-    assert out.shape == ref.shape
-    assert torch.equal(out, ref)
-
-
-def test_usp_merge_heads_unsupported_inputs_use_exact_fallback():
-    x = torch.randn(2, 4, 1, 4, 64, dtype=torch.bfloat16, device=DEVICE)
-    unsupported = [x.transpose(0, 1), x[:0]]
-
-    for value in unsupported:
-        assert not can_use_usp_merge_heads(value)
-        assert torch.equal(
-            usp_merge_heads(value), value.permute(2, 1, 0, 3, 4).contiguous()
-        )
-
-    with patch.object(torch.version, "hip", "6.3"):
-        assert not can_use_usp_merge_heads(x)
-        assert torch.equal(usp_merge_heads(x), x.permute(2, 1, 0, 3, 4).contiguous())
-
-
-def test_usp_merge_heads_fast_path_rejects_wrong_rank():
-    x = torch.randn(2, 4, 1, 4, 64, dtype=torch.bfloat16, device=DEVICE)
-    assert not can_use_usp_merge_heads(x[0])
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/diffusion/test_varlen_pack_pad.py

@@ -1,197 +0,0 @@
-"""Numerical correctness for fused varlen pack/scatter Triton kernels.
-
-Bit-exact comparison against the equivalent PyTorch ops (index_select,
-zeros + index_copy_) across bf16/fp16 and several shape/mask cases.
-"""
-
-import pytest
-import torch
-
-from sglang.kernels.jit.utils import get_ci_test_range
-from sglang.kernels.ops.diffusion.triton.varlen_pack_pad import (
-    build_inv_indices,
-    fused_pack_qkv,
-    fused_scatter_to_padded,
-)
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-
-register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-# Nightly is not redundant here: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1 to expand get_ci_test_range sweeps.
-register_cuda_ci(est_time=20, stage="nightly", runner_config="1-gpu-large")
-register_amd_ci(est_time=15, suite="nightly-amd-kernel-1-gpu", nightly=True)
-
-DEVICE = "cuda"
-DTYPES = get_ci_test_range([torch.bfloat16, torch.float16], [torch.bfloat16])
-# (bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens) tuples
-SHAPES = get_ci_test_range(
-    [
-        # name, bs, s_txt, s_img, H, D, valid_txt_lens
-        ("small_c2", 2, 64, 128, 4, 64, [32, 48]),
-        ("prod_c2", 2, 256, 1024, 24, 128, [128, 200]),
-        ("all_valid_b1", 1, 64, 128, 4, 64, [64]),
-        ("all_valid_b4", 4, 64, 128, 4, 64, [64, 64, 64, 64]),
-        ("c8_prod", 8, 256, 4096, 24, 128, [128, 200, 256, 100, 50, 256, 256, 50]),
-        # one batch with zero valid text tokens (image side still valid)
-        ("zero_txt_one_batch", 2, 64, 128, 4, 64, [0, 32]),
-        # bs=1 with no text validity (only image rows packed)
-        ("bs1_zero_txt", 1, 64, 128, 4, 64, [0]),
-    ],
-    [
-        ("small_c2", 2, 64, 128, 4, 64, [32, 48]),
-        ("prod_c2", 2, 256, 1024, 24, 128, [128, 200]),
-        ("all_valid_b4", 4, 64, 128, 4, 64, [64, 64, 64, 64]),
-    ],
-)
-
-
-def _build_mask(bs, s_txt, s_img, valid_txt_lens):
-    s = s_txt + s_img
-    mask = torch.zeros(bs, s, dtype=torch.bool, device=DEVICE)
-    for b, vt in enumerate(valid_txt_lens):
-        mask[b, :vt] = True
-        mask[b, s_txt:] = True
-    return mask
-
-
-def _ref_pack(q, k, v, indices):
-    bs, seq = q.shape[:2]
-
-    def flat(t):
-        return t.reshape(bs * seq, *t.shape[2:])
-
-    return (
-        flat(q).index_select(0, indices),
-        flat(k).index_select(0, indices),
-        flat(v).index_select(0, indices),
-    )
-
-
-def _ref_scatter(out_unpad, indices, bs, seq):
-    _, num_heads, head_dim = out_unpad.shape
-    flat = torch.zeros(
-        bs * seq, num_heads, head_dim, dtype=out_unpad.dtype, device=DEVICE
-    )
-    flat.index_copy_(0, indices, out_unpad)
-    return flat.view(bs, seq, num_heads, head_dim)
-
-
-def _build_meta(mask):
-    bs, seq = mask.shape
-    indices = mask.reshape(-1).nonzero(as_tuple=False).flatten()
-    inv_indices = build_inv_indices(indices, bs * seq)
-    return indices, inv_indices
-
-
-@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize(
-    "shape", SHAPES, ids=lambda s: s[0] if isinstance(s, tuple) else str(s)
-)
-def test_pack_matches_index_select(dtype, shape):
-    _, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens = shape
-    torch.manual_seed(0)
-    s = s_txt + s_img
-    mask = _build_mask(bs, s_txt, s_img, valid_txt_lens)
-    indices, _ = _build_meta(mask)
-
-    q = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
-    k = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
-    v = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
-
-    q_ref, k_ref, v_ref = _ref_pack(q, k, v, indices)
-    q_fused, k_fused, v_fused = fused_pack_qkv(q, k, v, indices)
-
-    # bit-exact: pack is pure gather, no math
-    assert torch.equal(q_ref, q_fused)
-    assert torch.equal(k_ref, k_fused)
-    assert torch.equal(v_ref, v_fused)
-
-
-@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize(
-    "shape", SHAPES, ids=lambda s: s[0] if isinstance(s, tuple) else str(s)
-)
-def test_scatter_matches_index_copy(dtype, shape):
-    _, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens = shape
-    torch.manual_seed(1)
-    s = s_txt + s_img
-    mask = _build_mask(bs, s_txt, s_img, valid_txt_lens)
-    indices, inv_indices = _build_meta(mask)
-    n_valid = indices.shape[0]
-
-    out_unpad = torch.randn(n_valid, num_heads, head_dim, dtype=dtype, device=DEVICE)
-    out_ref = _ref_scatter(out_unpad, indices, bs, s)
-    out_fused = fused_scatter_to_padded(out_unpad, inv_indices, bs, s)
-
-    # bit-exact: scatter is pure copy + zero-fill
-    assert torch.equal(out_ref, out_fused)
-    # Padding rows must be exactly zero
-    invalid = ~mask
-    if invalid.any():
-        assert out_fused[invalid].abs().max().item() == 0.0
-
-
-def test_pack_handles_non_contiguous_input():
-    """Helper must accept non-contiguous Q/K/V (auto .contiguous() inside)."""
-    torch.manual_seed(2)
-    bs, s_txt, s_img, num_heads, head_dim = 2, 64, 128, 4, 64
-    s = s_txt + s_img
-    mask = _build_mask(bs, s_txt, s_img, [32, 48])
-    indices, _ = _build_meta(mask)
-
-    # Build non-contiguous tensors via permute
-    qkv_pre = torch.randn(
-        bs, num_heads, s, head_dim, dtype=torch.bfloat16, device=DEVICE
-    )
-    q = qkv_pre.permute(0, 2, 1, 3)
-    k = torch.randn_like(qkv_pre).permute(0, 2, 1, 3)
-    v = torch.randn_like(qkv_pre).permute(0, 2, 1, 3)
-    assert not q.is_contiguous()
-
-    q_ref, k_ref, v_ref = _ref_pack(
-        q.contiguous(), k.contiguous(), v.contiguous(), indices
-    )
-    q_fused, k_fused, v_fused = fused_pack_qkv(q, k, v, indices)
-    assert torch.equal(q_ref, q_fused)
-    assert torch.equal(k_ref, k_fused)
-    assert torch.equal(v_ref, v_fused)
-
-
-def test_build_inv_indices_matches_manual():
-    """build_inv_indices output should match the manual full+scatter form."""
-    torch.manual_seed(3)
-    bs, s = 2, 32
-    mask = torch.bernoulli(torch.full((bs, s), 0.6, device=DEVICE)).to(torch.bool)
-    indices = mask.reshape(-1).nonzero(as_tuple=False).flatten()
-    n_valid = indices.shape[0]
-
-    manual = torch.full((bs * s,), -1, dtype=torch.int32, device=DEVICE)
-    if n_valid > 0:
-        manual[indices.long()] = torch.arange(n_valid, dtype=torch.int32, device=DEVICE)
-
-    built = build_inv_indices(indices, bs * s)
-    assert torch.equal(built, manual)
-
-
-def test_empty_valid_set_handled():
-    """All-False mask: pack returns empty tensors; scatter writes all zeros."""
-    bs, s, num_heads, head_dim = 2, 16, 4, 64
-    mask = torch.zeros(bs, s, dtype=torch.bool, device=DEVICE)
-    indices = mask.reshape(-1).nonzero(as_tuple=False).flatten()
-    inv_indices = build_inv_indices(indices, bs * s)
-    assert indices.numel() == 0
-
-    q = torch.randn(bs, s, num_heads, head_dim, dtype=torch.bfloat16, device=DEVICE)
-    q_unpad, k_unpad, v_unpad = fused_pack_qkv(q, q.clone(), q.clone(), indices)
-    assert q_unpad.shape == (0, num_heads, head_dim)
-    assert k_unpad.shape == (0, num_heads, head_dim)
-    assert v_unpad.shape == (0, num_heads, head_dim)
-
-    out_padded = fused_scatter_to_padded(q_unpad, inv_indices, bs, s)
-    assert out_padded.shape == (bs, s, num_heads, head_dim)
-    assert out_padded.abs().max().item() == 0.0
-
-
-if __name__ == "__main__":
-    import sys
-
-    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/diffusion/test_varlen_uspattn_equivalence.py

@@ -1,157 +0,0 @@
-"""End-to-end equivalence between USPAttention varlen path and SDPA reference.
-
-Compares the production varlen path (``build_varlen_mask_meta`` +
-``fused_pack_qkv`` + ``flash_attn_varlen_func`` + ``fused_scatter_to_padded``)
-against ``torch.nn.functional.scaled_dot_product_attention`` with a broadcast
-key mask, for inputs the gating in ``USPAttention.forward`` would accept.
-
-Verifies the documented contract:
-  * Valid (non-masked) query rows match SDPA within FA-vs-SDPA tolerance.
-  * Masked query rows are exactly zero in the varlen path (differs from
-    SDPA, which produces deterministic attention output at those rows).
-"""
-
-import pytest
-import torch
-import torch.nn.functional as F
-
-from sglang.kernels.jit.utils import get_ci_test_range
-from sglang.kernels.ops.attention.flash_attention import flash_attn_varlen_func
-from sglang.kernels.ops.diffusion.triton.varlen_pack_pad import (
-    fused_pack_qkv,
-    fused_scatter_to_padded,
-)
-from sglang.multimodal_gen.runtime.layers.attention.backends import (
-    flash_attn as _fa_backend,
-)
-from sglang.multimodal_gen.runtime.layers.attention.layer import (
-    build_varlen_mask_meta,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-# Nightly is not redundant here: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1 to expand get_ci_test_range sweeps.
-register_cuda_ci(est_time=20, stage="nightly", runner_config="1-gpu-large")
-
-DEVICE = "cuda"
-DTYPES = get_ci_test_range([torch.bfloat16, torch.float16], [torch.bfloat16])
-# (name, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens)
-SHAPES = get_ci_test_range(
-    [
-        ("small_c2", 2, 64, 128, 4, 64, [32, 48]),
-        ("prod_c2", 2, 256, 1024, 24, 128, [128, 200]),
-        ("all_valid_b1", 1, 64, 128, 4, 64, [64]),
-        ("zero_txt_one_batch", 2, 64, 128, 4, 64, [0, 32]),
-    ],
-    [
-        ("small_c2", 2, 64, 128, 4, 64, [32, 48]),
-    ],
-)
-
-
-def _build_mask(bs, s_txt, s_img, valid_txt_lens):
-    s = s_txt + s_img
-    mask = torch.zeros(bs, s, dtype=torch.bool, device=DEVICE)
-    for b, vt in enumerate(valid_txt_lens):
-        mask[b, :vt] = True
-        mask[b, s_txt:] = True
-    return mask
-
-
-def _sdpa_with_key_mask(q, k, v, key_mask, softmax_scale):
-    """Reference: SDPA with a ``[B, S]`` key mask broadcast to ``[B, 1, 1, S]``."""
-    q_ = q.transpose(1, 2)
-    k_ = k.transpose(1, 2)
-    v_ = v.transpose(1, 2)
-    mask = key_mask.to(dtype=q.dtype)[:, None, None, :]
-    mask = (mask - 1.0) * torch.finfo(q.dtype).max
-    out = F.scaled_dot_product_attention(
-        q_,
-        k_,
-        v_,
-        attn_mask=mask,
-        dropout_p=0.0,
-        is_causal=False,
-        scale=softmax_scale,
-    )
-    return out.transpose(1, 2)
-
-
-def _varlen_path(q, k, v, key_mask, softmax_scale):
-    """Production varlen path matching USPAttention.forward masked branch."""
-    bs, seq = q.shape[0], q.shape[1]
-    meta = build_varlen_mask_meta(key_mask)
-    indices = meta["indices"]
-    if indices.shape[0] == 0:
-        return torch.zeros_like(q)
-    q_unpad, k_unpad, v_unpad = fused_pack_qkv(q, k, v, indices)
-    out_unpad = flash_attn_varlen_func(
-        q=q_unpad,
-        k=k_unpad,
-        v=v_unpad,
-        cu_seqlens_q=meta["cu_seqlens"],
-        cu_seqlens_k=meta["cu_seqlens"],
-        max_seqlen_q=meta["max_seqlen"],
-        max_seqlen_k=meta["max_seqlen"],
-        softmax_scale=softmax_scale,
-        causal=False,
-        ver=_fa_backend.fa_ver,
-    )
-    return fused_scatter_to_padded(out_unpad, meta["inv_indices"], bs, seq)
-
-
-@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize(
-    "shape", SHAPES, ids=lambda s: s[0] if isinstance(s, tuple) else str(s)
-)
-def test_varlen_path_matches_sdpa_on_valid_rows(dtype, shape):
-    """Valid rows: varlen output ≈ SDPA output within FA tolerance."""
-    _, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens = shape
-    torch.manual_seed(0)
-    s = s_txt + s_img
-    softmax_scale = head_dim**-0.5
-    mask = _build_mask(bs, s_txt, s_img, valid_txt_lens)
-    q = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
-    k = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
-    v = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
-
-    out_sdpa = _sdpa_with_key_mask(q, k, v, mask, softmax_scale)
-    out_varlen = _varlen_path(q, k, v, mask, softmax_scale)
-
-    valid = mask[..., None, None].expand_as(out_sdpa)
-    rtol = 1e-2 if dtype == torch.bfloat16 else 5e-3
-    atol = 5e-2 if dtype == torch.bfloat16 else 1e-2
-    torch.testing.assert_close(
-        out_sdpa[valid],
-        out_varlen[valid],
-        rtol=rtol,
-        atol=atol,
-    )
-
-
-@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize(
-    "shape", SHAPES, ids=lambda s: s[0] if isinstance(s, tuple) else str(s)
-)
-def test_varlen_path_zeros_masked_rows(dtype, shape):
-    """Masked rows: varlen path produces exact zeros (documented contract)."""
-    _, bs, s_txt, s_img, num_heads, head_dim, valid_txt_lens = shape
-    torch.manual_seed(1)
-    s = s_txt + s_img
-    softmax_scale = head_dim**-0.5
-    mask = _build_mask(bs, s_txt, s_img, valid_txt_lens)
-    q = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
-    k = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
-    v = torch.randn(bs, s, num_heads, head_dim, dtype=dtype, device=DEVICE)
-
-    out_varlen = _varlen_path(q, k, v, mask, softmax_scale)
-
-    invalid = ~mask
-    if invalid.any():
-        assert (out_varlen[invalid] == 0).all(), "masked rows must be zero-filled"
-
-
-if __name__ == "__main__":
-    import sys
-
-    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/diffusion/test_wan_causal_cache.py

@@ -1,164 +0,0 @@
-"""Wan causal VAE data-movement kernels: the fused conv-input builder and the
-fused DupUp3D shortcut add must be bitwise identical to the aten op chains
-they replace (they are pure data movement plus zero fill / one fp32 add)."""
-
-import sys
-
-import pytest
-import torch
-import torch.nn.functional as F
-
-from sglang.kernels.ops.diffusion.triton.wan_causal_cache import (
-    cat_pad_channels_last_3d,
-    dup_up3d_add,
-)
-from sglang.multimodal_gen.runtime.models.vaes import wanvae
-from sglang.multimodal_gen.runtime.models.vaes.wanvae import (
-    CACHE_T,
-    WanCausalConv3d,
-    _cache_payload,
-    _run_cached_causal_conv,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-
-
-def _cl3d(shape, dtype):
-    return torch.randn(shape, device="cuda", dtype=dtype).contiguous(
-        memory_format=torch.channels_last_3d
-    )
-
-
-def _ref_cat_pad(x, cache, padding):
-    p = list(padding)
-    if cache is not None:
-        x = torch.cat([cache, x], dim=2)
-        p[4] -= cache.shape[2]
-    if any(p):
-        x = F.pad(x, p)
-    return x.contiguous(memory_format=torch.channels_last_3d)
-
-
-@torch.no_grad()
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
-@pytest.mark.parametrize(
-    "c,t,h,w,cache_t,pads",
-    [
-        (96, 1, 10, 14, 0, (1, 1, 1, 1, 2, 0)),  # first chunk, zero-fill front
-        (96, 1, 10, 14, 1, (1, 1, 1, 1, 2, 0)),  # legacy 1-frame cache
-        (96, 1, 10, 14, 2, (1, 1, 1, 1, 2, 0)),  # steady state k3 conv
-        (64, 1, 10, 14, 2, (0, 0, 0, 0, 2, 0)),  # time_conv (temporal only)
-        (48, 4, 10, 14, 2, (1, 1, 1, 1, 2, 0)),  # encoder-style T=4 chunk
-    ],
-)
-def test_cat_pad_bitwise(dtype, c, t, h, w, cache_t, pads) -> None:
-    torch.cuda.manual_seed(0)
-    x = _cl3d((1, c, t, h, w), dtype)
-    cache = None
-    if cache_t:
-        # Strided interior view: caches may arrive as non-contiguous slices.
-        ph, pw = pads[2], pads[0]
-        buf = _cl3d((1, c, cache_t, h + 2 * ph, w + 2 * pw), dtype)
-        cache = buf[:, :, :, ph : ph + h, pw : pw + w]
-    out = cat_pad_channels_last_3d(x, cache, pads)
-    ref = _ref_cat_pad(x, cache, pads)
-    assert out is not None and out.shape == ref.shape
-    assert out.is_contiguous(memory_format=torch.channels_last_3d)
-    assert torch.equal(out, ref)
-
-    # Dual-output mode: the same pass also emits the compact feature cache
-    # (unpadded interior of the last frames), bitwise equal to the slice.
-    pair = cat_pad_channels_last_3d(x, cache, pads, keep_cache_t=2)
-    assert pair is not None
-    out2, keep = pair
-    assert torch.equal(out2, ref)
-    ph, pw = pads[2], pads[0]
-    keep_t = min(2, ref.shape[2])
-    want = ref[:, :, ref.shape[2] - keep_t :, ph : ph + h, pw : pw + w]
-    assert keep.shape == want.shape
-    assert keep.is_contiguous(memory_format=torch.channels_last_3d)
-    assert torch.equal(keep, want)
-
-
-@torch.no_grad()
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
-@pytest.mark.parametrize(
-    "c_in,c_out,t,h,w,ft,fs,drop",
-    [
-        (128, 64, 1, 10, 14, 2, 2, False),
-        (128, 64, 1, 10, 14, 2, 2, True),  # first_chunk slicing
-        (64, 32, 2, 10, 14, 1, 2, False),
-    ],
-)
-def test_dup_up3d_add_bitwise(dtype, c_in, c_out, t, h, w, ft, fs, drop) -> None:
-    torch.cuda.manual_seed(0)
-    repeats = c_out * ft * fs * fs // c_in
-    src = _cl3d((1, c_in, t, h, w), dtype)
-    t_out = t * ft - (ft - 1 if drop else 0)
-    # Main arm as a permuted view, like the WanResample 2D output.
-    main = torch.randn(
-        (1, t_out, c_out, h * fs, w * fs), device="cuda", dtype=dtype
-    ).permute(0, 2, 1, 3, 4)
-
-    dup = src.repeat_interleave(repeats, dim=1)
-    dup = dup.view(1, c_out, ft, fs, fs, t, h, w)
-    dup = dup.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
-    dup = dup.view(1, c_out, t * ft, h * fs, w * fs)
-    if drop:
-        dup = dup[:, :, ft - 1 :, :, :]
-    ref = main + dup
-
-    out = dup_up3d_add(main, src, ft, fs, repeats, drop)
-    assert out is not None and out.shape == ref.shape
-    # Layout must match the aten add output exactly (downstream reductions
-    # are layout-sensitive), and every value must be bitwise identical.
-    assert out.stride() == ref.stride()
-    assert torch.equal(out, ref)
-
-
-@torch.no_grad()
-@pytest.mark.parametrize("pads_temporal_only", [False, True])
-def test_cached_conv_chunk_loop_bitwise(pads_temporal_only) -> None:
-    """The fused conv-input/compact-cache scheme must reproduce the original
-    clone/cat bookkeeping bitwise across a chunked decode, including the
-    first-chunk zero fill and the "Rep" marker start used by WanResample."""
-    torch.cuda.manual_seed(0)
-    c = 64
-    if pads_temporal_only:
-        conv = WanCausalConv3d(c, 2 * c, (3, 1, 1), padding=(1, 0, 0))
-    else:
-        conv = WanCausalConv3d(c, c, 3, padding=1)
-    conv = conv.to(device="cuda", dtype=torch.float32)
-    conv.weight.data = conv.weight.data.contiguous(memory_format=torch.channels_last_3d)
-    chunks = [_cl3d((1, c, 1, 10, 14), torch.float32) for _ in range(4)]
-
-    def run(force_fallback, start):
-        cache = [start]
-        outs = []
-        if force_fallback:
-            orig = wanvae.cat_pad_channels_last_3d
-            wanvae.cat_pad_channels_last_3d = None
-        try:
-            for x in chunks:
-                outs.append(_run_cached_causal_conv(conv, x, cache, 0))
-        finally:
-            if force_fallback:
-                wanvae.cat_pad_channels_last_3d = orig
-        return outs, cache[0]
-
-    for start in (None, "Rep"):
-        fused_outs, fused_cache = run(False, start)
-        ref_outs, ref_cache = run(True, start)
-        for got, want in zip(fused_outs, ref_outs, strict=True):
-            assert torch.equal(got, want)
-        got_payload = _cache_payload(fused_cache)
-        assert got_payload is not None and got_payload.shape[2] == CACHE_T
-        # Reference cache holds the last CACHE_T unpadded frames.
-        assert torch.equal(got_payload, ref_cache[:, :, -CACHE_T:])
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/diffusion/test_wan_vae_fastpath.py

@@ -1,79 +0,0 @@
-"""Wan VAE decoder fast path: fused-kernel numerics and gate dispatch
-(the lossless off-path must stay bit-exact)."""
-
-import sys
-
-import pytest
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from sglang.kernels.ops.diffusion.triton.wan_rmsnorm_silu import wan_rmsnorm_silu
-from sglang.multimodal_gen.runtime.models.vaes.wan_vae_cuda_opt import (
-    FusedWanRMSNormSiLU,
-    VaeFastPathGate,
-)
-from sglang.multimodal_gen.runtime.models.vaes.wanvae import WanRMS_norm
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-
-
-def _cl3d(shape, dtype):
-    return torch.randn(shape, device="cuda", dtype=dtype).contiguous(
-        memory_format=torch.channels_last_3d
-    )
-
-
-@torch.no_grad()
-@pytest.mark.parametrize(
-    "x_dtype,affine_dtype,atol,rtol",
-    [
-        (torch.float32, torch.float32, 1e-5, 1e-5),  # FastWan2.2 fp32 decode
-        (torch.bfloat16, torch.float32, 1.5e-1, 3e-2),  # Wan2.1 bf16 autocast
-    ],
-)
-def test_kernel_numerics(x_dtype, affine_dtype, atol, rtol) -> None:
-    torch.cuda.manual_seed(0)
-    x = _cl3d((1, 96, 3, 10, 14), x_dtype)
-    gamma = torch.randn((96, 1, 1, 1), device="cuda", dtype=affine_dtype)
-    for bias in (None, torch.randn_like(gamma)):
-        expected = F.silu(
-            F.normalize(x, dim=1) * 96**0.5 * gamma + (0 if bias is None else bias)
-        )
-        actual = wan_rmsnorm_silu(x, gamma, bias)
-        assert actual is not None and actual.dtype == expected.dtype
-        assert actual.stride() == x.stride()
-        torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
-
-
-@torch.no_grad()
-def test_fused_module_gate_dispatch() -> None:
-    # Gate off must stay bit-exact; gate on must route to the fused kernel.
-    torch.cuda.manual_seed(0)
-    norm = WanRMS_norm(96, images=False).to(device="cuda", dtype=torch.bfloat16)
-    norm.gamma.add_(torch.randn_like(norm.gamma))
-    gate = VaeFastPathGate()
-    fused = FusedWanRMSNormSiLU(norm, gate)
-    # Parameter names must not change (weight transfer matches by name).
-    assert [n for n, _ in fused.named_parameters()] == ["gamma"]
-    x = _cl3d((1, 96, 3, 10, 14), torch.bfloat16)
-    assert torch.equal(fused(x), nn.SiLU()(norm(x)))
-    gate.enabled = True
-    expected = wan_rmsnorm_silu(x, norm.gamma, rms_scale=float(norm.scale))
-    assert torch.equal(fused(x), expected)
-
-
-@torch.no_grad()
-def test_kernel_rejects_empty_input() -> None:
-    x = torch.empty(1, 96, 0, 2, 2, device="cuda", dtype=torch.bfloat16).to(
-        memory_format=torch.channels_last_3d
-    )
-    gamma = torch.ones(96, 1, 1, 1, device="cuda", dtype=torch.bfloat16)
-    assert wan_rmsnorm_silu(x, gamma) is None
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, "-v", "-s"]))