From 1ed9bfac2cd2d0de99271164e8f2a9b3b3a70fea Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Mon, 31 Aug 2026 18:19:42 +0800 Subject: [PATCH] [Diffusion] Fuse Qwen-Image FP8 norm and activation quantization (#37156) Co-authored-by: Cursor --- .../jit/csrc/diffusion/norm_scale_shift.cuh | 125 +++++++- .../sglang/kernels/ops/diffusion/__init__.py | 4 + .../diffusion/norm/norm_scale_shift_jit.py | 112 ++++++++ .../runtime/models/dits/qwen_image.py | 270 +++++++++++++++++- .../bench_qwen_image_norm_fp8_quant.py | 82 ++++++ .../test_qwen_image_norm_fp8_quant.py | 84 ++++++ .../models/test_qwen_image_fp8_norm_quant.py | 101 +++++++ 7 files changed, 761 insertions(+), 17 deletions(-) create mode 100644 test/registered/kernels/benchmark/diffusion/bench_qwen_image_norm_fp8_quant.py create mode 100644 test/registered/kernels/ops/diffusion/test_qwen_image_norm_fp8_quant.py create mode 100644 test/registered/unit/models/test_qwen_image_fp8_norm_quant.py diff --git a/python/sglang/kernels/jit/csrc/diffusion/norm_scale_shift.cuh b/python/sglang/kernels/jit/csrc/diffusion/norm_scale_shift.cuh index 3f19bddf8..e64adef87 100644 --- a/python/sglang/kernels/jit/csrc/diffusion/norm_scale_shift.cuh +++ b/python/sglang/kernels/jit/csrc/diffusion/norm_scale_shift.cuh @@ -17,6 +17,7 @@ #include // For TensorMatcher, SymbolicSize, SymbolicDevice #include // For device::math::rsqrt +#include // For DTypeTrait #include // For SGL_DEVICE, bf16_t, LaunchKernel #include // For AlignedVector #include // For warp::reduce_sum @@ -39,12 +40,14 @@ static_assert(kWarps == 6); struct NormScaleShiftParams { void* y; void* res_out; + void* quantized; const void* x; const void* input_bias; const void* residual; const void* gate; const void* scale; const void* shift; + const void* input_scale; float eps; }; @@ -66,7 +69,16 @@ SGL_DEVICE float cta_reduce_sum(float v, int warp, int lane, float* scratch) { return scratch[kWarps]; } -template +SGL_DEVICE float triton_scale_reciprocal(float scale) { + float reciprocal; + // Triton's static FP8 quantizer lowers `1.0 / scale` to div.full.f32. + // Match it exactly because a one-ULP difference at an E4M3 midpoint can + // change the quantized byte. + asm("div.full.f32 %0, %1, %2;" : "=f"(reciprocal) : "f"(1.0f), "f"(scale)); + return reciprocal; +} + +template __global__ void norm_scale_shift_kernel(const NormScaleShiftParams __grid_constant__ params) { using namespace device; using Vec = AlignedVector; @@ -135,15 +147,31 @@ __global__ void norm_scale_shift_kernel(const NormScaleShiftParams __grid_consta Vec scv; Vec shv; Vec yv; + AlignedVector qv; scv.load(static_cast(params.scale) + elem_offset); shv.load(static_cast(params.shift) + elem_offset); + float input_scale_inv = 0.0f; + if constexpr (kQuantizeFp8) { + input_scale_inv = triton_scale_reciprocal(*static_cast(params.input_scale)); + } + #pragma unroll for (int i = 0; i < kVecElems; ++i) { const float norm = static_cast(static_cast((v[i] - mean) * factor)); - yv[i] = static_cast(norm * (1.0f + static_cast(scv[i])) + static_cast(shv[i])); + const bf16_t rounded = static_cast(norm * (1.0f + static_cast(scv[i])) + static_cast(shv[i])); + yv[i] = rounded; + if constexpr (kQuantizeFp8) { + const float scaled = static_cast(rounded) * input_scale_inv; + const float clamped = + math::min(math::max(scaled, -DTypeTrait::kFloatMax), DTypeTrait::kFloatMax); + qv[i] = static_cast(clamped); + } } yv.store(static_cast(params.y) + row_offset + elem_offset); + if constexpr (kQuantizeFp8) { + qv.store(static_cast(params.quantized) + row_offset + elem_offset); + } } __global__ void bias_mul_add_kernel(const NormScaleShiftParams __grid_constant__ params) { @@ -197,12 +225,14 @@ struct NormScaleShiftKernel { const auto params = NormScaleShiftParams{ .y = y.data_ptr(), .res_out = nullptr, + .quantized = nullptr, .x = x.data_ptr(), .input_bias = nullptr, .residual = nullptr, .gate = nullptr, .scale = scale.data_ptr(), .shift = shift.data_ptr(), + .input_scale = nullptr, .eps = static_cast(eps), }; LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel, params); @@ -237,18 +267,105 @@ struct ScaleResidualNormScaleShiftKernel { const auto params = NormScaleShiftParams{ .y = y.data_ptr(), .res_out = res_out.data_ptr(), + .quantized = nullptr, .x = x.data_ptr(), .input_bias = nullptr, .residual = residual.data_ptr(), .gate = gate.data_ptr(), .scale = scale.data_ptr(), .shift = shift.data_ptr(), + .input_scale = nullptr, .eps = static_cast(eps), }; LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel, params); } }; +/** \brief Fuse Qwen LayerNorm/modulation with static E4M3 activation quantization. */ +struct NormScaleShiftFp8Kernel { + static void + run(tvm::ffi::TensorView y, + tvm::ffi::TensorView quantized, + tvm::ffi::TensorView x, + tvm::ffi::TensorView scale, + tvm::ffi::TensorView shift, + tvm::ffi::TensorView input_scale, + double eps) { + using namespace host; + auto N = SymbolicSize{"num_rows"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({N, kHidden}).with_dtype().with_device(device).verify(x).verify(y); + TensorMatcher({N, kHidden}).with_dtype().with_device(device).verify(quantized); + TensorMatcher({kHidden}).with_dtype().with_device(device).verify(scale).verify(shift); + TensorMatcher({1}).with_dtype().with_device(device).verify(input_scale); + + const uint32_t grid = verify_nss_geometry(N); + const auto params = NormScaleShiftParams{ + .y = y.data_ptr(), + .res_out = nullptr, + .quantized = quantized.data_ptr(), + .x = x.data_ptr(), + .input_bias = nullptr, + .residual = nullptr, + .gate = nullptr, + .scale = scale.data_ptr(), + .shift = shift.data_ptr(), + .input_scale = input_scale.data_ptr(), + .eps = static_cast(eps), + }; + LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel, params); + } +}; + +/** \brief Fuse Qwen residual LayerNorm/modulation with static E4M3 activation quantization. */ +struct ScaleResidualNormScaleShiftFp8Kernel { + static void + run(tvm::ffi::TensorView y, + tvm::ffi::TensorView quantized, + tvm::ffi::TensorView res_out, + tvm::ffi::TensorView residual, + tvm::ffi::TensorView x, + tvm::ffi::TensorView gate, + tvm::ffi::TensorView scale, + tvm::ffi::TensorView shift, + tvm::ffi::TensorView input_scale, + double eps) { + using namespace host; + auto N = SymbolicSize{"num_rows"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({N, kHidden}) + .with_dtype() + .with_device(device) + .verify(x) + .verify(residual) + .verify(y) + .verify(res_out); + TensorMatcher({N, kHidden}).with_dtype().with_device(device).verify(quantized); + TensorMatcher({kHidden}).with_dtype().with_device(device).verify(gate).verify(scale).verify(shift); + TensorMatcher({1}).with_dtype().with_device(device).verify(input_scale); + + const uint32_t grid = verify_nss_geometry(N); + const auto params = NormScaleShiftParams{ + .y = y.data_ptr(), + .res_out = res_out.data_ptr(), + .quantized = quantized.data_ptr(), + .x = x.data_ptr(), + .input_bias = nullptr, + .residual = residual.data_ptr(), + .gate = gate.data_ptr(), + .scale = scale.data_ptr(), + .shift = shift.data_ptr(), + .input_scale = input_scale.data_ptr(), + .eps = static_cast(eps), + }; + LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel, params); + } +}; + struct BiasScaleResidualNormScaleShiftKernel { static void run(tvm::ffi::TensorView y, @@ -284,12 +401,14 @@ struct BiasScaleResidualNormScaleShiftKernel { const auto params = NormScaleShiftParams{ .y = y.data_ptr(), .res_out = res_out.data_ptr(), + .quantized = nullptr, .x = x.data_ptr(), .input_bias = input_bias.data_ptr(), .residual = residual.data_ptr(), .gate = gate.data_ptr(), .scale = scale.data_ptr(), .shift = shift.data_ptr(), + .input_scale = nullptr, .eps = static_cast(eps), }; LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel, params); @@ -315,12 +434,14 @@ struct BiasMulAddKernel { const auto params = NormScaleShiftParams{ .y = y.data_ptr(), .res_out = nullptr, + .quantized = nullptr, .x = x.data_ptr(), .input_bias = input_bias.data_ptr(), .residual = residual.data_ptr(), .gate = gate.data_ptr(), .scale = nullptr, .shift = nullptr, + .input_scale = nullptr, .eps = 0.0f, }; LaunchKernel(grid, kThreads, device.unwrap())(bias_mul_add_kernel, params); diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index c13824a3e..540232d38 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/python/sglang/kernels/ops/diffusion/__init__.py @@ -382,6 +382,10 @@ _EXPORTS: dict[str, str] = { "fused_scale_residual_rmsnorm_scale_shift_bitexact": "norm.rmsnorm_scale_shift_bitexact", "fused_norm_scale_shift": "norm.scale_residual_norm_cutedsl", "fused_scale_residual_norm_scale_shift": "norm.scale_residual_norm_cutedsl", + "fused_norm_scale_shift_fp8": "norm.norm_scale_shift_jit", + "fused_scale_residual_norm_scale_shift_fp8": "norm.norm_scale_shift_jit", + "try_fused_norm_scale_shift_fp8": "norm.norm_scale_shift_jit", + "try_fused_scale_residual_norm_scale_shift_fp8": "norm.norm_scale_shift_jit", "validate_scale_shift": "norm.scale_residual_norm_cutedsl", "can_use_wan_rmsnorm_silu": "norm.wan_rmsnorm_silu_triton", "wan_rmsnorm_silu": "norm.wan_rmsnorm_silu_triton", diff --git a/python/sglang/kernels/ops/diffusion/norm/norm_scale_shift_jit.py b/python/sglang/kernels/ops/diffusion/norm/norm_scale_shift_jit.py index 5d0881a6e..8bd7607b8 100644 --- a/python/sglang/kernels/ops/diffusion/norm/norm_scale_shift_jit.py +++ b/python/sglang/kernels/ops/diffusion/norm/norm_scale_shift_jit.py @@ -67,6 +67,11 @@ def _row_bf16(t, device: torch.device): @cache_once def norm_scale_shift_module() -> Module: + device = torch.device("cuda", torch.cuda.current_device()) + if not _blackwell_or_newer(device): + raise RuntimeError( + "Qwen-Image norm-scale-shift JIT kernels require NVIDIA Blackwell or newer" + ) return load_jit( "norm_scale_shift_native", cuda_files=["diffusion/norm_scale_shift.cuh"], @@ -79,6 +84,14 @@ def norm_scale_shift_module() -> Module: "srnss_bf16_row", "norm_scale_shift::ScaleResidualNormScaleShiftKernel::run", ), + ( + "nss_fp8_row", + "norm_scale_shift::NormScaleShiftFp8Kernel::run", + ), + ( + "srnss_fp8_row", + "norm_scale_shift::ScaleResidualNormScaleShiftFp8Kernel::run", + ), ( "bias_srnss_bf16_row", "norm_scale_shift::BiasScaleResidualNormScaleShiftKernel::run", @@ -94,6 +107,44 @@ def norm_scale_shift_module() -> Module: _module = norm_scale_shift_module +def fused_norm_scale_shift_fp8(x, scale, shift, input_scale, eps): + """Return exact BF16 modulation output and its static E4M3 quantization.""" + normalized = torch.empty_like(x) + quantized = torch.empty_like(x, dtype=torch.float8_e4m3fn) + _module().nss_fp8_row( + normalized.view(-1, _HIDDEN), + quantized.view(-1, _HIDDEN), + x.view(-1, _HIDDEN), + scale, + shift, + input_scale.reshape(1), + float(eps), + ) + return normalized, quantized + + +def fused_scale_residual_norm_scale_shift_fp8( + residual, x, gate, scale, shift, input_scale, eps +): + """Return exact BF16 residual/modulation outputs and E4M3 quantization.""" + normalized = torch.empty_like(x) + quantized = torch.empty_like(x, dtype=torch.float8_e4m3fn) + residual_out = torch.empty_like(x) + _module().srnss_fp8_row( + normalized.view(-1, _HIDDEN), + quantized.view(-1, _HIDDEN), + residual_out.view(-1, _HIDDEN), + residual.view(-1, _HIDDEN), + x.view(-1, _HIDDEN), + gate, + scale, + shift, + input_scale.reshape(1), + float(eps), + ) + return normalized, quantized, residual_out + + 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 @@ -145,6 +196,67 @@ def try_fused_scale_residual_norm_scale_shift( return y, residual_out +def try_fused_norm_scale_shift_fp8( + x, weight, bias, scale, shift, input_scale, norm_type, eps +): + if norm_type != "layer" or weight is not None or bias is not None: + return None + if not _nss_activation(x) or not _blackwell_or_newer(x.device): + return None + + scale = _row_bf16(scale, x.device) + shift = _row_bf16(shift, x.device) + if scale is None or shift is None: + return None + if not _fp8_input_scale(input_scale, x.device): + return None + return fused_norm_scale_shift_fp8(x, scale, shift, input_scale, eps) + + +def try_fused_scale_residual_norm_scale_shift_fp8( + residual, + x, + gate, + weight, + bias, + scale, + shift, + input_scale, + norm_type, + eps, +): + if norm_type != "layer" or weight is not None or bias is not None: + return None + if not ( + _nss_activation(x) + and _nss_activation(residual, x) + and _blackwell_or_newer(x.device) + ): + return None + + gate = _row_bf16(gate, x.device) + scale = _row_bf16(scale, x.device) + shift = _row_bf16(shift, x.device) + if gate is None or scale is None or shift is None: + return None + if not _fp8_input_scale(input_scale, x.device): + return None + return fused_scale_residual_norm_scale_shift_fp8( + residual, x, gate, scale, shift, input_scale, eps + ) + + +def _fp8_input_scale(t, device: torch.device) -> bool: + return ( + isinstance(t, torch.Tensor) + and t.is_cuda + and t.device == device + and t.dtype == torch.float32 + and t.numel() == 1 + and t.is_contiguous() + ) + + def try_fused_bias_scale_residual_norm_scale_shift( residual, x, input_bias, gate, weight, bias, scale, shift, norm_type, eps ): 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 f2bf082d3..4a7eefcbf 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -22,6 +22,8 @@ from sglang.kernels.ops.diffusion import ( mark_fused_gelu_site, try_fused_bias_mul_add, try_fused_bias_scale_residual_norm_scale_shift, + try_fused_norm_scale_shift_fp8, + try_fused_scale_residual_norm_scale_shift_fp8, ) from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig from sglang.multimodal_gen.configs.models.fsdp import is_transformer_block @@ -70,6 +72,9 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config i NunchakuConfig, is_nunchaku_available, ) +from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( + ModelOptFp8LinearMethod, +) from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( apply_flashinfer_rope_qk_inplace, ) @@ -1125,6 +1130,65 @@ class QwenImageTransformerBlock(nn.Module): self.img_mlp = NunchakuFeedForward(self.img_mlp, **nunchaku_kwargs) self.txt_mlp = NunchakuFeedForward(self.txt_mlp, **nunchaku_kwargs) + self._fp8_img_attn_norm_quant = False + self._fp8_txt_attn_norm_quant = False + self._fp8_img_mlp_norm_quant = False + self._fp8_txt_mlp_norm_quant = False + + @staticmethod + def _valid_modelopt_fp8_linear(linear: nn.Module) -> bool: + input_scale = getattr(linear, "input_scale", None) + return ( + isinstance(getattr(linear, "quant_method", None), ModelOptFp8LinearMethod) + and isinstance(input_scale, torch.Tensor) + and input_scale.is_cuda + and input_scale.dtype == torch.float32 + and input_scale.numel() == 1 + and input_scale.is_contiguous() + and bool(torch.isfinite(input_scale).all().item()) + and bool((input_scale > 0).all().item()) + ) + + @classmethod + def _shared_modelopt_fp8_scale(cls, linears: list[nn.Module]) -> bool: + if not all(cls._valid_modelopt_fp8_linear(linear) for linear in linears): + return False + reference = linears[0].input_scale + return all(torch.equal(reference, linear.input_scale) for linear in linears[1:]) + + def configure_fp8_norm_quant(self) -> None: + """Enable exact norm+quant paths after checkpoint scales are materialized.""" + if not torch.cuda.is_available(): + return + capability = torch.cuda.get_device_capability() + if self.dim != 3072 or capability[0] < 10 or self.zero_cond_t: + return + if self.attn.use_fused_qkv: + self._fp8_img_attn_norm_quant = self._valid_modelopt_fp8_linear( + self.attn.to_qkv + ) + else: + self._fp8_img_attn_norm_quant = self._shared_modelopt_fp8_scale( + [self.attn.to_q, self.attn.to_k, self.attn.to_v] + ) + if self.attn.added_kv_proj_dim is not None: + if self.attn.use_fused_added_qkv: + self._fp8_txt_attn_norm_quant = self._valid_modelopt_fp8_linear( + self.attn.to_added_qkv + ) + else: + self._fp8_txt_attn_norm_quant = self._shared_modelopt_fp8_scale( + [self.attn.add_q_proj, self.attn.add_k_proj, self.attn.add_v_proj] + ) + if isinstance(self.img_mlp, QwenImageFeedForward): + self._fp8_img_mlp_norm_quant = self._valid_modelopt_fp8_linear( + self.img_mlp.net[0].proj + ) + if isinstance(self.txt_mlp, QwenImageFeedForward): + self._fp8_txt_mlp_norm_quant = self._valid_modelopt_fp8_linear( + self.txt_mlp.net[0].proj + ) + def _norm_scale_shift( self, norm_module: LayerNormScaleShift, @@ -1200,6 +1264,68 @@ class QwenImageTransformerBlock(nn.Module): ) return img_mod_params, txt_mod_params + def _try_fp8_norm_quant( + self, + norm_module: LayerNormScaleShift, + *, + x: torch.Tensor, + mod_params: torch.Tensor, + input_scale: Optional[torch.Tensor], + enabled: bool, + modulate_index: Optional[torch.Tensor], + use_bcg_helpers: bool, + ) -> Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + if not enabled or modulate_index is not None or use_bcg_helpers: + return None + shift, scale, gate = mod_params.chunk(3, dim=-1) + result = try_fused_norm_scale_shift_fp8( + x, + getattr(norm_module.norm, "weight", None), + getattr(norm_module.norm, "bias", None), + scale, + shift, + input_scale, + norm_module.norm_type, + norm_module.eps, + ) + if result is None: + return None + normalized, quantized = result + return quantized, gate.unsqueeze(1), normalized + + def _try_fp8_residual_norm_quant( + self, + norm_module: ScaleResidualLayerNormScaleShift, + *, + residual: torch.Tensor, + x: torch.Tensor, + residual_gate: torch.Tensor, + mod_params: torch.Tensor, + input_scale: Optional[torch.Tensor], + enabled: bool, + modulate_index: Optional[torch.Tensor], + use_bcg_helpers: bool, + ) -> Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]]: + if not enabled or modulate_index is not None or use_bcg_helpers: + return None + shift, scale, gate = mod_params.chunk(3, dim=-1) + result = try_fused_scale_residual_norm_scale_shift_fp8( + residual, + x, + residual_gate, + getattr(norm_module.norm, "weight", None), + getattr(norm_module.norm, "bias", None), + scale, + shift, + input_scale, + norm_module.norm_type, + norm_module.eps, + ) + if result is None: + return None + normalized, quantized, residual_out = result + return quantized, residual_out, gate.unsqueeze(1), normalized + def _modulate( self, x: torch.Tensor, @@ -1361,16 +1487,57 @@ class QwenImageTransformerBlock(nn.Module): use_bcg_helpers = is_in_breakable_cuda_graph() # Process image stream - norm1 + modulation - img_modulated, img_gate1 = self._modulate( - hidden_states, - img_mod1, + img_fp8 = self._try_fp8_norm_quant( self.img_norm1, - modulate_index, + x=hidden_states, + mod_params=img_mod1, + input_scale=( + ( + self.attn.to_qkv.input_scale + if self.attn.use_fused_qkv + else self.attn.to_q.input_scale + ) + if self._fp8_img_attn_norm_quant + else None + ), + enabled=self._fp8_img_attn_norm_quant, + modulate_index=modulate_index, use_bcg_helpers=use_bcg_helpers, ) + if img_fp8 is None: + img_modulated, img_gate1 = self._modulate( + hidden_states, + img_mod1, + self.img_norm1, + modulate_index, + use_bcg_helpers=use_bcg_helpers, + ) + img_modulated_bf16 = None + else: + img_modulated, img_gate1, img_modulated_bf16 = img_fp8 # Process text stream - norm1 + modulation + txt_fp8 = self._try_fp8_norm_quant( + self.txt_norm1, + x=encoder_hidden_states, + mod_params=txt_mod1, + input_scale=( + ( + self.attn.to_added_qkv.input_scale + if self.attn.use_fused_added_qkv + else self.attn.add_q_proj.input_scale + ) + if self._fp8_txt_attn_norm_quant + else None + ), + enabled=self._fp8_txt_attn_norm_quant, + modulate_index=modulate_index, + use_bcg_helpers=use_bcg_helpers, + ) txt_shift1, txt_scale1, txt_gate1_raw = txt_mod1.chunk(3, dim=-1) - if use_bcg_helpers: + if txt_fp8 is not None: + txt_modulated, txt_gate1, txt_modulated_bf16 = txt_fp8 + elif use_bcg_helpers: + txt_modulated_bf16 = None txt_modulated = self._norm_scale_shift( self.txt_norm1, encoder_hidden_states, @@ -1378,10 +1545,12 @@ class QwenImageTransformerBlock(nn.Module): scale=txt_scale1, ) else: + txt_modulated_bf16 = None txt_modulated = self.txt_norm1( encoder_hidden_states, shift=txt_shift1, scale=txt_scale1 ) - txt_gate1 = txt_gate1_raw.unsqueeze(1) + if txt_fp8 is None: + txt_gate1 = txt_gate1_raw.unsqueeze(1) # Use QwenAttnProcessor2_0 for joint attention computation # This directly implements the DoubleStreamLayerMegatron logic: @@ -1399,6 +1568,7 @@ class QwenImageTransformerBlock(nn.Module): image_rotary_emb=image_rotary_emb, **joint_attention_kwargs, ) + del img_modulated_bf16, txt_modulated_bf16 # QwenAttnProcessor2_0 returns (img_output, txt_output) when encoder_hidden_states is provided ( @@ -1408,16 +1578,40 @@ class QwenImageTransformerBlock(nn.Module): txt_attn_bias, ) = attn_output # Process image stream - norm2 + MLP - img_modulated2, hidden_states, img_gate2 = self._modulate( - img_attn_output, - img_mod2, + img_fp8_mlp = self._try_fp8_residual_norm_quant( self.img_norm2, - modulate_index, - gate_x=img_gate1, - residual_x=hidden_states, - x_bias=img_attn_bias, + residual=hidden_states, + x=img_attn_output, + residual_gate=img_gate1, + mod_params=img_mod2, + input_scale=( + self.img_mlp.net[0].proj.input_scale + if self._fp8_img_mlp_norm_quant + else None + ), + enabled=self._fp8_img_mlp_norm_quant, + modulate_index=modulate_index, use_bcg_helpers=use_bcg_helpers, ) + if img_fp8_mlp is None: + img_modulated2, hidden_states, img_gate2 = self._modulate( + img_attn_output, + img_mod2, + self.img_norm2, + modulate_index, + gate_x=img_gate1, + residual_x=hidden_states, + x_bias=img_attn_bias, + use_bcg_helpers=use_bcg_helpers, + ) + img_modulated2_bf16 = None + else: + ( + img_modulated2, + hidden_states, + img_gate2, + img_modulated2_bf16, + ) = img_fp8_mlp if isinstance(self.img_mlp, QwenImageFeedForward): img_mlp_output, img_mlp_bias = self.img_mlp.forward_with_bias( img_modulated2 @@ -1425,6 +1619,7 @@ class QwenImageTransformerBlock(nn.Module): else: img_mlp_output = self.img_mlp(img_modulated2) img_mlp_bias = None + del img_modulated2_bf16 if img_mlp_output.dim() == 2: img_mlp_output = img_mlp_output.unsqueeze(0) @@ -1438,7 +1633,30 @@ class QwenImageTransformerBlock(nn.Module): # Process text stream - norm2 + MLP txt_shift2, txt_scale2, txt_gate2_raw = txt_mod2.chunk(3, dim=-1) - if use_bcg_helpers: + txt_fp8_mlp = self._try_fp8_residual_norm_quant( + self.txt_norm2, + residual=encoder_hidden_states, + x=txt_attn_output, + residual_gate=txt_gate1, + mod_params=txt_mod2, + input_scale=( + self.txt_mlp.net[0].proj.input_scale + if self._fp8_txt_mlp_norm_quant + else None + ), + enabled=self._fp8_txt_mlp_norm_quant, + modulate_index=modulate_index, + use_bcg_helpers=use_bcg_helpers, + ) + if txt_fp8_mlp is not None: + ( + txt_modulated2, + encoder_hidden_states, + txt_gate2, + txt_modulated2_bf16, + ) = txt_fp8_mlp + elif use_bcg_helpers: + txt_modulated2_bf16 = None if txt_attn_bias is not None: txt_attn_output = txt_attn_output + txt_attn_bias ( @@ -1453,6 +1671,7 @@ class QwenImageTransformerBlock(nn.Module): scale=txt_scale2, ) elif txt_attn_bias is not None: + txt_modulated2_bf16 = None txt_modulated2, encoder_hidden_states, _ = self._modulate( txt_attn_output, txt_mod2, @@ -1462,6 +1681,7 @@ class QwenImageTransformerBlock(nn.Module): x_bias=txt_attn_bias, ) else: + txt_modulated2_bf16 = None txt_modulated2, encoder_hidden_states = self.txt_norm2( residual=encoder_hidden_states, x=txt_attn_output, @@ -1469,7 +1689,8 @@ class QwenImageTransformerBlock(nn.Module): shift=txt_shift2, scale=txt_scale2, ) - txt_gate2 = txt_gate2_raw.unsqueeze(1) + if txt_fp8_mlp is None: + txt_gate2 = txt_gate2_raw.unsqueeze(1) if isinstance(self.txt_mlp, QwenImageFeedForward): txt_mlp_output, txt_mlp_bias = self.txt_mlp.forward_with_bias( txt_modulated2 @@ -1477,6 +1698,7 @@ class QwenImageTransformerBlock(nn.Module): else: txt_mlp_output = self.txt_mlp(txt_modulated2) txt_mlp_bias = None + del txt_modulated2_bf16 if txt_mlp_output.dim() == 2: txt_mlp_output = txt_mlp_output.unsqueeze(0) @@ -1632,6 +1854,24 @@ class QwenImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin): self.layer_names = ["transformer_blocks"] + def post_load_weights(self) -> None: + super().post_load_weights() + for block in self.transformer_blocks: + block.configure_fp8_norm_quant() + enabled = sum( + block._fp8_img_attn_norm_quant + + block._fp8_txt_attn_norm_quant + + block._fp8_img_mlp_norm_quant + + block._fp8_txt_mlp_norm_quant + for block in self.transformer_blocks + ) + if enabled: + logger.info( + "Enabled Qwen FP8 norm+quant fusion for %d/%d block paths", + enabled, + 4 * len(self.transformer_blocks), + ) + @functools.lru_cache(maxsize=50) def build_modulate_index(self, img_shapes: tuple[int, int, int], device): sp_world_size = get_sp_world_size() diff --git a/test/registered/kernels/benchmark/diffusion/bench_qwen_image_norm_fp8_quant.py b/test/registered/kernels/benchmark/diffusion/bench_qwen_image_norm_fp8_quant.py new file mode 100644 index 000000000..2ff550f81 --- /dev/null +++ b/test/registered/kernels/benchmark/diffusion/bench_qwen_image_norm_fp8_quant.py @@ -0,0 +1,82 @@ +import torch + +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.ops.diffusion import ( + fused_norm_scale_shift_fp8, + fused_scale_residual_norm_scale_shift_fp8, +) +from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8 +from sglang.multimodal_gen.runtime.layers.layernorm import ( + LayerNormScaleShift, + ScaleResidualLayerNormScaleShift, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=12, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" +) + +DEVICE = "cuda" +DTYPE = torch.bfloat16 +HIDDEN = 3072 +EPS = 1e-6 + + +@marker.parametrize("rows", [128, 1024, 4096], [128]) +@marker.parametrize("residual_path", [False, True], [False, True]) +@marker.benchmark("impl", ["split", "fused"], unit="us") +def benchmark(rows: int, residual_path: bool, impl: str): + if impl == "fused" and torch.cuda.get_device_capability()[0] < 10: + marker.skip("Fused Qwen-Image norm+FP8 quant requires NVIDIA Blackwell") + + generator = torch.Generator(device=DEVICE) + generator.manual_seed(20260831 + rows + int(residual_path)) + x = torch.randn((1, rows, HIDDEN), dtype=DTYPE, device=DEVICE, generator=generator) + residual = torch.randn_like(x) + gate = torch.randn((HIDDEN,), dtype=DTYPE, device=DEVICE, generator=generator) + scale = torch.randn((HIDDEN,), dtype=DTYPE, device=DEVICE, generator=generator) + shift = torch.randn((HIDDEN,), dtype=DTYPE, device=DEVICE, generator=generator) + input_scale = torch.tensor(0.03125, dtype=torch.float32, device=DEVICE) + + if residual_path: + layer = ScaleResidualLayerNormScaleShift( + HIDDEN, eps=EPS, elementwise_affine=False, dtype=DTYPE + ).to(DEVICE) + + if impl == "split": + + def fn(): + normalized, residual_out = layer.forward_cuda( + residual, x, gate, shift, scale + ) + quantized, _ = static_quant_fp8(normalized, input_scale) + return quantized, residual_out + + else: + + def fn(): + return fused_scale_residual_norm_scale_shift_fp8( + residual, x, gate, scale, shift, input_scale, EPS + ) + + else: + layer = LayerNormScaleShift( + HIDDEN, eps=EPS, elementwise_affine=False, dtype=DTYPE + ).to(DEVICE) + + if impl == "split": + + def fn(): + normalized = layer.forward_cuda(x, shift, scale) + return static_quant_fp8(normalized, input_scale)[0] + + else: + + def fn(): + return fused_norm_scale_shift_fp8(x, scale, shift, input_scale, EPS) + + return marker.do_bench(fn, disable_log_bandwidth=True) + + +if __name__ == "__main__": + benchmark.run() diff --git a/test/registered/kernels/ops/diffusion/test_qwen_image_norm_fp8_quant.py b/test/registered/kernels/ops/diffusion/test_qwen_image_norm_fp8_quant.py new file mode 100644 index 000000000..ce6617070 --- /dev/null +++ b/test/registered/kernels/ops/diffusion/test_qwen_image_norm_fp8_quant.py @@ -0,0 +1,84 @@ +import sys + +import pytest +import torch + +from sglang.kernels.ops.diffusion import ( + fused_norm_scale_shift_fp8, + fused_scale_residual_norm_scale_shift_fp8, +) +from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8 +from sglang.multimodal_gen.runtime.layers.layernorm import ( + LayerNormScaleShift, + ScaleResidualLayerNormScaleShift, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + +DEVICE = "cuda" +DTYPE = torch.bfloat16 +HIDDEN = 3072 +EPS = 1e-6 + + +def _make_inputs(rows: int): + generator = torch.Generator(device=DEVICE) + generator.manual_seed(20260831 + rows) + x = torch.randn((1, rows, HIDDEN), dtype=DTYPE, device=DEVICE, generator=generator) + residual = torch.randn_like(x) + gate = torch.randn((HIDDEN,), dtype=DTYPE, device=DEVICE, generator=generator) + scale = torch.randn((HIDDEN,), dtype=DTYPE, device=DEVICE, generator=generator) + shift = torch.randn((HIDDEN,), dtype=DTYPE, device=DEVICE, generator=generator) + return x, residual, gate, scale, shift + + +@pytest.mark.parametrize("rows", [1, 127, 1024]) +@pytest.mark.parametrize( + "input_scale_value", [0.005, 0.03125, 0.4263392984867096, 0.4754464328289032, 1.0] +) +def test_norm_scale_shift_fp8_is_bit_exact(rows: int, input_scale_value: float) -> None: + x, _, _, scale, shift = _make_inputs(rows) + input_scale = torch.tensor(input_scale_value, dtype=torch.float32, device=DEVICE) + layer = LayerNormScaleShift( + HIDDEN, eps=EPS, elementwise_affine=False, dtype=DTYPE + ).to(DEVICE) + + normalized = layer.forward_cuda(x, shift, scale) + expected, _ = static_quant_fp8(normalized, input_scale) + actual_normalized, actual = fused_norm_scale_shift_fp8( + x, scale, shift, input_scale, EPS + ) + + assert torch.equal(actual_normalized, normalized) + assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) + + +@pytest.mark.parametrize("rows", [1, 127, 1024]) +@pytest.mark.parametrize( + "input_scale_value", [0.005, 0.03125, 0.4263392984867096, 0.4754464328289032, 1.0] +) +def test_residual_norm_scale_shift_fp8_is_bit_exact( + rows: int, input_scale_value: float +) -> None: + x, residual, gate, scale, shift = _make_inputs(rows) + input_scale = torch.tensor(input_scale_value, dtype=torch.float32, device=DEVICE) + layer = ScaleResidualLayerNormScaleShift( + HIDDEN, eps=EPS, elementwise_affine=False, dtype=DTYPE + ).to(DEVICE) + + normalized, expected_residual = layer.forward_cuda(residual, x, gate, shift, scale) + expected, _ = static_quant_fp8(normalized, input_scale) + actual_normalized, actual, actual_residual = ( + fused_scale_residual_norm_scale_shift_fp8( + residual, x, gate, scale, shift, input_scale, EPS + ) + ) + + assert torch.equal(actual_normalized, normalized) + assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) + assert torch.equal(actual_residual, expected_residual) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/test/registered/unit/models/test_qwen_image_fp8_norm_quant.py b/test/registered/unit/models/test_qwen_image_fp8_norm_quant.py new file mode 100644 index 000000000..18d33803f --- /dev/null +++ b/test/registered/unit/models/test_qwen_image_fp8_norm_quant.py @@ -0,0 +1,101 @@ +"""Unit tests for Qwen-Image ModelOpt FP8 norm+quant activation gates.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch +import torch.nn as nn + +from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( + ModelOptFp8LinearMethod, +) +from sglang.multimodal_gen.runtime.models.dits.qwen_image import ( + QwenImageTransformerBlock, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=5, stage="base-b", runner_config="1-gpu-small") + + +def _fp8_linear(input_scale: float) -> nn.Module: + linear = nn.Module() + linear.quant_method = object.__new__(ModelOptFp8LinearMethod) + linear.register_parameter( + "input_scale", + nn.Parameter( + torch.tensor(input_scale, dtype=torch.float32, device="cuda"), + requires_grad=False, + ), + ) + return linear + + +def _attention(*, fused: bool, scales: tuple[float, ...]) -> SimpleNamespace: + if fused: + return SimpleNamespace( + use_fused_qkv=True, + to_qkv=_fp8_linear(scales[0]), + added_kv_proj_dim=None, + ) + return SimpleNamespace( + use_fused_qkv=False, + to_q=_fp8_linear(scales[0]), + to_k=_fp8_linear(scales[1]), + to_v=_fp8_linear(scales[2]), + added_kv_proj_dim=None, + ) + + +def _block(attn: SimpleNamespace) -> QwenImageTransformerBlock: + block = object.__new__(QwenImageTransformerBlock) + nn.Module.__init__(block) + block.dim = 3072 + block.zero_cond_t = False + block.attn = attn + block.img_mlp = None + block.txt_mlp = None + block._fp8_img_attn_norm_quant = False + block._fp8_txt_attn_norm_quant = False + block._fp8_img_mlp_norm_quant = False + block._fp8_txt_mlp_norm_quant = False + return block + + +@patch("torch.cuda.get_device_capability", return_value=(10, 0)) +@patch("torch.cuda.is_available", return_value=True) +class TestQwenImageFp8NormQuantGate(CustomTestCase): + def test_separate_qkv_requires_identical_input_scales( + self, _is_available, _capability + ) -> None: + matching = _block(_attention(fused=False, scales=(0.25, 0.25, 0.25))) + mismatched = _block(_attention(fused=False, scales=(0.25, 0.5, 0.25))) + + matching.configure_fp8_norm_quant() + mismatched.configure_fp8_norm_quant() + + self.assertTrue(matching._fp8_img_attn_norm_quant) + self.assertFalse(mismatched._fp8_img_attn_norm_quant) + + def test_merged_qkv_uses_its_materialized_input_scale( + self, _is_available, _capability + ) -> None: + block = _block(_attention(fused=True, scales=(0.25,))) + + block.configure_fp8_norm_quant() + + self.assertTrue(block._fp8_img_attn_norm_quant) + + def test_nonpositive_scale_keeps_fusion_disabled( + self, _is_available, _capability + ) -> None: + block = _block(_attention(fused=True, scales=(0.0,))) + + block.configure_fp8_norm_quant() + + self.assertFalse(block._fp8_img_attn_norm_quant) + + +if __name__ == "__main__": + unittest.main()