From 644c16fe3523e3d36c101dc06e64b5d48783b42c Mon Sep 17 00:00:00 2001 From: Hanjun Cho Date: Wed, 14 Oct 2026 10:47:07 +0900 Subject: [PATCH 1/4] [Model] Fuse SiLU*mul with the NVFP4 quantization of down_proj in Qwen2MLP Qwen3.5's dense MLP runs SiLU*mul and the down_proj NVFP4 input quantization as one FlashInfer kernel; Qwen3, whose decoder layers use Qwen2MLP, ran them as two. Move the fused call and its enabling check from Qwen2MoeMLP and qwen3_5.py to models/utils.py, give Qwen2MLP the same path, and enable it for Qwen3 layers whose MLP projections are both NVFP4 W4A4. Keep it off for NVFP4 AWQ checkpoints, whose per-input-channel scale the fused producer does not apply, and skip it in a forward where sequence parallelism is active or down_proj is wrapped by LoRA, since both read the activation tensor. Signed-off-by: Hanjun Cho --- python/sglang/srt/models/qwen2.py | 13 +++- python/sglang/srt/models/qwen2_moe.py | 34 ++------ python/sglang/srt/models/qwen3.py | 6 +- python/sglang/srt/models/qwen3_5.py | 31 +------- python/sglang/srt/models/utils.py | 78 ++++++++++++++++++- .../unit/models/test_qwen3_5_modelopt_fp4.py | 45 ++++++++++- 6 files changed, 146 insertions(+), 61 deletions(-) diff --git a/python/sglang/srt/models/qwen2.py b/python/sglang/srt/models/qwen2.py index a32957785919..836dc352c6d9 100644 --- a/python/sglang/srt/models/qwen2.py +++ b/python/sglang/srt/models/qwen2.py @@ -60,6 +60,10 @@ kv_cache_scales_loader, supports_quantized_rl_reload, ) +from sglang.srt.models.utils import ( + can_run_silu_fp4_quant_fusion, + silu_and_mul_fp4_quant, +) from sglang.srt.platforms import current_platform from sglang.srt.runtime_context import get_exec, get_parallel from sglang.srt.utils import add_prefix, make_pp_layers @@ -104,6 +108,11 @@ def __init__( self.act_fn = SiluAndMul() self.dual_gemm = DualGemm(self.gate_up_proj, self.down_proj, hidden_size) self.activation_quant = DenseSiluAndMulQuant(self.down_proj) + # maybe_enable_silu_fp4_quant_fusion turns this on for NVFP4 projections; + # SiLU+mul and the down_proj input quantization then run as one kernel. + self._enable_silu_fp4_quant_fusion = False + self._masked_m_cache: dict = {} + self._down_input_scale_inv_1d = None def forward( self, @@ -117,7 +126,9 @@ def forward( x = self.dual_gemm(x, self.gate_up_proj) else: gate_up, _ = self.gate_up_proj(x) - if forward_batch is not None and forward_batch.forward_mode.is_decode(): + if can_run_silu_fp4_quant_fusion(self, gate_up): + x = silu_and_mul_fp4_quant(self, gate_up) + elif forward_batch is not None and forward_batch.forward_mode.is_decode(): x = self.activation_quant.forward_decode( gate_up, self.act_fn, self.down_proj ) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 8e9664446ddb..4a46077f34fb 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -92,6 +92,10 @@ is_in_breakable_cuda_graph, ) from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.utils import ( + can_run_silu_fp4_quant_fusion, + silu_and_mul_fp4_quant, +) from sglang.srt.runtime_context import ( get_exec, get_parallel, @@ -212,39 +216,13 @@ def __init__( # at construction time); the fused kernel requires a 1-D global scale. self._down_input_scale_inv_1d = None - def _silu_fp4_quant_fused(self, gate_up: torch.Tensor) -> tuple: - from flashinfer import silu_and_mul_scaled_nvfp4_experts_quantize - - if self._down_input_scale_inv_1d is None: - self._down_input_scale_inv_1d = self.down_proj.input_scale_inv.reshape(1) - num_tokens = gate_up.shape[0] - masked_m = self._masked_m_cache.get(num_tokens) - if masked_m is None: - masked_m = torch.tensor( - [num_tokens], dtype=torch.int32, device=gate_up.device - ) - self._masked_m_cache[num_tokens] = masked_m - y_fp4, y_sf = silu_and_mul_scaled_nvfp4_experts_quantize( - gate_up.unsqueeze(0), - masked_m, - self._down_input_scale_inv_1d, - ) - # [M, K/2, 1] -> [M, K/2]; scale: expert-grouped 6-D swizzle - # (32, 4, m_blocks, 4, K/64, 1) -> the dense swizzled layout - # fp4_gemm expects (verified bit-exact vs fp4_quantize up to FP4 - # rounding ties for M in 64..8192). - y_fp4 = y_fp4.squeeze(-1).view(torch.uint8) - m_padded = y_sf.shape[2] * y_sf.shape[0] * y_sf.shape[3] # m_blocks*32*4 - y_sf = y_sf.view(torch.uint8).permute(2, 4, 0, 1, 3, 5).reshape(m_padded, -1) - return y_fp4, y_sf - def forward( self, x, ): gate_up, _ = self.gate_up_proj(x) - if self._enable_silu_fp4_quant_fusion and not isinstance(gate_up, tuple): - x, _ = self.down_proj(self._silu_fp4_quant_fused(gate_up)) + if can_run_silu_fp4_quant_fusion(self, gate_up): + x, _ = self.down_proj(silu_and_mul_fp4_quant(self, gate_up)) return x x = self.act_fn(gate_up) x, _ = self.down_proj(x) diff --git a/python/sglang/srt/models/qwen3.py b/python/sglang/srt/models/qwen3.py index aa63a5a2419b..53d7185fa116 100644 --- a/python/sglang/srt/models/qwen3.py +++ b/python/sglang/srt/models/qwen3.py @@ -38,7 +38,10 @@ ) from sglang.srt.models.qwen2 import Qwen2MLP as Qwen3MLP from sglang.srt.models.qwen2 import Qwen2Model -from sglang.srt.models.utils import apply_qk_norm +from sglang.srt.models.utils import ( + apply_qk_norm, + maybe_enable_silu_fp4_quant_fusion, +) from sglang.srt.platforms import current_platform from sglang.srt.runtime_context import get_exec, get_parallel, get_stream from sglang.srt.utils import add_prefix, get_bool_env_var, is_cuda, is_hip, is_npu @@ -476,6 +479,7 @@ def __init__( prefix=add_prefix("mlp", prefix), reduce_results=False, ) + maybe_enable_silu_fp4_quant_fusion(self.mlp) norm_kwargs = ( dict( diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 2191ddce8038..6ea2e3daac89 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -15,7 +15,6 @@ """Inference-only Qwen3.5 model and Qwen3.5 MoE model compatible with HuggingFace weights.""" import logging -import os from functools import lru_cache from typing import Iterable, Optional, Set, Tuple, Union @@ -127,6 +126,7 @@ WeightsMapper, fused_qk_gemma_rmsnorm, fused_qk_gemma_rmsnorm_with_gate, + maybe_enable_silu_fp4_quant_fusion, ) from sglang.srt.runtime_context import ( get_exec, @@ -215,33 +215,6 @@ def _disable_shared_experts_fusion() -> bool: return bool(uses_cutedsl_ar_fusion() or is_shared_experts_fusion_disabled()) -def _maybe_enable_silu_fp4_quant_fusion(mlp: nn.Module) -> None: - """Fuse SiLU+mul with the down_proj NVFP4 input quantization. - - Replaces the separate act_and_mul and per-token FP4 quantize kernels with - one FlashInfer kernel and feeds down_proj a prequantized (fp4, scale) - tuple. Enabled when both dense-MLP projections use the NVFP4 (W4A4) - linear method; kill switch: SGLANG_DISABLE_SILU_FP4_QUANT_FUSION=1. - """ - if os.environ.get("SGLANG_DISABLE_SILU_FP4_QUANT_FUSION", "0") == "1": - return - from sglang.srt.layers.quantization.modelopt_quant import ModelOptFp4LinearMethod - - if not ( - isinstance(mlp.gate_up_proj.quant_method, ModelOptFp4LinearMethod) - and mlp.gate_up_proj.quant_method.quant_mode == "w4a4" - and isinstance(mlp.down_proj.quant_method, ModelOptFp4LinearMethod) - ): - return - try: - from flashinfer import silu_and_mul_scaled_nvfp4_experts_quantize # noqa: F401 - except ImportError: - return - mlp._enable_silu_fp4_quant_fusion = True - mlp.down_proj._accepts_prequantized_fp4 = True - logger.info("Enabled fused SiLU+mul+FP4-quant for dense MLP down_proj input.") - - def _use_mnnvl_cutedsl_fusion(config: Qwen3_5TextConfig, is_nextn: bool) -> bool: return bool( not is_nextn @@ -1191,7 +1164,7 @@ def __init__( # The stage boundary completes this stage's sum. reduce_results=False, ) - _maybe_enable_silu_fp4_quant_fusion(self.mlp) + maybe_enable_silu_fp4_quant_fusion(self.mlp) else: raise ValueError(f"Invalid model type: {config.model_type}") diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index 16df5d01b60a..73e73b9a8ca0 100644 --- a/python/sglang/srt/models/utils.py +++ b/python/sglang/srt/models/utils.py @@ -14,6 +14,8 @@ from __future__ import annotations import itertools +import logging +import os from collections.abc import Iterable, Mapping from dataclasses import dataclass, field from functools import lru_cache @@ -41,10 +43,12 @@ ) from sglang.srt.model_executor.runner import get_is_capture_mode from sglang.srt.model_loader.weight_utils import default_weight_loader -from sglang.srt.runtime_context import get_exec, get_parallel +from sglang.srt.runtime_context import get_exec, get_forward, get_parallel from sglang.srt.utils import get_current_device_stream_fast, is_cpu, is_cuda, is_hip from sglang.srt.utils.custom_op import register_custom_op +logger = logging.getLogger(__name__) + if TYPE_CHECKING: from sglang.srt.layers.layernorm import RMSNorm @@ -385,6 +389,78 @@ def create_fused_set_kv_buffer_arg( } +def maybe_enable_silu_fp4_quant_fusion(mlp: torch.nn.Module) -> None: + """Fuse SiLU+mul with the down_proj NVFP4 input quantization. + + Replaces the separate act_and_mul and per-token FP4 quantize kernels with + one FlashInfer kernel and feeds down_proj a prequantized (fp4, scale) + tuple. Enabled when both dense-MLP projections use the NVFP4 (W4A4) + linear method without AWQ; kill switch: SGLANG_DISABLE_SILU_FP4_QUANT_FUSION=1. + """ + if os.environ.get("SGLANG_DISABLE_SILU_FP4_QUANT_FUSION", "0") == "1": + return + from sglang.srt.layers.quantization.modelopt_quant import ModelOptFp4LinearMethod + + if not ( + isinstance(mlp.gate_up_proj.quant_method, ModelOptFp4LinearMethod) + and mlp.gate_up_proj.quant_method.quant_mode == "w4a4" + and isinstance(mlp.down_proj.quant_method, ModelOptFp4LinearMethod) + ): + return + # The fused producer does not apply AWQ's per-input-channel scale. + if mlp.down_proj.quant_method.quant_config.is_awq: + return + try: + from flashinfer import silu_and_mul_scaled_nvfp4_experts_quantize # noqa: F401 + except ImportError: + return + mlp._enable_silu_fp4_quant_fusion = True + mlp.down_proj._accepts_prequantized_fp4 = True + logger.info("Enabled fused SiLU+mul+FP4-quant for dense MLP down_proj input.") + + +def can_run_silu_fp4_quant_fusion( + mlp: torch.nn.Module, gate_up: torch.Tensor | tuple +) -> bool: + """Whether this forward can hand down_proj the fused (fp4, scale) tuple.""" + return ( + mlp._enable_silu_fp4_quant_fusion + and not isinstance(gate_up, tuple) + and not get_forward().sp_active # the SP g-bar matmul needs a tensor + and not hasattr(mlp.down_proj, "base_layer") # LoRA reads the activation + ) + + +def silu_and_mul_fp4_quant( + mlp: torch.nn.Module, gate_up: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + """SiLU(gate) * up of a dense MLP, quantized to NVFP4 for its down_proj.""" + from flashinfer import silu_and_mul_scaled_nvfp4_experts_quantize + + # input_scale_inv only exists after weight loading, so derive it lazily; + # the fused kernel requires a 1-D global scale. + if mlp._down_input_scale_inv_1d is None: + mlp._down_input_scale_inv_1d = mlp.down_proj.input_scale_inv.reshape(1) + num_tokens = gate_up.shape[0] + masked_m = mlp._masked_m_cache.get(num_tokens) + if masked_m is None: + masked_m = torch.tensor([num_tokens], dtype=torch.int32, device=gate_up.device) + mlp._masked_m_cache[num_tokens] = masked_m + y_fp4, y_sf = silu_and_mul_scaled_nvfp4_experts_quantize( + gate_up.unsqueeze(0), + masked_m, + mlp._down_input_scale_inv_1d, + ) + # [M, K/2, 1] -> [M, K/2]; scale: expert-grouped 6-D swizzle + # (32, 4, m_blocks, 4, K/64, 1) -> the dense swizzled layout + # fp4_gemm expects (verified bit-exact vs fp4_quantize up to FP4 + # rounding ties for M in 64..8192). + y_fp4 = y_fp4.squeeze(-1).view(torch.uint8) + m_padded = y_sf.shape[2] * y_sf.shape[0] * y_sf.shape[3] # m_blocks*32*4 + y_sf = y_sf.view(torch.uint8).permute(2, 4, 0, 1, 3, 5).reshape(m_padded, -1) + return y_fp4, y_sf + + def permute_inv(perm: torch.Tensor) -> torch.Tensor: inv_perm = torch.empty_like(perm) inv_perm[perm] = torch.arange(perm.numel(), device=perm.device, dtype=perm.dtype) diff --git a/test/registered/unit/models/test_qwen3_5_modelopt_fp4.py b/test/registered/unit/models/test_qwen3_5_modelopt_fp4.py index d0e3f60ad949..f36057c90456 100644 --- a/test/registered/unit/models/test_qwen3_5_modelopt_fp4.py +++ b/test/registered/unit/models/test_qwen3_5_modelopt_fp4.py @@ -8,16 +8,28 @@ load into and silently fall back to 1.0. 3. QWEN3_5_KV_SCALE_MAPPER remaps the checkpoint's baked KV-scale names onto the RadixAttention parameter names. + 4. The dense-MLP SiLU+FP4 fusion stays off where down_proj cannot take its + (fp4, scale) tuple: NVFP4 AWQ, sequence parallelism, a LoRA-wrapped down_proj. """ import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock import torch -from sglang.srt.layers.quantization.modelopt_quant import ModelOptFp4Config +from sglang.srt.layers.quantization.modelopt_quant import ( + ModelOptFp4Config, + ModelOptFp4LinearMethod, +) from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.qwen3_5 import QWEN3_5_KV_SCALE_MAPPER +from sglang.srt.models.utils import ( + can_run_silu_fp4_quant_fusion, + maybe_enable_silu_fp4_quant_fusion, +) +from sglang.srt.runtime_context import get_forward from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -153,5 +165,36 @@ def test_mapped_scale_loads_via_default_weight_loader(self): self.assertAlmostEqual(scale_param.item(), 0.0347, places=6) +class TestSiluFp4QuantFusionGates(CustomTestCase): + @staticmethod + def _enabled_mlp(down_proj) -> SimpleNamespace: + return SimpleNamespace(_enable_silu_fp4_quant_fusion=True, down_proj=down_proj) + + def test_runs_for_a_plain_down_proj(self): + mlp = self._enabled_mlp(SimpleNamespace()) + self.assertTrue(can_run_silu_fp4_quant_fusion(mlp, torch.empty(2, 32))) + + def test_skips_a_lora_wrapped_down_proj(self): + mlp = self._enabled_mlp(SimpleNamespace(base_layer=SimpleNamespace())) + self.assertFalse(can_run_silu_fp4_quant_fusion(mlp, torch.empty(2, 32))) + + def test_skips_under_sequence_parallelism(self): + mlp = self._enabled_mlp(SimpleNamespace()) + with get_forward().scoped(sp_active=True): + self.assertFalse(can_run_silu_fp4_quant_fusion(mlp, torch.empty(2, 32))) + + def test_stays_off_for_nvfp4_awq(self): + # The fused producer would skip AWQ's per-input-channel pre_quant_scale. + method = MagicMock(spec=ModelOptFp4LinearMethod, quant_mode="w4a4") + method.quant_config = SimpleNamespace(is_awq=True) + mlp = SimpleNamespace( + _enable_silu_fp4_quant_fusion=False, + gate_up_proj=SimpleNamespace(quant_method=method), + down_proj=SimpleNamespace(quant_method=method), + ) + maybe_enable_silu_fp4_quant_fusion(mlp) + self.assertFalse(mlp._enable_silu_fp4_quant_fusion) + + if __name__ == "__main__": unittest.main() From f5c2129dddc21918cdae68ec554f1a4417b666a3 Mon Sep 17 00:00:00 2001 From: Hanjun Cho Date: Tue, 3 Nov 2026 18:03:14 +0900 Subject: [PATCH 2/4] [Model] Keep the SiLU+FP4 fusion off on the Marlin FP4 backend Marlin, the default FP4 GEMM backend on SM80-SM90, takes a floating-point activation, so it cannot consume the prequantized (fp4, scale) tuple and a Qwen3 NVFP4 model crashed on the first forward. Leave the fusion off when the FP4 GEMM backend is Marlin. The enabling-check tests now use a stand-in FlashInfer, so the AWQ and Marlin gates are exercised whether or not FlashInfer is installed, and a positive case shows the fusion turns on for W4A4 on a native FP4 GEMM backend. Signed-off-by: Hanjun Cho --- python/sglang/srt/models/utils.py | 7 ++- .../unit/models/test_qwen3_5_modelopt_fp4.py | 50 +++++++++++++++---- 2 files changed, 47 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index 73e73b9a8ca0..fab5eda12a07 100644 --- a/python/sglang/srt/models/utils.py +++ b/python/sglang/srt/models/utils.py @@ -395,10 +395,12 @@ def maybe_enable_silu_fp4_quant_fusion(mlp: torch.nn.Module) -> None: Replaces the separate act_and_mul and per-token FP4 quantize kernels with one FlashInfer kernel and feeds down_proj a prequantized (fp4, scale) tuple. Enabled when both dense-MLP projections use the NVFP4 (W4A4) - linear method without AWQ; kill switch: SGLANG_DISABLE_SILU_FP4_QUANT_FUSION=1. + linear method without AWQ and the FP4 GEMM backend is not Marlin; + kill switch: SGLANG_DISABLE_SILU_FP4_QUANT_FUSION=1. """ if os.environ.get("SGLANG_DISABLE_SILU_FP4_QUANT_FUSION", "0") == "1": return + from sglang.srt.layers.quantization.fp4_utils import get_fp4_gemm_runner_backend from sglang.srt.layers.quantization.modelopt_quant import ModelOptFp4LinearMethod if not ( @@ -410,6 +412,9 @@ def maybe_enable_silu_fp4_quant_fusion(mlp: torch.nn.Module) -> None: # The fused producer does not apply AWQ's per-input-channel scale. if mlp.down_proj.quant_method.quant_config.is_awq: return + # Marlin, the SM80-SM90 default, takes a floating-point input, not the tuple. + if get_fp4_gemm_runner_backend().is_marlin(): + return try: from flashinfer import silu_and_mul_scaled_nvfp4_experts_quantize # noqa: F401 except ImportError: diff --git a/test/registered/unit/models/test_qwen3_5_modelopt_fp4.py b/test/registered/unit/models/test_qwen3_5_modelopt_fp4.py index f36057c90456..6beed915fe96 100644 --- a/test/registered/unit/models/test_qwen3_5_modelopt_fp4.py +++ b/test/registered/unit/models/test_qwen3_5_modelopt_fp4.py @@ -1,6 +1,6 @@ """Unit tests for modelopt_fp4 checkpoints that quantize Qwen3.5 attention. -Covers three things: +Covers four things: 1. ModelOptFp4Config.is_layer_excluded() decides per prefix whether attention is quantized or kept in BF16. 2. RadixAttention registers k_scale/v_scale when built with a quant_config that @@ -8,16 +8,20 @@ load into and silently fall back to 1.0. 3. QWEN3_5_KV_SCALE_MAPPER remaps the checkpoint's baked KV-scale names onto the RadixAttention parameter names. - 4. The dense-MLP SiLU+FP4 fusion stays off where down_proj cannot take its - (fp4, scale) tuple: NVFP4 AWQ, sequence parallelism, a LoRA-wrapped down_proj. + 4. The dense-MLP SiLU+FP4 fusion turns on for NVFP4 W4A4 and stays off where + down_proj cannot take its (fp4, scale) tuple: NVFP4 AWQ, the Marlin FP4 GEMM + backend, sequence parallelism, a LoRA-wrapped down_proj. """ +import sys import unittest from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import torch +from sglang.srt.layers.quantization import fp4_utils +from sglang.srt.layers.quantization.fp4_utils import Fp4GemmRunnerBackend from sglang.srt.layers.quantization.modelopt_quant import ( ModelOptFp4Config, ModelOptFp4LinearMethod, @@ -183,17 +187,45 @@ def test_skips_under_sequence_parallelism(self): with get_forward().scoped(sp_active=True): self.assertFalse(can_run_silu_fp4_quant_fusion(mlp, torch.empty(2, 32))) - def test_stays_off_for_nvfp4_awq(self): - # The fused producer would skip AWQ's per-input-channel pre_quant_scale. + @staticmethod + def _enables_fusion(*, is_awq: bool, backend: Fp4GemmRunnerBackend) -> bool: method = MagicMock(spec=ModelOptFp4LinearMethod, quant_mode="w4a4") - method.quant_config = SimpleNamespace(is_awq=True) + method.quant_config = SimpleNamespace(is_awq=is_awq) mlp = SimpleNamespace( _enable_silu_fp4_quant_fusion=False, gate_up_proj=SimpleNamespace(quant_method=method), down_proj=SimpleNamespace(quant_method=method), ) - maybe_enable_silu_fp4_quant_fusion(mlp) - self.assertFalse(mlp._enable_silu_fp4_quant_fusion) + # A stand-in FlashInfer keeps the result independent of the installed one. + flashinfer = SimpleNamespace(silu_and_mul_scaled_nvfp4_experts_quantize=None) + with ( + patch.dict(sys.modules, {"flashinfer": flashinfer}), + patch.object(fp4_utils, "FP4_GEMM_RUNNER_BACKEND", backend), + ): + maybe_enable_silu_fp4_quant_fusion(mlp) + return mlp._enable_silu_fp4_quant_fusion + + def test_turns_on_for_w4a4_on_a_native_fp4_backend(self): + self.assertTrue( + self._enables_fusion( + is_awq=False, backend=Fp4GemmRunnerBackend.FLASHINFER_CUTLASS + ) + ) + + def test_stays_off_for_nvfp4_awq(self): + # The fused producer would skip AWQ's per-input-channel pre_quant_scale. + self.assertFalse( + self._enables_fusion( + is_awq=True, backend=Fp4GemmRunnerBackend.FLASHINFER_CUTLASS + ) + ) + + def test_stays_off_for_the_marlin_backend(self): + # Marlin, the SM80-SM90 default, takes a floating-point input; + # handing it the (fp4, scale) tuple crashes the first forward. + self.assertFalse( + self._enables_fusion(is_awq=False, backend=Fp4GemmRunnerBackend.MARLIN) + ) if __name__ == "__main__": From 47dcb3e45b14aad502efe99c6cd1fb04baff95be Mon Sep 17 00:00:00 2001 From: Hanjun Cho Date: Tue, 3 Nov 2026 19:29:58 +0900 Subject: [PATCH 3/4] [Model] Keep the SiLU+FP4 fusion off when K is not a multiple of 64 The NVFP4 block scales of a K-wide input are padded to a multiple of four columns. fp4_quantize zeroes that padding, but the fused SiLU*mul+FP4 kernel never writes it, and the FP4 GEMM reads it. On an RTX PRO 6000, with K = 1376 or 608 and freed memory holding a NaN byte, every output of the down_proj GEMM became NaN (CUTLASS and cuDNN backends); K = 1216 was unaffected. Leave the fusion off unless down_proj's input width is a multiple of 64. Every Qwen3 intermediate size at TP <= 8 is. Signed-off-by: Hanjun Cho --- python/sglang/srt/models/utils.py | 3 +++ .../unit/models/test_qwen3_5_modelopt_fp4.py | 24 +++++++++++++++---- 2 files changed, 23 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index fab5eda12a07..a99b0533781f 100644 --- a/python/sglang/srt/models/utils.py +++ b/python/sglang/srt/models/utils.py @@ -415,6 +415,9 @@ def maybe_enable_silu_fp4_quant_fusion(mlp: torch.nn.Module) -> None: # Marlin, the SM80-SM90 default, takes a floating-point input, not the tuple. if get_fp4_gemm_runner_backend().is_marlin(): return + # The fused kernel never writes the scale padding, which exists when K % 64 != 0. + if mlp.down_proj.input_size_per_partition % 64 != 0: + return try: from flashinfer import silu_and_mul_scaled_nvfp4_experts_quantize # noqa: F401 except ImportError: diff --git a/test/registered/unit/models/test_qwen3_5_modelopt_fp4.py b/test/registered/unit/models/test_qwen3_5_modelopt_fp4.py index 6beed915fe96..622eaa7d82b6 100644 --- a/test/registered/unit/models/test_qwen3_5_modelopt_fp4.py +++ b/test/registered/unit/models/test_qwen3_5_modelopt_fp4.py @@ -9,8 +9,9 @@ 3. QWEN3_5_KV_SCALE_MAPPER remaps the checkpoint's baked KV-scale names onto the RadixAttention parameter names. 4. The dense-MLP SiLU+FP4 fusion turns on for NVFP4 W4A4 and stays off where - down_proj cannot take its (fp4, scale) tuple: NVFP4 AWQ, the Marlin FP4 GEMM - backend, sequence parallelism, a LoRA-wrapped down_proj. + its (fp4, scale) tuple would be wrong or unusable: NVFP4 AWQ, the Marlin FP4 + GEMM backend, K not a multiple of 64, sequence parallelism, a LoRA-wrapped + down_proj. """ import sys @@ -188,13 +189,17 @@ def test_skips_under_sequence_parallelism(self): self.assertFalse(can_run_silu_fp4_quant_fusion(mlp, torch.empty(2, 32))) @staticmethod - def _enables_fusion(*, is_awq: bool, backend: Fp4GemmRunnerBackend) -> bool: + def _enables_fusion( + *, is_awq: bool, backend: Fp4GemmRunnerBackend, down_proj_k: int = 128 + ) -> bool: method = MagicMock(spec=ModelOptFp4LinearMethod, quant_mode="w4a4") method.quant_config = SimpleNamespace(is_awq=is_awq) mlp = SimpleNamespace( _enable_silu_fp4_quant_fusion=False, gate_up_proj=SimpleNamespace(quant_method=method), - down_proj=SimpleNamespace(quant_method=method), + down_proj=SimpleNamespace( + quant_method=method, input_size_per_partition=down_proj_k + ), ) # A stand-in FlashInfer keeps the result independent of the installed one. flashinfer = SimpleNamespace(silu_and_mul_scaled_nvfp4_experts_quantize=None) @@ -220,6 +225,17 @@ def test_stays_off_for_nvfp4_awq(self): ) ) + def test_stays_off_when_k_is_not_a_multiple_of_64(self): + """K = 1376 pads 86 scale columns to 88; the fused kernel never writes the + padding, so leftover NaN bytes there turned the whole GEMM output into NaN.""" + self.assertFalse( + self._enables_fusion( + is_awq=False, + backend=Fp4GemmRunnerBackend.FLASHINFER_CUTLASS, + down_proj_k=1376, + ) + ) + def test_stays_off_for_the_marlin_backend(self): # Marlin, the SM80-SM90 default, takes a floating-point input; # handing it the (fp4, scale) tuple crashes the first forward. From 8b512342aa749690aa53efc9276e3390ec4f6ef3 Mon Sep 17 00:00:00 2001 From: Hanjun Cho Date: Tue, 3 Nov 2026 19:30:11 +0900 Subject: [PATCH 4/4] [Model] Read the SiLU+FP4 fusion kill switch through envs Register SGLANG_DISABLE_SILU_FP4_QUANT_FUSION in Envs instead of reading os.environ in models/utils.py, log the enable message once per process rather than once per decoder layer, trim comments that restated the gates or named the caller, and state in each regression test's docstring the behavior it guards. Signed-off-by: Hanjun Cho --- python/sglang/srt/environ.py | 2 ++ python/sglang/srt/models/qwen2.py | 3 +-- python/sglang/srt/models/utils.py | 20 +++++++++---------- .../unit/models/test_qwen3_5_modelopt_fp4.py | 20 +++++++++---------- 4 files changed, 23 insertions(+), 22 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 711e6ae59860..716697e4f65d 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1255,6 +1255,8 @@ class Envs: SGLANG_MOE_NVFP4_DISPATCH = EnvBool(False) SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN = EnvBool(False) SGLANG_NVFP4_CKPT_FP8_NEXTN_MOE = EnvBool(False) + # Kill switch for fusing dense-MLP SiLU*mul with the NVFP4 down_proj input quant. + SGLANG_DISABLE_SILU_FP4_QUANT_FUSION = EnvBool(False) # GLM NextN (MTP): cast the draft layer's bf16 fused MoE to per-channel FP8 # on load. Unrelated to the NVFP4 block-FP8 NextN path above. SGLANG_GLM_NEXTN_MOE_PTPC = EnvBool(False) diff --git a/python/sglang/srt/models/qwen2.py b/python/sglang/srt/models/qwen2.py index 836dc352c6d9..6cef42c423c8 100644 --- a/python/sglang/srt/models/qwen2.py +++ b/python/sglang/srt/models/qwen2.py @@ -108,8 +108,7 @@ def __init__( self.act_fn = SiluAndMul() self.dual_gemm = DualGemm(self.gate_up_proj, self.down_proj, hidden_size) self.activation_quant = DenseSiluAndMulQuant(self.down_proj) - # maybe_enable_silu_fp4_quant_fusion turns this on for NVFP4 projections; - # SiLU+mul and the down_proj input quantization then run as one kernel. + # Off unless the decoder layer enables the fused SiLU+mul+NVFP4 quant path. self._enable_silu_fp4_quant_fusion = False self._masked_m_cache: dict = {} self._down_input_scale_inv_1d = None diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index a99b0533781f..c8bbb960614b 100644 --- a/python/sglang/srt/models/utils.py +++ b/python/sglang/srt/models/utils.py @@ -14,8 +14,6 @@ from __future__ import annotations import itertools -import logging -import os from collections.abc import Iterable, Mapping from dataclasses import dataclass, field from functools import lru_cache @@ -44,11 +42,15 @@ from sglang.srt.model_executor.runner import get_is_capture_mode from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.runtime_context import get_exec, get_forward, get_parallel -from sglang.srt.utils import get_current_device_stream_fast, is_cpu, is_cuda, is_hip +from sglang.srt.utils import ( + get_current_device_stream_fast, + is_cpu, + is_cuda, + is_hip, + print_info_once, +) from sglang.srt.utils.custom_op import register_custom_op -logger = logging.getLogger(__name__) - if TYPE_CHECKING: from sglang.srt.layers.layernorm import RMSNorm @@ -394,11 +396,9 @@ def maybe_enable_silu_fp4_quant_fusion(mlp: torch.nn.Module) -> None: Replaces the separate act_and_mul and per-token FP4 quantize kernels with one FlashInfer kernel and feeds down_proj a prequantized (fp4, scale) - tuple. Enabled when both dense-MLP projections use the NVFP4 (W4A4) - linear method without AWQ and the FP4 GEMM backend is not Marlin; - kill switch: SGLANG_DISABLE_SILU_FP4_QUANT_FUSION=1. + tuple. Kill switch: SGLANG_DISABLE_SILU_FP4_QUANT_FUSION=1. """ - if os.environ.get("SGLANG_DISABLE_SILU_FP4_QUANT_FUSION", "0") == "1": + if envs.SGLANG_DISABLE_SILU_FP4_QUANT_FUSION.get(): return from sglang.srt.layers.quantization.fp4_utils import get_fp4_gemm_runner_backend from sglang.srt.layers.quantization.modelopt_quant import ModelOptFp4LinearMethod @@ -424,7 +424,7 @@ def maybe_enable_silu_fp4_quant_fusion(mlp: torch.nn.Module) -> None: return mlp._enable_silu_fp4_quant_fusion = True mlp.down_proj._accepts_prequantized_fp4 = True - logger.info("Enabled fused SiLU+mul+FP4-quant for dense MLP down_proj input.") + print_info_once("Enabled fused SiLU+mul+FP4-quant for dense MLP down_proj input.") def can_run_silu_fp4_quant_fusion( diff --git a/test/registered/unit/models/test_qwen3_5_modelopt_fp4.py b/test/registered/unit/models/test_qwen3_5_modelopt_fp4.py index 622eaa7d82b6..1b3cdb39b7c0 100644 --- a/test/registered/unit/models/test_qwen3_5_modelopt_fp4.py +++ b/test/registered/unit/models/test_qwen3_5_modelopt_fp4.py @@ -189,7 +189,7 @@ def test_skips_under_sequence_parallelism(self): self.assertFalse(can_run_silu_fp4_quant_fusion(mlp, torch.empty(2, 32))) @staticmethod - def _enables_fusion( + def _fusion_turns_on( *, is_awq: bool, backend: Fp4GemmRunnerBackend, down_proj_k: int = 128 ) -> bool: method = MagicMock(spec=ModelOptFp4LinearMethod, quant_mode="w4a4") @@ -212,24 +212,24 @@ def _enables_fusion( def test_turns_on_for_w4a4_on_a_native_fp4_backend(self): self.assertTrue( - self._enables_fusion( + self._fusion_turns_on( is_awq=False, backend=Fp4GemmRunnerBackend.FLASHINFER_CUTLASS ) ) def test_stays_off_for_nvfp4_awq(self): - # The fused producer would skip AWQ's per-input-channel pre_quant_scale. + """The fused producer would skip AWQ's per-input-channel pre_quant_scale.""" self.assertFalse( - self._enables_fusion( + self._fusion_turns_on( is_awq=True, backend=Fp4GemmRunnerBackend.FLASHINFER_CUTLASS ) ) def test_stays_off_when_k_is_not_a_multiple_of_64(self): - """K = 1376 pads 86 scale columns to 88; the fused kernel never writes the - padding, so leftover NaN bytes there turned the whole GEMM output into NaN.""" + """K = 1376 pads 86 scale columns to 88, which the fused kernel never writes; + a NaN byte left there makes the whole GEMM output NaN.""" self.assertFalse( - self._enables_fusion( + self._fusion_turns_on( is_awq=False, backend=Fp4GemmRunnerBackend.FLASHINFER_CUTLASS, down_proj_k=1376, @@ -237,10 +237,10 @@ def test_stays_off_when_k_is_not_a_multiple_of_64(self): ) def test_stays_off_for_the_marlin_backend(self): - # Marlin, the SM80-SM90 default, takes a floating-point input; - # handing it the (fp4, scale) tuple crashes the first forward. + """Marlin, the SM80-SM90 default, takes a floating-point input; the + (fp4, scale) tuple crashes its first forward.""" self.assertFalse( - self._enables_fusion(is_awq=False, backend=Fp4GemmRunnerBackend.MARLIN) + self._fusion_turns_on(is_awq=False, backend=Fp4GemmRunnerBackend.MARLIN) )