[Perf] qwen3: Fuse SiLU*mul with the NVFP4 quantization of down_proj in Qwen2MLP

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)

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,10 @@ 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)
+        # 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
 
     def forward(
         self,
@@ -117,7 +125,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
                 )

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)

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(

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}")
 

python/sglang/srt/models/utils.py

@@ -41,8 +41,14 @@
 )
 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.utils import get_current_device_stream_fast, is_cpu, is_cuda, is_hip
+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,
+    print_info_once,
+)
 from sglang.srt.utils.custom_op import register_custom_op
 
 if TYPE_CHECKING:
@@ -385,6 +391,84 @@ 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. Kill switch: SGLANG_DISABLE_SILU_FP4_QUANT_FUSION=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
+
+    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
+    # 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:
+        return
+    mlp._enable_silu_fp4_quant_fusion = True
+    mlp.down_proj._accepts_prequantized_fp4 = True
+    print_info_once("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)

test/registered/unit/models/test_qwen3_5_modelopt_fp4.py

@@ -1,23 +1,40 @@
 """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
      declares kv_cache_quant_algo; without them, baked KV scales have nowhere to
      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 turns on for NVFP4 W4A4 and stays off where
+     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
 import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
 
 import torch
 
-from sglang.srt.layers.quantization.modelopt_quant import ModelOptFp4Config
+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,
+)
 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 +170,79 @@ 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)))
+
+    @staticmethod
+    def _fusion_turns_on(
+        *, 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, 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)
+        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._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."""
+        self.assertFalse(
+            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, which the fused kernel never writes;
+        a NaN byte left there makes the whole GEMM output NaN."""
+        self.assertFalse(
+            self._fusion_turns_on(
+                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; the
+        (fp4, scale) tuple crashes its first forward."""
+        self.assertFalse(
+            self._fusion_turns_on(is_awq=False, backend=Fp4GemmRunnerBackend.MARLIN)
+        )
+
+
 if __name__ == "__main__":
     unittest.main()