[diffusion] refactor: remove stale kernels and dead code (#32651)

2026-07-29 06:23:23 +08:00
parent 7f438a6031
commit 7778dd23ea
29 changed files with 301 additions and 763 deletions
@@ -1,6 +1,6 @@
-"""Diffusion-model kernels (group-norm+silu, residual-gate-add, qk-norm+rope).
+"""Registered diffusion-model kernels and their public wrappers.
-These are JIT CUDA kernels; the wrappers forward to ``sglang.kernels.ops.diffusion``.
+Implementations use the backend recorded by each kernel specification.
 """
 from __future__ import annotations
@@ -25,11 +25,11 @@ _CUDA = frozenset({CapabilityRequirement.CUDA})
 register_kernel(
    KernelSpec(
        op="diffusion.apply_group_norm_silu",
-        backend=KernelBackend.JIT,
+        backend=KernelBackend.TRITON,
        target="sglang.kernels.ops.diffusion.group_norm_silu:apply_group_norm_silu",
        capabilities=_CUDA,
        format_signature=FormatSignature(description="fused GroupNorm + SiLU"),
-        description="Fused group-norm + SiLU (sglang.kernels.jit).",
+        description="Fused group-norm + SiLU (Triton).",
    )
 )
 register_kernel(
@@ -60,7 +60,7 @@ def apply_group_norm_silu(
    x: torch.Tensor, norm: nn.Module, activation: nn.Module
 ) -> torch.Tensor:
    """Fused GroupNorm + SiLU (falls back to eager when unsupported)."""
-    return get_kernel("diffusion.apply_group_norm_silu", KernelBackend.JIT)(
+    return get_kernel("diffusion.apply_group_norm_silu", KernelBackend.TRITON)(
        x, norm, activation
    )
@@ -1,344 +0,0 @@
 from typing import Optional, Tuple
 import cuda.bindings.driver as cuda
 import cutlass
 import cutlass.cute as cute
 import torch
 from sglang.kernels.ops.diffusion.cutedsl.common.norm_fusion import (
    apply_norm_cta,
    broadcast_tensor_for_bsfd,
    tensor_slice_for_bsfd,
 )
 from sglang.kernels.ops.diffusion.cutedsl.utils import (
    WARP_SIZE,
    to_cute_arg,
    to_fake_cute_args,
 )
 _COMPILE_CACHE = {}
 class NormTanhMulAddNormScale:
    @classmethod
    def make_hash_key(cls, *inputs):
        """
        Compile-time values:
          - D: hidden dimension (size of the last dimension)
          - norm_type: layer norm or RMS norm
          - tensor dtype
          - tensor rank (i.e., tensor.ndim)
        Runtime values:
          - all other inputs
        This hash key defines the compile-time specialization boundary for
        NormTanhMulAddNormScale kernels.
        """
        def _sig(val):
            if isinstance(val, torch.Tensor):
                return (val.dtype, val.ndim, val.shape[-1])
            return val
        return tuple(_sig(val) for val in inputs)
    def __init__(self, D: int, norm_type: str, is_norm2: bool):
        self.D = D
        self.norm_type = norm_type  # "layer" or "rms"
        self.is_norm2 = is_norm2  # single norm or double norm
        self.num_warps = self.D // 256  # num of warps per cta
        self.num_threads = self.num_warps * WARP_SIZE  # num of threads per cta
    @cute.jit
    def __call__(
        self,
        mY,
        mY2,
        mX,
        mWeight,
        mBias,
        mScale,
        mShift,
        mWeight2,
        mBias2,
        mScale2,
        eps: cutlass.Float32 = cutlass.Float32(1e-5),
        stream: cuda.CUstream = cuda.CUstream(cuda.CUstream_flags.CU_STREAM_DEFAULT),
    ):
        # Tensor shapes
        B, S, _ = mX.shape  # (batch, seq_len, hidden_dim)
        # Vectorized copy configuration
        num_vectorized = 8  # maximum num of elem per copy
        atom_copy = cute.make_copy_atom(
            cute.nvgpu.CopyUniversalOp(),
            mX.element_type,
            num_bits_per_copy=128,
        )
        # Thread/value layouts for tiled copy
        t_layout = cute.make_layout(self.num_threads)  # thread layout within a CTA
        v_layout = cute.make_layout(num_vectorized)  # per-thread vector layout
        tiled_copy = cute.make_tiled_copy_tv(atom_copy, t_layout, v_layout)
        self.kernel(
            mY,
            mY2,
            mX,
            mWeight,
            mBias,
            mScale,
            mShift,
            mWeight2,
            mBias2,
            mScale2,
            tiled_copy,
            eps,
        ).launch(
            grid=[B * S, 1, 1],
            block=[self.num_threads, 1, 1],
            stream=stream,
        )
    @cute.kernel
    def kernel(
        self,
        mY,
        mY2,
        mX,
        mWeight,
        mBias,
        mScale,
        mShift,
        mWeight2,
        mBias2,
        mScale2,
        tiled_copy: cute.TiledCopy,
        eps: cutlass.Float32,
    ):
        _, S, _ = mX.shape
        tidx, _, _ = cute.arch.thread_idx()  # thread index
        bid, _, _ = cute.arch.block_idx()  # cta index
        bidx = cutlass.Int32(bid // S)  # batch index
        bidy = cutlass.Int32(bid % S)  # seq_len index
        thr_copy = tiled_copy.get_slice(tidx)
        @cute.jit
        def slice_if(mV):
            if cutlass.const_expr(isinstance(mV, cute.Tensor)):
                return tensor_slice_for_bsfd(mV, thr_copy, bidx, bidy, S, self.D)
            return mV, mV
        @cute.jit
        def copy_if(src, dst):
            if cutlass.const_expr(
                isinstance(src, cute.Tensor) and isinstance(dst, cute.Tensor)
            ):
                cute.autovec_copy(src, dst)  # LDG.128
        @cute.jit
        def norm(x, weight, bias):
            return apply_norm_cta(
                self.norm_type, self.num_warps, tidx, x, weight, bias, self.D, eps
            )
        # Slice: retrieve the per-thread data slices for both global memory (gmem)
        tXgX, tXrX = slice_if(mX)  # x
        tWgW, tWrW = slice_if(mWeight)  # weight
        tBgB, tBrB = slice_if(mBias)  # bias
        tSCgSC, tSCrSC = slice_if(mScale)  # scale
        tSHgSH, tSHrSH = slice_if(mShift)  # shift
        tYgY, tYrY = slice_if(mY)  # y
        if cutlass.const_expr(self.is_norm2):
            tYgY2, tYrY2 = slice_if(mY2)  # y2
            tWgW2, tWrW2 = slice_if(mWeight2)  # weight2
            tBgB2, tBrB2 = slice_if(mBias2)  # bias2
            tSCgSC2, tSCrSC2 = slice_if(mScale2)  # scale2
        # Load: load tensor from global memory to registers
        copy_if(tXgX, tXrX)  # gmem -> rmem
        copy_if(tWgW, tWrW)  # gmem -> rmem
        copy_if(tBgB, tBrB)  # gmem -> rmem
        tNrN = norm(tXrX, tWrW, tBrB)
        # Compute: value = value * tanh(<scale>) + <shift>
        copy_if(tSCgSC, tSCrSC)  # gmem -> rmem
        copy_if(tSHgSH, tSHrSH)  # gmem -> rmem
        value = tNrN.load() * cute.tanh(tSCrSC.load()) + tSHrSH.load()
        # Store: y
        tYrY.store(value.to(tYrY.element_type))
        copy_if(tYrY, tYgY)  # rmem -> gmem
        if cutlass.const_expr(self.is_norm2):
            copy_if(tWgW2, tWrW2)  # gmem -> rmem
            copy_if(tBgB2, tBrB2)  # gmem -> rmem
            tNrN2 = norm(tYrY, tWrW2, tBrB2)
            # Compute: value2 = value2 * (1 + <scale2>)
            copy_if(tSCgSC2, tSCrSC2)  # gmem -> rmem
            value2 = tNrN2.load() * (1 + tSCrSC2.load())
            # Store: y2
            tYrY2.store(value2.to(tYrY2.element_type))
            copy_if(tYrY2, tYgY2)  # rmem -> gmem
 def validate_3d(t: torch.Tensor, B: int, S: int, D: int):
    if t.dtype not in (torch.float16, torch.bfloat16, torch.float32):
        raise ValueError(f"Validate failed: unsupported dtype: {t.dtype}")
    if (
        t.ndim != 3
        or (t.shape[0] not in (1, B))
        or (t.shape[1] not in (1, S) or t.shape[2] != D)
    ):
        raise ValueError(f"Validate failed: unsupported 3d-tensor: {t.shape}.")
    if t.stride()[-1] != 1:
        raise ValueError(f"Validate failed: not contiguous on dim D.")
 def validate_weight_bias(t: Optional[torch.Tensor], D: int):
    if t is None:
        return
    if t.dtype not in (torch.float16, torch.bfloat16, torch.float32):
        raise ValueError(f"Validate failed: unsupported dtype: {t.dtype}")
    if t.shape != (D,):
        raise ValueError(f"Validate failed: unsupported tensor shape: {t.shape}.")
    if t.stride()[-1] != 1:
        raise ValueError(f"Validate failed: not contiguous on dim D.")
@torch.library.custom_op("sglang::fused_norm_tanh_mul_add", mutates_args=())
 def fused_norm_tanh_mul_add(
    x: torch.Tensor,
    weight: Optional[torch.Tensor],
    bias: Optional[torch.Tensor],
    scale: torch.Tensor,
    shift: torch.Tensor,
    norm_type: str,
    eps: float = 1e-5,
 ) -> torch.Tensor:
    """
    Fuse: norm(x) * tanh(scale) + shift
      where norm is either layernorm or rmsnorm.
    Expects:
      - x: [B, S, D]
      - weight/bias: None, [D]
      - scale/shift: [1/B, 1/S, D]
      - norm_type: str, "layer" or "rms"
      - eps: Optional[float], default: 1e-5
    D must be a multiple of 256 and <= 8192 to enable LDG.128 vectorized loads per
    thread and avoid predicated loads (e.g., bounds checks such as `index < D`).
    """
    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
    # Tensor Validation
    BSD = x.shape
    validate_3d(x, *BSD)
    validate_weight_bias(weight, BSD[2])
    validate_weight_bias(bias, BSD[2])
    validate_3d(scale, *BSD)
    validate_3d(shift, *BSD)
    if norm_type == "layer" or norm_type == "rms":
        D = x.shape[-1]
        if D % 256 != 0 or D > 8192:
            raise ValueError(
                f"D={D} not supported, must be multiple of 256 and <= 8192"
            )
        y = torch.empty_like(x)  # create output tensor
        scale = broadcast_tensor_for_bsfd(scale, *x.shape)  # handle various shapes
        shift = broadcast_tensor_for_bsfd(shift, *x.shape)  # handle various shapes
        # y2, weight2, bias2, scale2 is None
        torch_tensors = [y, None, x, weight, bias, scale, shift, None, None, None]
        cute_tensor_args = [to_cute_arg(t) for t in torch_tensors]
        # Compile cache
        hash_key = NormTanhMulAddNormScale.make_hash_key(norm_type, *torch_tensors)
        compiled_fn = _COMPILE_CACHE.get(hash_key)
        if compiled_fn is None:
            kernel = NormTanhMulAddNormScale(D, norm_type, is_norm2=False)
            fake_sig_args = [to_fake_cute_args(t) for t in torch_tensors]
            compiled_fn = cute.compile(
                kernel, *fake_sig_args, options="--enable-tvm-ffi"
            )
            _COMPILE_CACHE[hash_key] = compiled_fn
        # Execute
        compiled_fn(*cute_tensor_args, eps, stream)
        return y
    else:
        raise ValueError(f'norm_type must be one of "layer" and "rms"')
@fused_norm_tanh_mul_add.register_fake
 def _fused_norm_tanh_mul_add_fake(x, weight, bias, scale, shift, norm_type, eps=1e-5):
    return x.new_empty(x.shape)
@torch.library.custom_op("sglang::fused_norm_tanh_mul_add_norm_scale", mutates_args=())
 def fused_norm_tanh_mul_add_norm_scale(
    x: torch.Tensor,
    weight: Optional[torch.Tensor],
    bias: Optional[torch.Tensor],
    scale: torch.Tensor,
    shift: torch.Tensor,
    weight2: Optional[torch.Tensor],
    bias2: Optional[torch.Tensor],
    scale2: torch.Tensor,
    norm_type: str,
    eps: float = 1e-5,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Fuse:
      y = norm(x) * tanh(scale) + shift
      y2 = norm(y) * (1 + scale2)
      where norm is either layernorm or rmsnorm.
    Expects:
      - x: [B, S, D]
      - weight/bia/weight2/bias2: None, [D]
      - scale/shift/scale2: [1/B, 1/S, D]
      - norm_type: str, "layer" or "rms"
      - eps: Optional[float], default: 1e-5
    D must be a multiple of 256 and <= 8192 to enable LDG.128 vectorized loads per
    thread and avoid predicated loads (e.g., bounds checks such as `index < D`).
    """
    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
    # Tensor Validation
    BSD = x.shape
    validate_3d(x, *BSD)
    validate_weight_bias(weight, BSD[2])
    validate_weight_bias(bias, BSD[2])
    validate_3d(scale, *BSD)
    validate_3d(shift, *BSD)
    validate_weight_bias(weight2, BSD[2])
    validate_weight_bias(bias2, BSD[2])
    validate_3d(scale2, *BSD)
    if norm_type == "layer" or norm_type == "rms":
        D = x.shape[-1]
        if D % 256 != 0 or D > 8192:
            raise ValueError(
                f"D={D} not supported, must be multiple of 256 and <= 8192"
            )
        y = torch.empty_like(x)  # create output tensor
        y2 = torch.empty_like(x)  # create output tensor
        scale = broadcast_tensor_for_bsfd(scale, *x.shape)  # handle various shapes
        shift = broadcast_tensor_for_bsfd(shift, *x.shape)  # handle various shapes
        scale2 = broadcast_tensor_for_bsfd(scale2, *x.shape)  # handle various shapes
        torch_tensors = [y, y2, x, weight, bias, scale, shift, weight2, bias2, scale2]
        cute_tensor_args = [to_cute_arg(t) for t in torch_tensors]
        # Compile cache
        hash_key = NormTanhMulAddNormScale.make_hash_key(norm_type, *torch_tensors)
        compiled_fn = _COMPILE_CACHE.get(hash_key)
        if compiled_fn is None:
            kernel = NormTanhMulAddNormScale(D, norm_type, is_norm2=True)
            fake_sig_args = [to_fake_cute_args(t) for t in torch_tensors]
            compiled_fn = cute.compile(
                kernel, *fake_sig_args, options="--enable-tvm-ffi"
            )
            _COMPILE_CACHE[hash_key] = compiled_fn
        # Execute
        compiled_fn(*cute_tensor_args, eps, stream)
        return y, y2
    else:
        raise ValueError(f'norm_type must be one of "layer" and "rms"')
@fused_norm_tanh_mul_add_norm_scale.register_fake
 def _fused_norm_tanh_mul_add_norm_scale_fake(
    x, weight, bias, scale, shift, weight2, bias2, scale2, norm_type, eps=1e-5
 ):
    return x.new_empty(x.shape), x.new_empty(x.shape)
@@ -621,6 +621,6 @@ Before submitting, verify:
 After the model produces non-noise output, read
 [references/testing-and-accuracy.md](references/testing-and-accuracy.md) before
 adding GPU cases, component-accuracy skips/hooks, suite entries, or benchmark
-claims. That reference tracks the current `gpu_cases.py` / `testcase_configs.py`
+claims. That reference tracks the current `gpu_cases.py`,
-/ `accuracy_testcase_configs.py` / `run_suite.py` split and the component-accuracy
+`DiffusionTestCase.run_component_accuracy_check`,
-decision rules.
+`single_test_file/component_accuracy/`, and `run_suite.py` split.
@@ -7,9 +7,10 @@ produce a non-noise image or video.
 - Add concrete GPU integration cases in `python/sglang/multimodal_gen/test/server/gpu_cases.py`.
 - Keep reusable dataclasses, constants, thresholds, and testcase factory helpers in `python/sglang/multimodal_gen/test/server/testcase_configs.py`.
- Add the case id to `python/sglang/multimodal_gen/test/server/accuracy_testcase_configs.py`
+- Set `DiffusionTestCase.run_component_accuracy_check=False` only when the case
-  only when it should be part of component-accuracy coverage. Adding a GPU case
+  should not enter component-accuracy coverage. Eligible cases default to
-  alone does not enroll it there.
+  `True`; `python/sglang/multimodal_gen/test/single_test_file/component_accuracy/testcase_configs.py`
  enrolls them automatically.
 - Let `python/sglang/multimodal_gen/test/run_suite.py` own suite selection, runtime-based partitioning, and standalone test files. Do not hard-code CI shard lists elsewhere.
 - If a new standalone test file is added to a suite, update `STANDALONE_FILE_EST_TIMES` after the first measured CI/runtime value is known.
@@ -23,38 +24,53 @@ PYTHONPATH=python python3 python/sglang/multimodal_gen/test/run_suite.py --suite
 ## Component Accuracy When Adding A GPU Case
-If you add a new entry to `ONE_GPU_CASES`, `TWO_GPU_CASES`, or a B200-specific
+If you add a new entry to `ONE_GPU_CASES` or `TWO_GPU_CASES`, treat component
-case group in `gpu_cases.py`, treat component accuracy as part of the
+accuracy as part of the model-adding workflow. Cases with
-model-adding workflow. Do not assume the new testcase will automatically fit or
+`run_component_accuracy_check=True` are selected automatically. The selector
-enter the existing component-accuracy harness.
+deduplicates each component by source model, component override, and GPU
 topology; later equivalent cases receive an automatic duplicate skip reason.
 B200-only groups are not currently inputs to the component-accuracy selector;
 add or identify a representative regular GPU case when that coverage is
 required.
 The component-accuracy harness compares SGLang components against Diffusers/HF
 reference components. This is stricter than pipeline-level inference. New GPU
 cases commonly fail here for one of three reasons:
-1. The model family needs explicit hook wiring in `python/sglang/multimodal_gen/test/server/accuracy_hooks.py`.
+1. The model family needs explicit hook wiring in `python/sglang/multimodal_gen/test/single_test_file/component_accuracy/hooks.py`.
   - Add hook logic only when the harness cannot call the raw component correctly without it.
   - Valid reasons include missing required forward arguments, required autocast/runtime context, or family-specific input preparation for the same component contract.
   - Do not change the compared output mode or add harness-side behavior that changes the component contract just to make the test pass.
 2. The component is already covered by another testcase with the same source component and topology.
   - Do not add redundant component-accuracy coverage.
-   - Add a skip entry in `python/sglang/multimodal_gen/test/server/accuracy_config.py` with a concrete reason such as `Representative VAE accuracy is already covered by ... for the same source component and topology`.
+   - Let `_select_accuracy_cases` in
     `python/sglang/multimodal_gen/test/single_test_file/component_accuracy/testcase_configs.py`
     deduplicate the component automatically. Do not add a manual skip for a
     duplicate that the selector can identify.
   - This is the preferred path for variant-only cases such as LoRA, Cache-DiT, upscaling, or other cases that reuse the same underlying component weights and topology.
 3. The HF/Diffusers reference component cannot be loaded or compared faithfully in the harness.
-   - Add a skip entry in `accuracy_config.py` with the exact technical failure.
+   - Add a skip entry in
     `python/sglang/multimodal_gen/test/single_test_file/component_accuracy/config.py`
     with the exact technical failure.
   - Good reasons include missing/unsupported HF component layout, incomplete checkpoints, unsupported raw component contract, or proven divergence after matched weight transfer and matching output shape.
   - Keep the skip reason concrete and technical. Do not write vague reasons like "component accuracy flaky" or "needs investigation."
 When adding a new GPU case, make this decision explicitly:
- if the case should have component-accuracy coverage, add its case id to
+- if the case should have component-accuracy coverage, leave
-  `accuracy_testcase_configs.py`
+  `run_component_accuracy_check=True`
- if the family needs minimal harness wiring, add the smallest possible change in `accuracy_hooks.py`
+- if the family needs minimal harness wiring, add the smallest possible change
- if the case is only a variant of an already covered source component and topology, add a skip in `accuracy_config.py`
+  in
- if the HF/Diffusers reference component cannot be compared faithfully, add a skip in `accuracy_config.py`
+  `python/sglang/multimodal_gen/test/single_test_file/component_accuracy/hooks.py`
- if the case is intentionally GPU-smoke-only, leave it out of `accuracy_testcase_configs.py` and keep that choice explicit in the PR notes
+- if the case is only a variant of an already covered source component and
  topology, rely on automatic per-component deduplication
 - if the HF/Diffusers reference component cannot be compared faithfully, add a
  concrete skip in
  `python/sglang/multimodal_gen/test/single_test_file/component_accuracy/config.py`
 - if the case is intentionally GPU-smoke-only, set
  `run_component_accuracy_check=False` and explain the choice in the PR notes
 Do not add a new GPU case and wait for CI to discover missing component-accuracy
 wiring.
@@ -54,7 +54,7 @@ Before calling a diffusion hotspot "new", first classify it with `existing-fast-
 Always rule out these existing families first:
 - HunyuanVideo VAE GroupNorm+SiLU
 - LTX upsampler GroupNorm+SiLU
- Z-Image residual-form modulation
+- Z-Image bf16-native Triton RMSNorm scale/tanh-residual modulation
 - SANA packed self-attention Q/K/V and cross-attention K/V GEMMs
 - fused diffusion `QK norm + RoPE`
 - LTX2 split RoPE
@@ -198,9 +198,11 @@ Use the preset categories this way:
 | `zimage` | `Tongyi-MAI/Z-Image-Turbo` | Yes: `zimage_turbo_t2i_1024` | Prompt, 1024x1024, seed 42, 2 GPUs, TP size 2; no explicit steps/guidance override |
 | `wan-t2v` | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | Yes: `wan22_t2v_a14b_720p` | 1280x720, 81 frames, 4 GPUs, CFG parallel, Ulysses degree 2, text encoder CPU offload and pinned CPU memory |
 | `wan-ti2v` | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | Yes: `wan22_ti2v_5b_720p` | Nightly cat image and motion prompt, 1280x720, 81 frames, seed 42 |
 | `ltx2` | `Lightricks/LTX-2` | Yes: `ltx2_twostage_t2v` | `LTX2TwoStagePipeline`, 2 GPUs, CFG parallel, 768x512, 121 frames, seed 42 |
 | `ltx23-ti2v-two-stage` | `Lightricks/LTX-2.3` | Yes: `ltx2.3_twostage_ti2v_2gpus` | Nightly cat image, motion prompt, `LTX2TwoStagePipeline`, 2 GPUs, `--cfg-parallel-size 2`, 768x512, 121 frames, seed 42 |
 | `ideogram4-fp8` | `ideogram-ai/ideogram-4-fp8` | Yes: `ideogram4_fp8_t2i_2gpu` | Prompt, 1024x1024, seed 42, 2 GPUs, TP size 2, FlashAttention backend; sampling preset owns steps/guidance |
 | `cosmos3-super-t2v` | `nvidia/Cosmos3-Super` | Yes: `cosmos3_super_t2v_2gpu` | Prompt, 1280x720, 81 frames, seed 42, 2 GPUs, TP size 2, guardrails disabled for benchmark isolation |
 | `wan-i2v` | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | Yes: `wan22_i2v_a14b_720p` | Nightly cat image and motion prompt, 1280x720, 81 frames, 4 GPUs, CFG parallel, Ulysses degree 2, text encoder CPU offload and pinned CPU memory |
 | `ltx2` | `Lightricks/LTX-2` | No | Current-source two-stage LTX-2 preset with 2 GPUs, CFG parallel, 768x512, 121 frames |
 | `qwen-image` | `Qwen/Qwen-Image` | No | Current-source extra covering the base Qwen-Image native path, separate from the nightly `Qwen-Image-2512` case |
 | `qwen-edit-2509` | `Qwen/Qwen-Image-Edit-2509` | No | Current-source extra for the pre-2511 edit-plus path; uses the cat image, 1024x1024 |
 | `zimage-base` | `Tongyi-MAI/Z-Image` | No | Current-source extra for non-turbo Z-Image; keep it separate from `zimage` / `Z-Image-Turbo` |
@@ -208,7 +210,6 @@ Use the preset categories this way:
 | `flux2-klein-base` | `black-forest-labs/FLUX.2-klein-base-4B` | No | Current-source extra for the undistilled FLUX.2 Klein Base path; gated repo, 1024x1024, DiT layerwise offload disabled |
 | `cosmos3-nano-t2i` | `nvidia/Cosmos3-Nano` | No | Current-source extra for the single-frame Cosmos3 image path; sets `SGLANG_DISABLE_COSMOS3_GUARDRAILS=1` in the helper environment |
 | `cosmos3-nano-t2v` | `nvidia/Cosmos3-Nano` | No | Current-source extra for a short Cosmos3 video path; sets `SGLANG_DISABLE_COSMOS3_GUARDRAILS=1` in the helper environment |
 | `ideogram4-fp8` | `ideogram-ai/ideogram-4-fp8` | No | Current-source extra matching the native Ideogram 4 FP8 pipeline; do not override steps/guidance directly because the sampling preset owns them |
 | `ernie-image-turbo` | `baidu/ERNIE-Image-Turbo` | No | Current-source extra for ERNIE-Image Turbo |
 | `glm-image` | `zai-org/GLM-Image` | No | Current-source extra for GLM-Image |
 | `sana-1.5-1.6b` | `Efficient-Large-Model/SANA1.5_1.6B_1024px_diffusers` | No | Current-source extra for a SANA native image path |
@@ -488,7 +489,7 @@ the known mainline families.
 | --- | --- |
 | `fused_inplace_qknorm_rope` missing, but separate qk norm plus rope show up | Check whether the fused diffusion `QK norm + RoPE` path should have engaged |
 | `to_q -> to_k -> to_v` on NVFP4 or Nunchaku FLUX-family checkpoints | Treat as a packed-QKV fast-path miss or checkpoint-format mismatch |
-| `fused_norm_tanh_mul_add*` missing on Z-Image | Treat as a missing mainline modulation path, not a new fusion request |
+| `zimage_rmsnorm_scale` or `zimage_rmsnorm_tanh_residual` missing on Z-Image | Check the bf16-native Triton eligibility guards before proposing a new fusion |
 | LTX-2 split RoPE appears as a long PyTorch elementwise chain | Check the `apply_ltx2_split_rotary_emb` Triton path and its shape guards |
 | masked attention spends time packing/unpacking Q/K/V | Check whether fused varlen USP pack/scatter should have engaged |
 | `all_to_all`, ring attention, or async A2A dominate | Classify against Ulysses, USP, or turbo-layer overlap first |
@@ -15,20 +15,22 @@ framework-specific optimization workflow.
 - `python/sglang/kernels/ops/diffusion/triton/group_norm_silu.py`
 - `python/sglang/kernels/ops/diffusion/triton/norm.py`
 - `python/sglang/kernels/ops/diffusion/triton/rmsnorm_onepass.py`
 - `python/sglang/kernels/ops/diffusion/triton/zimage_native_norm.py`
 - `python/sglang/kernels/ops/diffusion/triton/rotary.py`
 - `python/sglang/kernels/ops/diffusion/triton/ltx2_rotary.py`
 - `python/sglang/kernels/ops/diffusion/residual_gate_add.py`
 - `python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh`
 - `python/sglang/kernels/ops/diffusion/triton/varlen_pack_pad.py`
 - `python/sglang/kernels/ops/diffusion/cutedsl/scale_residual_norm_scale_shift.py`
- `test/registered/jit/diffusion/test_qwen_image_modulation.py`
+- `test/registered/kernels/ops/diffusion/test_qwen_image_modulation.py`
- `test/registered/jit/diffusion/test_group_norm_silu.py`
+- `test/registered/kernels/ops/diffusion/test_group_norm_silu.py`
- `test/registered/jit/diffusion/test_residual_gate_add.py`
+- `test/registered/kernels/ops/diffusion/test_residual_gate_add.py`
- `test/registered/jit/diffusion/test_varlen_pack_pad.py`
+- `test/registered/kernels/ops/diffusion/test_varlen_pack_pad.py`
- `test/registered/jit/diffusion/test_varlen_uspattn_equivalence.py`
+- `test/registered/kernels/ops/diffusion/test_varlen_uspattn_equivalence.py`
- `test/registered/jit/benchmark/diffusion/bench_qwen_image_modulation.py`
+- `test/registered/kernels/ops/diffusion/test_zimage_native_norm.py`
- `test/registered/jit/benchmark/diffusion/bench_group_norm_silu.py`
+- `test/registered/kernels/benchmark/diffusion/bench_qwen_image_modulation.py`
- `test/registered/jit/benchmark/diffusion/bench_residual_gate_add.py`
+- `test/registered/kernels/benchmark/diffusion/bench_group_norm_silu.py`
 - `test/registered/kernels/benchmark/diffusion/bench_residual_gate_add.py`
 - `python/sglang/kernels/ops/layernorm/norm.py`
 - `python/sglang/multimodal_gen/runtime/platforms/cuda.py`
 - `python/sglang/multimodal_gen/runtime/layers/attention/selector.py`
@@ -42,7 +44,7 @@ framework-specific optimization workflow.
 - Use cases: `x * (1 + scale) + shift`, `a * (k + b) + c`, and Qwen-style `(layernorm/residual layernorm) + scale/shift + gate select`.
 - Constraints: `x` must be CUDA and contiguous. `scale/shift` support 0D/1D/2D/3D/4D broadcast. 4D `[B, F, 1, C]` requires `L % F == 0`.
 - NPU fallback: `scale_shift.py` swaps to `npu_fallback` native path.
- Validation: `test/registered/jit/diffusion/test_qwen_image_modulation.py`.
+- Validation: `test/registered/kernels/ops/diffusion/test_qwen_image_modulation.py`.
 2. Norm + Scale/Shift fusion (CuTe DSL)
 - Kernels: `fused_norm_scale_shift`, `fused_scale_residual_norm_scale_shift`
@@ -53,22 +55,25 @@ framework-specific optimization workflow.
 - Constraints: `D % 256 == 0` and `D <= 8192`. `x/residual/gate/scale/shift` must pass shape and stride validation. Dtypes limited to fp16/bf16/fp32.
 - Behavior: CuTe DSL compilation cached by `(dtype, ndim, D, norm_type)`. `None` tensors replaced by scalar placeholders. If constraints fail, `layernorm.py` warns and falls back to native PyTorch.
-3. Z-Image fused tanh/gate modulation
+3. Z-Image bf16-native RMSNorm modulation (Triton)
- Kernels: `fused_norm_tanh_mul_add`, `fused_norm_tanh_mul_add_norm_scale`
+- Kernels: `zimage_rmsnorm_scale`, `zimage_rmsnorm_tanh_residual`
- Locations: `layernorm.py`, `cutedsl/norm_tanh_mul_add_norm_scale.py`, `zimage.py`
+- Locations: `triton/zimage_native_norm.py`, `zimage.py`
 - Use cases:
-  - `y = tanh(gate) * norm(x) + shift`
+  - `y = rmsnorm(x) * scale`
-  - `y, y2 = tanh(gate) * norm(x) + shift`, then `y2 = norm(y) * (1 + scale)`
+  - `y = residual + tanh(gate) * rmsnorm(x)`
- Constraints: same CuTe DSL envelope as the norm+scale/shift family in practice: contiguous last dim, fp16/bf16/fp32, and `D % 256 == 0`, `D <= 8192`.
+- Constraints: CUDA bf16 tensors, contiguous weights, flattenable row strides,
- Validation: `test/registered/jit/diffusion/test_norm_tanh_mul_add_norm_scale.py`
+  compatible modulation row counts, and `D <= 8192`.
- Behavior: this is already a mainline fast path, so if Z-Image traces show the unfused chain, treat it as a missing or regressed existing optimization before proposing a new kernel.
+- Validation: `test/registered/kernels/ops/diffusion/test_zimage_native_norm.py`
 - Behavior: the kernels preserve Z-Image's native bf16 arithmetic. They return
  `None` when an eligibility guard fails, and the runtime wrapper executes the
  native PyTorch formula.
 4. Triton LayerNorm/RMSNorm fusion
 - Kernels: `rms_norm_fn`, `layer_norm_fn`, `norm_infer`
 - Locations: `triton/norm.py`, `layernorm.py`
 - Use cases: fp32 RMSNorm with residual/dropout/rowscale/x1 branches, and inference-friendly `norm_infer`.
 - Constraints: last dim must be contiguous, and `N * element_size < 64KB`.
- Validation: `test/registered/jit/test_rmsnorm.py`.
+- Validation: `test/registered/kernels/ops/layernorm/test_rmsnorm.py`.
 5. Triton one-pass RMSNorm (small hidden size fast path)
 - Kernel: `triton_one_pass_rms_norm`
@@ -82,7 +87,7 @@ framework-specific optimization workflow.
 - Use case: GPT-J style RoPE when not Neox.
 - Constraints: `head_size` must be even.
 - NPU fallback: `npu_fallback.apply_rotary_embedding_native`.
- Validation: `test/registered/jit/test_rope.py`.
+- Validation: `test/registered/kernels/ops/attention/test_rope.py`.
 7. LTX2 split RoPE fusion
 - Kernel: `apply_ltx2_split_rotary_emb`
@@ -97,8 +102,8 @@ framework-specific optimization workflow.
 - Use case: `residual + update * gate` in LTX2 self-attention, prompt cross-attention, audio/video cross-attention, and feed-forward residual updates.
 - Constraints: `residual`, `update`, and `gate` must be CUDA tensors on the same device, contiguous, same dtype (`fp16`, `bf16`, or `fp32`), with `update.shape == residual.shape`; `gate` can match `residual` or be row-broadcast with the last dimension matching.
 - Behavior: `_ltx2_residual_gate_add(...)` uses the CUDA custom op while guards pass. On a runtime exception outside `torch.compile`, it logs once, disables the fast path for the process, and falls back to `residual + update * gate`.
- Validation: `test/registered/jit/diffusion/test_residual_gate_add.py`.
+- Validation: `test/registered/kernels/ops/diffusion/test_residual_gate_add.py`.
- Microbench: `test/registered/jit/benchmark/diffusion/bench_residual_gate_add.py`.
+- Microbench: `test/registered/kernels/benchmark/diffusion/bench_residual_gate_add.py`.
 - Workflow rule: if LTX2 traces show repeated elementwise `mul` + `add` ladders around attention or MLP residuals, check whether this existing CUDA path was disabled by shape, dtype, contiguity, or a prior runtime failure before proposing another elementwise fusion.
 9. HunyuanVideo / LTX upsampler GroupNorm + SiLU fusion
@@ -107,8 +112,8 @@ framework-specific optimization workflow.
 - Use case: `activation(group_norm(x))` when the activation is non-inplace `nn.SiLU` and the GroupNorm is affine.
 - Enablement: mainline uses `apply_group_norm_silu(...)` in HunyuanVideo VAE paths and LTX latent upsampler paths by default; there is no env toggle. The wrapper dispatches to Triton only when guards pass.
 - Constraints: CUDA inference path only; no grad, `x.requires_grad == False`, `nn.GroupNorm`, `nn.SiLU(inplace=False)`, affine norm with weight and bias. Unsupported cases fall back to native `activation(norm(x))`.
- Validation: `test/registered/jit/diffusion/test_group_norm_silu.py`.
+- Validation: `test/registered/kernels/ops/diffusion/test_group_norm_silu.py`.
- Microbench: `test/registered/jit/benchmark/diffusion/bench_group_norm_silu.py`.
+- Microbench: `test/registered/kernels/benchmark/diffusion/bench_group_norm_silu.py`.
 **Faster CUDA Kernel Usage Points**
@@ -116,7 +121,8 @@ framework-specific optimization workflow.
 - Location: `layernorm.py`
 - Behavior:
 - Standard `bf16`/`fp16` CUDA paths use `sgl_kernel.fused_add_rmsnorm` and `sgl_kernel.rmsnorm`.
- The Z-Image `fp32` `32x2560` path under `torch.compile` avoids `wrap_triton` and uses the native fp32 path.
+- Z-Image keeps bf16 arithmetic and uses its dedicated Triton native-norm
  kernels when their guards pass.
 - `hidden_size <= 128` uses Triton one-pass.
 - ROCm falls back to native.
@@ -131,7 +137,7 @@ framework-specific optimization workflow.
 4. Varlen USP attention pack/scatter
 - Locations: `runtime/layers/attention/layer.py`, `triton/varlen_pack_pad.py`
 - Behavior: masked `USPAttention.forward` can gather dense Q/K/V into packed `[total_valid, H, D]` rows with `fused_pack_qkv`, run varlen attention, then scatter back with `fused_scatter_to_padded`.
- Validation: `test/registered/jit/diffusion/test_varlen_pack_pad.py` and `test_varlen_uspattn_equivalence.py`.
+- Validation: `test/registered/kernels/ops/diffusion/test_varlen_pack_pad.py` and `test/registered/kernels/ops/diffusion/test_varlen_uspattn_equivalence.py`.
 - Workflow rule: if a masked attention trace spends time in Python/advanced indexing pack or scatter, first check whether this fused varlen path should have engaged.
 **QK Norm Optimization**
@@ -144,7 +150,7 @@ framework-specific optimization workflow.
  - `can_use_fused_inplace_qknorm(head_dim, dtype)` returns true.
  - Supported head dims: `64, 128, 256, 512, 1024`.
 - Behavior: Fused path operates on `q` and `k` in place after reshaping to `[B, -1, head_dim]`. If preconditions fail, fall back to per-tensor RMSNorm.
- Validation: `test/registered/jit/test_qknorm.py` and `test/registered/jit/test_qknorm_across_heads.py`.
+- Validation: `test/registered/kernels/ops/layernorm/test_qknorm.py` and `test/registered/kernels/ops/layernorm/test_qknorm_across_heads.py`.
 **QK Norm + RoPE Optimization**
@@ -159,7 +165,7 @@ framework-specific optimization workflow.
  - `can_use_fused_inplace_qknorm_rope(head_dim, rope_dim, is_neox, dtype)` returns true.
  - Supported head dims: `64, 128, 256`.
 - Behavior: `apply_qk_norm_rope` prefers the fused JIT kernel when all guards pass; otherwise it falls back to `apply_qk_norm(...)` plus `apply_flashinfer_rope_qk_inplace(...)`.
- Validation: `test/registered/jit/diffusion/test_qknorm_rope.py`.
+- Validation: `test/registered/kernels/ops/diffusion/test_qknorm_rope.py`.
 - Workflow rule: treat LTX2 traces that miss the generic fused path as an enablement/shape-guard issue first, and check the separate LTX2 split-RoPE path before proposing new attention-prep kernels.
 **Nunchaku Fused GELU MLP**
@@ -187,7 +193,9 @@ framework-specific optimization workflow.
 **Common Entry Points in Diffusion Models**
 - AdaLN modulation: `LayerNormScaleShift`, `RMSNormScaleShift`, `ScaleResidual*` in `layernorm.py`.
 - Qwen-Image gating: `fuse_layernorm_scale_shift_gate_select01_kernel` and `fuse_residual_layernorm_scale_shift_gate_select01_kernel` through `fused_scale_shift_gate.py` and `qwen_image.py`.
- Z-Image residual-form modulation: `fused_norm_tanh_mul_add` and `fused_norm_tanh_mul_add_norm_scale` in `zimage.py`.
+- Z-Image native norm modulation: `zimage_rmsnorm_scale` and
  `zimage_rmsnorm_tanh_mul_add` in `zimage.py`, backed by
  `triton/zimage_native_norm.py`.
 - HunyuanVideo VAE and LTX upsampler GroupNorm+SiLU: `apply_group_norm_silu` in `hunyuanvae.py` and `latent_upsampler.py`; default-eligible when wrapper guards pass.
 - QK norm: `apply_qk_norm` used in `flux.py`, `flux_2.py`, `qwen_image.py`, `zimage.py`, `wanvideo.py`, `ltx_2.py`, `hunyuanvideo.py`.
 - QK norm + RoPE: `apply_qk_norm_rope` in `layernorm.py`; use this path when the model wants fused attention prep instead of separate QK norm and RoPE calls.
@@ -77,8 +77,9 @@ NIGHTLY_PRESET_ORDER = (
    "zimage",
    "wan-t2v",
    "wan-ti2v",
    "ltx2",
    "ltx23-ti2v-two-stage",
    "ideogram4-fp8",
    "cosmos3-super-t2v",
    "wan-i2v",
 )
@@ -184,21 +185,7 @@ MODELS = {
            "--num-frames=81",
        ],
    },
-    # 8. Nightly: ltx2_twostage_t2v
+    # 8. Nightly: ltx2.3_twostage_ti2v_2gpus
    "ltx2": {
        "nightly_case_id": "ltx2_twostage_t2v",
        "path": "Lightricks/LTX-2",
        "prompt": "A cat and a dog baking a cake together in a kitchen.",
        "extra_args": [
            "--pipeline-class-name=LTX2TwoStagePipeline",
            "--width=768",
            "--height=512",
            "--num-frames=121",
            "--num-gpus=2",
            "--enable-cfg-parallel",
        ],
    },
    # 9. Nightly: ltx2.3_twostage_ti2v_2gpus
    # Requires: <repo>/inputs/diffusion_benchmark/figs/cat.png
    "ltx23-ti2v-two-stage": {
        "nightly_case_id": "ltx2.3_twostage_ti2v_2gpus",
@@ -214,7 +201,36 @@ MODELS = {
            "--cfg-parallel-size=2",
        ],
    },
-    # 10. Nightly: wan22_i2v_a14b_720p
+    # 9. Nightly: ideogram4_fp8_t2i_2gpu
    "ideogram4-fp8": {
        "nightly_case_id": "ideogram4_fp8_t2i_2gpu",
        "path": "ideogram-ai/ideogram-4-fp8",
        "prompt": "A futuristic cyberpunk city at night, neon lights reflecting on wet streets",
        "extra_args": [
            "--width=1024",
            "--height=1024",
            "--num-gpus=2",
            "--tp-size=2",
            "--attention-backend=fa",
        ],
    },
    # 10. Nightly: cosmos3_super_t2v_2gpu
    "cosmos3-super-t2v": {
        "nightly_case_id": "cosmos3_super_t2v_2gpu",
        "path": "nvidia/Cosmos3-Super",
        "prompt": "A cat and a dog baking a cake together in a kitchen.",
        "env": {
            "SGLANG_DISABLE_COSMOS3_GUARDRAILS": "1",
        },
        "extra_args": [
            "--width=1280",
            "--height=720",
            "--num-frames=81",
            "--num-gpus=2",
            "--tp-size=2",
        ],
    },
    # 11. Nightly: wan22_i2v_a14b_720p
    # Requires: <repo>/inputs/diffusion_benchmark/figs/cat.png
    "wan-i2v": {
        "nightly_case_id": "wan22_i2v_a14b_720p",
@@ -233,6 +249,18 @@ MODELS = {
        ],
    },
    # Source-tracked extras from current registry / GPU test coverage.
    "ltx2": {
        "path": "Lightricks/LTX-2",
        "prompt": "A cat and a dog baking a cake together in a kitchen.",
        "extra_args": [
            "--pipeline-class-name=LTX2TwoStagePipeline",
            "--width=768",
            "--height=512",
            "--num-frames=121",
            "--num-gpus=2",
            "--enable-cfg-parallel",
        ],
    },
    "qwen-image": {
        "path": "Qwen/Qwen-Image",
        "prompt": "A futuristic cyberpunk city at night, neon lights reflecting on wet streets",
@@ -305,14 +333,6 @@ MODELS = {
            "--num-inference-steps=4",
        ],
    },
    "ideogram4-fp8": {
        "path": "ideogram-ai/ideogram-4-fp8",
        "prompt": "A clean product poster for a new open-source inference engine",
        "extra_args": [
            "--width=1024",
            "--height=1024",
        ],
    },
    "ernie-image-turbo": {
        "path": "baidu/ERNIE-Image-Turbo",
        "prompt": "A futuristic cyberpunk city at night, neon lights reflecting on wet streets",
@@ -657,6 +677,8 @@ def validate_nightly_alignment() -> int:
            errors.append(f"{model_key}: reference image presence differs")
        if preset.get("seed", 42) != case.get("seed"):
            errors.append(f"{model_key}: seed differs")
        if preset.get("env", {}) != case["frameworks"]["sglang"].get("extra_env", {}):
            errors.append(f"{model_key}: environment differs")
        actual_args = {
            key: _normalize_cli_value(value)
@@ -403,5 +403,5 @@ When documenting results:
 | `tools/build_modelopt_nvfp4_transformer.py` | Build mixed BF16+NVFP4 transformer directories when a family needs preserved BF16 layers |
 | `tools/compare_diffusion_trajectory_similarity.py` | reduced deterministic BF16-vs-quantized validation |
 | `docs_new/docs/sglang-diffusion/quantization.mdx` | public ModelOpt support matrix and CLI examples |
-| `test/server/testcase_configs.py` | reusable ModelOpt testcase constants, thresholds, and helpers |
+| `python/sglang/multimodal_gen/test/server/testcase_configs.py` | reusable ModelOpt testcase constants, thresholds, and helpers |
-| `test/server/gpu_cases.py` | concrete GPU and B200 ModelOpt CI case lists |
+| `python/sglang/multimodal_gen/test/server/gpu_cases.py` | concrete GPU and B200 ModelOpt CI case lists |
@@ -82,7 +82,7 @@ For Wan2.2 specifically:
 - for pure latency tuning, benchmark pure Ulysses too, for example `--ulysses-degree=4 --ring-degree=1` on 4 GPUs
 - on 8 GPUs, compare pure `--ulysses-degree=8` against `--enable-cfg-parallel --ulysses-degree=4`
-### Nightly-aligned model, 2 GPUs: LTX-2 two-stage
+### Current-source model, 2 GPUs: LTX-2 two-stage
 ```bash
 sglang generate --model-path Lightricks/LTX-2 \
@@ -94,7 +94,11 @@ sglang generate --model-path Lightricks/LTX-2 \
  --enable-torch-compile --warmup --save-output
 ```
-Note: this generate recipe is aligned with the nightly comparison case `ltx2_twostage_t2v`. The nightly config omits explicit steps and guidance, so this command omits them too and uses runtime defaults. `LTX2TwoStagePipeline` is a native path and auto-resolves the spatial upsampler plus distilled LoRA from the same model snapshot unless you override them.
+Note: LTX-2 is a current-source benchmark preset rather than a nightly
 comparison case. The command uses runtime-default steps and guidance.
 `LTX2TwoStagePipeline` is a native path and auto-resolves the spatial
 upsampler plus distilled LoRA from the same model snapshot unless you override
 them.
 ### Nightly-aligned model, 2 GPUs: LTX-2.3 TI2V two-stage
@@ -254,7 +258,7 @@ Use these as first commands to benchmark, not as universal winners.
 | FLUX.1 / FLUX.2 image | 1024x1024, runtime-default steps/guidance, 1 GPU | `--enable-torch-compile --warmup --dit-layerwise-offload false` | `black-forest-labs/FLUX.*` repos are gated; for FP8/NVFP4 use validated `--transformer-path` or `--transformer-weights-path` flows from the quant skill. |
 | FLUX.2 Klein / Klein Base | 1024x1024, runtime-default steps/guidance, 1 GPU | `--enable-torch-compile --warmup --dit-layerwise-offload false` | Current registry has `black-forest-labs/FLUX.2-klein-4B`, `FLUX.2-klein-9B`, and base variants. Klein is step-distilled; Klein Base is not. |
 | Qwen-Image / Qwen-Image-Edit | 1024x1024, runtime-default steps/guidance, 1 GPU | `--enable-torch-compile --warmup`; optionally native `SGLANG_CACHE_DIT_ENABLED=true` | Cache-DiT is lossy. For edit tasks, keep reference image, seed, and output size fixed. |
-| Z-Image / Z-Image-Turbo | 1024x1024, runtime-default steps/guidance, 1 GPU | `--enable-torch-compile --warmup` | Keep base Z-Image separate from Turbo: base uses 50-step CFG defaults, Turbo uses 9-step zero-CFG defaults. Mainline has Z-Image tanh/gate norm fusions. |
+| Z-Image / Z-Image-Turbo | 1024x1024, runtime-default steps/guidance, 1 GPU | `--enable-torch-compile --warmup` | Keep base Z-Image separate from Turbo: base uses 50-step CFG defaults, Turbo uses 9-step zero-CFG defaults. Mainline has bf16-native Triton RMSNorm scale and tanh-residual fusions. |
 | Wan2.2 A14B T2V/I2V | 1280x720, 81 frames | Nightly: `--num-gpus 4 --enable-cfg-parallel --ulysses-degree 2 --text-encoder-cpu-offload --pin-cpu-memory` | For lowest latency, also benchmark pure Ulysses on the same GPUs. |
 | Wan2.2 TI2V 5B | 1280x720, 81 frames, 1 GPU | `--enable-torch-compile --warmup` | Keep the input image and motion prompt fixed when comparing sparse attention or Cache-DiT. |
 | Wan2.1 / FastWan / TurboWan variants | 480p or 720p video, family defaults | `--enable-torch-compile --warmup`; add `--ulysses-degree` / CFG parallel only after measuring | Current registry includes Wan2.1, FastWan2.1, FastWan2.2 TI2V, TurboWan2.1, TurboWan2.2 I2V, and Wan2.1-Fun InP. Use the compatibility matrix and benchmark presets before choosing topology. |
@@ -285,6 +289,6 @@ about whether the work has merged:
 - **Offload tuning**: after the first request, the runtime logs peak GPU memory and which components could stay resident. Use this to decide which `--*-cpu-offload` flags to disable.
 - **Backend selection**: `--backend sglang` (default, auto-detected) enables native optimizations (fused kernels, SP, native Cache-DiT env knobs, etc.). `--backend diffusers` falls back to Diffusers pipelines and is the path that accepts `--cache-dit-config` plus diffusers attention backend names.
 - **Wan2.2-I2V sizing**: explicit `--width/--height` on `Wan2.2-I2V-A14B` control the target area while preserving the condition-image aspect ratio.
- **Mainline diffusion fast paths**: before proposing a new kernel or overlap scheme, check `sglang-diffusion-benchmark-profile/existing-fast-paths.md`. It covers GroupNorm+SiLU, Z-Image residual-form modulation, fused diffusion `QK norm + RoPE`, LTX2 split RoPE, LTX2 residual-gate add, varlen USP pack/scatter, packed QKV/NVFP4 expectations, and existing multi-GPU overlap families such as Ulysses / USP and turbo-layer async all-to-all.
+- **Mainline diffusion fast paths**: before proposing a new kernel or overlap scheme, check `sglang-diffusion-benchmark-profile/existing-fast-paths.md`. It covers GroupNorm+SiLU, Z-Image bf16-native Triton norm modulation, fused diffusion `QK norm + RoPE`, LTX2 split RoPE, LTX2 residual-gate add, varlen USP pack/scatter, packed QKV/NVFP4 expectations, and existing multi-GPU overlap families such as Ulysses / USP and turbo-layer async all-to-all.
 - **NVFP4 trace interpretation**: on FLUX.2 NVFP4 and Nunchaku-style checkpoints, packed QKV is expected. SGLang intentionally uses fused projection modules such as `to_qkv` / `to_added_qkv` instead of separate `to_q` / `to_k` / `to_v`, so a split-QKV trace usually means the quantized path did not engage rather than a brand new fusion opportunity.
 - **Hotspot workflow split**: use `sglang-diffusion-benchmark-profile` to prove and classify a slowdown with perf dumps plus `torch.profiler`; hand concrete kernel work off with the perf/profile evidence attached instead of expanding the benchmark skill.
@@ -716,8 +716,3 @@ class LTX2PipelineConfig(PipelineConfig):
@dataclasses.dataclass
 class LTX23PipelineConfig(LTX2PipelineConfig):
    """Configuration overrides for LTX-2.3."""
@dataclasses.dataclass
 class LTX2I2VPipelineConfig(LTX2PipelineConfig):
    task_type: ModelTaskType = ModelTaskType.TI2V
@@ -57,10 +57,6 @@ def zimage_postprocess_text(
    return pad_text_embeddings_with_mask(split_hidden_states)
 class TransformersModelConfig(EncoderConfig):
    tokenizer_kwargs: dict = field(default_factory=lambda: {})
@dataclass
 class ZImagePipelineConfig(ZImageRolloutPipelineMixin, ImagePipelineConfig):
    should_use_guidance: bool = False
@@ -19,11 +19,3 @@ class Krea2SamplingParams(SamplingParams):
    width: int = 1024
    guidance_scale: float = 1.0
    num_inference_steps: int = 8
@dataclass
 class Krea2RawSamplingParams(Krea2SamplingParams):
    """Base `oss_raw` defaults: full sampler with CFG."""
    guidance_scale: float = 4.5
    num_inference_steps: int = 52
@@ -12,17 +12,14 @@ from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var
 logger = logging.getLogger(__name__)
 if TYPE_CHECKING:
    SGLANG_DIFFUSION_RINGBUFFER_WARNING_INTERVAL: int = 60
    SGLANG_DIFFUSION_NCCL_SO_PATH: str | None = None
    LD_LIBRARY_PATH: str | None = None
    LOCAL_RANK: int = 0
    CUDA_VISIBLE_DEVICES: str | None = None
    SGLANG_DIFFUSION_CACHE_ROOT: str = os.path.expanduser("~/.cache/sgl_diffusion")
    SGLANG_DIFFUSION_CONFIG_ROOT: str = os.path.expanduser("~/.config/sgl_diffusion")
    SGLANG_DIFFUSION_CONFIGURE_LOGGING: int = 1
    SGLANG_DIFFUSION_LOGGING_LEVEL: str = "INFO"
    SGLANG_DIFFUSION_LOGGING_PREFIX: str = ""
    SGLANG_DIFFUSION_LOGGING_CONFIG_PATH: str | None = None
    SGLANG_DIFFUSION_TRACE_FUNCTION: int = 0
    SGLANG_DIFFUSION_WORKER_MULTIPROC_METHOD: str = "fork"
    SGLANG_DIFFUSION_TARGET_DEVICE: str = "cuda"
@@ -62,7 +59,6 @@ if TYPE_CHECKING:
    SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND: str | None = None
    SGLANG_DIFFUSION_ENABLE_W8A8_FP8_GEMM: bool = False
    SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D: str = "auto"
    SGLANG_USE_CUDA_HUNYUANVIDEO_GROUP_NORM_SILU: bool = False
    SGLANG_USE_ROCM_VAE: bool = False
    SGLANG_USE_ROCM_CUDNN_BENCHMARK: bool = False
    SGLANG_USE_ROCM_VAE_CONV2D: bool = False
@@ -83,10 +79,6 @@ def get_default_config_root() -> str:
    )
 def maybe_convert_int(value: str | None) -> int | None:
    return int(value) if value is not None else None
 # helpers for environment variable definitions
 def _lazy_str(key: str, default: str | None = None) -> Callable[[], str | None]:
    return lambda: os.getenv(key, default)
@@ -110,20 +102,6 @@ def _lazy_bool(key: str, default: str = "false") -> Callable[[], bool]:
    return lambda: get_bool_env_var(key, default)
 def _lazy_bool_any(keys: list[str], default: str = "false") -> Callable[[], bool]:
    def _getter():
        for key in keys:
            if get_bool_env_var(key, "false"):
                return True
        return (
            get_bool_env_var("", default)
            if not keys
            else get_bool_env_var(keys[0], default)
        )
    return _getter
 def _lazy_path(
    key: str, default_func: Callable[[], str] | None = None
 ) -> Callable[[], str | None]:
@@ -157,13 +135,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
    # By default this is 1.
    # If set, `MAX_JOBS` will be reduced to avoid oversubscribing the CPU.
    "NVCC_THREADS": _lazy_str("NVCC_THREADS"),
    # If set, sgl_diffusion will use precompiled binaries (*.so)
    "SGLANG_DIFFUSION_USE_PRECOMPILED": _lazy_bool_any(
        [
            "SGLANG_DIFFUSION_USE_PRECOMPILED",
            "SGLANG_DIFFUSION_PRECOMPILED_WHEEL_LOCATION",
        ]
    ),
    # CMake build type
    # If not set, defaults to "Debug" or "RelWithDebInfo"
    # Available options: "Debug", "Release", "RelWithDebInfo"
@@ -186,39 +157,17 @@ environment_variables: dict[str, Callable[[], Any]] = {
        "SGLANG_DIFFUSION_CACHE_ROOT",
        lambda: os.path.join(get_default_cache_root(), "sgl_diffusion"),
    ),
    # Interval in seconds to log a warning message when the ring buffer is full
    "SGLANG_DIFFUSION_RINGBUFFER_WARNING_INTERVAL": _lazy_int(
        "SGLANG_DIFFUSION_RINGBUFFER_WARNING_INTERVAL", 60
    ),
    # Path to the NCCL library file. It is needed because nccl>=2.19 brought
    # by PyTorch contains a bug: https://github.com/NVIDIA/nccl/issues/1234
    "SGLANG_DIFFUSION_NCCL_SO_PATH": _lazy_str("SGLANG_DIFFUSION_NCCL_SO_PATH"),
    # when `SGLANG_DIFFUSION_NCCL_SO_PATH` is not set, sgl_diffusion will try to find the nccl
    # library file in the locations specified by `LD_LIBRARY_PATH`
    "LD_LIBRARY_PATH": _lazy_str("LD_LIBRARY_PATH"),
    # Internal flag to enable Dynamo fullgraph capture
    "SGLANG_DIFFUSION_TEST_DYNAMO_FULLGRAPH_CAPTURE": _lazy_bool(
        "SGLANG_DIFFUSION_TEST_DYNAMO_FULLGRAPH_CAPTURE", "1"
    ),
    # local rank of the process in the distributed setting, used to determine
    # the GPU device id
    "LOCAL_RANK": _lazy_int("LOCAL_RANK", 0),
    # used to control the visible devices in the distributed setting
    "CUDA_VISIBLE_DEVICES": _lazy_str("CUDA_VISIBLE_DEVICES"),
    # timeout for each iteration in the engine
    "SGLANG_DIFFUSION_ENGINE_ITERATION_TIMEOUT_S": _lazy_int(
        "SGLANG_DIFFUSION_ENGINE_ITERATION_TIMEOUT_S", 60
    ),
    # Logging configuration
    # If set to 0, sgl_diffusion will not configure logging
    # If set to 1, sgl_diffusion will configure logging using the default configuration
    #    or the configuration file specified by SGLANG_DIFFUSION_LOGGING_CONFIG_PATH
    "SGLANG_DIFFUSION_CONFIGURE_LOGGING": _lazy_int(
        "SGLANG_DIFFUSION_CONFIGURE_LOGGING", 1
    ),
    "SGLANG_DIFFUSION_LOGGING_CONFIG_PATH": _lazy_str(
        "SGLANG_DIFFUSION_LOGGING_CONFIG_PATH"
    ),
    # this is used for configuring the default logging level
    "SGLANG_DIFFUSION_LOGGING_LEVEL": _lazy_str(
        "SGLANG_DIFFUSION_LOGGING_LEVEL", "INFO"
@@ -111,7 +111,6 @@ class GeluAndMul(CustomOp):
@CustomOp.register("gelu_new")
 class NewGELU(CustomOp):
    def __init__(self):
        super().__init__()
@@ -161,18 +160,3 @@ def get_act_fn(act_fn_name: str) -> nn.Module:
        raise ValueError(f"Activation function {act_fn_name!r} is not supported.")
    return _ACTIVATION_REGISTRY[act_fn_name]()
 _ACTIVATION_AND_MUL_REGISTRY = {
    "gelu": GeluAndMul,
    "silu": SiluAndMul,
 }
 def get_act_and_mul_fn(act_fn_name: str) -> nn.Module:
    """Get an activation-and-mul (i.e. SiluAndMul) function by name."""
    act_fn_name = act_fn_name.lower()
    if act_fn_name not in _ACTIVATION_AND_MUL_REGISTRY:
        raise ValueError(f"Activation function {act_fn_name!r} is not supported.")
    return _ACTIVATION_AND_MUL_REGISTRY[act_fn_name]()
@@ -541,7 +541,7 @@ class _ScaleResidualNormScaleShift(CustomOp):
        if residual.numel() == 0 or x.numel() == 0:
            return self.forward_native(residual, x, gate, shift, scale)
-        if x.shape[-1] % 256 != 0 and x.shape[-1] <= 8192:
+        if x.shape[-1] % 256 != 0 or x.shape[-1] > 8192:
            import warnings
            warnings.warn(
@@ -719,7 +719,7 @@ class _NormScaleShift(CustomOp):
    def forward_cuda(
        self, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor
    ) -> torch.Tensor:
-        if x.shape[-1] % 256 != 0 and x.shape[-1] <= 8192:
+        if x.shape[-1] % 256 != 0 or x.shape[-1] > 8192:
            import warnings
            warnings.warn(
@@ -810,81 +810,6 @@ class RMSNormScaleShift(_NormScaleShift):
    norm_type = "rms"
 ################################################################################
 # NormTanhMulAdd
 # y = norm(x) * tanh(scale) + shift (where norm is layernorm or rmsnorm)
 # See details in norm_tanh_mul_add_norm_scale.py
 ################################################################################
 class _NormTanhMulAdd(CustomOp):
    norm_type: str
    def __init__(
        self,
        hidden_size: int,
        eps: float = 1e-6,
        affine: bool = False,
        dtype: torch.dtype = torch.float32,
    ):
        super().__init__()
        self.eps = eps
        if self.norm_type == "rms":
            self.norm = RMSNorm(hidden_size, eps=eps, dtype=dtype)
        elif self.norm_type == "layer":
            self.norm = FP32LayerNorm(
                hidden_size, elementwise_affine=affine, eps=eps, dtype=dtype
            )
        else:
            raise NotImplementedError(f"Norm type {self.norm_type} not implemented")
    def forward_cuda(
        self, x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor
    ) -> torch.Tensor:
        if x.shape[-1] % 256 != 0 and x.shape[-1] <= 8192:
            import warnings
            warnings.warn(
                "FusedNormScaleShift cuda not available, using native fallback",
                stacklevel=2,
            )
            return self.forward_native(x, scale, shift)
        from sglang.kernels.ops.diffusion.cutedsl.norm_tanh_mul_add_norm_scale import (
            fused_norm_tanh_mul_add,
        )
        x, scale, shift = x.contiguous(), scale.contiguous(), shift.contiguous()
        weight = _ensure_contiguous(getattr(self.norm, "weight", None))
        bias = _ensure_contiguous(getattr(self.norm, "bias", None))
        return fused_norm_tanh_mul_add(
            x,
            weight,
            bias,
            scale,
            shift,
            self.norm_type,
            self.eps,
        )
    def forward_hip(self, *args, **kwargs):
        # Fallback to native because ROCm does not support CuTeDSL.
        return self.forward_native(*args, **kwargs)
    @torch.compile(disable=current_platform.is_npu() or current_platform.is_rocm())
    def forward_native(
        self, x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor
    ) -> torch.Tensor:
        y = self.norm(x) * torch.tanh(scale) + shift
        return y.to(x.dtype)
 class LayerNormTanhMulAdd(_NormTanhMulAdd):
    norm_type = "layer"
 class RMSNormTanhMulAdd(_NormTanhMulAdd):
    norm_type = "rms"
 def apply_qk_norm(
    q: torch.Tensor,
    k: torch.Tensor,
@@ -1078,34 +1003,6 @@ def apply_qk_norm_rope(
    )
 def apply_rmsnorm_tanh_mul_add(
    x: torch.Tensor,
    gate: torch.Tensor,
    residual: torch.Tensor,
    norm: "RMSNorm",
 ) -> torch.Tensor:
    """Compute residual + tanh(gate) * rmsnorm(x), with a fused CUDA fast path."""
    if get_bool_env_var("SGLANG_ENABLE_DETERMINISTIC_INFERENCE"):
        return residual + torch.tanh(gate) * norm(x)
    if _is_cuda and x.is_cuda and x.shape[-1] % 256 == 0 and x.shape[-1] <= 8192:
        from sglang.kernels.ops.diffusion.cutedsl.norm_tanh_mul_add_norm_scale import (
            fused_norm_tanh_mul_add,
        )
        return fused_norm_tanh_mul_add(
            x.contiguous(),
            norm.weight.data.contiguous(),
            None,
            gate.contiguous(),
            residual.contiguous(),
            "rms",
            norm.variance_epsilon,
        )
    return residual + torch.tanh(gate) * norm(x)
 def tensor_parallel_rms_norm(x: torch.Tensor, norm: "RMSNorm") -> torch.Tensor:
    tp_rank = get_tensor_model_parallel_rank()
    tp_size = get_tensor_model_parallel_world_size()
@@ -115,7 +115,7 @@ def norm_scale_shift(
    ``weight`` is the effective RMSNorm weight (K2 stores ``scale``, so callers
    pass ``scale + 1``), kept off the checkpoint so the identity load is unaffected.
    """
-    if x.is_cuda and x.shape[-1] % 256 == 0:
+    if x.is_cuda and x.shape[-1] % 256 == 0 and x.shape[-1] <= 8192:
        from sglang.kernels.ops.diffusion.cutedsl.scale_residual_norm_scale_shift import (
            fused_norm_scale_shift,
        )
@@ -1,8 +1,8 @@
 # Copied and adapted from: mossVG/mova/diffusion/models/wan_audio_dit.py
 # SPDX-License-Identifier: Apache-2.0
 #
-# NOTE: This module reuses common functions from mova_video_dit.py to reduce code duplication.
+# NOTE: This module reuses common functions from mova_video_dit.py to reduce
-# Audio-specific functions (precompute_freqs_cis_1d, legacy_precompute_freqs_cis_1d) are kept here.
+# code duplication. Audio-specific precompute_freqs_cis_1d is kept here.
 import math
 from typing import Any, Optional, Tuple
@@ -27,23 +27,6 @@ from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from .mova_video_dit import DiTBlock, precompute_freqs_cis, sinusoidal_embedding_1d
 # Audio-specific positional encoding functions
 def legacy_precompute_freqs_cis_1d(
    dim: int,
    end: int = 16384,
    theta: float = 10000.0,
    base_tps=4.0,
    target_tps=44100 / 2048,
 ):
    s = float(base_tps) / float(target_tps)
    # 1d rope precompute
    f_freqs_cis = precompute_freqs_cis(dim - 2 * (dim // 3), end, theta, s)
    # No positional encoding is applied to the remaining dimensions
    no_freqs_cis = precompute_freqs_cis(dim // 3, end, theta, s)
    no_freqs_cis = torch.ones_like(no_freqs_cis)
    return f_freqs_cis, no_freqs_cis, no_freqs_cis
 def precompute_freqs_cis_1d(dim: int, end: int = 16384, theta: float = 10000.0):
    f_freqs_cis = precompute_freqs_cis(dim, end, theta)
    return f_freqs_cis.chunk(3, dim=-1)
@@ -82,15 +82,6 @@ def precompute_freqs_cis(
    return freqs_cis
 def rope_apply(x, freqs, num_heads):
    x = rearrange(x, "b s (n d) -> b s n d", n=num_heads)
    x_out = torch.view_as_complex(
        x.to(torch.float64).reshape(x.shape[0], x.shape[1], x.shape[2], -1, 2)
    )
    x_out = torch.view_as_real(x_out * freqs).flatten(2)
    return x_out.to(x.dtype)
 def rope_apply_head_dim(x, freqs, head_dim):
    x = rearrange(x, "b s (n d) -> b s n d", d=head_dim)
    x_out = torch.view_as_complex(
@@ -127,14 +127,6 @@ def zimage_rmsnorm_scale(
    return norm(x) * scale
 class SelectFirstElement(nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x):
        return x[0]
 class TimestepEmbedder(nn.Module):
    def __init__(self, out_size, mid_size=None, frequency_embedding_size=256):
        super().__init__()
@@ -1,46 +0,0 @@
 # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
 # SPDX-License-Identifier: Apache-2.0
 # type: ignore
 import os
 import torch
 import torch.nn as nn
 from transformers import BertModel, BertTokenizer
 class HunyuanClip(nn.Module):
    """
    Hunyuan clip code copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py
    hunyuan's clip used BertModel and BertTokenizer, so we copy it.
    """
    def __init__(self, model_dir, max_length=77):
        super().__init__()
        self.max_length = max_length
        self.tokenizer = BertTokenizer.from_pretrained(
            os.path.join(model_dir, "tokenizer")
        )
        self.text_encoder = BertModel.from_pretrained(
            os.path.join(model_dir, "clip_text_encoder")
        )
    @torch.no_grad
    def forward(self, prompts, with_mask=True):
        self.device = next(self.text_encoder.parameters()).device
        text_inputs = self.tokenizer(
            prompts,
            padding="max_length",
            max_length=self.max_length,
            truncation=True,
            return_attention_mask=True,
            return_tensors="pt",
        )
        prompt_embeds = self.text_encoder(
            text_inputs.input_ids.to(self.device),
            attention_mask=(
                text_inputs.attention_mask.to(self.device) if with_mask else None
            ),
        )
        return prompt_embeds.last_hidden_state, prompt_embeds.pooler_output
@@ -126,21 +126,6 @@ def _can_use_unmasked_causal_attention(
    return bool(torch.all(attention_mask > 0).item())
 def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    """
    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep).
    The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to
    (batch, num_attention_heads, seqlen, head_dim)
    """
    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(
        batch, num_key_value_heads, n_rep, slen, head_dim
    )
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
 class MistralAttention(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""
@@ -71,7 +71,6 @@ from transformers.activations import ACT2FN
 from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
    Qwen2_5_VisionRotaryEmbedding,
    Qwen2_5_VisionTransformerPretrainedModel,
    Qwen2_5_VLAttention,
    Qwen2_5_VLCausalLMOutputWithPast,
    Qwen2_5_VLModelOutputWithPast,
    Qwen2_5_VLRotaryEmbedding,
@@ -3,51 +3,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/vision.py
 from abc import ABC, abstractmethod
 from typing import Generic, TypeVar
 import torch
 from transformers import PretrainedConfig
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 logger = init_logger(__name__)
 _C = TypeVar("_C", bound=PretrainedConfig)
 class VisionEncoderInfo(ABC, Generic[_C]):
    def __init__(self, vision_config: _C) -> None:
        super().__init__()
        self.vision_config = vision_config
    @abstractmethod
    def get_num_image_tokens(
        self,
        *,
        image_width: int,
        image_height: int,
    ) -> int:
        raise NotImplementedError
    @abstractmethod
    def get_max_image_tokens(self) -> int:
        raise NotImplementedError
    @abstractmethod
    def get_image_size(self) -> int:
        raise NotImplementedError
    @abstractmethod
    def get_patch_size(self) -> int:
        raise NotImplementedError
    @abstractmethod
    def get_patch_grid_length(self) -> int:
        raise NotImplementedError
 def resolve_visual_encoder_outputs(
    encoder_outputs: torch.Tensor | list[torch.Tensor],
@@ -123,10 +123,6 @@ class Platform:
    def is_cuda(self) -> bool:
        return self.is_cuda_static()
    @lru_cache(maxsize=1)
    def is_npu(self) -> bool:
        return self._enum == PlatformEnum.NPU
    @lru_cache(maxsize=1)
    def is_rocm(self) -> bool:
        return self.is_rocm_static()
@@ -163,7 +163,7 @@ async def run_async_client_warmup(
            response = await forward(req)
            if response.error is not None:
                raise RuntimeError(response.error)
-    except Exception as e:
+    except Exception:
        if fail_open:
            logger.warning(
                "Synthetic server warmup failed; continuing startup", exc_info=True
@@ -0,0 +1,71 @@
 import sys
 from types import SimpleNamespace
 from unittest.mock import patch
 import pytest
 import torch
 from sglang.multimodal_gen.runtime.layers.layernorm import (
    RMSNormScaleShift,
    ScaleResidualRMSNormScaleShift,
 )
 _CUTEDSL_MODULE = "sglang.kernels.ops.diffusion.cutedsl.scale_residual_norm_scale_shift"
@pytest.mark.parametrize("hidden_size", [257, 8448])
 def test_norm_scale_shift_cuda_falls_back_for_unsupported_hidden_size(hidden_size):
    layer = RMSNormScaleShift(hidden_size)
    x = torch.empty(1, 1, hidden_size)
    shift = torch.empty(1, 1, hidden_size)
    scale = torch.empty(1, 1, hidden_size)
    expected = object()
    with (
        patch.object(layer, "forward_native", return_value=expected) as native,
        pytest.warns(UserWarning, match="native fallback"),
    ):
        actual = layer.forward_cuda(x, shift, scale)
    assert actual is expected
    native.assert_called_once_with(x, shift, scale)
@pytest.mark.parametrize("hidden_size", [257, 8448])
 def test_scale_residual_cuda_falls_back_for_unsupported_hidden_size(hidden_size):
    layer = ScaleResidualRMSNormScaleShift(hidden_size)
    residual = torch.empty(1, 1, hidden_size)
    x = torch.empty(1, 1, hidden_size)
    gate = torch.empty(1, 1, hidden_size)
    shift = torch.empty(1, 1, hidden_size)
    scale = torch.empty(1, 1, hidden_size)
    expected = object()
    with (
        patch.object(layer, "forward_native", return_value=expected) as native,
        pytest.warns(UserWarning, match="native fallback"),
    ):
        actual = layer.forward_cuda(residual, x, gate, shift, scale)
    assert actual is expected
    native.assert_called_once_with(residual, x, gate, shift, scale)
 def test_norm_scale_shift_cuda_uses_cutedsl_for_supported_hidden_size(monkeypatch):
    hidden_size = 256
    layer = RMSNormScaleShift(hidden_size)
    x = torch.empty(1, 1, hidden_size)
    shift = torch.empty(1, 1, hidden_size)
    scale = torch.empty(1, 1, hidden_size)
    expected = object()
    def fused_norm_scale_shift(*args):
        return expected
    monkeypatch.setitem(
        sys.modules,
        _CUTEDSL_MODULE,
        SimpleNamespace(fused_norm_scale_shift=fused_norm_scale_shift),
    )
    assert layer.forward_cuda(x, shift, scale) is expected
@@ -0,0 +1,85 @@
 import pytest
 import torch
 from sglang.kernels.ops.diffusion.triton.zimage_native_norm import (
    zimage_rmsnorm_scale,
    zimage_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_zimage_native_norm_rejects_cpu_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 zimage_rmsnorm_scale(x, weight, modulation, EPS) is None
    assert zimage_rmsnorm_tanh_residual(x, modulation, residual, 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_zimage_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 = zimage_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_zimage_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 = zimage_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_zimage_native_norm_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 zimage_rmsnorm_scale(x, weight, modulation, EPS) is None
    assert zimage_rmsnorm_tanh_residual(x, modulation, residual, weight, EPS) is None
 if __name__ == "__main__":
    import sys
    sys.exit(pytest.main([__file__]))
@@ -46,6 +46,7 @@ EXPECTED = {
    "moe.moe_align_block_size": {"aot", "jit"},
    "quantization.nvfp4_gemm_swiglu_nvfp4_quant": {"cute_dsl"},
    "kvcache.reshape_and_cache_flash": {"triton"},
    "diffusion.apply_group_norm_silu": {"triton"},
 }
 _CPU = PlatformInfo(device_type="cpu")