From 771e613d96de0ee89631bc308a2525aaeae9f13e Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Mon, 31 Aug 2026 18:22:20 +0800 Subject: [PATCH] [Diffusion] Fuse Qwen-Image final adaptive LayerNorm (#37144) --- .../runtime/models/dits/qwen_image.py | 72 ++++++++++++++++- .../diffusion/bench_qwen_image_modulation.py | 48 +++++++++++- .../ops/diffusion/test_model_fast_paths.py | 77 +++++++++++++++++++ 3 files changed, 195 insertions(+), 2 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 4a7eefcbf..2b6fa6345 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -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 diff --git a/test/registered/kernels/benchmark/diffusion/bench_qwen_image_modulation.py b/test/registered/kernels/benchmark/diffusion/bench_qwen_image_modulation.py index 5d7ce04b2..83d00aebd 100644 --- a/test/registered/kernels/benchmark/diffusion/bench_qwen_image_modulation.py +++ b/test/registered/kernels/benchmark/diffusion/bench_qwen_image_modulation.py @@ -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") diff --git a/test/registered/kernels/ops/diffusion/test_model_fast_paths.py b/test/registered/kernels/ops/diffusion/test_model_fast_paths.py index 9e5bd6f5c..2bf36dd64 100644 --- a/test/registered/kernels/ops/diffusion/test_model_fast_paths.py +++ b/test/registered/kernels/ops/diffusion/test_model_fast_paths.py @@ -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 # -------------------------------------------------------------------------