[perf] Fuse NVFP4 gate_up_gemm + swiglu + output FP4 quant (#26626)

.codespellrc

@@ -1,3 +1,3 @@
 [codespell]
-ignore-words-list = ans, als, hel, boostrap, childs, te, vas, hsa, ment, cann, thi, makro, wil, rouge, PRIS, ather, MIS, medias, allready, inout, nd, fo, visibles, nothink, renderD, ond, tbe, CopyIn, notin
+ignore-words-list = ans, als, hel, boostrap, childs, te, vas, hsa, ment, cann, thi, makro, wil, rouge, PRIS, ather, MIS, medias, allready, inout, nd, fo, visibles, nothink, renderD, ond, tbe, CopyIn, notin, subtile, subtiles
 skip = *.json, *.jsonl, *.patch, *.txt, *.lock

python/sglang/srt/environ.py

@@ -639,6 +639,7 @@ class Envs:
     SGLANG_OPT_USE_FUSED_COMPRESS_TRITON = EnvBool(False)
     SGLANG_OPT_USE_FUSED_QK_NORM_ROPE = EnvBool(True)
     SGLANG_OPT_USE_FUSED_CLAMP_ACT_MUL = EnvBool(True)
+    SGLANG_ENABLE_NVFP4_GEMM_SWIGLU_FUSION = EnvBool(True)
     SGLANG_FIX_MTP_HC_HIDDEN = EnvBool(False)
     # ====================================================================
 

python/sglang/srt/layers/quantization/modelopt_quant.py

@@ -1503,6 +1503,15 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
         K_padded = round_up_to_multiple(K, 4)
         padded_scales = torch.zeros((B, M_padded, K_padded), dtype=scales.dtype)
         padded_scales[:B, :M, :K] = scales
+
+        # Snapshot the raw (pre-swizzle) scale BEFORE alias_or_bind_derived_param
+        # overwrites layer.weight_scale.data in-place via .copy_() on the broadcast
+        # path. Without this, the swiglu side-channel below would read the swizzled
+        # bytes when it later re-reads layer.weight_scale.
+        raw_scale_snapshot = (
+            (scales.squeeze(0) if scale_ndim == 2 else scales).detach().clone()
+        )
+
         batches, rows, cols = padded_scales.shape
         assert rows % 128 == 0
         assert cols % 4 == 0
@@ -1518,23 +1527,73 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
             layer, "weight_scale", "weight_scale_interleaved", padded_scales
         )
 
+        if getattr(layer, "_interleave_for_swiglu_fusion", False):
+            from sglang.srt.layers.quantization.nvfp4_gemm_swiglu_nvfp4_quant import (
+                interleave_linear_and_gate,
+                swizzle_blockscale_2d,
+            )
+
+            w = layer.weight.data
+            assert weights_padding_cols == 0, (
+                "_interleave_for_swiglu_fusion does not support K-padded weights; "
+                f"got weights_padding_cols={weights_padding_cols}."
+            )
+            assert raw_scale_snapshot.shape[0] == w.shape[0], (
+                "_interleave_for_swiglu_fusion requires no N-padding; "
+                f"raw_scale rows={raw_scale_snapshot.shape[0]} vs weight rows={w.shape[0]}."
+            )
+            assert w.shape[0] % 128 == 0, (
+                "_interleave_for_swiglu_fusion requires N % 128 == 0 (group_size=64 "
+                f"with gate+up halves); got N={w.shape[0]}."
+            )
+
+            gate_w, up_w = w.chunk(2, dim=0)
+            w_swiglu = interleave_linear_and_gate(
+                torch.cat((up_w, gate_w), dim=0), group_size=64, dim=0
+            )
+
+            gate_s, up_s = raw_scale_snapshot.chunk(2, dim=0)
+            w_scale_swiglu = swizzle_blockscale_2d(
+                interleave_linear_and_gate(
+                    torch.cat((up_s, gate_s), dim=0), group_size=64, dim=0
+                )
+            )
+
+            layer.weight_swiglu_interleaved = w_swiglu
+            layer.weight_scale_swiglu_interleaved = w_scale_swiglu
+
+            # Keep the Parameter objects alive so weight reload can refill
+            # them and re-run this hook; free their storage in the meantime.
+            layer.weight.data = torch.empty(
+                0, dtype=layer.weight.dtype, device=layer.weight.device
+            )
+            layer.weight_scale_interleaved.data = torch.empty(
+                0,
+                dtype=layer.weight_scale_interleaved.dtype,
+                device=layer.weight_scale_interleaved.device,
+            )
+
     def apply(
         self,
         layer: torch.nn.Module,
         x: torch.Tensor,
         bias: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-        output_dtype = x.dtype
+        # `_accepts_prequantized_fp4` is the explicit opt-in so an accidental
-        x_m, _ = x.shape
+        # tuple from unrelated code can't silently bypass quantization.
+        if getattr(layer, "_accepts_prequantized_fp4", False) and isinstance(x, tuple):
+            x_fp4, x_scale_interleaved = x
+            x_m = x_fp4.shape[0]
+            output_dtype = layer.params_dtype
+        else:
+            x_fp4, x_scale_interleaved = fp4_quantize(x, layer.input_scale_inv)
+            x_m, _ = x.shape
+            output_dtype = x.dtype
 
-        # Get original output size (before padding) and padded weight size
         output_size = layer.output_size_per_partition
         w_n, _ = layer.weight.shape
         output_shape = [x_m, output_size]
 
-        # Quantize BF16 or FP16 to (FP4 and interleaved block scale)
-        x_fp4, x_scale_interleaved = fp4_quantize(x, layer.input_scale_inv)
-
         assert x_fp4.dtype == torch.uint8
         assert layer.weight.dtype == torch.uint8
         assert layer.weight_scale_interleaved.dtype == torch.float8_e4m3fn

python/sglang/srt/models/deepseek_v2.py

@@ -275,6 +275,35 @@ class DeepseekV2MLP(nn.Module):
         if (self.tp_size == 1) and x.shape[0] == 0:
             return x
 
+        if (
+            getattr(self, "_enable_nvfp4_gemm_swiglu_fusion", False)
+            and self.swiglu_limit is None
+            and not isinstance(x, tuple)
+        ):
+            from flashinfer import fp4_quantize
+
+            from sglang.srt.layers.quantization.nvfp4_gemm_swiglu_nvfp4_quant import (
+                nvfp4_gemm_swiglu_nvfp4_quant,
+            )
+
+            x_fp4, x_scale = fp4_quantize(
+                x, self.gate_up_proj.input_scale_inv, enable_pdl=True
+            )
+            out_fp4, out_scale = nvfp4_gemm_swiglu_nvfp4_quant(
+                x_fp4,
+                x_scale,
+                self.gate_up_proj.weight_swiglu_interleaved,
+                self.gate_up_proj.weight_scale_swiglu_interleaved,
+                self.gate_up_proj.alpha,
+                self.down_proj.input_scale_inv,
+                enable_pdl=True,
+            )
+            out, _ = self.down_proj(
+                (out_fp4, out_scale),
+                skip_all_reduce=should_allreduce_fusion or use_reduce_scatter,
+            )
+            return out
+
         if (
             gemm_output_zero_allocator is not None
             and x.shape[0] <= 256
@@ -673,6 +702,32 @@ class DeepseekV2MoE(nn.Module):
                 prefix=add_prefix("shared_experts", prefix),
                 **(dict(tp_rank=0, tp_size=1) if _shared_expert_use_tp1 else {}),
             )
+            # Flags must be set before weight load so
+            # process_weights_after_loading sees them and builds the
+            # [Up, Gate]-interleaved weight + scale.
+            from sglang.srt.layers.quantization.modelopt_quant import (
+                ModelOptFp4LinearMethod,
+            )
+            from sglang.srt.utils.common import is_sm100_supported
+
+            fc1_n = self.shared_experts.gate_up_proj.output_size_per_partition
+            if (
+                envs.SGLANG_ENABLE_NVFP4_GEMM_SWIGLU_FUSION.get()
+                and is_sm100_supported()
+                and isinstance(
+                    self.shared_experts.gate_up_proj.quant_method,
+                    ModelOptFp4LinearMethod,
+                )
+                and isinstance(
+                    self.shared_experts.down_proj.quant_method,
+                    ModelOptFp4LinearMethod,
+                )
+                and fc1_n % 128 == 0
+                and get_global_server_args().disable_piecewise_cuda_graph
+            ):
+                self.shared_experts.gate_up_proj._interleave_for_swiglu_fusion = True
+                self.shared_experts._enable_nvfp4_gemm_swiglu_fusion = True
+                self.shared_experts.down_proj._accepts_prequantized_fp4 = True
             self._shared_expert_tp1 = _shared_expert_use_tp1
             is_packed_weight = hasattr(
                 self.shared_experts.gate_up_proj.quant_method, "quant_config"