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 e64adef87..648bff847 100644 --- a/python/sglang/kernels/jit/csrc/diffusion/norm_scale_shift.cuh +++ b/python/sglang/kernels/jit/csrc/diffusion/norm_scale_shift.cuh @@ -22,6 +22,14 @@ #include // For AlignedVector #include // For warp::reduce_sum +#if defined(ENABLE_FP4) && ENABLE_FP4 +// FlashInfer's TensorRT-LLM quantization helper uses the C macro directly. +#ifndef FLT_MAX +#define FLT_MAX __FLT_MAX__ +#endif +#include +#endif + #include namespace sglang { @@ -41,6 +49,7 @@ struct NormScaleShiftParams { void* y; void* res_out; void* quantized; + void* quant_scales; const void* x; const void* input_bias; const void* residual; @@ -48,6 +57,8 @@ struct NormScaleShiftParams { const void* scale; const void* shift; const void* input_scale; + const void* global_scale; + uint32_t num_rows; float eps; }; @@ -78,13 +89,24 @@ SGL_DEVICE float triton_scale_reciprocal(float scale) { return reciprocal; } -template +template __global__ void norm_scale_shift_kernel(const NormScaleShiftParams __grid_constant__ params) { + static_assert(!(kQuantizeFp8 && kQuantizeNvfp4)); using namespace device; using Vec = AlignedVector; const int row = blockIdx.x; const int tid = threadIdx.x; +#if defined(ENABLE_FP4) && ENABLE_FP4 + if constexpr (kQuantizeNvfp4) { + if (row >= params.num_rows) { + auto* scales = static_cast(params.quant_scales); + const int64_t scale_offset = tensorrt_llm::kernels::get_sf_out_offset_128x4(row, tid, kThreads); + scales[scale_offset] = 0; + return; + } + } +#endif const int lane = tid & int(kWarpThreads - 1); const int warp = tid >> 5; const int row_offset = row * kHidden; @@ -168,9 +190,29 @@ __global__ void norm_scale_shift_kernel(const NormScaleShiftParams __grid_consta 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); + if constexpr (kQuantizeNvfp4) { +#if defined(ENABLE_FP4) && ENABLE_FP4 + tensorrt_llm::kernels::PackedVec<__nv_bfloat16, kVecElems> quant_vec; + auto* quant_values = reinterpret_cast<__nv_bfloat16*>(&quant_vec); +#pragma unroll + for (int i = 0; i < kVecElems; ++i) { + quant_values[i] = static_cast<__nv_bfloat16>(yv[i]); + } + + auto* scales = static_cast(params.quant_scales); + const int64_t scale_offset = tensorrt_llm::kernels::get_sf_out_offset_128x4(row, tid, kThreads); + const float global_scale = *static_cast(params.global_scale); + const uint64_t packed = tensorrt_llm::kernels::cvt_warp_fp16_to_fp4<__nv_bfloat16, kVecElems, kVecElems, false>( + quant_vec, global_scale, scales + scale_offset); + static_cast(params.quantized)[int64_t(row) * kThreads + tid] = packed; +#else + static_assert(!kQuantizeNvfp4); +#endif + } else { + yv.store(static_cast(params.y) + row_offset + elem_offset); + if constexpr (kQuantizeFp8) { + qv.store(static_cast(params.quantized) + row_offset + elem_offset); + } } } @@ -226,6 +268,7 @@ struct NormScaleShiftKernel { .y = y.data_ptr(), .res_out = nullptr, .quantized = nullptr, + .quant_scales = nullptr, .x = x.data_ptr(), .input_bias = nullptr, .residual = nullptr, @@ -233,6 +276,8 @@ struct NormScaleShiftKernel { .scale = scale.data_ptr(), .shift = shift.data_ptr(), .input_scale = nullptr, + .global_scale = nullptr, + .num_rows = grid, .eps = static_cast(eps), }; LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel, params); @@ -268,6 +313,7 @@ struct ScaleResidualNormScaleShiftKernel { .y = y.data_ptr(), .res_out = res_out.data_ptr(), .quantized = nullptr, + .quant_scales = nullptr, .x = x.data_ptr(), .input_bias = nullptr, .residual = residual.data_ptr(), @@ -275,6 +321,8 @@ struct ScaleResidualNormScaleShiftKernel { .scale = scale.data_ptr(), .shift = shift.data_ptr(), .input_scale = nullptr, + .global_scale = nullptr, + .num_rows = grid, .eps = static_cast(eps), }; LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel, params); @@ -306,6 +354,7 @@ struct NormScaleShiftFp8Kernel { .y = y.data_ptr(), .res_out = nullptr, .quantized = quantized.data_ptr(), + .quant_scales = nullptr, .x = x.data_ptr(), .input_bias = nullptr, .residual = nullptr, @@ -313,6 +362,8 @@ struct NormScaleShiftFp8Kernel { .scale = scale.data_ptr(), .shift = shift.data_ptr(), .input_scale = input_scale.data_ptr(), + .global_scale = nullptr, + .num_rows = grid, .eps = static_cast(eps), }; LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel, params); @@ -353,6 +404,7 @@ struct ScaleResidualNormScaleShiftFp8Kernel { .y = y.data_ptr(), .res_out = res_out.data_ptr(), .quantized = quantized.data_ptr(), + .quant_scales = nullptr, .x = x.data_ptr(), .input_bias = nullptr, .residual = residual.data_ptr(), @@ -360,6 +412,8 @@ struct ScaleResidualNormScaleShiftFp8Kernel { .scale = scale.data_ptr(), .shift = shift.data_ptr(), .input_scale = input_scale.data_ptr(), + .global_scale = nullptr, + .num_rows = grid, .eps = static_cast(eps), }; LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel, params); @@ -402,6 +456,7 @@ struct BiasScaleResidualNormScaleShiftKernel { .y = y.data_ptr(), .res_out = res_out.data_ptr(), .quantized = nullptr, + .quant_scales = nullptr, .x = x.data_ptr(), .input_bias = input_bias.data_ptr(), .residual = residual.data_ptr(), @@ -409,6 +464,8 @@ struct BiasScaleResidualNormScaleShiftKernel { .scale = scale.data_ptr(), .shift = shift.data_ptr(), .input_scale = nullptr, + .global_scale = nullptr, + .num_rows = grid, .eps = static_cast(eps), }; LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel, params); @@ -435,6 +492,7 @@ struct BiasMulAddKernel { .y = y.data_ptr(), .res_out = nullptr, .quantized = nullptr, + .quant_scales = nullptr, .x = x.data_ptr(), .input_bias = input_bias.data_ptr(), .residual = residual.data_ptr(), @@ -442,12 +500,70 @@ struct BiasMulAddKernel { .scale = nullptr, .shift = nullptr, .input_scale = nullptr, + .global_scale = nullptr, + .num_rows = grid, .eps = 0.0f, }; LaunchKernel(grid, kThreads, device.unwrap())(bias_mul_add_kernel, params); } }; +#if defined(ENABLE_FP4) && ENABLE_FP4 +struct ScaleResidualNormScaleShiftNvfp4Kernel { + static void + run(tvm::ffi::TensorView quantized, + tvm::ffi::TensorView quant_scales, + tvm::ffi::TensorView res_out, + tvm::ffi::TensorView residual, + tvm::ffi::TensorView x, + tvm::ffi::TensorView input_bias, + tvm::ffi::TensorView gate, + tvm::ffi::TensorView scale, + tvm::ffi::TensorView shift, + tvm::ffi::TensorView global_scale, + double eps) { + using namespace host; + auto N = SymbolicSize{"num_rows"}; + auto NP = SymbolicSize{"num_rows_padded"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({N, kHidden}).with_dtype().with_device(device).verify(x).verify(residual).verify(res_out); + TensorMatcher({kHidden}) + .with_dtype() + .with_device(device) + .verify(input_bias) + .verify(gate) + .verify(scale) + .verify(shift); + TensorMatcher({N, kHidden / 2}).with_dtype().with_device(device).verify(quantized); + TensorMatcher({NP, kThreads}).with_dtype().with_device(device).verify(quant_scales); + TensorMatcher({1}).with_dtype().with_device(device).verify(global_scale); + + const uint32_t num_rows = verify_nss_geometry(N); + const uint32_t num_rows_padded = div_ceil(num_rows, uint32_t(128)) * 128; + RuntimeCheck(NP.unwrap() == num_rows_padded, "quant scale rows must be padded to 128"); + const auto params = NormScaleShiftParams{ + .y = nullptr, + .res_out = res_out.data_ptr(), + .quantized = quantized.data_ptr(), + .quant_scales = quant_scales.data_ptr(), + .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, + .global_scale = global_scale.data_ptr(), + .num_rows = num_rows, + .eps = static_cast(eps), + }; + LaunchKernel(num_rows_padded, kThreads, device.unwrap())(norm_scale_shift_kernel, params); + } +}; +#endif + } // namespace norm_scale_shift } // namespace sglang diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index 411ec5131..fb55b7221 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/python/sglang/kernels/ops/diffusion/__init__.py @@ -103,6 +103,13 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = ( _HIP, "FlyDSL (ROCm gfx950) residual + norm + scale/shift.", ), + ( + "diffusion.scale_residual_norm_scale_shift_nvfp4", + KernelBackend.JIT, + "norm.norm_scale_shift_jit:try_fused_scale_residual_norm_scale_shift_nvfp4", + _CUDA, + "Qwen residual LayerNorm/modulation + NVFP4 quantization.", + ), ( "diffusion.norm_scale_shift", KernelBackend.CUTE_DSL, @@ -405,6 +412,7 @@ _EXPORTS: dict[str, str] = { "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", + "try_fused_scale_residual_norm_scale_shift_nvfp4": "norm.norm_scale_shift_jit", "can_use_wan_rmsnorm_silu": "norm.wan_rmsnorm_silu_triton", "wan_rmsnorm_silu": "norm.wan_rmsnorm_silu_triton", "can_use_qk_rmsnorm_native": "norm.zimage_qk_rmsnorm_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 8bd7607b8..94b08a387 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 @@ -1,5 +1,6 @@ from __future__ import annotations +import os from typing import TYPE_CHECKING import torch @@ -31,6 +32,12 @@ def _sm103(device: torch.device) -> bool: ) +def _blackwell_sm10x(device: torch.device) -> bool: + return ( + torch.cuda.is_available() and torch.cuda.get_device_capability(device)[0] == 10 + ) + + def _nss_activation(t, like=None) -> bool: return ( isinstance(t, torch.Tensor) @@ -107,6 +114,25 @@ def norm_scale_shift_module() -> Module: _module = norm_scale_shift_module +@cache_once +def norm_scale_shift_nvfp4_module() -> Module: + return load_jit( + "norm_scale_shift_nvfp4_native", + cuda_files=["diffusion/norm_scale_shift.cuh"], + cuda_wrappers=[ + ( + "srnss_nvfp4_row", + "norm_scale_shift::ScaleResidualNormScaleShiftNvfp4Kernel::run", + ), + ], + extra_cuda_cflags=["-DENABLE_BF16", "-DENABLE_FP4"], + extra_dependencies=["flashinfer", "flashinfer_nv_internal"], + ) + + +_nvfp4_module = norm_scale_shift_nvfp4_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) @@ -257,6 +283,81 @@ def _fp8_input_scale(t, device: torch.device) -> bool: ) +def _env_enabled(name: str) -> bool: + return os.getenv(name, "").strip().lower() not in {"", "0", "false", "off", "no"} + + +def try_fused_scale_residual_norm_scale_shift_nvfp4( + residual, + x, + input_bias, + gate, + weight, + bias, + scale, + shift, + global_scale, + norm_type, + eps, +): + """Fuse Qwen residual LayerNorm/modulation with FC1-input NVFP4 quantization.""" + if ( + torch.compiler.is_compiling() + or _env_enabled("FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH") + or _env_enabled("TRTLLM_DISABLE_FP4_QUANT_FAST_MATH") + or _env_enabled("FLASHINFER_NVFP4_4OVER6") + ): + return None + 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_sm10x(x.device) + ): + return None + if torch.cuda.is_current_stream_capturing(): + return None + + gate = _row_bf16(gate, x.device) + input_bias = _row_bf16(input_bias, x.device) + scale = _row_bf16(scale, x.device) + shift = _row_bf16(shift, x.device) + if input_bias is None or gate is None or scale is None or shift is None: + return None + if not ( + isinstance(global_scale, torch.Tensor) + and global_scale.is_cuda + and global_scale.device == x.device + and global_scale.dtype == torch.float32 + and global_scale.numel() == 1 + and global_scale.is_contiguous() + ): + return None + + rows = x.numel() // _HIDDEN + padded_rows = (rows + 127) // 128 * 128 + quantized = torch.empty((rows, _HIDDEN // 2), dtype=torch.uint8, device=x.device) + quant_scales = torch.empty( + (padded_rows, _HIDDEN // 16), dtype=torch.uint8, device=x.device + ) + residual_out = torch.empty_like(x) + _nvfp4_module().srnss_nvfp4_row( + quantized, + quant_scales, + residual_out.view(-1, _HIDDEN), + residual.view(-1, _HIDDEN), + x.view(-1, _HIDDEN), + input_bias, + gate, + scale, + shift, + global_scale.reshape(1), + float(eps), + ) + return (quantized, quant_scales), residual_out + + 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 eefec49fc..3a6a97f05 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -29,6 +29,7 @@ from sglang.kernels.ops.diffusion import ( try_fused_norm_scale_shift_fp8, try_fused_qwen_qkv_epilogue, try_fused_scale_residual_norm_scale_shift_fp8, + try_fused_scale_residual_norm_scale_shift_nvfp4, ) from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig from sglang.multimodal_gen.configs.models.fsdp import is_transformer_block @@ -78,7 +79,9 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config i is_nunchaku_available, ) from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( + ModelOptFp4LinearMethod, ModelOptFp8LinearMethod, + apply_nvfp4_gemm_prequantized, ) from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( apply_flashinfer_rope_qk_inplace, @@ -1095,12 +1098,23 @@ class QwenImageGELU(nn.Module): # epilogue. Off by default; mounted per batch by the denoising stage. mark_fused_gelu_site(self, "proj") - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if fused_gelu_active(self) and can_use_linear_gelu(self.proj, hidden_states): + def forward( + self, + hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + if isinstance(hidden_states, tuple): + hidden_states = apply_nvfp4_gemm_prequantized( + self.proj, + *hidden_states, + output_dtype=self.proj.params_dtype, + bias=self.proj.bias, + ) + elif fused_gelu_active(self) and can_use_linear_gelu(self.proj, hidden_states): return fused_linear_gelu_tanh( hidden_states, self.proj.weight, self.proj.bias ) - hidden_states, _ = self.proj(hidden_states) + else: + hidden_states, _ = self.proj(hidden_states) return F.gelu(hidden_states, approximate="tanh") @@ -1154,13 +1168,15 @@ class QwenImageFeedForward(nn.Module): ) def forward_with_bias( - self, hidden_states: torch.Tensor + self, hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor] ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: hidden_states = self.net[0](hidden_states) hidden_states = self.net[1](hidden_states) return self.net[2](hidden_states) - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + def forward( + self, hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor] + ) -> torch.Tensor: hidden_states, bias = self.forward_with_bias(hidden_states) return hidden_states if bias is None else hidden_states + bias @@ -1294,6 +1310,20 @@ class QwenImageTransformerBlock(nn.Module): self.img_mlp = NunchakuFeedForward(self.img_mlp, **nunchaku_kwargs) self.txt_mlp = NunchakuFeedForward(self.txt_mlp, **nunchaku_kwargs) + self._enable_nvfp4_resnorm_quant = False + capability = current_platform.get_device_capability() + if ( + not nunchaku_enabled + and dim == 3072 + and capability is not None + and capability.major == 10 + ): + img_fc1 = self.img_mlp.net[0].proj + txt_fc1 = self.txt_mlp.net[0].proj + self._enable_nvfp4_resnorm_quant = isinstance( + img_fc1.quant_method, ModelOptFp4LinearMethod + ) and isinstance(txt_fc1.quant_method, ModelOptFp4LinearMethod) + self._fp8_img_attn_norm_quant = False self._fp8_txt_attn_norm_quant = False self._fp8_img_mlp_norm_quant = False @@ -1428,6 +1458,52 @@ class QwenImageTransformerBlock(nn.Module): ) return img_mod_params, txt_mod_params + def _try_nvfp4_resnorm_quant( + self, + norm_module: ScaleResidualLayerNormScaleShift, + mlp: QwenImageFeedForward, + *, + residual: torch.Tensor, + x: torch.Tensor, + x_bias: Optional[torch.Tensor], + residual_gate: torch.Tensor, + mod_params: torch.Tensor, + modulate_index: Optional[torch.Tensor], + use_bcg_helpers: bool, + ) -> Optional[ + tuple[ + tuple[torch.Tensor, torch.Tensor], + torch.Tensor, + torch.Tensor, + ] + ]: + if ( + not self._enable_nvfp4_resnorm_quant + or modulate_index is not None + or use_bcg_helpers + ): + return None + + shift, scale, gate = mod_params.chunk(3, dim=-1) + fc1 = mlp.net[0].proj + result = try_fused_scale_residual_norm_scale_shift_nvfp4( + residual, + x, + x_bias, + residual_gate, + getattr(norm_module.norm, "weight", None), + getattr(norm_module.norm, "bias", None), + scale.unsqueeze(1), + shift.unsqueeze(1), + fc1.input_scale_inv, + norm_module.norm_type, + norm_module.eps, + ) + if result is None: + return None + packed, residual_out = result + return packed, residual_out, gate.unsqueeze(1) + def _try_fp8_norm_quant( self, norm_module: LayerNormScaleShift, @@ -1742,40 +1818,55 @@ class QwenImageTransformerBlock(nn.Module): txt_attn_bias, ) = attn_output # Process image stream - norm2 + MLP - img_fp8_mlp = self._try_fp8_residual_norm_quant( + img_nvfp4 = self._try_nvfp4_resnorm_quant( self.img_norm2, + self.img_mlp, residual=hidden_states, x=img_attn_output, + x_bias=img_attn_bias, 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, - ) + if img_nvfp4 is not None: + img_modulated2, hidden_states, img_gate2 = img_nvfp4 img_modulated2_bf16 = None else: - ( - img_modulated2, - hidden_states, - img_gate2, - img_modulated2_bf16, - ) = img_fp8_mlp + img_fp8_mlp = self._try_fp8_residual_norm_quant( + self.img_norm2, + 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 @@ -1797,63 +1888,79 @@ class QwenImageTransformerBlock(nn.Module): # Process text stream - norm2 + MLP txt_shift2, txt_scale2, txt_gate2_raw = txt_mod2.chunk(3, dim=-1) - txt_fp8_mlp = self._try_fp8_residual_norm_quant( + txt_nvfp4 = self._try_nvfp4_resnorm_quant( self.txt_norm2, + self.txt_mlp, residual=encoder_hidden_states, x=txt_attn_output, + x_bias=txt_attn_bias, 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_fp8_mlp = None + if txt_nvfp4 is not None: + txt_modulated2, encoder_hidden_states, txt_gate2 = txt_nvfp4 txt_modulated2_bf16 = None - if txt_attn_bias is not None: - txt_attn_output = txt_attn_output + txt_attn_bias - ( - txt_modulated2, - encoder_hidden_states, - ) = self._scale_residual_norm_scale_shift( - self.txt_norm2, - residual=encoder_hidden_states, - x=txt_attn_output, - gate=txt_gate1, - shift=txt_shift2, - 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, - self.txt_norm2, - gate_x=txt_gate1, - residual_x=encoder_hidden_states, - x_bias=txt_attn_bias, - ) else: - txt_modulated2_bf16 = None - txt_modulated2, encoder_hidden_states = self.txt_norm2( + txt_fp8_mlp = self._try_fp8_residual_norm_quant( + self.txt_norm2, residual=encoder_hidden_states, x=txt_attn_output, - gate=txt_gate1, - shift=txt_shift2, - scale=txt_scale2, + 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 None: + 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 + ( + txt_modulated2, + encoder_hidden_states, + ) = self._scale_residual_norm_scale_shift( + self.txt_norm2, + residual=encoder_hidden_states, + x=txt_attn_output, + gate=txt_gate1, + shift=txt_shift2, + 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, + self.txt_norm2, + gate_x=txt_gate1, + residual_x=encoder_hidden_states, + 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, + gate=txt_gate1, + shift=txt_shift2, + scale=txt_scale2, + ) + if txt_nvfp4 is None and 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( diff --git a/test/registered/kernels/benchmark/diffusion/bench_qwen_image_resnorm_nvfp4_quant.py b/test/registered/kernels/benchmark/diffusion/bench_qwen_image_resnorm_nvfp4_quant.py new file mode 100644 index 000000000..98863e27b --- /dev/null +++ b/test/registered/kernels/benchmark/diffusion/bench_qwen_image_resnorm_nvfp4_quant.py @@ -0,0 +1,97 @@ +import flashinfer +import torch + +from sglang.kernels.ops.diffusion import ( + fused_scale_residual_norm_scale_shift, + try_fused_scale_residual_norm_scale_shift_nvfp4, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=30, + stage="base-b-kernel-benchmark", + runner_config="1-gpu-large", + disabled="standalone Qwen-Image NVFP4 residual-norm benchmark", +) + + +def _benchmark(fn, iterations: int = 100) -> float: + for _ in range(10): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) * 1000 / iterations + + +def _run_case(token_count: int) -> None: + hidden_size = 3072 + generator = torch.Generator(device="cuda") + generator.manual_seed(20260830 + token_count) + + def randn(shape): + return torch.randn( + shape, device="cuda", dtype=torch.bfloat16, generator=generator + ).contiguous() + + x = randn((1, token_count, hidden_size)) + residual = randn(x.shape) + input_bias = randn((hidden_size,)) + gate = randn((1, 1, hidden_size)) + scale = randn((1, 1, hidden_size)) + shift = randn((1, 1, hidden_size)) + global_scale = torch.tensor(0.625, device="cuda", dtype=torch.float32) + + def baseline(): + modulated, residual_out = fused_scale_residual_norm_scale_shift( + residual, x + input_bias, gate, None, None, scale, shift, "layer", 1e-6 + ) + quantized, quant_scales = flashinfer.fp4_quantize( + modulated.view(-1, hidden_size), global_scale + ) + return quantized, quant_scales, residual_out + + def fused(): + result = try_fused_scale_residual_norm_scale_shift_nvfp4( + residual, + x, + input_bias, + gate, + None, + None, + scale, + shift, + global_scale, + "layer", + 1e-6, + ) + assert result is not None + (quantized, quant_scales), residual_out = result + return quantized, quant_scales, residual_out + + expected = baseline() + actual = fused() + exact = [torch.equal(lhs, rhs) for lhs, rhs in zip(actual, expected)] + baseline_us = _benchmark(baseline) + fused_us = _benchmark(fused) + print( + { + "tokens": token_count, + "baseline_us": baseline_us, + "fused_us": fused_us, + "speedup": baseline_us / fused_us, + "exact": exact, + } + ) + + +if __name__ == "__main__": + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10: + raise RuntimeError("This benchmark requires an NVIDIA Blackwell SM10x GPU") + for tokens in (17, 1024, 4096, 4608): + _run_case(tokens) diff --git a/test/registered/kernels/ops/quantization/test_diffusion_nvfp4_scaled_mm.py b/test/registered/kernels/ops/quantization/test_diffusion_nvfp4_scaled_mm.py index 873a90954..d99be8193 100644 --- a/test/registered/kernels/ops/quantization/test_diffusion_nvfp4_scaled_mm.py +++ b/test/registered/kernels/ops/quantization/test_diffusion_nvfp4_scaled_mm.py @@ -4,6 +4,10 @@ import flashinfer import pytest import torch +from sglang.kernels.ops.diffusion import ( + fused_scale_residual_norm_scale_shift, + try_fused_scale_residual_norm_scale_shift_nvfp4, +) from sglang.multimodal_gen.runtime.layers.quantization import ( modelopt_quant as diffusion_modelopt_quant, ) @@ -37,6 +41,62 @@ def _nvfp4_supported() -> bool: return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0) +def _qwen_resnorm_nvfp4_supported() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10 + + +@pytest.mark.skipif( + not _qwen_resnorm_nvfp4_supported(), + reason="Qwen-Image fused residual norm + NVFP4 quantization requires SM10x", +) +@pytest.mark.parametrize("token_count", [17, 1024]) +def test_qwen_image_fused_resnorm_nvfp4_quant_is_exact(token_count: int) -> None: + hidden_size = 3072 + generator = torch.Generator(device=DEVICE) + generator.manual_seed(20260830 + token_count) + + def randn(shape): + return torch.randn( + shape, device=DEVICE, dtype=DTYPE, generator=generator + ).contiguous() + + residual = randn((1, token_count, hidden_size)) + x = randn((1, token_count, hidden_size)) + input_bias = randn((hidden_size,)) + gate = randn((1, 1, hidden_size)) + scale = randn((1, 1, hidden_size)) + shift = randn((1, 1, hidden_size)) + global_scale = torch.tensor(512.0, device=DEVICE, dtype=torch.float32) + + expected_modulated, expected_residual = fused_scale_residual_norm_scale_shift( + residual, x + input_bias, gate, None, None, scale, shift, "layer", 1e-6 + ) + expected_quantized, expected_scales = flashinfer.fp4_quantize( + expected_modulated.view(-1, hidden_size), global_scale + ) + + actual = try_fused_scale_residual_norm_scale_shift_nvfp4( + residual, + x, + input_bias, + gate, + None, + None, + scale, + shift, + global_scale, + "layer", + 1e-6, + ) + assert actual is not None + (actual_quantized, actual_scales), actual_residual = actual + assert torch.equal(actual_quantized, expected_quantized) + assert torch.equal( + actual_scales.view(torch.uint8), expected_scales.view(torch.uint8) + ) + assert torch.equal(actual_residual, expected_residual) + + def _make_global_scale(x: torch.Tensor) -> torch.Tensor: max_abs = torch.amax(x.abs()).clamp_min_(1e-6) return (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / max_abs).to(torch.float32)