minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
2026-08-10 22:14:09 +08:00
committed by GitHub
co-authored by Cursor Claude Fable 5
parent 77c90e7e54
commit fd3036523a
24 changed files with 496 additions and 288 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/diffusion/norm_scale_shift.cuh
+12 -12
View File
@@ -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 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<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 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<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
+14 -11
View File
@@ -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
+118
View File
@@ -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
+11 -11
View File
@@ -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),
python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md
+1 -1
View File
@@ -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
+23 -8
View File
@@ -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`.
python/sglang/multimodal_gen/runtime/models/dits/common.py
+47
View File
@@ -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
+33 -39
View File
@@ -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
python/sglang/multimodal_gen/runtime/models/dits/flux.py
+19 -46
View File
@@ -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,
python/sglang/multimodal_gen/runtime/models/dits/flux_2.py
+2 -25
View File
@@ -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):
python/sglang/multimodal_gen/runtime/models/dits/glm_image.py
+39 -43
View File
@@ -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)
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
+2 -26
View File
@@ -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):
python/sglang/multimodal_gen/runtime/models/dits/sana.py
+23 -18
View File
@@ -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)
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
+7 -14
View File
@@ -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)
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/shape.py
+7 -11
View File
@@ -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)
python/sglang/multimodal_gen/runtime/utils/component_load.py
+29
View File
@@ -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
+1 -1
View File
@@ -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
+11
View File
@@ -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
+3 -8
View File
@@ -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"