From fd3036523a1af7ba0775bc785d8063f4d7ee5094 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Mon, 10 Aug 2026 22:14:09 +0800 Subject: [PATCH] [diffusion] Clean up shared bitexact gates, helpers, and stale naming (#34180) Co-authored-by: Cursor Co-authored-by: Claude Fable 5 --- .../jit/csrc/diffusion/norm_scale_shift.cuh | 24 ++-- .../sglang/kernels/ops/diffusion/__init__.py | 25 ++-- .../kernels/ops/diffusion/bitexact_gate.py | 118 ++++++++++++++++++ .../ops/diffusion/norm_scale_shift_native.py | 22 ++-- .../benchmark-and-profile.md | 2 +- .../existing-fast-paths.md | 31 +++-- .../runtime/models/dits/common.py | 47 +++++++ .../runtime/models/dits/ernie_image.py | 72 +++++------ .../runtime/models/dits/flux.py | 65 +++------- .../runtime/models/dits/flux_2.py | 27 +--- .../runtime/models/dits/glm_image.py | 82 ++++++------ .../runtime/models/dits/qwen_image.py | 28 +---- .../runtime/models/dits/sana.py | 41 +++--- .../pipelines_core/stages/denoising.py | 21 ++-- .../model_specific_stages/hunyuan3d/shape.py | 18 ++- .../runtime/utils/component_load.py | 29 +++++ .../multimodal_gen/runtime/utils/precision.py | 2 +- .../runtime/utils/precision_types.py | 11 ++ python/sglang/multimodal_gen/utils.py | 11 +- .../ops/diffusion/test_bitexact_gate.py | 80 ++++++++++++ .../diffusion/test_ernie_norm_scale_shift.py | 8 +- .../ops/diffusion/test_flux_ln_modulate.py | 4 +- .../diffusion/test_glm_image_ln_modulate.py | 8 +- .../ops/diffusion/test_sana_ln_modulate.py | 8 +- 24 files changed, 496 insertions(+), 288 deletions(-) create mode 100644 python/sglang/kernels/ops/diffusion/bitexact_gate.py create mode 100644 python/sglang/multimodal_gen/runtime/utils/component_load.py create mode 100644 python/sglang/multimodal_gen/runtime/utils/precision_types.py create mode 100644 test/registered/kernels/ops/diffusion/test_bitexact_gate.py diff --git a/python/sglang/kernels/jit/csrc/diffusion/norm_scale_shift.cuh b/python/sglang/kernels/jit/csrc/diffusion/norm_scale_shift.cuh index bbc8be168..ec3673203 100644 --- a/python/sglang/kernels/jit/csrc/diffusion/norm_scale_shift.cuh +++ b/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 -__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; @@ -135,14 +135,14 @@ __global__ void qwen_image_norm_scale_shift_kernel(const QwenImageNormParams __g yv.store(static_cast(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(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().with_device(device).verify(x).verify(y); TensorMatcher({kHidden}).with_dtype().with_device(device).verify(scale).verify(shift); - const uint32_t grid = verify_qwen_geometry(N); - const auto params = QwenImageNormParams{ + const uint32_t grid = verify_nss_geometry(N); + 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(eps), }; - LaunchKernel(grid, kThreads, device.unwrap())(qwen_image_norm_scale_shift_kernel, params); + LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel, 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().with_device(device).verify(gate).verify(scale).verify(shift); - const uint32_t grid = verify_qwen_geometry(N); - const auto params = QwenImageNormParams{ + const uint32_t grid = verify_nss_geometry(N); + 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(eps), }; - LaunchKernel(grid, kThreads, device.unwrap())(qwen_image_norm_scale_shift_kernel, params); + LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel, params); } }; diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index 6c97dd8aa..02c56db51 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/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", - ) -) diff --git a/python/sglang/kernels/ops/diffusion/bitexact_gate.py b/python/sglang/kernels/ops/diffusion/bitexact_gate.py new file mode 100644 index 000000000..d0ac7badd --- /dev/null +++ b/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)) diff --git a/python/sglang/kernels/ops/diffusion/norm_scale_shift_native.py b/python/sglang/kernels/ops/diffusion/norm_scale_shift_native.py index 30e58766b..5b1712112 100644 --- a/python/sglang/kernels/ops/diffusion/norm_scale_shift_native.py +++ b/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", - "norm_scale_shift::QwenImageNormScaleShiftKernel::run", + "nss_bf16_row", + "norm_scale_shift::NormScaleShiftKernel::run", ), ( - "qwen_image_srnss_bf16_row", - "norm_scale_shift::" "QwenImageScaleResidualNormScaleShiftKernel::run", + "srnss_bf16_row", + "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) - and _qwen_activation(residual, x) + _nss_activation(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), diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md index 2a1aba3d9..98df19ca5 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md +++ b/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 | diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md index 7f8eee21a..c558acf51 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md +++ b/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` -- Locations: `triton/native_bf16_rmsnorm.py`, compatibility exports in - `triton/zimage_native_norm.py`, and `zimage.py` +- Kernels: `rmsnorm_scale`, `rmsnorm_tanh_residual` +- Locations: `triton/native_bf16_rmsnorm.py`, with wrappers in `zimage.py` and + `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 - `zimage_rmsnorm_tanh_mul_add` in `zimage.py`, backed by the shared - `triton/native_bf16_rmsnorm.py` kernels. +- Z-Image native norm modulation: `rmsnorm_scale` and `rmsnorm_tanh_residual` + in `triton/native_bf16_rmsnorm.py`, with wrappers in `zimage.py` / + `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`. diff --git a/python/sglang/multimodal_gen/runtime/models/dits/common.py b/python/sglang/multimodal_gen/runtime/models/dits/common.py index 4f98ef813..7c05fe98b 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/common.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py b/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py index dce4aa256..f3db9c5c9 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py +++ b/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_FUSED_NORM_VERIFIED = False -_ERNIE_FUSED_GATED_NORM_DISABLED = False -_ERNIE_FUSED_GATED_NORM_VERIFIED = False +_ERNIE_NORM = BitExactFusionGate("ERNIE fused-norm") +_ERNIE_GATED_NORM = BitExactFusionGate("ERNIE fused gated-norm") 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(): - raise - logger.warning_once(f"Disabling ERNIE fused-norm fast path: {exc}") - _ERNIE_FUSED_NORM_DISABLED = True + _ERNIE_NORM.on_exception(exc, logger=logger) else: - if _ERNIE_FUSED_NORM_VERIFIED: + if verified: return out - ref = _eager_norm_scale_shift(norm, x, scale, shift) - if torch.equal(out, ref): - _ERNIE_FUSED_NORM_VERIFIED = True - return out - logger.warning_once( - "ERNIE fused-norm fast path is not bit-exact against this " - "platform's rmsnorm dispatch; falling back to eager" + return _ERNIE_NORM.accept_or_fallback( + out, + _eager_norm_scale_shift(norm, x, scale, shift), + logger=logger, + 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(): - raise - logger.warning_once(f"Disabling ERNIE fused gated-norm fast path: {exc}") - _ERNIE_FUSED_GATED_NORM_DISABLED = True + _ERNIE_GATED_NORM.on_exception(exc, logger=logger) 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): - _ERNIE_FUSED_GATED_NORM_VERIFIED = True - return out, res - logger.warning_once( - "ERNIE fused gated-norm fast path is not bit-exact against " - "this platform's rmsnorm dispatch; falling back to eager" + return _ERNIE_GATED_NORM.accept_or_fallback( + (out, res), + (ref, res_ref), + 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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux.py b/python/sglang/multimodal_gen/runtime/models/dits/flux.py index 12d671532..0035ca821 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux.py +++ b/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 -# (shape, stride, eps) signatures whose fused output has been verified -# ``torch.equal`` against the live eager chain. -_FLUX_FUSED_LN_MOD_VERIFIED: set = set() +_FLUX_LN_MOD = BitExactFusionGate("FLUX fused LN+modulate", per_signature=True) +# Keep the pre-refactor direct set lookup in this launch-sensitive hot path. +_FLUX_LN_MOD_SIGS = _FLUX_LN_MOD.verified_sigs +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(): - raise - logger.warning_once(f"Disabling FLUX fused LN+modulate fast path: {exc}") - _FLUX_FUSED_LN_MOD_DISABLED = True + _FLUX_LN_MOD.on_exception(exc, logger=logger) return None if verified: return out ref = modulate_scale_shift(norm(x), scale, shift) - if torch.equal(out, ref): - _FLUX_FUSED_LN_MOD_VERIFIED.add(sig) - return out - logger.warning_once( - "FLUX fused LN+modulate fast path is not bit-exact against this " - "platform's LayerNorm dispatch; falling back to eager" + return _FLUX_LN_MOD.accept_or_fallback( + out, + ref, + 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, diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py index ddf9e5c44..d3d6a7432 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/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 - -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 +_get_qkv_projections = get_qkv_projections class Flux2SwiGLU(nn.Module): diff --git a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py index b26065c62..a9bb5a919 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py +++ b/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_FUSED_LN_MOD_VERIFIED = False -_GLM_FUSED_QK_LN_DISABLED = False -_GLM_FUSED_QK_LN_VERIFIED = False +_GLM_LN_MOD = BitExactFusionGate("GLM fused LN+modulate") +_GLM_QK_LN = BitExactFusionGate("GLM fused qk-LayerNorm") 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(): - raise - logger.warning_once(f"Disabling GLM fused LN+modulate fast path: {exc}") - _GLM_FUSED_LN_MOD_DISABLED = True + _GLM_LN_MOD.on_exception(exc, logger=logger) else: - if _GLM_FUSED_LN_MOD_VERIFIED: + if verified: return out - ref = _eager_ln_modulate(norm, x, scale, shift, dtype) - if torch.equal(out, ref): - _GLM_FUSED_LN_MOD_VERIFIED = True - return out - logger.warning_once( - "GLM fused LN+modulate fast path is not bit-exact against " - "this platform's LayerNorm dispatch; falling back to eager" + return _GLM_LN_MOD.accept_or_fallback( + out, + _eager_ln_modulate(norm, x, scale, shift, dtype), + logger=logger, + 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(): - raise - logger.warning_once(f"Disabling GLM fused qk-LayerNorm fast path: {exc}") - _GLM_FUSED_QK_LN_DISABLED = True + _GLM_QK_LN.on_exception(exc, logger=logger) else: - if _GLM_FUSED_QK_LN_VERIFIED: - return q_out, k_out - q_ref = norm_q(query).to(dtype=dtype) - k_ref = norm_k(key).to(dtype=dtype) - if torch.equal(q_out, q_ref) and torch.equal(k_out, k_ref): - _GLM_FUSED_QK_LN_VERIFIED = True - return q_out, k_out - logger.warning_once( - "GLM fused qk-LayerNorm fast path is not bit-exact against " - "this platform's LayerNorm dispatch; falling back to eager" + if verified: + return out + ref = ( + norm_q(query).to(dtype=dtype), + norm_k(key).to(dtype=dtype), + ) + return _GLM_QK_LN.accept_or_fallback( + 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) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 6a1d68a17..56ae6bc9e 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/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( - 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 +_get_qkv_projections = get_qkv_projections class QwenTimestepProjEmbeddings(nn.Module): diff --git a/python/sglang/multimodal_gen/runtime/models/dits/sana.py b/python/sglang/multimodal_gen/runtime/models/dits/sana.py index 98c3a693f..303e561ef 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/sana.py +++ b/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 -# (shape/stride/dtype/eps) signatures whose fused output was torch.equal- -# verified against the live eager chain. -_SANA_FUSED_LN_MOD_OK_SIGS: set = set() +_SANA_LN_MOD = BitExactFusionGate("Sana fused LN+modulate", per_signature=True) +# Direct module-level state keeps BCG warmup launch overhead equal to the +# pre-refactor path; the gate still owns first-sight verification transitions. +_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_FUSED_LN_MOD_DISABLED = True + _SANA_LN_MOD.on_exception(exc, logger=logger) + _SANA_LN_MOD_DISABLED = True else: - ref = _eager_ln_modulate(norm, x, scale, shift) - if torch.equal(out, ref): - _SANA_FUSED_LN_MOD_OK_SIGS.add(sig) - return out - logger.warning_once( - "Sana fused LN+modulate fast path is not bit-exact against " - "this platform's LayerNorm dispatch; falling back to eager" + result = _SANA_LN_MOD.accept_or_fallback( + out, + _eager_ln_modulate(norm, x, scale, shift), + sig=sig, + 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 - return ref + _SANA_LN_MOD_DISABLED = _SANA_LN_MOD.disabled + return result return _eager_ln_modulate(norm, x, scale, shift) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 9e59fe530..891a9c533 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/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"] - if not transformer_was_loaded: - # FIXME: reuse more code - loader = TransformerLoader() - self.transformer = loader.load( - server_args.model_paths["transformer"], server_args, "transformer" - ) + freshly_loaded = load_transformer_if_needed(self, server_args) self._maybe_enable_cache_dit_and_torch_compile( cache_dit_num_inference_steps, batch ) - if not transformer_was_loaded: - if pipeline: - pipeline.add_module("transformer", self.transformer) - server_args.model_loaded["transformer"] = True + if freshly_loaded: + register_loaded_transformer(self, server_args, pipeline) if batch.rollout: self._maybe_prepare_rollout(batch) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/shape.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/shape.py index 35961d2a8..7eabec492 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/shape.py +++ b/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"]: - loader = TransformerLoader() - self.transformer = loader.load( - server_args.model_paths["transformer"], server_args, "transformer" - ) + freshly_loaded = load_transformer_if_needed(self, server_args) + if freshly_loaded: self._maybe_enable_cache_dit(cache_dit_num_inference_steps, batch) self._maybe_torch_compile(self.transformer) - if pipeline: - pipeline.add_module("transformer", self.transformer) - server_args.model_loaded["transformer"] = True + register_loaded_transformer(self, server_args, pipeline) else: self._maybe_enable_cache_dit(cache_dit_num_inference_steps, batch) diff --git a/python/sglang/multimodal_gen/runtime/utils/component_load.py b/python/sglang/multimodal_gen/runtime/utils/component_load.py new file mode 100644 index 000000000..84cbcc918 --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/utils/precision.py b/python/sglang/multimodal_gen/runtime/utils/precision.py index 238dacdf5..6ac31c787 100644 --- a/python/sglang/multimodal_gen/runtime/utils/precision.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/utils/precision_types.py b/python/sglang/multimodal_gen/runtime/utils/precision_types.py new file mode 100644 index 000000000..84a3a581c --- /dev/null +++ b/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"] diff --git a/python/sglang/multimodal_gen/utils.py b/python/sglang/multimodal_gen/utils.py index df543fdf4..6729138ee 100644 --- a/python/sglang/multimodal_gen/utils.py +++ b/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" diff --git a/test/registered/kernels/ops/diffusion/test_bitexact_gate.py b/test/registered/kernels/ops/diffusion/test_bitexact_gate.py new file mode 100644 index 000000000..aa30fcfa7 --- /dev/null +++ b/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"])) diff --git a/test/registered/kernels/ops/diffusion/test_ernie_norm_scale_shift.py b/test/registered/kernels/ops/diffusion/test_ernie_norm_scale_shift.py index aeca75111..703ed0794 100644 --- a/test/registered/kernels/ops/diffusion/test_ernie_norm_scale_shift.py +++ b/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_FUSED_GATED_NORM_VERIFIED - assert not ernie_image._ERNIE_FUSED_NORM_DISABLED - assert not ernie_image._ERNIE_FUSED_GATED_NORM_DISABLED + assert ernie_image._ERNIE_NORM.verified + assert ernie_image._ERNIE_GATED_NORM.verified + assert not ernie_image._ERNIE_NORM.disabled + assert not ernie_image._ERNIE_GATED_NORM.disabled if __name__ == "__main__": diff --git a/test/registered/kernels/ops/diffusion/test_flux_ln_modulate.py b/test/registered/kernels/ops/diffusion/test_flux_ln_modulate.py index 867591212..7e42c7e31 100644 --- a/test/registered/kernels/ops/diffusion/test_flux_ln_modulate.py +++ b/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 flux._FLUX_FUSED_LN_MOD_VERIFIED + assert not flux._FLUX_LN_MOD.disabled + assert flux._FLUX_LN_MOD.verified def test_flux_norm_modulate_bitexact_supersedes_high_fold(): diff --git a/test/registered/kernels/ops/diffusion/test_glm_image_ln_modulate.py b/test/registered/kernels/ops/diffusion/test_glm_image_ln_modulate.py index d2bb151d6..a3d177218 100644 --- a/test/registered/kernels/ops/diffusion/test_glm_image_ln_modulate.py +++ b/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 not glm_image._GLM_FUSED_LN_MOD_DISABLED + assert glm_image._GLM_LN_MOD.verified + assert not glm_image._GLM_LN_MOD.disabled @pytest.mark.parametrize("shape", [(1, 4360, 32, 128), (2, 37, 3, 40), (1, 129, 5, 64)]) @@ -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 not glm_image._GLM_FUSED_QK_LN_DISABLED + assert glm_image._GLM_QK_LN.verified + assert not glm_image._GLM_QK_LN.disabled if __name__ == "__main__": diff --git a/test/registered/kernels/ops/diffusion/test_sana_ln_modulate.py b/test/registered/kernels/ops/diffusion/test_sana_ln_modulate.py index eb0a01eac..36405ac8c 100644 --- a/test/registered/kernels/ops/diffusion/test_sana_ln_modulate.py +++ b/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__":