[Kernel] Reclassify kernel tests by ops group + move helpers out of the package (RFC #29630) (#32128)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

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

@@ -0,0 +1,89 @@
+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_diffusion_modelopt_fp8_scaled_mm.py

@@ -0,0 +1,141 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8
+from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
+    ModelOptFp8Config,
+    ModelOptFp8LinearMethod,
+)
+from sglang.srt.layers.quantization.fp8_utils import (
+    cutlass_fp8_supported,
+    input_to_float8,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+MAX_FP8_DIFF = 5e-4
+TEST_CASES = [
+    pytest.param(19, 150, 80, id="misaligned_projection_shape"),
+    pytest.param(512, 3072, 4096, id="flux2_added_kv_projection_shape"),
+]
+
+
+def _modelopt_fp8_supported() -> bool:
+    return torch.cuda.is_available() and cutlass_fp8_supported()
+
+
+def _calc_diff(x: torch.Tensor, y: torch.Tensor) -> float:
+    x, y = x.double(), y.double()
+    denominator = (x * x + y * y).sum()
+    if denominator == 0:
+        return 0.0
+    sim = 2 * (x * y).sum() / denominator
+    return (1 - sim).item()
+
+
+def _dequantize_fp8_input(qinput: torch.Tensor, x_scale: torch.Tensor) -> torch.Tensor:
+    return qinput.to(torch.float32) * x_scale.to(torch.float32)
+
+
+def _dequantize_fp8_weight(
+    weight: torch.Tensor, weight_scale: torch.Tensor
+) -> torch.Tensor:
+    if weight_scale.ndim == 0 or weight_scale.numel() == 1:
+        scale = weight_scale.to(torch.float32)
+    else:
+        scale = weight_scale.to(torch.float32).reshape(-1, 1).t()
+    return weight.to(torch.float32) * scale
+
+
+def _build_layer(
+    weight_q: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_scale: torch.Tensor,
+) -> tuple[torch.nn.Module, ModelOptFp8LinearMethod]:
+    output_size, input_size = weight_q.shape
+    method = ModelOptFp8LinearMethod(
+        ModelOptFp8Config(is_checkpoint_fp8_serialized=True)
+    )
+    layer = torch.nn.Module()
+    method.create_weights(
+        layer=layer,
+        input_size_per_partition=input_size,
+        output_partition_sizes=[output_size],
+        input_size=input_size,
+        output_size=output_size,
+        params_dtype=DTYPE,
+        weight_loader=lambda *args, **kwargs: None,
+    )
+    layer = layer.to(device=DEVICE)
+
+    layer.weight.data.copy_(weight_q)
+    layer.weight_scale.data.copy_(weight_scale.reshape_as(layer.weight_scale))
+    layer.input_scale.data.copy_(input_scale.reshape_as(layer.input_scale))
+    method.process_weights_after_loading(layer)
+    return layer, method
+
+
+@pytest.mark.skipif(
+    not _modelopt_fp8_supported(),
+    reason="Diffusion ModelOpt FP8 scaled mm correctness requires CUDA FP8 support",
+)
+@pytest.mark.parametrize("m,n,k", TEST_CASES)
+def test_checkpoint_processing(m: int, n: int, k: int) -> None:
+    generator = torch.Generator(device=DEVICE)
+    generator.manual_seed(20260410 + m + n + k)
+
+    weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
+    weight_q, weight_scale = input_to_float8(weight)
+    input_scale = torch.tensor(1.0, device=DEVICE, dtype=torch.float32)
+
+    layer, _ = _build_layer(weight_q, weight_scale, input_scale)
+
+    assert tuple(layer.weight.shape) == (k, n)
+    assert tuple(layer.weight.stride()) == (1, k)
+    assert layer.weight.dtype == torch.float8_e4m3fn
+    assert layer.input_scale.ndim == 0
+    assert tuple(layer.weight_scale.shape) == (n, 1)
+
+    expected_weight = weight_q.t().to(torch.float32) * weight_scale.to(torch.float32)
+    actual_weight = _dequantize_fp8_weight(layer.weight, layer.weight_scale)
+    torch.testing.assert_close(actual_weight, expected_weight, atol=0.0, rtol=0.0)
+
+
+@pytest.mark.skipif(
+    not _modelopt_fp8_supported(),
+    reason="Diffusion ModelOpt FP8 scaled mm correctness requires CUDA FP8 support",
+)
+@pytest.mark.parametrize("m,n,k", TEST_CASES)
+def test_shape_correctness(m: int, n: int, k: int) -> None:
+    generator = torch.Generator(device=DEVICE)
+    generator.manual_seed(20260410 + m + n + k)
+
+    x = torch.randn((m, k), device=DEVICE, dtype=DTYPE, generator=generator)
+    weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
+    weight_q, weight_scale = input_to_float8(weight)
+    _, input_scale = input_to_float8(x)
+
+    layer, method = _build_layer(weight_q, weight_scale, input_scale)
+
+    qinput, x_scale = static_quant_fp8(
+        x.contiguous(),
+        layer.input_scale,
+        repeat_scale=method.cutlass_fp8_supported,
+    )
+    expected = torch.matmul(
+        _dequantize_fp8_input(qinput, x_scale),
+        _dequantize_fp8_weight(layer.weight, layer.weight_scale),
+    )
+
+    actual = method.apply(layer, x)
+    diff = _calc_diff(actual, expected.to(dtype=DTYPE))
+    assert diff < MAX_FP8_DIFF, f"{m=}, {n=}, {k=}, {diff=:.6f}"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

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

@@ -0,0 +1,447 @@
+import sys
+
+import flashinfer
+import pytest
+import torch
+
+from sglang.multimodal_gen.runtime.layers.quantization import (
+    modelopt_quant as diffusion_modelopt_quant,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
+    ModelOptFp4Config,
+    ModelOptFp4LinearMethod,
+)
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.srt.layers.quantization.modelopt_quant import pad_nvfp4_weight
+from sglang.test.ci.ci_register import register_cuda_ci
+
+# B200-only correctness coverage for diffusion NVFP4 scaled mm.
+register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+BLOCK_SIZE = 16
+FLOAT4_E2M1_MAX = 6.0
+FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
+FP4_VALUE_LUT = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
+DEEPGEMM_FP4_MAX_DIFF = 0.02
+TEST_CASES = [
+    pytest.param(19, 150, 80, id="padding_regression"),
+    pytest.param(512, 6144, 128, id="flux2_projection_shape"),
+]
+FLUX2_PROJECTION_SHAPE = (512, 6144, 128)
+
+
+def _nvfp4_supported() -> bool:
+    return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0)
+
+
+def _make_global_scale(x: torch.Tensor) -> torch.Tensor:
+    max_abs = torch.amax(x.abs()).clamp_min_(1e-6)
+    return (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / max_abs).to(torch.float32)
+
+
+def _calc_diff(x: torch.Tensor, y: torch.Tensor) -> float:
+    x, y = x.double(), y.double()
+    denominator = (x * x + y * y).sum()
+    if denominator == 0:
+        return 0.0
+    sim = 2 * (x * y).sum() / denominator
+    return (1 - sim).item()
+
+
+def _swap_fp4_nibbles(packed: torch.Tensor) -> torch.Tensor:
+    return ((packed >> 4) | (packed << 4)).contiguous()
+
+
+def _fp4_lut(device: torch.device) -> torch.Tensor:
+    return torch.tensor(FP4_VALUE_LUT, dtype=torch.float32, device=device)
+
+
+def _unpack_fp4_bytes(packed: torch.Tensor) -> torch.Tensor:
+    assert packed.dtype == torch.uint8
+    lut = _fp4_lut(packed.device)
+
+    def _decode(nibbles: torch.Tensor) -> torch.Tensor:
+        values = lut[(nibbles & 0x7).to(torch.long)]
+        return torch.where((nibbles & 0x8) != 0, -values, values)
+
+    low = _decode(packed & 0x0F)
+    high = _decode((packed & 0xF0) >> 4)
+    return torch.stack((low, high), dim=-1).reshape(
+        packed.shape[0], packed.shape[1] * 2
+    )
+
+
+def _swizzled_to_linear(
+    scales_swizzled: torch.Tensor,
+    rows: int,
+    cols: int,
+) -> torch.Tensor:
+    scales_swizzled = scales_swizzled.view(torch.float8_e4m3fn)
+    row_tiles = (rows + 128 - 1) // 128
+    tile_cols = BLOCK_SIZE * 4
+    col_tiles = (cols + tile_cols - 1) // tile_cols
+    tmp = scales_swizzled.reshape(1, row_tiles, col_tiles, 32, 4, 4)
+    tmp = tmp.permute(0, 1, 4, 3, 2, 5)
+    linear = tmp.reshape(row_tiles * 128, col_tiles * tile_cols // BLOCK_SIZE)
+    return linear[:rows, : cols // BLOCK_SIZE]
+
+
+def _dequantize_nvfp4(
+    packed: torch.Tensor,
+    scales_swizzled: torch.Tensor,
+    global_scale: torch.Tensor,
+) -> torch.Tensor:
+    rows, packed_cols = packed.shape
+    cols = packed_cols * 2
+    unpacked = _unpack_fp4_bytes(packed).reshape(rows, cols // BLOCK_SIZE, BLOCK_SIZE)
+    scales_linear = _swizzled_to_linear(scales_swizzled, rows, cols).to(torch.float32)
+    return (unpacked * (scales_linear / global_scale).unsqueeze(-1)).reshape(rows, cols)
+
+
+def _quantize_weight_for_checkpoint(
+    weight: torch.Tensor, weight_global_scale: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    weight_fp4, weight_scale_linear = flashinfer.fp4_quantize(
+        weight,
+        weight_global_scale,
+        is_sf_swizzled_layout=False,
+    )
+    if weight_scale_linear.dtype == torch.uint8:
+        weight_scale_linear = weight_scale_linear.view(torch.float8_e4m3fn)
+    return weight_fp4, weight_scale_linear.contiguous()
+
+
+def _set_diffusion_fp4_backend(
+    monkeypatch: pytest.MonkeyPatch, backend: str | None
+) -> None:
+    if backend is None:
+        monkeypatch.delenv(
+            "SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND", raising=False
+        )
+    else:
+        monkeypatch.setenv("SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND", backend)
+
+    current_platform.__class__.get_modelopt_flashinfer_fp4_backend.cache_clear()
+    current_platform.__class__.get_modelopt_fp4_gemm_op.cache_clear()
+    diffusion_modelopt_quant._get_fp4_gemm_op.cache_clear()
+
+
+def _build_layer(
+    weight_fp4: torch.Tensor,
+    weight_scale_linear: torch.Tensor,
+    input_global_scale: torch.Tensor,
+    weight_global_scale: torch.Tensor,
+    *,
+    weight_scale_device: torch.device | str | None = None,
+    checkpoint_weight_scale_layout: str = "linear",
+) -> tuple[ModelOptFp4LinearMethod, torch.nn.Module]:
+    output_size, input_size_half = weight_fp4.shape
+    input_size = input_size_half * 2
+    method = ModelOptFp4LinearMethod(
+        ModelOptFp4Config(
+            is_checkpoint_nvfp4_serialized=True,
+            group_size=BLOCK_SIZE,
+            swap_weight_nibbles=True,
+            checkpoint_weight_scale_layout=checkpoint_weight_scale_layout,
+        )
+    )
+    layer = torch.nn.Module()
+    method.create_weights(
+        layer,
+        input_size_per_partition=input_size,
+        output_partition_sizes=[output_size],
+        input_size=input_size,
+        output_size=output_size,
+        params_dtype=DTYPE,
+        weight_loader=lambda *args, **kwargs: None,
+    )
+    layer = layer.to(device=DEVICE)
+
+    checkpoint_weight = _swap_fp4_nibbles(weight_fp4)
+    layer.weight.data.copy_(checkpoint_weight)
+    layer.input_scale.data.copy_(
+        (1.0 / input_global_scale).reshape_as(layer.input_scale)
+    )
+    layer.weight_scale_2.data.copy_(
+        (1.0 / weight_global_scale).reshape_as(layer.weight_scale_2)
+    )
+    layer.weight_scale.data.copy_(weight_scale_linear)
+    if weight_scale_device is not None:
+        layer.weight_scale = torch.nn.Parameter(
+            layer.weight_scale.detach().to(weight_scale_device), requires_grad=False
+        )
+
+    method.process_weights_after_loading(layer)
+
+    _, flashinfer_backend = current_platform.get_modelopt_fp4_gemm_op()
+    if flashinfer_backend == "trtllm":
+        expected_weight, _ = pad_nvfp4_weight(
+            weight_fp4, n_alignment=128, k_alignment=0
+        )
+        expected_scale = (
+            _swizzled_to_linear(weight_scale_linear, output_size, input_size)
+            if checkpoint_weight_scale_layout == "swizzled"
+            else weight_scale_linear
+        )
+        if expected_scale.shape[0] != expected_weight.shape[0]:
+            pad_n = expected_weight.shape[0] - expected_scale.shape[0]
+            expected_scale = torch.nn.functional.pad(expected_scale, (0, 0, 0, pad_n))
+
+        expected_padding_cols = 0
+        if expected_scale.shape[1] % 4 != 0:
+            padded_scale_k = ((expected_scale.shape[1] + 4 - 1) // 4) * 4
+            pad_scale_k = padded_scale_k - expected_scale.shape[1]
+            expected_scale = torch.nn.functional.pad(
+                expected_scale, (0, pad_scale_k, 0, 0)
+            )
+            pad_weight_k = pad_scale_k * 8
+            expected_weight = torch.nn.functional.pad(
+                expected_weight, (0, pad_weight_k, 0, 0)
+            )
+            expected_padding_cols = pad_weight_k
+
+        expected_weight = flashinfer.shuffle_matrix_a(
+            expected_weight.view(torch.uint8), 128
+        )
+        expected_scale = (
+            flashinfer.shuffle_matrix_sf_a(expected_scale.view(torch.uint8), 128)
+            .reshape(expected_scale.shape)
+            .view(torch.float8_e4m3fn)
+        )
+
+        assert torch.equal(layer.weight, expected_weight)
+        assert torch.equal(
+            layer.weight_scale_interleaved.view(torch.uint8),
+            expected_scale.view(torch.uint8),
+        )
+        assert layer.weights_padding_cols == expected_padding_cols
+    else:
+        expected_weight, expected_padding_cols = pad_nvfp4_weight(weight_fp4)
+        expected_scale_shape = (
+            ((output_size + 128 - 1) // 128) * 128,
+            (((input_size // BLOCK_SIZE) + 4 - 1) // 4) * 4,
+        )
+
+        assert torch.equal(layer.weight, expected_weight)
+        assert layer.weight_scale_interleaved.shape == expected_scale_shape
+        assert layer.weight_scale_interleaved.dtype == torch.float8_e4m3fn
+        assert layer.weights_padding_cols == expected_padding_cols
+    torch.testing.assert_close(
+        layer.alpha,
+        (1.0 / (input_global_scale * weight_global_scale)).to(torch.float32),
+    )
+    torch.testing.assert_close(
+        layer.input_scale_inv,
+        input_global_scale.to(torch.float32),
+    )
+    return method, layer
+
+
+def _resolve_mode(mode: str):
+    if mode == "flashinfer2":
+        return flashinfer.fp4_quantize, flashinfer.mm_fp4, "cudnn"
+    if mode == "flashinfer_trtllm":
+        return flashinfer.fp4_quantize, flashinfer.mm_fp4, "trtllm"
+    raise ValueError(f"Unknown mode: {mode}")
+
+
+@pytest.mark.skipif(
+    not _nvfp4_supported(),
+    reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
+)
+@pytest.mark.parametrize(
+    "backend", [None, "flashinfer_trtllm"], ids=["default", "flashinfer_trtllm"]
+)
+@pytest.mark.parametrize("m,n,k", TEST_CASES)
+def test_checkpoint_processing(
+    monkeypatch: pytest.MonkeyPatch, backend: str | None, m: int, n: int, k: int
+) -> None:
+    _set_diffusion_fp4_backend(monkeypatch, backend)
+    generator = torch.Generator(device=DEVICE)
+    generator.manual_seed(20260404 + m + n + k)
+
+    weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
+    input_global_scale = torch.tensor(512.0, device=DEVICE, dtype=torch.float32)
+    weight_global_scale = _make_global_scale(weight)
+    weight_fp4, weight_scale_linear = _quantize_weight_for_checkpoint(
+        weight, weight_global_scale
+    )
+
+    _build_layer(
+        weight_fp4, weight_scale_linear, input_global_scale, weight_global_scale
+    )
+
+
+@pytest.mark.skipif(
+    not _nvfp4_supported(),
+    reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
+)
+@pytest.mark.parametrize("mode", ["flashinfer2"])
+def test_flux2_shape_correctness(mode: str) -> None:
+    m, n, k = FLUX2_PROJECTION_SHAPE
+    quantize_op, gemm_op, gemm_backend = _resolve_mode(mode)
+    generator = torch.Generator(device=DEVICE)
+    generator.manual_seed(20260404 + m + n + k)
+
+    x = torch.randn((m, k), device=DEVICE, dtype=DTYPE, generator=generator)
+    weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
+    input_global_scale = _make_global_scale(x)
+    weight_global_scale = _make_global_scale(weight)
+    alpha = (1.0 / (input_global_scale * weight_global_scale)).to(torch.float32)
+
+    x_fp4, x_scale_swizzled = quantize_op(x, input_global_scale)
+    weight_fp4, weight_scale_swizzled = quantize_op(weight, weight_global_scale)
+    if x_scale_swizzled.dtype == torch.uint8:
+        x_scale_swizzled = x_scale_swizzled.view(torch.float8_e4m3fn)
+    if weight_scale_swizzled.dtype == torch.uint8:
+        weight_scale_swizzled = weight_scale_swizzled.view(torch.float8_e4m3fn)
+
+    expected = torch.matmul(
+        _dequantize_nvfp4(x_fp4, x_scale_swizzled, input_global_scale),
+        _dequantize_nvfp4(weight_fp4, weight_scale_swizzled, weight_global_scale).t(),
+    )
+
+    actual = gemm_op(
+        x_fp4,
+        weight_fp4.t(),
+        x_scale_swizzled,
+        weight_scale_swizzled.t(),
+        alpha,
+        DTYPE,
+        backend=gemm_backend,
+    )
+
+    diff = _calc_diff(actual, expected.to(dtype=DTYPE))
+    assert diff < DEEPGEMM_FP4_MAX_DIFF, f"{mode=}, {m=}, {n=}, {k=}, {diff=:.6f}"
+
+
+@pytest.mark.skipif(
+    not _nvfp4_supported(),
+    reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
+)
+def test_flux2_shape_correctness_flashinfer_trtllm(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _set_diffusion_fp4_backend(monkeypatch, "flashinfer_trtllm")
+
+    m, n, k = FLUX2_PROJECTION_SHAPE
+    generator = torch.Generator(device=DEVICE)
+    generator.manual_seed(20260404 + m + n + k + 17)
+
+    x = torch.randn((m, k), device=DEVICE, dtype=DTYPE, generator=generator)
+    weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
+    input_global_scale = _make_global_scale(x)
+    weight_global_scale = _make_global_scale(weight)
+    weight_fp4, weight_scale_linear = _quantize_weight_for_checkpoint(
+        weight, weight_global_scale
+    )
+
+    method, layer = _build_layer(
+        weight_fp4, weight_scale_linear, input_global_scale, weight_global_scale
+    )
+    actual = method.apply(layer, x)
+
+    x_fp4, x_scale_swizzled = flashinfer.fp4_quantize(x, input_global_scale)
+    weight_fp4_ref, weight_scale_swizzled = flashinfer.fp4_quantize(
+        weight, weight_global_scale
+    )
+    if x_scale_swizzled.dtype == torch.uint8:
+        x_scale_swizzled = x_scale_swizzled.view(torch.float8_e4m3fn)
+    if weight_scale_swizzled.dtype == torch.uint8:
+        weight_scale_swizzled = weight_scale_swizzled.view(torch.float8_e4m3fn)
+
+    expected = torch.matmul(
+        _dequantize_nvfp4(x_fp4, x_scale_swizzled, input_global_scale),
+        _dequantize_nvfp4(
+            weight_fp4_ref, weight_scale_swizzled, weight_global_scale
+        ).t(),
+    )
+
+    diff = _calc_diff(actual, expected.to(dtype=DTYPE))
+    assert diff < DEEPGEMM_FP4_MAX_DIFF, f"{m=}, {n=}, {k=}, {diff=:.6f}"
+
+
+@pytest.mark.skipif(
+    not _nvfp4_supported(),
+    reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
+)
+def test_flux2_swizzled_scale_checkpoint_flashinfer_trtllm_matches_cudnn(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _set_diffusion_fp4_backend(monkeypatch, "flashinfer_trtllm")
+
+    m, n, k = FLUX2_PROJECTION_SHAPE
+    generator = torch.Generator(device=DEVICE)
+    generator.manual_seed(20260517 + m + n + k)
+
+    x = torch.randn((m, k), device=DEVICE, dtype=DTYPE, generator=generator)
+    weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
+    input_global_scale = _make_global_scale(x)
+    weight_global_scale = _make_global_scale(weight)
+    alpha = (1.0 / (input_global_scale * weight_global_scale)).to(torch.float32)
+
+    x_fp4, x_scale_swizzled = flashinfer.fp4_quantize(x, input_global_scale)
+    weight_fp4, weight_scale_swizzled = flashinfer.fp4_quantize(
+        weight, weight_global_scale
+    )
+    if x_scale_swizzled.dtype == torch.uint8:
+        x_scale_swizzled = x_scale_swizzled.view(torch.float8_e4m3fn)
+    if weight_scale_swizzled.dtype == torch.uint8:
+        weight_scale_swizzled = weight_scale_swizzled.view(torch.float8_e4m3fn)
+
+    method, layer = _build_layer(
+        weight_fp4,
+        weight_scale_swizzled,
+        input_global_scale,
+        weight_global_scale,
+        checkpoint_weight_scale_layout="swizzled",
+    )
+    actual = method.apply(layer, x)
+
+    expected = flashinfer.mm_fp4(
+        x_fp4,
+        weight_fp4.t(),
+        x_scale_swizzled,
+        weight_scale_swizzled.t(),
+        alpha,
+        DTYPE,
+        backend="cudnn",
+    )
+
+    diff = _calc_diff(actual, expected)
+    assert diff < DEEPGEMM_FP4_MAX_DIFF, f"{m=}, {n=}, {k=}, {diff=:.6f}"
+
+
+@pytest.mark.skipif(
+    not _nvfp4_supported(),
+    reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
+)
+def test_checkpoint_processing_flashinfer_trtllm_cpu_weight_scale(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _set_diffusion_fp4_backend(monkeypatch, "flashinfer_trtllm")
+
+    m, n, k = FLUX2_PROJECTION_SHAPE
+    generator = torch.Generator(device=DEVICE)
+    generator.manual_seed(20260413 + m + n + k)
+
+    weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
+    input_global_scale = torch.tensor(512.0, device=DEVICE, dtype=torch.float32)
+    weight_global_scale = _make_global_scale(weight)
+    weight_fp4, weight_scale_linear = _quantize_weight_for_checkpoint(
+        weight, weight_global_scale
+    )
+
+    _build_layer(
+        weight_fp4,
+        weight_scale_linear,
+        input_global_scale,
+        weight_global_scale,
+        weight_scale_device="cpu",
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

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

@@ -0,0 +1,133 @@
+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_norm_scale_shift.py

@@ -0,0 +1,251 @@
+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_group_norm_silu.py

@@ -0,0 +1,104 @@
+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_ltx2_ada_values.py

@@ -0,0 +1,70 @@
+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):
+        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_qknorm_split_rope.py

@@ -0,0 +1,221 @@
+import sys
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from sglang.kernels.ops.diffusion.ltx2_qknorm_split_rope import (
+    can_use_ltx2_qknorm_split_rope_cuda,
+    ltx2_qknorm_split_rope_cuda,
+)
+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")
+
+BF16_FUSED_ATOL = 1.6e-1
+
+
+def _require_cuda_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(
+    batch: int, seq_len: int, num_heads: int, head_dim: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    half_dim = head_dim // 2
+    cos = torch.randn(
+        batch, seq_len, num_heads, half_dim, device="cuda", dtype=torch.bfloat16
+    ).transpose(1, 2)
+    sin = torch.randn(
+        batch, seq_len, num_heads, half_dim, device="cuda", dtype=torch.bfloat16
+    ).transpose(1, 2)
+    return cos, sin
+
+
+def _apply_split_rotary_ref(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> torch.Tensor:
+    x_dtype = x.dtype
+    batch = x.shape[0]
+    _, num_heads, seq_len, _ = cos.shape
+    x = x.reshape(batch, seq_len, num_heads, -1).swapaxes(1, 2)
+    last = x.shape[-1]
+    half = last // 2
+
+    split_x = x.reshape(*x.shape[:-1], 2, half)
+    first_x = split_x[..., :1, :]
+    second_x = split_x[..., 1:, :]
+    cos_u = cos.unsqueeze(-2)
+    sin_u = sin.unsqueeze(-2)
+
+    out = split_x * cos_u
+    out[..., :1, :].addcmul_(-sin_u, second_x)
+    out[..., 1:, :].addcmul_(sin_u, first_x)
+    out = out.reshape(*out.shape[:-2], last)
+    return out.swapaxes(1, 2).reshape(batch, seq_len, -1).to(dtype=x_dtype)
+
+
+def _reference(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_cos: torch.Tensor,
+    q_sin: torch.Tensor,
+    k_cos: torch.Tensor,
+    k_sin: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    # rms_norm isn't autocast fp32-preserving, so feed fp32 inputs directly
+    # to keep the normalized value unrounded until the final RoPE output.
+    q_norm = F.rms_norm(q.float(), (q.shape[-1],), q_weight.float(), eps)
+    k_norm = F.rms_norm(k.float(), (k.shape[-1],), k_weight.float(), eps)
+    q_ref = _apply_split_rotary_ref(q_norm, q_cos, q_sin)
+    k_ref = _apply_split_rotary_ref(k_norm, k_cos, k_sin)
+    return q_ref.to(dtype=torch.bfloat16), k_ref.to(dtype=torch.bfloat16)
+
+
+@pytest.mark.parametrize(
+    "batch,q_seq,k_seq,num_heads,head_dim",
+    [
+        (1, 3, 3, 32, 128),
+        (1, 5, 2, 32, 64),
+        (2, 4, 3, 32, 64),
+    ],
+)
+def test_ltx2_qknorm_split_rope_matches_torch_exactly(
+    batch: int, q_seq: int, k_seq: int, num_heads: int, head_dim: int
+) -> None:
+    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_weight = torch.randn(hidden, device="cuda", dtype=torch.bfloat16)
+    k_weight = torch.randn(hidden, device="cuda", dtype=torch.bfloat16)
+
+    assert can_use_ltx2_qknorm_split_rope_cuda(
+        q,
+        q_cos,
+        q_sin,
+        q_weight,
+        k,
+        k_cos,
+        k_sin,
+        k_weight,
+        num_heads=num_heads,
+        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_out, k_out = ltx2_qknorm_split_rope_cuda(
+        q,
+        q_cos,
+        q_sin,
+        q_weight,
+        k,
+        k_cos,
+        k_sin,
+        k_weight,
+        eps=eps,
+        num_heads=num_heads,
+        head_dim=head_dim,
+    )
+    torch.cuda.synchronize()
+
+    torch.testing.assert_close(q_out, q_ref, rtol=0, atol=BF16_FUSED_ATOL)
+    torch.testing.assert_close(k_out, k_ref, rtol=0, atol=BF16_FUSED_ATOL)
+
+
+def test_ltx2_qknorm_split_rope_rejects_unsupported_inputs() -> None:
+    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_weight = torch.randn(4096, device="cuda", dtype=torch.bfloat16)
+    k_weight = torch.randn(4096, device="cuda", dtype=torch.bfloat16)
+
+    assert can_use_ltx2_qknorm_split_rope_cuda(
+        q,
+        q_cos,
+        q_sin,
+        q_weight,
+        k,
+        q_cos,
+        q_sin,
+        k_weight,
+        num_heads=32,
+        head_dim=128,
+    )
+    assert not can_use_ltx2_qknorm_split_rope_cuda(
+        q.float(),
+        q_cos,
+        q_sin,
+        q_weight,
+        k,
+        q_cos,
+        q_sin,
+        k_weight,
+        num_heads=32,
+        head_dim=128,
+    )
+    assert not can_use_ltx2_qknorm_split_rope_cuda(
+        q,
+        q_cos,
+        q_sin,
+        q_weight,
+        k,
+        q_cos.transpose(-1, -2),
+        q_sin,
+        k_weight,
+        num_heads=32,
+        head_dim=128,
+    )
+
+
+def test_ltx2_qknorm_split_rope_custom_op_torch_compile_fullgraph() -> None:
+    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_weight = torch.randn(hidden, device="cuda", dtype=torch.bfloat16)
+    k_weight = torch.randn(hidden, device="cuda", dtype=torch.bfloat16)
+
+    def fn(q, k, q_cos, q_sin, k_cos, k_sin, q_weight, k_weight):
+        return ltx2_qknorm_split_rope_cuda(
+            q,
+            q_cos,
+            q_sin,
+            q_weight,
+            k,
+            k_cos,
+            k_sin,
+            k_weight,
+            eps=1e-6,
+            num_heads=num_heads,
+            head_dim=head_dim,
+        )
+
+    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, k, q_cos, q_sin, k_cos, k_sin, q_weight, k_weight, 1e-6
+    )
+    torch.cuda.synchronize()
+    torch.testing.assert_close(q_out, q_ref, rtol=0, atol=BF16_FUSED_ATOL)
+    torch.testing.assert_close(k_out, k_ref, rtol=0, atol=BF16_FUSED_ATOL)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

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

@@ -0,0 +1,154 @@
+import itertools
+import sys
+
+import pytest
+import torch
+import triton
+
+from sglang.kernels.jit.utils import get_ci_test_range
+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.
+register_cuda_ci(est_time=176, suite="nightly-kernel-1-gpu", nightly=True)
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+MAX_SEQ_LEN = 131072
+ROPE_BASE = 10000.0
+ATOL = 8e-2
+RTOL = 1e-2
+
+
+def create_cos_sin_cache(
+    rotary_dim: int,
+    max_position: int = MAX_SEQ_LEN,
+    base: float = ROPE_BASE,
+) -> torch.Tensor:
+    inv_freq = 1.0 / (
+        base
+        ** (
+            torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=DEVICE)
+            / rotary_dim
+        )
+    )
+    t = torch.arange(max_position, dtype=torch.float32, device=DEVICE)
+    freqs = torch.einsum("i,j->ij", t, inv_freq)
+    return torch.cat((freqs.cos(), freqs.sin()), dim=-1)
+
+
+def split_qknorm_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    positions: torch.Tensor,
+    is_neox: bool,
+) -> None:
+    from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace
+
+    from sglang.kernels.ops.layernorm._jit_norm import fused_inplace_qknorm
+
+    fused_inplace_qknorm(q, k, q_weight, k_weight)
+    apply_rope_with_cos_sin_cache_inplace(
+        positions=positions.long(),
+        query=q.view(q.shape[0], -1),
+        key=k.view(k.shape[0], -1),
+        head_size=q.shape[-1],
+        cos_sin_cache=cos_sin_cache,
+        is_neox=is_neox,
+    )
+
+
+def fused_qknorm_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    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,
+        q_weight,
+        k_weight,
+        cos_sin_cache,
+        positions,
+        is_neox=is_neox,
+        rope_dim=cos_sin_cache.shape[-1],
+    )
+
+
+BS_LIST = [2**n for n in range(13)]
+BS_LIST += [x + 1 for x in BS_LIST]
+BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 129, 257, 2049, 4097])
+HEADS_LIST = get_ci_test_range([8, 16, 24, 32], [8, 24])
+HEAD_DIM_LIST = get_ci_test_range([64, 128, 256], [64, 128, 256])
+IS_NEOX_LIST = [False, True]
+POSITION_DTYPES = [torch.int32, torch.int64]
+ROPE_DIM_CHOICES = {
+    64: [64],
+    128: [64, 128],
+    256: [64, 128, 256],
+}
+
+
+@pytest.mark.parametrize(
+    "batch_size,num_heads,head_dim,is_neox,position_dtype",
+    list(
+        itertools.product(
+            BS_LIST,
+            HEADS_LIST,
+            HEAD_DIM_LIST,
+            IS_NEOX_LIST,
+            POSITION_DTYPES,
+        )
+    ),
+)
+def test_qknorm_rope(
+    batch_size: int,
+    num_heads: int,
+    head_dim: int,
+    is_neox: bool,
+    position_dtype: torch.dtype,
+) -> None:
+    rope_dims = ROPE_DIM_CHOICES[head_dim]
+    for rope_dim in rope_dims:
+        if is_neox:
+            elems_per_thread = head_dim // 32
+            rotary_lanes = rope_dim // elems_per_thread
+            if rotary_lanes < 2 or rotary_lanes & (rotary_lanes - 1):
+                continue
+
+        q = torch.randn(batch_size, num_heads, head_dim, device=DEVICE, dtype=DTYPE)
+        k = torch.randn(batch_size, num_heads, head_dim, device=DEVICE, dtype=DTYPE)
+        q_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
+        k_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
+        positions = torch.randint(
+            0, MAX_SEQ_LEN, (batch_size,), device=DEVICE, dtype=position_dtype
+        )
+        cos_sin_cache = create_cos_sin_cache(rope_dim)
+
+        q_ref, k_ref = q.clone(), k.clone()
+        q_fused, k_fused = q.clone(), k.clone()
+
+        split_qknorm_rope(
+            q_ref, k_ref, q_weight, k_weight, cos_sin_cache, positions, is_neox
+        )
+        fused_qknorm_rope(
+            q_fused, k_fused, q_weight, k_weight, cos_sin_cache, positions, is_neox
+        )
+
+        # The split baseline mixes a separate BF16 qknorm kernel with FlashInfer RoPE,
+        # which differs from the fused path by about one BF16 rounding step on H200.
+        triton.testing.assert_close(q_ref, q_fused, atol=ATOL, rtol=RTOL)
+        triton.testing.assert_close(k_ref, k_fused, atol=ATOL, rtol=RTOL)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

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

@@ -0,0 +1,228 @@
+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=120, suite="nightly-kernel-1-gpu", nightly=True)
+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

@@ -0,0 +1,101 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion.residual_gate_add import (
+    can_use_residual_gate_add_cuda,
+    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)),
+]
+
+
+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)
+
+
+@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])
+
+    # 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
+    )
+
+
+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_cuda(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_timestep_embedding.py

@@ -0,0 +1,184 @@
+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=120, suite="nightly-kernel-1-gpu", nightly=True)
+
+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):
+            output_fn = kernel_fn(*args, **kwargs)
+        torch.cuda.synchronize()
+
+        start.record()
+        for _ in range(repeat_times):
+            output_fn = 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_varlen_pack_pad.py

@@ -0,0 +1,195 @@
+"""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=60, suite="nightly-kernel-1-gpu", nightly=True)
+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]
+    flat = lambda t: 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):
+    n_valid = indices.shape[0]
+    _, 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

@@ -0,0 +1,157 @@
+"""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=60, suite="nightly-kernel-1-gpu", nightly=True)
+
+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"]))