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

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

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