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 a32957785919..6cef42c423c8 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,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 ) 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..c8bbb960614b 100644 --- a/python/sglang/srt/models/utils.py +++ b/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) 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..1b3cdb39b7c0 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,33 @@ 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()