From d0913fca8dc5be24c27221a06ff8193954f54846 Mon Sep 17 00:00:00 2001 From: Yuhao Yang <47235274+yhyang201@users.noreply.github.com> Date: Wed, 13 May 2026 22:36:44 +0800 Subject: [PATCH] Port fused SiLU+clamp+FP8 quant from DSV4 dev branch (#24897) Co-authored-by: Cheng Wan Co-authored-by: fzyzcjy Co-authored-by: zcnrex --- python/sglang/srt/models/deepseek_v2.py | 57 ++++++++++++++++++++++--- 1 file changed, 51 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index cd02b9c0b..2d623b83d 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -27,6 +27,10 @@ import torch.nn.functional as F from torch import nn from transformers import PretrainedConfig +from sglang.jit_kernel.deepseek_v4 import ( + silu_and_mul_clamp, + silu_and_mul_contig_post_quant, +) from sglang.srt.batch_overlap.single_batch_overlap import SboFlags, compute_overlap_args from sglang.srt.batch_overlap.two_batch_overlap import ( MaybeTboDeepEPDispatcher, @@ -107,6 +111,9 @@ from sglang.srt.layers.moe.utils import ( ) from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.quantization.fp8 import Fp8Config +from sglang.srt.layers.quantization.fp8_kernel import ( + create_per_token_group_quant_fp8_output_scale, +) from sglang.srt.layers.quantization.mxfp4_flashinfer_trtllm_moe import ( maybe_fuse_routed_scale_and_shared_add, ) @@ -269,13 +276,51 @@ class DeepseekV2MLP(nn.Module): x = (x, None, y) gate_up, _ = self.gate_up_proj(x) - if self.swiglu_limit is not None: - _g, _u = gate_up.chunk(2, dim=-1) - _lim = float(self.swiglu_limit) - gate_up = torch.cat( - [_g.clamp(max=_lim), _u.clamp(min=-_lim, max=_lim)], dim=-1 + # Fast path: fused silu+clamp+fp8_quant+deepgemm when conditions met. + # Only valid when down_proj does NOT need an all-reduce and its weights + # are fp8 (uint8 storage with weight_scale_inv). + if ( + self.swiglu_limit is not None + and not self.down_proj.reduce_results + and self.down_proj.weight.dtype == torch.uint8 + and hasattr(self.down_proj, "weight_scale_inv") + ): + M, N = gate_up.shape + down_input_fp8 = gate_up.new_empty((M, N // 2), dtype=torch.float8_e4m3fn) + scale_block_size = 128 + down_input_scale = create_per_token_group_quant_fp8_output_scale( + x_shape=(M, N // 2), + device=gate_up.device, + group_size=scale_block_size, + column_major_scales=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + scale_tma_aligned=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, ) - x = self.act_fn(gate_up) + silu_and_mul_contig_post_quant( + input=gate_up, + output=down_input_fp8, + output_scale=down_input_scale, + quant_group_size=scale_block_size, + scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + transposed=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + swiglu_limit=float(self.swiglu_limit), + ) + down_output = gate_up.new_empty( + (M, self.down_proj.output_size), dtype=torch.bfloat16 + ) + deep_gemm_wrapper.gemm_nt_f8f8bf16( + (down_input_fp8, down_input_scale), + (self.down_proj.weight, self.down_proj.weight_scale_inv), + down_output, + ) + return down_output + # Fallback: fused silu+clamp kernel (still faster than unfused) + if self.swiglu_limit is not None: + M, N = gate_up.shape + x = gate_up.new_empty((M, N // 2)) + silu_and_mul_clamp(gate_up, x, float(self.swiglu_limit)) + else: + x = self.act_fn(gate_up) x, _ = self.down_proj( x, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter,