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

[Diffusion] Fuse Qwen-Image final adaptive LayerNorm (#37144)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-08-31 18:22:20 +08:00
committed by GitHub
parent 1ed9bfac2c
commit 771e613d96
3 changed files with 195 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
+71 -1
View File
@@ -16,9 +16,13 @@ from diffusers.models.modeling_outputs import Transformer2DModelOutput
from diffusers.models.normalization import AdaLayerNormContinuous
from sglang.kernels.ops.diffusion import (
BitExactFusionGate,
can_use_fused_layernorm_modulate,
can_use_linear_gelu,
fused_gelu_active,
fused_layernorm_modulate_raw,
fused_linear_gelu_tanh,
is_plain_layer_norm,
mark_fused_gelu_site,
try_fused_bias_mul_add,
try_fused_bias_scale_residual_norm_scale_shift,
@@ -93,6 +97,72 @@ from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import
)
logger = init_logger(__name__) # pylint: disable=invalid-name
_QWEN_NORM_OUT = BitExactFusionGate("Qwen-Image fused norm_out", per_signature=True)
_QWEN_NORM_OUT_SIGS = _QWEN_NORM_OUT.verified_sigs
assert _QWEN_NORM_OUT_SIGS is not None
def _qwen_norm_out(
norm_out: AdaLayerNormContinuous,
hidden_states: torch.Tensor,
conditioning_embedding: torch.Tensor,
) -> torch.Tensor:
"""Bit-exact final LayerNorm and adaLN scale/shift fusion."""
if torch.compiler.is_compiling():
return norm_out(hidden_states, conditioning_embedding)
# Keep the projected modulation available for direct eager dispatch. The
# public custom-op wrapper costs more than this small final norm itself.
emb = norm_out.linear(norm_out.silu(conditioning_embedding).to(hidden_states.dtype))
scale, shift = torch.chunk(emb, 2, dim=1)
if (
_QWEN_NORM_OUT.disabled
or not is_plain_layer_norm(norm_out.norm, hidden_states.shape[-1])
or not can_use_fused_layernorm_modulate(hidden_states, scale, shift)
):
return (
norm_out.norm(hidden_states) * (1 + scale)[:, None, :] + shift[:, None, :]
)
sig = (
hidden_states.dtype,
hidden_states.device,
hidden_states.shape[0],
hidden_states.shape[-1],
hidden_states.stride(-1),
scale.stride(0) if scale.shape[0] > 1 else hidden_states.shape[-1],
shift.stride(0) if shift.shape[0] > 1 else hidden_states.shape[-1],
norm_out.norm.eps,
)
verified = sig in _QWEN_NORM_OUT_SIGS
if not verified and torch.cuda.is_current_stream_capturing():
return (
norm_out.norm(hidden_states) * (1 + scale)[:, None, :] + shift[:, None, :]
)
try:
fused = fused_layernorm_modulate_raw(
hidden_states, scale, shift, norm_out.norm.eps
)
except Exception as exc:
_QWEN_NORM_OUT.on_exception(exc, logger=logger)
return (
norm_out.norm(hidden_states) * (1 + scale)[:, None, :] + shift[:, None, :]
)
if verified:
return fused
reference = (
norm_out.norm(hidden_states) * (1 + scale)[:, None, :] + shift[:, None, :]
)
return _QWEN_NORM_OUT.accept_or_fallback(
fused,
reference,
sig=sig,
logger=logger,
mismatch_msg=(
"Qwen-Image fused norm_out is not bit-exact on this platform; "
"falling back to eager"
),
)
def _attn_mask_meta_local_pad(attn_mask_meta) -> int:
@@ -2069,7 +2139,7 @@ class QwenImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
+ controlnet_block_samples[index_block // interval_control]
)
# Use only the image part (hidden_states) from the dual-stream blocks
hidden_states = self.norm_out(hidden_states, temb_txt)
hidden_states = _qwen_norm_out(self.norm_out, hidden_states, temb_txt)
output, _ = self.proj_out(hidden_states)
return output
test/registered/kernels/benchmark/diffusion/bench_qwen_image_modulation.py
+47 -1
View File
@@ -7,20 +7,23 @@ from sglang.kernels.jit.benchmark.utils import run_benchmark_no_cudagraph
from sglang.kernels.ops.diffusion import (
fuse_layernorm_scale_shift_gate_select01_kernel,
fuse_residual_layernorm_scale_shift_gate_select01_kernel,
fused_layernorm_modulate_raw,
norm_infer,
)
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.utils import is_in_ci
register_cuda_ci(
est_time=13, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
est_time=15, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
)
register_amd_ci(est_time=13, stage="jit-kernel-benchmark", runner_config="amd")
if is_in_ci():
B_RANGE, S_RANGE, D_RANGE = [1], [128], [3072]
NORM_OUT_S_RANGE = [128]
else:
B_RANGE, S_RANGE, D_RANGE = [1, 2], [128, 512, 2048], [1024, 1536, 3072]
NORM_OUT_S_RANGE = [128, 512, 2048, 4096, 4608]
DTYPE = torch.bfloat16
DEVICE = "cuda"
@@ -29,6 +32,7 @@ LINE_VALS = ["split", "fused"]
LINE_NAMES = ["Triton Norm + Torch Select", "Fused Triton"]
STYLES = [("red", "-"), ("blue", "--")]
CONFIG = [(b, s, d) for b in B_RANGE for s in S_RANGE for d in D_RANGE]
NORM_OUT_CONFIG = [(1, s, 3072) for s in NORM_OUT_S_RANGE]
def _make_common_inputs(batch_size: int, seq_len: int, hidden_size: int):
@@ -62,6 +66,40 @@ def _apply_select01_modulation(
return x * (1 + scale) + shift, gate
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["B", "S", "D"],
x_vals=NORM_OUT_CONFIG,
line_arg="provider",
line_vals=LINE_VALS,
line_names=["Torch LayerNorm + Modulate", "Fused Triton"],
styles=STYLES,
ylabel="us",
plot_name="qwen_image_norm_out",
args={},
)
)
def bench_norm_out(B: int, S: int, D: int, provider: str) -> Tuple[float, float, float]:
x = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
scale = torch.randn(B, D, dtype=DTYPE, device=DEVICE)
shift = torch.randn(B, D, dtype=DTYPE, device=DEVICE)
norm = torch.nn.LayerNorm(
D, eps=EPS, elementwise_affine=False, device=DEVICE, dtype=DTYPE
)
if provider == "split":
def fn():
return norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
else:
def fn():
return fused_layernorm_modulate_raw(x, scale, shift, EPS)
return run_benchmark_no_cudagraph(fn)
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["B", "S", "D"],
@@ -176,6 +214,14 @@ def bench_residual_layernorm_scale_shift_gate_select01(
if __name__ == "__main__":
# The bit-exact LayerNorm kernel uses CUDA inline PTX; keep the shared
# Qwen modulation benchmark runnable on its registered ROCm lane.
if torch.version.hip is None:
print(f"\n{'=' * 80}")
print("Benchmark: qwen_image norm_out")
print(f"{'=' * 80}\n")
bench_norm_out.run(print_data=True)
print(f"\n{'=' * 80}")
print("Benchmark: qwen_image layernorm + scale_shift_gate_select01")
print(f"{'=' * 80}\n")
test/registered/kernels/ops/diffusion/test_model_fast_paths.py
+77
View File
@@ -36,6 +36,7 @@ import sglang.multimodal_gen.runtime.models.dits.flux_2 as flux2
import sglang.multimodal_gen.runtime.models.dits.glm_image as glm_image
import sglang.multimodal_gen.runtime.models.dits.longcat_image as longcat_image
import sglang.multimodal_gen.runtime.models.dits.ltx_2 as ltx2_module
import sglang.multimodal_gen.runtime.models.dits.qwen_image as qwen_image
import sglang.multimodal_gen.runtime.models.dits.sana as sana
from sglang.kernels.ops.diffusion import (
can_use_fused_layernorm_modulate,
@@ -98,6 +99,7 @@ from sglang.multimodal_gen.runtime.models.dits.ltx_2 import _ltx2_rms_norm_modul
from sglang.multimodal_gen.runtime.models.dits.qwen_image import (
QwenImageTransformerBlock,
_qwen_modulation_cache_key,
_qwen_norm_out,
)
from sglang.multimodal_gen.runtime.models.dits.sana import (
_eager_ln_modulate as _sana_eager_ln_modulate,
@@ -151,6 +153,81 @@ def test_bitexact_norm_guards_follow_platform():
assert can_use_fused_rmsnorm_scale_shift(x, weight, vec, vec) is is_cuda()
# -------------------------------------------------------------------------
# Qwen-Image -- final LayerNorm + adaLN scale/shift
# -------------------------------------------------------------------------
@requires_inline_ptx
def test_qwen_norm_out_matches_adaln_reference():
qwen_image._QWEN_NORM_OUT.disabled = False
qwen_image._QWEN_NORM_OUT.verified = False
qwen_image._QWEN_NORM_OUT_SIGS.clear()
torch.manual_seed(0)
norm_out = (
qwen_image.AdaLayerNormContinuous(
3072, 3072, elementwise_affine=False, eps=1e-6
)
.cuda()
.bfloat16()
)
hidden_states = torch.randn(1, 257, 3072, device="cuda", dtype=torch.bfloat16)
conditioning = torch.randn(1, 3072, device="cuda", dtype=torch.bfloat16)
expected = norm_out(hidden_states, conditioning)
actual = _qwen_norm_out(norm_out, hidden_states, conditioning)
assert torch.equal(actual, expected)
assert qwen_image._QWEN_NORM_OUT.verified
assert not qwen_image._QWEN_NORM_OUT.disabled
def test_qwen_norm_out_preserves_compile_path(monkeypatch):
norm_out = (
qwen_image.AdaLayerNormContinuous(16, 16, elementwise_affine=False, eps=1e-6)
.cuda()
.bfloat16()
)
hidden_states = torch.randn(1, 3, 16, device="cuda", dtype=torch.bfloat16)
conditioning = torch.randn(1, 16, device="cuda", dtype=torch.bfloat16)
expected = norm_out(hidden_states, conditioning)
monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
monkeypatch.setattr(
qwen_image,
"fused_layernorm_modulate_raw",
lambda *args, **kwargs: pytest.fail("compile path must not dispatch kernel"),
)
assert torch.equal(_qwen_norm_out(norm_out, hidden_states, conditioning), expected)
def test_qwen_norm_out_does_not_verify_during_graph_capture(monkeypatch):
qwen_image._QWEN_NORM_OUT.disabled = False
qwen_image._QWEN_NORM_OUT.verified = False
qwen_image._QWEN_NORM_OUT_SIGS.clear()
norm_out = (
qwen_image.AdaLayerNormContinuous(
3072, 3072, elementwise_affine=False, eps=1e-6
)
.cuda()
.bfloat16()
)
hidden_states = torch.randn(1, 17, 3072, device="cuda", dtype=torch.bfloat16)
conditioning = torch.randn(1, 3072, device="cuda", dtype=torch.bfloat16)
expected = norm_out(hidden_states, conditioning)
monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
monkeypatch.setattr(
qwen_image,
"fused_layernorm_modulate_raw",
lambda *args, **kwargs: pytest.fail("capture must not verify a new layout"),
)
assert torch.equal(_qwen_norm_out(norm_out, hidden_states, conditioning), expected)
assert not qwen_image._QWEN_NORM_OUT_SIGS
# -------------------------------------------------------------------------
# FLUX.1 -- LayerNorm + adaLN modulate, and the shared-FF GELU site
# -------------------------------------------------------------------------