[diffusion] Clean up shared bitexact gates, helpers, and stale naming (#34180)

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

python/sglang/kernels/jit/csrc/diffusion/norm_scale_shift.cuh

@@ -1,4 +1,4 @@
-// Minimal native-CUDA fast path for Qwen-Image diffusion norm-scale-shift.
+// Minimal native-CUDA fast path for generic bf16 hidden=3072 norm-scale-shift.
 //
 // Supported shape family:
 //   - bf16 activations, B == 1, hidden dim == 3072
@@ -36,7 +36,7 @@ constexpr float kInvHidden = 1.0f / float(kHidden);
 static_assert(kThreads == 192);
 static_assert(kWarps == 6);
 
-struct QwenImageNormParams {
+struct NormScaleShiftParams {
   void* y;
   void* res_out;
   const void* x;
@@ -66,7 +66,7 @@ SGL_DEVICE float cta_reduce_sum(float v, int warp, int lane, float* scratch) {
 }
 
 template <bool kHasResidual>
-__global__ void qwen_image_norm_scale_shift_kernel(const QwenImageNormParams __grid_constant__ params) {
+__global__ void norm_scale_shift_kernel(const NormScaleShiftParams __grid_constant__ params) {
   using namespace device;
   using Vec = AlignedVector<bf16_t, kVecElems>;
 
@@ -135,14 +135,14 @@ __global__ void qwen_image_norm_scale_shift_kernel(const QwenImageNormParams __g
   yv.store(static_cast<bf16_t*>(params.y) + row_offset + elem_offset);
 }
 
-inline uint32_t verify_qwen_geometry(host::SymbolicSize& num_rows) {
+inline uint32_t verify_nss_geometry(host::SymbolicSize& num_rows) {
   using namespace host;
   RuntimeCheck(num_rows.unwrap() > 0, "num_rows must be positive");
   RuntimeCheck(num_rows.unwrap() <= int64_t(UINT32_MAX), "num_rows out of range");
   return static_cast<uint32_t>(num_rows.unwrap());
 }
 
-struct QwenImageNormScaleShiftKernel {
+struct NormScaleShiftKernel {
   static void
   run(tvm::ffi::TensorView y,
       tvm::ffi::TensorView x,
@@ -157,8 +157,8 @@ struct QwenImageNormScaleShiftKernel {
     TensorMatcher({N, kHidden}).with_dtype<bf16_t>().with_device(device).verify(x).verify(y);
     TensorMatcher({kHidden}).with_dtype<bf16_t>().with_device(device).verify(scale).verify(shift);
 
-    const uint32_t grid = verify_qwen_geometry(N);
+    const uint32_t grid = verify_nss_geometry(N);
-    const auto params = QwenImageNormParams{
+    const auto params = NormScaleShiftParams{
         .y = y.data_ptr(),
         .res_out = nullptr,
         .x = x.data_ptr(),
@@ -168,11 +168,11 @@ struct QwenImageNormScaleShiftKernel {
         .shift = shift.data_ptr(),
         .eps = static_cast<float>(eps),
     };
-    LaunchKernel(grid, kThreads, device.unwrap())(qwen_image_norm_scale_shift_kernel<false>, params);
+    LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel<false>, params);
   }
 };
 
-struct QwenImageScaleResidualNormScaleShiftKernel {
+struct ScaleResidualNormScaleShiftKernel {
   static void
   run(tvm::ffi::TensorView y,
       tvm::ffi::TensorView res_out,
@@ -196,8 +196,8 @@ struct QwenImageScaleResidualNormScaleShiftKernel {
         .verify(res_out);
     TensorMatcher({kHidden}).with_dtype<bf16_t>().with_device(device).verify(gate).verify(scale).verify(shift);
 
-    const uint32_t grid = verify_qwen_geometry(N);
+    const uint32_t grid = verify_nss_geometry(N);
-    const auto params = QwenImageNormParams{
+    const auto params = NormScaleShiftParams{
         .y = y.data_ptr(),
         .res_out = res_out.data_ptr(),
         .x = x.data_ptr(),
@@ -207,7 +207,7 @@ struct QwenImageScaleResidualNormScaleShiftKernel {
         .shift = shift.data_ptr(),
         .eps = static_cast<float>(eps),
     };
-    LaunchKernel(grid, kThreads, device.unwrap())(qwen_image_norm_scale_shift_kernel<true>, params);
+    LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel<true>, params);
   }
 };
 

python/sglang/kernels/ops/diffusion/__init__.py

@@ -1,6 +1,7 @@
 """Registered diffusion-model kernels and their public wrappers.
 
-Implementations use the backend recorded by each kernel specification.
+Hot paths import concrete implementations from submodules. The package-level
+wrappers remain available for backward compatibility.
 """
 
 from __future__ import annotations
@@ -54,6 +55,18 @@ register_kernel(
         description="Fused QK-norm + RoPE (sglang.kernels.jit).",
     )
 )
+# Migrated from multimodal_gen (RFC #29630, Phase 2.5). Hot paths import the
+# Triton symbol directly; the registry entry remains for namespace discovery.
+register_kernel(
+    KernelSpec(
+        op="diffusion.sparse_linear_attn_fwd",
+        backend=KernelBackend.TRITON,
+        target="sglang.kernels.ops.diffusion.sparse_linear_attn_kernels:_attn_fwd",
+        capabilities=_CUDA,
+        format_signature=FormatSignature(description="sparse linear attention fwd"),
+        description="Sparse linear attention forward (Triton).",
+    )
+)
 
 
 def apply_group_norm_silu(
@@ -107,13 +120,3 @@ __all__ = [
     "residual_gate_add",
     "fused_inplace_qknorm_rope",
 ]
-
-
-# Migrated from multimodal_gen (RFC #29630, Phase 2.5).
-register_kernel(
-    KernelSpec(
-        op="diffusion.sparse_linear_attn_fwd",
-        backend=KernelBackend.TRITON,
-        target="sglang.kernels.ops.diffusion.sparse_linear_attn_kernels:_attn_fwd",
-    )
-)

python/sglang/kernels/ops/diffusion/bitexact_gate.py

@@ -0,0 +1,118 @@
+"""Shared first-sight verification for bit-exact diffusion fast paths."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+import torch
+
+T = TypeVar("T")
+EqualFn = Callable[[Any, Any], bool]
+
+
+class BitExactFusionGate:
+    """Track permanent disable + first-sight ``torch.equal`` verification.
+
+    Two modes:
+
+    * **once-for-all** (default): the first successful equal-check enables the
+      fused path for every later call (GLM / Ernie).
+    * **per-signature**: each distinct ``sig`` is verified independently
+      (FLUX / Sana), matching aten LayerNorm dispatch that can vary by shape.
+
+    The first-sight check runs the eager reference chain plus a host sync, so
+    it must never happen inside ``torch.compile`` tracing or CUDA graph
+    capture. Once-for-all callers get both guards from
+    :meth:`can_attempt_once`; per-signature callers must keep their own
+    compile/capture checks next to the signature lookup (see FLUX / Sana).
+    """
+
+    __slots__ = ("name", "disabled", "verified", "verified_sigs")
+
+    def __init__(self, name: str, *, per_signature: bool = False) -> None:
+        self.name = name
+        self.disabled = False
+        # This is a plain field because steady-state DiT blocks read it on
+        # every invocation; a property descriptor is measurable at this scale.
+        self.verified = False
+        self.verified_sigs: set[Any] | None = set() if per_signature else None
+
+    def is_verified(self, sig: Any = None) -> bool:
+        if self.verified_sigs is not None:
+            return sig in self.verified_sigs
+        return self.verified
+
+    def mark_verified(self, sig: Any = None) -> None:
+        if self.verified_sigs is not None:
+            assert sig is not None
+            self.verified_sigs.add(sig)
+        self.verified = True
+
+    def disable(self) -> None:
+        self.disabled = True
+
+    def can_attempt_once(self) -> bool:
+        """Once-for-all mode: may we launch the fused kernel right now?"""
+        if self.disabled:
+            return False
+        if self.verified:
+            return True
+        # First-sight verify runs the eager reference chain and a host sync:
+        # attempt neither inside compile tracing nor CUDA graph capture (the
+        # sync would abort the capture; BCG then blocks the signature). Once
+        # verified, the fused kernel runs alone and is compile/capture-safe.
+        if torch.compiler.is_compiling():
+            return False
+        return not (
+            torch.cuda.is_available() and torch.cuda.is_current_stream_capturing()
+        )
+
+    def on_exception(
+        self,
+        exc: BaseException,
+        *,
+        logger: logging.Logger | None = None,
+        re_raise_if_compiling: bool = True,
+    ) -> None:
+        if re_raise_if_compiling and torch.compiler.is_compiling():
+            raise exc
+        if logger is not None:
+            logger.warning_once(f"Disabling {self.name} fast path: {exc}")
+        self.disable()
+
+    def accept_or_fallback(
+        self,
+        out: T,
+        ref: T,
+        *,
+        sig: Any = None,
+        equal: EqualFn | None = None,
+        logger: logging.Logger | None = None,
+        mismatch_msg: str | None = None,
+    ) -> T:
+        """Return ``out`` when bit-exact; otherwise disable and return ``ref``."""
+        if self.is_verified(sig):
+            return out
+        eq = equal or torch.equal
+        if eq(out, ref):
+            self.mark_verified(sig)
+            return out
+        if logger is not None:
+            logger.warning_once(
+                mismatch_msg
+                or (
+                    f"{self.name} fast path is not bit-exact against this "
+                    "platform's reference dispatch; falling back to eager"
+                )
+            )
+        self.disable()
+        return ref
+
+
+def tensors_equal(a: Any, b: Any) -> bool:
+    """``torch.equal`` for a tensor or a sequence of tensors."""
+    if isinstance(a, torch.Tensor):
+        return torch.equal(a, b)
+    return all(torch.equal(x, y) for x, y in zip(a, b, strict=True))

python/sglang/kernels/ops/diffusion/norm_scale_shift_native.py

@@ -24,7 +24,7 @@ def _blackwell_or_newer(device: torch.device) -> bool:
     )
 
 
-def _qwen_activation(t, like=None) -> bool:
+def _nss_activation(t, like=None) -> bool:
     return (
         isinstance(t, torch.Tensor)
         and t.is_cuda
@@ -61,16 +61,16 @@ def _row_bf16(t, device: torch.device):
 @cache_once
 def norm_scale_shift_module() -> Module:
     return load_jit(
-        "qwen_image_norm_scale_shift_native",
+        "norm_scale_shift_native",
         cuda_files=["diffusion/norm_scale_shift.cuh"],
         cuda_wrappers=[
             (
-                "qwen_image_nss_bf16_row",
+                "nss_bf16_row",
-                "norm_scale_shift::QwenImageNormScaleShiftKernel::run",
+                "norm_scale_shift::NormScaleShiftKernel::run",
             ),
             (
-                "qwen_image_srnss_bf16_row",
+                "srnss_bf16_row",
-                "norm_scale_shift::" "QwenImageScaleResidualNormScaleShiftKernel::run",
+                "norm_scale_shift::ScaleResidualNormScaleShiftKernel::run",
             ),
         ],
     )
@@ -82,7 +82,7 @@ _module = norm_scale_shift_module
 def try_fused_norm_scale_shift(x, weight, bias, scale, shift, norm_type, eps):
     if norm_type != "layer" or weight is not None or bias is not None:
         return None
-    if not _qwen_activation(x) or not _blackwell_or_newer(x.device):
+    if not _nss_activation(x) or not _blackwell_or_newer(x.device):
         return None
 
     scale = _row_bf16(scale, x.device)
@@ -91,7 +91,7 @@ def try_fused_norm_scale_shift(x, weight, bias, scale, shift, norm_type, eps):
         return None
 
     y = torch.empty_like(x)
-    _module().qwen_image_nss_bf16_row(
+    _module().nss_bf16_row(
         y.view(-1, _HIDDEN), x.view(-1, _HIDDEN), scale, shift, float(eps)
     )
     return y
@@ -103,8 +103,8 @@ def try_fused_scale_residual_norm_scale_shift(
     if norm_type != "layer" or weight is not None or bias is not None:
         return None
     if not (
-        _qwen_activation(x)
+        _nss_activation(x)
-        and _qwen_activation(residual, x)
+        and _nss_activation(residual, x)
         and _blackwell_or_newer(x.device)
     ):
         return None
@@ -117,7 +117,7 @@ def try_fused_scale_residual_norm_scale_shift(
 
     y = torch.empty_like(x)
     residual_out = torch.empty_like(x)
-    _module().qwen_image_srnss_bf16_row(
+    _module().srnss_bf16_row(
         y.view(-1, _HIDDEN),
         residual_out.view(-1, _HIDDEN),
         residual.view(-1, _HIDDEN),

python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md

@@ -619,7 +619,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 |
-| `zimage_rmsnorm_scale` or `zimage_rmsnorm_tanh_residual` missing on Z-Image | Check the bf16-native Triton eligibility guards before proposing a new fusion |
+| `rmsnorm_scale` or `rmsnorm_tanh_residual` missing on Z-Image | Check the bf16-native Triton eligibility guards before proposing a new fusion |
 | FLUX.1, GLM-Image, or SANA shows separate LayerNorm plus adaLN elementwise kernels | Check the bit-exact `modulate_scale_shift` and `fused_layernorm_modulate` guards/self-test before proposing another norm fusion |
 | `quality=high` shows the same FLUX/GLM DiT or FLUX-family/Wan VAE chain as `lossless` | Check whether the request-scoped quality gate mounted and whether every site passed its all-or-nothing compatibility checks |
 | LTX-2 split RoPE appears as a long PyTorch elementwise chain | Check the `apply_ltx2_split_rotary_emb` Triton path and its shape guards |

python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md

@@ -14,6 +14,7 @@ framework-specific optimization workflow.
 - `python/sglang/kernels/ops/diffusion/modulate_scale_shift.py`
 - `python/sglang/kernels/ops/diffusion/fused_ln_modulate.py`
 - `python/sglang/kernels/ops/diffusion/quality_gate.py`
+- `python/sglang/kernels/ops/diffusion/bitexact_gate.py`
 - `python/sglang/kernels/ops/diffusion/group_norm_silu.py`
 - `python/sglang/kernels/ops/diffusion/triton/group_norm_silu.py`
 - `python/sglang/kernels/ops/diffusion/triton/group_norm_silu_twopass.py`
@@ -24,6 +25,8 @@ framework-specific optimization workflow.
 - `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/ltx2_qknorm_split_rope.py`
+- `python/sglang/kernels/ops/diffusion/ltx2_rmsnorm_modulate.py`
 - `python/sglang/kernels/ops/diffusion/triton/indexed_modulation.py`
 - `python/sglang/kernels/ops/diffusion/triton/ulysses_qkv.py`
 - `python/sglang/kernels/ops/diffusion/usp_relayout.py`
@@ -48,6 +51,8 @@ framework-specific optimization workflow.
 - `test/registered/kernels/ops/diffusion/test_glm_image_ln_modulate.py`
 - `test/registered/kernels/ops/diffusion/test_sana_ln_modulate.py`
 - `test/registered/kernels/ops/diffusion/test_quality_gate.py`
+- `test/registered/kernels/ops/diffusion/test_ltx2_rms_norm_modulate.py`
+- `test/registered/kernels/ops/diffusion/test_bitexact_gate.py`
 - `test/registered/kernels/ops/diffusion/test_wan_causal_cache.py`
 - `test/registered/kernels/ops/diffusion/test_stage_profiler_sync.py`
 - `test/registered/kernels/benchmark/diffusion/bench_qwen_image_modulation.py`
@@ -116,9 +121,9 @@ framework-specific optimization workflow.
   `test_vae_fast_path_gate.py`.
 
 5. Z-Image bf16-native RMSNorm modulation (Triton)
-- Kernels: `zimage_rmsnorm_scale`, `zimage_rmsnorm_tanh_residual`
+- Kernels: `rmsnorm_scale`, `rmsnorm_tanh_residual`
-- Locations: `triton/native_bf16_rmsnorm.py`, compatibility exports in
+- Locations: `triton/native_bf16_rmsnorm.py`, with wrappers in `zimage.py` and
-  `triton/zimage_native_norm.py`, and `zimage.py`
+  `fused_gate_rmsnorm.py`. Note: `triton/zimage_native_norm.py` is QK-only.
 - Use cases:
   - `y = rmsnorm(x) * scale`
   - `y = residual + tanh(gate) * rmsnorm(x)`
@@ -162,7 +167,7 @@ framework-specific optimization workflow.
 - Locations: `diffusion/residual_gate_add.py`, `csrc/diffusion/residual_gate_add.cuh`, `runtime/models/dits/ltx_2.py`
 - 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`.
+- Behavior: LTX2 calls `residual_gate_add(...)` from the kernels package directly. The CUDA custom op is used 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/kernels/ops/diffusion/test_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.
@@ -291,17 +296,27 @@ framework-specific optimization workflow.
 - Request-scoped high-quality acceleration: `QualityGatedFusion` in
   `quality_gate.py`, `_maybe_toggle_quality_fusions` in `denoising.py`, and
   `use_vae_fast_path` in `decoding.py`.
+- Bit-exact first-sight verify/disable: `BitExactFusionGate` in
+  `bitexact_gate.py`, used by FLUX / GLM / Sana / Ernie fused norm sites.
 - 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 native norm modulation: `zimage_rmsnorm_scale` and
+- Z-Image native norm modulation: `rmsnorm_scale` and `rmsnorm_tanh_residual`
-  `zimage_rmsnorm_tanh_mul_add` in `zimage.py`, backed by the shared
+  in `triton/native_bf16_rmsnorm.py`, with wrappers in `zimage.py` /
-  `triton/native_bf16_rmsnorm.py` kernels.
+  `fused_gate_rmsnorm.py`. `zimage_native_norm.py` is QK-only.
 - HunyuanVideo VAE and LTX upsampler GroupNorm+SiLU: `apply_group_norm_silu` in `hunyuanvae.py` and `latent_upsampler.py`; default-eligible when wrapper guards pass.
 - MiniMax-H3 indexed modulation: `_modulate_scale_shift` and `_modulate_gate` in `minimax_h3.py`, backed by `triton/indexed_modulation.py`.
 - MiniMax-H3 Ulysses relayout: `_usp_input_all_to_all_packed_qkv` and `usp_merge_heads` through `runtime/layers/usp.py`.
 - 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.
 - LTX2 split RoPE: `apply_ltx2_split_rotary_emb` in `ltx_2.py`.
-- LTX2 residual-gate add: `_ltx2_residual_gate_add` in `ltx_2.py` wraps the CUDA `diffusion_residual_gate_add` custom op for attention, cross-attention, and MLP residual updates.
+- LTX2 RMSNorm+modulate and FFN GELU epilogue under `quality="high"`:
+  `mark_ltx2_rms_norm_modulate_site` / `fused_ltx2_rms_norm_modulate` in
+  `kernels/ops/diffusion/ltx2_rmsnorm_modulate.py` (mount-based
+  `QualityGatedFusion`, not a first-sight `BitExactFusionGate` — the fused
+  kernel is <=1 ULP off aten, so it is request-gated instead of verified),
+  wired at the six `LTX2TransformerBlock` adaLN sites in `ltx_2.py`.
+- LTX2 residual-gate add: `ltx_2.py` calls `residual_gate_add` from
+  `kernels/ops/diffusion/residual_gate_add.py` directly for attention,
+  cross-attention, and MLP residual updates.
 - Wan causal VAE: `cat_pad_channels_last_3d` and `dup_up3d_add` in
   `wanvae.py`, backed by `triton/wan_causal_cache.py`.
 - Varlen USP attention: `fused_pack_qkv` and `fused_scatter_to_padded` in `attention/layer.py`.

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

@@ -1,5 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
+from typing import Any
+
 import torch
 
 
@@ -16,3 +20,46 @@ def modulate(
     if scale is None:
         return x + shift.unsqueeze(1)  # type: ignore[union-attr]
     return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+def get_qkv_projections(
+    attn: Any,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor | None = None,
+) -> tuple[
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor | None,
+    torch.Tensor | None,
+    torch.Tensor | None,
+]:
+    """Shared fused/unfused QKV (+ optional added-KV) projection helper.
+
+    Used by FLUX / FLUX.2 / Qwen-Image attention blocks that expose the same
+    ``to_qkv`` / ``to_added_qkv`` packing flags. ``use_fused_qkv`` is always
+    set by those blocks' constructors, and ``use_fused_added_qkv`` whenever
+    ``added_kv_proj_dim`` is not ``None`` — direct attribute access so a
+    renamed flag fails loudly instead of silently unfusing.
+    """
+    if attn.use_fused_qkv:
+        qkv, _ = attn.to_qkv(hidden_states)
+        query, key, value = [t.contiguous() for t in qkv.chunk(3, dim=-1)]
+    else:
+        query, _ = attn.to_q(hidden_states)
+        key, _ = attn.to_k(hidden_states)
+        value, _ = attn.to_v(hidden_states)
+
+    encoder_query = encoder_key = encoder_value = None
+    if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None:
+        if attn.use_fused_added_qkv:
+            added_qkv, _ = attn.to_added_qkv(encoder_hidden_states)
+            encoder_query, encoder_key, encoder_value = [
+                t.contiguous() for t in added_qkv.chunk(3, dim=-1)
+            ]
+        else:
+            encoder_query, _ = attn.add_q_proj(encoder_hidden_states)
+            encoder_key, _ = attn.add_k_proj(encoder_hidden_states)
+            encoder_value, _ = attn.add_v_proj(encoder_hidden_states)
+
+    return query, key, value, encoder_query, encoder_key, encoder_value

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

@@ -19,6 +19,10 @@ import torch.nn as nn
 import torch.nn.functional as F
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 
+from sglang.kernels.ops.diffusion.bitexact_gate import (
+    BitExactFusionGate,
+    tensors_equal,
+)
 from sglang.kernels.ops.diffusion.residual_gate_add import residual_gate_add
 from sglang.kernels.ops.diffusion.triton.rmsnorm_scale_shift_bitexact import (
     can_use_fused_rmsnorm_scale_shift,
@@ -52,10 +56,8 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 logger = init_logger(__name__)
 
 
-_ERNIE_FUSED_NORM_DISABLED = False
+_ERNIE_NORM = BitExactFusionGate("ERNIE fused-norm")
-_ERNIE_FUSED_NORM_VERIFIED = False
+_ERNIE_GATED_NORM = BitExactFusionGate("ERNIE fused gated-norm")
-_ERNIE_FUSED_GATED_NORM_DISABLED = False
-_ERNIE_FUSED_GATED_NORM_VERIFIED = False
 
 
 def _eager_norm_scale_shift(
@@ -75,36 +77,31 @@ def _ernie_norm_scale_shift(
     first call verifies ``torch.equal`` against the eager chain and disables
     the fast path permanently on any mismatch.
     """
-    global _ERNIE_FUSED_NORM_DISABLED, _ERNIE_FUSED_NORM_VERIFIED
+    verified = _ERNIE_NORM.verified
-
     if (
-        not _ERNIE_FUSED_NORM_DISABLED
+        not _ERNIE_NORM.disabled
         and norm.variance_size_override is None
         and can_use_fused_rmsnorm_scale_shift(x, norm.weight, scale, shift)
-        and (_ERNIE_FUSED_NORM_VERIFIED or not torch.compiler.is_compiling())
+        and (verified or _ERNIE_NORM.can_attempt_once())
     ):
         try:
             out = fused_rmsnorm_scale_shift_bitexact(
                 x, norm.weight, scale, shift, norm.variance_epsilon
             )
         except Exception as exc:
-            if torch.compiler.is_compiling():
+            _ERNIE_NORM.on_exception(exc, logger=logger)
-                raise
-            logger.warning_once(f"Disabling ERNIE fused-norm fast path: {exc}")
-            _ERNIE_FUSED_NORM_DISABLED = True
         else:
-            if _ERNIE_FUSED_NORM_VERIFIED:
+            if verified:
                 return out
-            ref = _eager_norm_scale_shift(norm, x, scale, shift)
+            return _ERNIE_NORM.accept_or_fallback(
-            if torch.equal(out, ref):
+                out,
-                _ERNIE_FUSED_NORM_VERIFIED = True
+                _eager_norm_scale_shift(norm, x, scale, shift),
-                return out
+                logger=logger,
-            logger.warning_once(
+                mismatch_msg=(
                     "ERNIE fused-norm fast path is not bit-exact against this "
                     "platform's rmsnorm dispatch; falling back to eager"
+                ),
             )
-            _ERNIE_FUSED_NORM_DISABLED = True
-            return ref
 
     return _eager_norm_scale_shift(norm, x, scale, shift)
 
@@ -123,15 +120,14 @@ def _ernie_gated_norm_scale_shift(
     (and the ``residual_gate_add_cuda`` fast path) + norm chain; first call
     self-verifies like :func:`_ernie_norm_scale_shift`.
     """
-    global _ERNIE_FUSED_GATED_NORM_DISABLED, _ERNIE_FUSED_GATED_NORM_VERIFIED
+    verified = _ERNIE_GATED_NORM.verified
-
     if (
-        not _ERNIE_FUSED_GATED_NORM_DISABLED
+        not _ERNIE_GATED_NORM.disabled
         and norm.variance_size_override is None
         and can_use_fused_scale_residual_rmsnorm_scale_shift(
             residual, update, gate, norm.weight, scale, shift
         )
-        and (_ERNIE_FUSED_GATED_NORM_VERIFIED or not torch.compiler.is_compiling())
+        and (verified or _ERNIE_GATED_NORM.can_attempt_once())
     ):
         try:
             out, res = fused_scale_residual_rmsnorm_scale_shift_bitexact(
@@ -144,24 +140,22 @@ def _ernie_gated_norm_scale_shift(
                 norm.variance_epsilon,
             )
         except Exception as exc:
-            if torch.compiler.is_compiling():
+            _ERNIE_GATED_NORM.on_exception(exc, logger=logger)
-                raise
-            logger.warning_once(f"Disabling ERNIE fused gated-norm fast path: {exc}")
-            _ERNIE_FUSED_GATED_NORM_DISABLED = True
         else:
-            if _ERNIE_FUSED_GATED_NORM_VERIFIED:
+            if verified:
                 return out, res
             res_ref = residual + gate * update
             ref = _eager_norm_scale_shift(norm, res_ref, scale, shift)
-            if torch.equal(out, ref) and torch.equal(res, res_ref):
+            return _ERNIE_GATED_NORM.accept_or_fallback(
-                _ERNIE_FUSED_GATED_NORM_VERIFIED = True
+                (out, res),
-                return out, res
+                (ref, res_ref),
-            logger.warning_once(
+                equal=tensors_equal,
+                logger=logger,
+                mismatch_msg=(
                     "ERNIE fused gated-norm fast path is not bit-exact against "
                     "this platform's rmsnorm dispatch; falling back to eager"
+                ),
             )
-            _ERNIE_FUSED_GATED_NORM_DISABLED = True
-            return ref, res_ref
 
     res = residual_gate_add(residual, update, gate)
     return _eager_norm_scale_shift(norm, res, scale, shift), res

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

@@ -28,6 +28,7 @@ from diffusers.models.normalization import (
 )
 from torch.nn import LayerNorm as LayerNorm
 
+from sglang.kernels.ops.diffusion.bitexact_gate import BitExactFusionGate
 from sglang.kernels.ops.diffusion.fused_linear_gelu import (
     can_fuse_linear_gelu,
     fused_gelu_active,
@@ -91,16 +92,18 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload im
     LayerwiseOffloadableModuleMixin,
 )
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+from sglang.multimodal_gen.runtime.models.dits.common import get_qkv_projections
 from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)  # pylint: disable=invalid-name
 
+_get_qkv_projections = get_qkv_projections
 
-_FLUX_FUSED_LN_MOD_DISABLED = False
+_FLUX_LN_MOD = BitExactFusionGate("FLUX fused LN+modulate", per_signature=True)
-# (shape, stride, eps) signatures whose fused output has been verified
+# Keep the pre-refactor direct set lookup in this launch-sensitive hot path.
-# ``torch.equal`` against the live eager chain.
+_FLUX_LN_MOD_SIGS = _FLUX_LN_MOD.verified_sigs
-_FLUX_FUSED_LN_MOD_VERIFIED: set = set()
+assert _FLUX_LN_MOD_SIGS is not None
 
 
 def _flux_fused_ln_modulate(
@@ -118,10 +121,8 @@ def _flux_fused_ln_modulate(
     verified ``torch.equal`` against the eager chain on first sight, and any
     mismatch disables the fast path permanently.
     """
-    global _FLUX_FUSED_LN_MOD_DISABLED
-
     if (
-        _FLUX_FUSED_LN_MOD_DISABLED
+        _FLUX_LN_MOD.disabled
         or not is_plain_layer_norm(norm, x.shape[-1])
         or not can_use_fused_layernorm_modulate(x, scale, shift)
     ):
@@ -135,7 +136,7 @@ def _flux_fused_ln_modulate(
         shift.stride(),
         norm.eps,
     )
-    verified = sig in _FLUX_FUSED_LN_MOD_VERIFIED
+    verified = sig in _FLUX_LN_MOD_SIGS
     if not verified and (
         torch.compiler.is_compiling() or torch.cuda.is_current_stream_capturing()
     ):
@@ -145,23 +146,21 @@ def _flux_fused_ln_modulate(
     try:
         out = fused_layernorm_modulate(x, scale, shift, norm.eps)
     except Exception as exc:
-        if torch.compiler.is_compiling():
+        _FLUX_LN_MOD.on_exception(exc, logger=logger)
-            raise
-        logger.warning_once(f"Disabling FLUX fused LN+modulate fast path: {exc}")
-        _FLUX_FUSED_LN_MOD_DISABLED = True
         return None
     if verified:
         return out
     ref = modulate_scale_shift(norm(x), scale, shift)
-    if torch.equal(out, ref):
+    return _FLUX_LN_MOD.accept_or_fallback(
-        _FLUX_FUSED_LN_MOD_VERIFIED.add(sig)
+        out,
-        return out
+        ref,
-    logger.warning_once(
+        sig=sig,
+        logger=logger,
+        mismatch_msg=(
             "FLUX fused LN+modulate fast path is not bit-exact against this "
             "platform's LayerNorm dispatch; falling back to eager"
+        ),
     )
-    _FLUX_FUSED_LN_MOD_DISABLED = True
-    return ref
 
 
 def _flux_norm_modulate(
@@ -373,32 +372,6 @@ def _fused_gelu_mlp(
     return output.view(batch_size, seq_len, -1)
 
 
-def _get_qkv_projections(
-    attn: "FluxAttention", hidden_states, encoder_hidden_states=None
-):
-    if getattr(attn, "use_fused_qkv", False):
-        qkv, _ = attn.to_qkv(hidden_states)
-        query, key, value = [x.contiguous() for x in qkv.chunk(3, dim=-1)]
-    else:
-        query, _ = attn.to_q(hidden_states)
-        key, _ = attn.to_k(hidden_states)
-        value, _ = attn.to_v(hidden_states)
-
-    encoder_query = encoder_key = encoder_value = None
-    if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None:
-        if attn.use_fused_added_qkv:
-            added_qkv, _ = attn.to_added_qkv(encoder_hidden_states)
-            encoder_query, encoder_key, encoder_value = [
-                x.contiguous() for x in added_qkv.chunk(3, dim=-1)
-            ]
-        else:
-            encoder_query, _ = attn.add_q_proj(encoder_hidden_states)
-            encoder_key, _ = attn.add_k_proj(encoder_hidden_states)
-            encoder_value, _ = attn.add_v_proj(encoder_hidden_states)
-
-    return query, key, value, encoder_query, encoder_key, encoder_value
-
-
 class FluxGELU(nn.Module):
     def __init__(
         self,

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

@@ -58,6 +58,7 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload im
     LayerwiseOffloadableModuleMixin,
 )
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+from sglang.multimodal_gen.runtime.models.dits.common import get_qkv_projections
 from sglang.multimodal_gen.runtime.platforms import (
     AttentionBackendEnum,
     current_platform,
@@ -66,31 +67,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)  # pylint: disable=invalid-name
 
-
+_get_qkv_projections = get_qkv_projections
-def _get_qkv_projections(
-    attn: "Flux2Attention", hidden_states, encoder_hidden_states=None
-):
-    if attn.use_fused_qkv:
-        qkv, _ = attn.to_qkv(hidden_states)
-        query, key, value = [t.contiguous() for t in qkv.chunk(3, dim=-1)]
-    else:
-        query, _ = attn.to_q(hidden_states)
-        key, _ = attn.to_k(hidden_states)
-        value, _ = attn.to_v(hidden_states)
-
-    encoder_query = encoder_key = encoder_value = None
-    if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None:
-        if attn.use_fused_added_qkv:
-            added_qkv, _ = attn.to_added_qkv(encoder_hidden_states)
-            encoder_query, encoder_key, encoder_value = [
-                t.contiguous() for t in added_qkv.chunk(3, dim=-1)
-            ]
-        else:
-            encoder_query, _ = attn.add_q_proj(encoder_hidden_states)
-            encoder_key, _ = attn.add_k_proj(encoder_hidden_states)
-            encoder_value, _ = attn.add_v_proj(encoder_hidden_states)
-
-    return query, key, value, encoder_query, encoder_key, encoder_value
 
 
 class Flux2SwiGLU(nn.Module):

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

@@ -18,6 +18,10 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from sglang.kernels.ops.diffusion.bitexact_gate import (
+    BitExactFusionGate,
+    tensors_equal,
+)
 from sglang.kernels.ops.diffusion.fused_linear_gelu import (
     can_fuse_linear_gelu,
     fused_gelu_active,
@@ -73,10 +77,8 @@ logger = init_logger(__name__)
 
 _is_cuda = current_platform.is_cuda()
 
-_GLM_FUSED_LN_MOD_DISABLED = False
+_GLM_LN_MOD = BitExactFusionGate("GLM fused LN+modulate")
-_GLM_FUSED_LN_MOD_VERIFIED = False
+_GLM_QK_LN = BitExactFusionGate("GLM fused qk-LayerNorm")
-_GLM_FUSED_QK_LN_DISABLED = False
-_GLM_FUSED_QK_LN_VERIFIED = False
 
 
 def _eager_ln_modulate(
@@ -102,36 +104,31 @@ def _glm_ln_modulate(
     the first call verifies ``torch.equal`` against the eager chain and
     disables the fast path permanently on any mismatch.
     """
-    global _GLM_FUSED_LN_MOD_DISABLED, _GLM_FUSED_LN_MOD_VERIFIED
+    verified = _GLM_LN_MOD.verified
-
     if (
-        not _GLM_FUSED_LN_MOD_DISABLED
+        not _GLM_LN_MOD.disabled
         and _is_cuda
         and dtype is x.dtype
         and is_plain_layer_norm(norm, x.shape[-1])
         and can_use_fused_layernorm_modulate(x, scale, shift)
-        and (_GLM_FUSED_LN_MOD_VERIFIED or not torch.compiler.is_compiling())
+        and (verified or _GLM_LN_MOD.can_attempt_once())
     ):
         try:
             out = fused_layernorm_modulate(x, scale, shift, norm.eps)
         except Exception as exc:
-            if torch.compiler.is_compiling():
+            _GLM_LN_MOD.on_exception(exc, logger=logger)
-                raise
-            logger.warning_once(f"Disabling GLM fused LN+modulate fast path: {exc}")
-            _GLM_FUSED_LN_MOD_DISABLED = True
         else:
-            if _GLM_FUSED_LN_MOD_VERIFIED:
+            if verified:
                 return out
-            ref = _eager_ln_modulate(norm, x, scale, shift, dtype)
+            return _GLM_LN_MOD.accept_or_fallback(
-            if torch.equal(out, ref):
+                out,
-                _GLM_FUSED_LN_MOD_VERIFIED = True
+                _eager_ln_modulate(norm, x, scale, shift, dtype),
-                return out
+                logger=logger,
-            logger.warning_once(
+                mismatch_msg=(
                     "GLM fused LN+modulate fast path is not bit-exact against "
                     "this platform's LayerNorm dispatch; falling back to eager"
+                ),
             )
-            _GLM_FUSED_LN_MOD_DISABLED = True
-            return ref
 
     return _eager_ln_modulate(norm, x, scale, shift, dtype)
 
@@ -148,10 +145,9 @@ def _glm_qk_layernorm(
     First call verifies ``torch.equal`` against the eager pair and falls
     back permanently on any mismatch.
     """
-    global _GLM_FUSED_QK_LN_DISABLED, _GLM_FUSED_QK_LN_VERIFIED
+    verified = _GLM_QK_LN.verified
-
     if (
-        not _GLM_FUSED_QK_LN_DISABLED
+        not _GLM_QK_LN.disabled
         and _is_cuda
         and dtype is query.dtype
         and dtype is key.dtype
@@ -159,29 +155,29 @@ def _glm_qk_layernorm(
         and is_plain_layer_norm(norm_k, key.shape[-1])
         and norm_q.eps == norm_k.eps
         and can_use_fused_qk_head_layernorm(query, key)
-        and (_GLM_FUSED_QK_LN_VERIFIED or not torch.compiler.is_compiling())
+        and (verified or _GLM_QK_LN.can_attempt_once())
     ):
         try:
-            q_out, k_out = fused_qk_head_layernorm(query, key, norm_q.eps)
+            out = fused_qk_head_layernorm(query, key, norm_q.eps)
         except Exception as exc:
-            if torch.compiler.is_compiling():
+            _GLM_QK_LN.on_exception(exc, logger=logger)
-                raise
-            logger.warning_once(f"Disabling GLM fused qk-LayerNorm fast path: {exc}")
-            _GLM_FUSED_QK_LN_DISABLED = True
         else:
-            if _GLM_FUSED_QK_LN_VERIFIED:
+            if verified:
-                return q_out, k_out
+                return out
-            q_ref = norm_q(query).to(dtype=dtype)
+            ref = (
-            k_ref = norm_k(key).to(dtype=dtype)
+                norm_q(query).to(dtype=dtype),
-            if torch.equal(q_out, q_ref) and torch.equal(k_out, k_ref):
+                norm_k(key).to(dtype=dtype),
-                _GLM_FUSED_QK_LN_VERIFIED = True
+            )
-                return q_out, k_out
+            return _GLM_QK_LN.accept_or_fallback(
-            logger.warning_once(
+                out,
+                ref,
+                equal=tensors_equal,
+                logger=logger,
+                mismatch_msg=(
                     "GLM fused qk-LayerNorm fast path is not bit-exact against "
                     "this platform's LayerNorm dispatch; falling back to eager"
+                ),
             )
-            _GLM_FUSED_QK_LN_DISABLED = True
-            return q_ref, k_ref
 
     return norm_q(query).to(dtype=dtype), norm_k(key).to(dtype=dtype)
 

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

@@ -73,6 +73,7 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload im
     LayerwiseOffloadableModuleMixin,
 )
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+from sglang.multimodal_gen.runtime.models.dits.common import get_qkv_projections
 from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
@@ -104,32 +105,7 @@ def _local_seq_len(seq_len: int, sp_world_size: int) -> int:
     return padded_len // sp_world_size
 
 
-def _get_qkv_projections(
+_get_qkv_projections = get_qkv_projections
-    attn: "QwenImageCrossAttention", hidden_states, encoder_hidden_states=None
-):
-    if attn.use_fused_qkv:
-        img_qkv, _ = attn.to_qkv(hidden_states)
-        img_query, img_key, img_value = [
-            x.contiguous() for x in img_qkv.chunk(3, dim=-1)
-        ]
-    else:
-        img_query, _ = attn.to_q(hidden_states)
-        img_key, _ = attn.to_k(hidden_states)
-        img_value, _ = attn.to_v(hidden_states)
-
-    txt_query = txt_key = txt_value = None
-    if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None:
-        if attn.use_fused_added_qkv:
-            txt_qkv, _ = attn.to_added_qkv(encoder_hidden_states)
-            txt_query, txt_key, txt_value = [
-                x.contiguous() for x in txt_qkv.chunk(3, dim=-1)
-            ]
-        else:
-            txt_query, _ = attn.add_q_proj(encoder_hidden_states)
-            txt_key, _ = attn.add_k_proj(encoder_hidden_states)
-            txt_value, _ = attn.add_v_proj(encoder_hidden_states)
-
-    return img_query, img_key, img_value, txt_query, txt_key, txt_value
 
 
 class QwenTimestepProjEmbeddings(nn.Module):

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

@@ -5,6 +5,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding
 
+from sglang.kernels.ops.diffusion.bitexact_gate import BitExactFusionGate
 from sglang.kernels.ops.diffusion.triton.layernorm_modulate import (
     can_use_fused_layernorm_modulate,
     fused_layernorm_modulate_raw,
@@ -22,10 +23,12 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)
 
-_SANA_FUSED_LN_MOD_DISABLED = False
+_SANA_LN_MOD = BitExactFusionGate("Sana fused LN+modulate", per_signature=True)
-# (shape/stride/dtype/eps) signatures whose fused output was torch.equal-
+# Direct module-level state keeps BCG warmup launch overhead equal to the
-# verified against the live eager chain.
+# pre-refactor path; the gate still owns first-sight verification transitions.
-_SANA_FUSED_LN_MOD_OK_SIGS: set = set()
+_SANA_LN_MOD_SIGS = _SANA_LN_MOD.verified_sigs
+assert _SANA_LN_MOD_SIGS is not None
+_SANA_LN_MOD_DISABLED = False
 
 
 def _eager_ln_modulate(
@@ -61,9 +64,9 @@ def _sana_ln_modulate(
     layout); aten's LayerNorm contiguizes internally, and the fast path
     issues the same copy explicitly.
     """
-    global _SANA_FUSED_LN_MOD_DISABLED
+    global _SANA_LN_MOD_DISABLED
 
-    if _SANA_FUSED_LN_MOD_DISABLED or torch.compiler.is_compiling() or not x.is_cuda:
+    if _SANA_LN_MOD_DISABLED or torch.compiler.is_compiling() or not x.is_cuda:
         return _eager_ln_modulate(norm, x, scale, shift)
 
     capturing = torch.cuda.is_current_stream_capturing()
@@ -79,7 +82,7 @@ def _sana_ln_modulate(
         shift.stride(),
         norm.eps,
     )
-    if sig in _SANA_FUSED_LN_MOD_OK_SIGS:
+    if sig in _SANA_LN_MOD_SIGS:
         return fused_layernorm_modulate_raw(
             x.contiguous(), scale[:, 0], shift[:, 0], norm.eps
         )
@@ -101,19 +104,21 @@ def _sana_ln_modulate(
         try:
             out = fused_layernorm_modulate_raw(x_c, scale[:, 0], shift[:, 0], norm.eps)
         except Exception as exc:
-            logger.warning_once(f"Disabling Sana fused LN+modulate fast path: {exc}")
+            _SANA_LN_MOD.on_exception(exc, logger=logger)
-            _SANA_FUSED_LN_MOD_DISABLED = True
+            _SANA_LN_MOD_DISABLED = True
         else:
-            ref = _eager_ln_modulate(norm, x, scale, shift)
+            result = _SANA_LN_MOD.accept_or_fallback(
-            if torch.equal(out, ref):
+                out,
-                _SANA_FUSED_LN_MOD_OK_SIGS.add(sig)
+                _eager_ln_modulate(norm, x, scale, shift),
-                return out
+                sig=sig,
-            logger.warning_once(
+                logger=logger,
+                mismatch_msg=(
                     "Sana fused LN+modulate fast path is not bit-exact against "
                     "this platform's LayerNorm dispatch; falling back to eager"
+                ),
             )
-            _SANA_FUSED_LN_MOD_DISABLED = True
+            _SANA_LN_MOD_DISABLED = _SANA_LN_MOD.disabled
-            return ref
+            return result
 
     return _eager_ln_modulate(norm, x, scale, shift)
 

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -85,9 +85,6 @@ from sglang.multimodal_gen.runtime.layers.attention.STA_configuration import (
     configure_sta,
     save_mask_search_results,
 )
-from sglang.multimodal_gen.runtime.loader.component_loaders.transformer_loader import (
-    TransformerLoader,
-)
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
 from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
     ComponentUse,
@@ -125,6 +122,10 @@ from sglang.multimodal_gen.runtime.post_training.rollout_denoising_mixin import
     RolloutDenoisingMixin,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.component_load import (
+    load_transformer_if_needed,
+    register_loaded_transformer,
+)
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.runtime.utils.nvtx_pytorch_hooks import maybe_nvtx_range
 from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler
@@ -886,22 +887,14 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         else:
             cache_dit_num_inference_steps = num_inference_steps
 
-        transformer_was_loaded = server_args.model_loaded["transformer"]
+        freshly_loaded = load_transformer_if_needed(self, server_args)
-        if not transformer_was_loaded:
-            # FIXME: reuse more code
-            loader = TransformerLoader()
-            self.transformer = loader.load(
-                server_args.model_paths["transformer"], server_args, "transformer"
-            )
 
         self._maybe_enable_cache_dit_and_torch_compile(
             cache_dit_num_inference_steps, batch
         )
 
-        if not transformer_was_loaded:
+        if freshly_loaded:
-            if pipeline:
+            register_loaded_transformer(self, server_args, pipeline)
-                pipeline.add_module("transformer", self.transformer)
-            server_args.model_loaded["transformer"] = True
 
         if batch.rollout:
             self._maybe_prepare_rollout(batch)

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/shape.py

@@ -16,9 +16,6 @@ import torch
 from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import (
     Hunyuan3D2PipelineConfig,
 )
-from sglang.multimodal_gen.runtime.loader.component_loaders.transformer_loader import (
-    TransformerLoader,
-)
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
@@ -33,6 +30,10 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
     VerificationResult,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.component_load import (
+    load_transformer_if_needed,
+    register_loaded_transformer,
+)
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.runtime.utils.mesh3d_utils import export_to_trimesh
 
@@ -293,16 +294,11 @@ class Hunyuan3DShapeDenoisingStage(DenoisingStage):
         cache_dit_num_inference_steps = batch.extra.get(
             "cache_dit_num_inference_steps", batch.num_inference_steps
         )
-        if not server_args.model_loaded["transformer"]:
+        freshly_loaded = load_transformer_if_needed(self, server_args)
-            loader = TransformerLoader()
+        if freshly_loaded:
-            self.transformer = loader.load(
-                server_args.model_paths["transformer"], server_args, "transformer"
-            )
             self._maybe_enable_cache_dit(cache_dit_num_inference_steps, batch)
             self._maybe_torch_compile(self.transformer)
-            if pipeline:
+            register_loaded_transformer(self, server_args, pipeline)
-                pipeline.add_module("transformer", self.transformer)
-            server_args.model_loaded["transformer"] = True
         else:
             self._maybe_enable_cache_dit(cache_dit_num_inference_steps, batch)
 

python/sglang/multimodal_gen/runtime/utils/component_load.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Shared helpers for lazily loading pipeline components onto stages."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from sglang.multimodal_gen.runtime.loader.component_loaders.transformer_loader import (
+    TransformerLoader,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+
+def load_transformer_if_needed(stage: Any, server_args: ServerArgs) -> bool:
+    """Load transformer onto stage if not marked loaded. Returns True if freshly loaded."""
+    if server_args.model_loaded["transformer"]:
+        return False
+    stage.transformer = TransformerLoader().load(
+        server_args.model_paths["transformer"], server_args, "transformer"
+    )
+    return True
+
+
+def register_loaded_transformer(
+    stage: Any, server_args: ServerArgs, pipeline: Any
+) -> None:
+    if pipeline is not None:
+        pipeline.add_module("transformer", stage.transformer)
+    server_args.model_loaded["transformer"] = True

python/sglang/multimodal_gen/runtime/utils/precision.py

@@ -4,7 +4,7 @@ from typing import Iterator, Optional, Union
 import torch
 
 from sglang.multimodal_gen.runtime.platforms import current_platform
-from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
+from sglang.multimodal_gen.runtime.utils.precision_types import PRECISION_TO_TYPE
 
 
 def precision_to_dtype(precision: str, field_name: str = "precision") -> torch.dtype:

python/sglang/multimodal_gen/runtime/utils/precision_types.py

@@ -0,0 +1,11 @@
+"""Canonical string→dtype map for diffusion precision configs."""
+
+import torch
+
+PRECISION_TO_TYPE = {
+    "fp32": torch.float32,
+    "fp16": torch.float16,
+    "bf16": torch.bfloat16,
+}
+
+__all__ = ["PRECISION_TO_TYPE"]

python/sglang/multimodal_gen/utils.py

@@ -29,6 +29,9 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import (
     SortedHelpFormatter,
     init_logger,
 )
+from sglang.multimodal_gen.runtime.utils.precision_types import (
+    PRECISION_TO_TYPE as PRECISION_TO_TYPE,
+)
 
 logger = init_logger(__name__)
 
@@ -52,14 +55,6 @@ def expand_path_fields(obj) -> None:
             )
 
 
-# TODO(will): used to convert server_args.precision to torch.dtype. Find a
-# cleaner way to do this.
-PRECISION_TO_TYPE = {
-    "fp32": torch.float32,
-    "fp16": torch.float16,
-    "bf16": torch.bfloat16,
-}
-
 STR_BACKEND_ENV_VAR: str = "SGLANG_DIFFUSION_ATTENTION_BACKEND"
 STR_ATTN_CONFIG_ENV_VAR: str = "SGLANG_DIFFUSION_ATTENTION_CONFIG"
 

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

@@ -0,0 +1,80 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion.bitexact_gate import (
+    BitExactFusionGate,
+    tensors_equal,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+
+def test_bitexact_gate_once_mode_verifies_then_reuses():
+    gate = BitExactFusionGate("once")
+    calls = {"fused": 0, "ref": 0}
+
+    def fused():
+        calls["fused"] += 1
+        return torch.tensor([1.0])
+
+    def ref():
+        calls["ref"] += 1
+        return torch.tensor([1.0])
+
+    assert torch.equal(gate.accept_or_fallback(fused(), ref()), torch.tensor([1.0]))
+    assert gate.verified and not gate.disabled and calls == {"fused": 1, "ref": 1}
+    assert torch.equal(fused(), torch.tensor([1.0]))
+    assert calls == {"fused": 2, "ref": 1}
+
+
+def test_bitexact_gate_mismatch_disables_permanently():
+    gate = BitExactFusionGate("mismatch")
+
+    out = gate.accept_or_fallback(
+        torch.tensor([1.0]),
+        torch.tensor([2.0]),
+        mismatch_msg="mismatch",
+    )
+    assert torch.equal(out, torch.tensor([2.0]))
+    assert gate.disabled and not gate.verified
+
+
+def test_bitexact_gate_per_signature_tracks_each_sig():
+    gate = BitExactFusionGate("sig", per_signature=True)
+    a = torch.tensor([1.0])
+    assert torch.equal(gate.accept_or_fallback(a, a, sig=("a",)), a)
+    assert gate.is_verified(("a",))
+    assert not gate.is_verified(("b",))
+    assert torch.equal(gate.accept_or_fallback(a, a, sig=("b",)), a)
+    assert gate.verified_sigs == {("a",), ("b",)}
+
+
+def test_bitexact_gate_skips_first_sight_during_graph_capture(monkeypatch):
+    # Negative-branch contract: an unverified gate must not attempt first-sight
+    # verification inside CUDA graph capture — the eager-reference host sync
+    # would abort the capture (and BCG would permanently block the signature).
+    gate = BitExactFusionGate("capture")
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
+    assert not gate.can_attempt_once()
+    # A verified gate replays the fused kernel alone, which is capture-safe.
+    gate.mark_verified()
+    assert gate.can_attempt_once()
+
+
+def test_tensors_equal_supports_sequences():
+    assert tensors_equal(
+        (torch.tensor([1.0]), torch.tensor([2.0])),
+        (torch.tensor([1.0]), torch.tensor([2.0])),
+    )
+    assert not tensors_equal(
+        (torch.tensor([1.0]), torch.tensor([2.0])),
+        (torch.tensor([1.0]), torch.tensor([3.0])),
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

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

@@ -47,10 +47,10 @@ def test_fused_norm_scale_shift_is_bit_exact(shape):
         assert torch.equal(out2, ref2)
 
     # the fast paths must actually be in use (not silently disabled)
-    assert ernie_image._ERNIE_FUSED_NORM_VERIFIED
+    assert ernie_image._ERNIE_NORM.verified
-    assert ernie_image._ERNIE_FUSED_GATED_NORM_VERIFIED
+    assert ernie_image._ERNIE_GATED_NORM.verified
-    assert not ernie_image._ERNIE_FUSED_NORM_DISABLED
+    assert not ernie_image._ERNIE_NORM.disabled
-    assert not ernie_image._ERNIE_FUSED_GATED_NORM_DISABLED
+    assert not ernie_image._ERNIE_GATED_NORM.disabled
 
 
 if __name__ == "__main__":

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

@@ -48,8 +48,8 @@ def test_flux_fused_ln_modulate_is_bit_exact(shape, chunks):
     out = _flux_fused_ln_modulate(norm, x, scale, shift)
     assert out is not None
     assert torch.equal(out, _eager(norm, x, scale, shift))
-    assert not flux._FLUX_FUSED_LN_MOD_DISABLED
+    assert not flux._FLUX_LN_MOD.disabled
-    assert flux._FLUX_FUSED_LN_MOD_VERIFIED
+    assert flux._FLUX_LN_MOD.verified
 
 
 def test_flux_norm_modulate_bitexact_supersedes_high_fold():

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

@@ -28,8 +28,8 @@ def test_fused_ln_modulate_is_bit_exact(shape):
     shift, scale = chunks[0], chunks[2]
     out = _glm_ln_modulate(norm, x, scale, shift, x.dtype)
     assert torch.equal(out, _eager_ln_modulate(norm, x, scale, shift, x.dtype))
-    assert glm_image._GLM_FUSED_LN_MOD_VERIFIED
+    assert glm_image._GLM_LN_MOD.verified
-    assert not glm_image._GLM_FUSED_LN_MOD_DISABLED
+    assert not glm_image._GLM_LN_MOD.disabled
 
 
 @pytest.mark.parametrize("shape", [(1, 4360, 32, 128), (2, 37, 3, 40), (1, 129, 5, 64)])
@@ -45,8 +45,8 @@ def test_fused_qk_head_layernorm_is_bit_exact(shape):
     q_out, k_out = _glm_qk_layernorm(norm_q, norm_k, q, k, q.dtype)
     assert torch.equal(q_out, norm_q(q).to(q.dtype))
     assert torch.equal(k_out, norm_k(k).to(k.dtype))
-    assert glm_image._GLM_FUSED_QK_LN_VERIFIED
+    assert glm_image._GLM_QK_LN.verified
-    assert not glm_image._GLM_FUSED_QK_LN_DISABLED
+    assert not glm_image._GLM_QK_LN.disabled
 
 
 if __name__ == "__main__":

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

@@ -36,17 +36,17 @@ def test_sana_fused_ln_modulate_is_bit_exact(shape, nmod, transposed):
     emb = torch.randn(batch, nmod, hidden, device="cuda").bfloat16()
     shift, scale = emb.chunk(nmod, dim=1)[0], emb.chunk(nmod, dim=1)[-1]
     # default-stream eager serving must stay on the untouched eager chain
-    n_sigs = len(sana._SANA_FUSED_LN_MOD_OK_SIGS)
+    n_sigs = len(sana._SANA_LN_MOD.verified_sigs)
     _sana_ln_modulate(norm, x, scale, shift)
-    assert len(sana._SANA_FUSED_LN_MOD_OK_SIGS) == n_sigs
+    assert len(sana._SANA_LN_MOD.verified_sigs) == n_sigs
     # the fusion engages on non-default streams (the BCG warmup/capture path)
     with torch.cuda.stream(torch.cuda.Stream()):
         out = _sana_ln_modulate(norm, x, scale, shift)
-        assert len(sana._SANA_FUSED_LN_MOD_OK_SIGS) == n_sigs + 1  # verified
+        assert len(sana._SANA_LN_MOD.verified_sigs) == n_sigs + 1  # verified
         out2 = _sana_ln_modulate(norm, x, scale, shift)  # verified-sig lane
     torch.cuda.synchronize()
     assert torch.equal(out, _eager_ln_modulate(norm, x, scale, shift))
-    assert torch.equal(out2, out) and not sana._SANA_FUSED_LN_MOD_DISABLED
+    assert torch.equal(out2, out) and not sana._SANA_LN_MOD.disabled
 
 
 if __name__ == "__main__":