chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[1/2] Refactor DeepGeem requant for FP8 Linear on Blackwell (#13601)

Browse Source 
Co-authored-by: fy1214
This commit is contained in:
Baizhou Zhang
2025-11-23 16:07:56 -08:00
committed by GitHub
parent 9054e844ea
commit 4683e244fe
4 changed files with 39 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
python/sglang/srt/layers/quantization/fp8.py
View File

@@ -44,6 +44,7 @@ from sglang.srt.layers.quantization.fp8_utils import (
dispatch_w8a8_block_fp8_linear,
input_to_float8,
normalize_e4m3fn_to_e4m3fnuz,
requant_weight_ue8m0_inplace,
)
from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod
from sglang.srt.layers.quantization.marlin_utils_fp8 import (
@@ -261,6 +262,7 @@ class Fp8LinearMethod(LinearMethodBase):
layer.input_size_per_partition = input_size_per_partition
layer.output_size_per_partition = output_size_per_partition
layer.orig_dtype = params_dtype
layer.executed_weight_requant_ue8m0 = False
# WEIGHT
weight_dtype = (
@@ -347,7 +349,34 @@ class Fp8LinearMethod(LinearMethodBase):
)
return
else:
# For fp8 linear weights run with deepgemm, the weights and scales need be requantized to ue8m0
from sglang.srt.layers.quantization.fp8_utils import (
deepgemm_w8a8_block_fp8_linear_with_fallback,
)
from sglang.srt.model_loader.utils import (
should_deepgemm_weight_requant_ue8m0,
)
if (
should_deepgemm_weight_requant_ue8m0(
weight_block_size=getattr(
self.quant_config, "weight_block_size", None
),
)
and (
self.w8a8_block_fp8_linear
is deepgemm_w8a8_block_fp8_linear_with_fallback
)
and (not layer.executed_weight_requant_ue8m0)
):
requant_weight_ue8m0_inplace(
layer.weight,
layer.weight_scale_inv,
self.quant_config.weight_block_size,
)
layer.executed_weight_requant_ue8m0 = True
weight, weight_scale = layer.weight.data, layer.weight_scale_inv.data
layer.weight.data = weight.data
layer.weight_scale_inv.data = weight_scale.data
else:

1
python/sglang/srt/models/deepseek_nextn.py
View File

@@ -220,7 +220,6 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM):
use_attn_tp_group=get_global_server_args().enable_dp_lm_head,
)
self.logits_processor = LogitsProcessor(config)
self._executed_weight_requant_ue8m0 = False
@torch.no_grad()
def forward(

86
python/sglang/srt/models/deepseek_v2.py
View File

@@ -108,7 +108,6 @@ from sglang.srt.layers.quantization.fp8_kernel import (
per_token_group_quant_mla_deep_gemm_masked_fp8,
)
from sglang.srt.layers.quantization.fp8_utils import (
ENABLE_FLASHINFER_FP8_GEMM,
block_quant_dequant,
block_quant_to_tensor_quant,
channel_quant_to_tensor_quant,
@@ -3419,7 +3418,6 @@ class DeepseekV2ForCausalLM(nn.Module):
}
)
self.capture_aux_hidden_states = False
self._executed_weight_requant_ue8m0 = False
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if self.nsa_enable_prefill_cp:
@@ -3594,13 +3592,14 @@ class DeepseekV2ForCausalLM(nn.Module):
weight = w
weight_scale = self_attn.kv_b_proj.weight_scale_inv
# In multiple weight loading scenarios (e.g. RL), we need to inverse the scale of the weights after the requantization happened at the first loading.
if (
should_deepgemm_weight_requant_ue8m0(
weight_block_size=getattr(
self.quant_config, "weight_block_size", None
)
)
and self._executed_weight_requant_ue8m0
and self_attn.kv_b_proj.executed_weight_requant_ue8m0
):
weight_scale = inverse_transform_scale_ue8m0(
weight_scale, mn=weight.shape[-2]
@@ -3716,27 +3715,15 @@ class DeepseekV2ForCausalLM(nn.Module):
self_attn.w_vc = bind_or_assign(self_attn.w_vc, w_vc.contiguous())
self_attn.use_deep_gemm_bmm = True
if (
not ENABLE_FLASHINFER_FP8_GEMM
and should_deepgemm_weight_requant_ue8m0(
weight_block_size=getattr(self.quant_config, "weight_block_size", None)
)
and not self._executed_weight_requant_ue8m0
):
self._executed_weight_requant_ue8m0 = True
self._weight_requant_ue8m0(is_nextn)
# TODO can move weight_requant_ue8m0 and transform_scale_ue8m0 into Fp8LinearMethod.process_weights_after_loading
if (
deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
and get_bool_env_var("SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN")
):
self._transform_scale_ue8m0(is_nextn)
# Requant the weights and scales of MoE layers
if get_moe_runner_backend().is_deep_gemm():
self._maybe_moe_weight_requant_ue8m0(is_nextn)
if is_nextn and enable_nextn_moe_bf16_cast_to_fp8(self.quant_config):
self._transform_scale_nextn_moe_ue8m0()
def _weight_requant_ue8m0(self, is_nextn=False):
def _maybe_moe_weight_requant_ue8m0(self, is_nextn=False):
# Dense fp8 layers will be processed in Fp8LinearMethod.process_weights_after_loading
# So we only need to process sparse MoE layers here
weight_block_size = self.quant_config.weight_block_size
moe_layers = list(
@@ -3755,70 +3742,15 @@ class DeepseekV2ForCausalLM(nn.Module):
else:
layer = self.model.layers[layer_id]
module_list = [
layer.self_attn.kv_b_proj,
layer.self_attn.o_proj,
]
if self.config.q_lora_rank is not None:
module_list.append(layer.self_attn.fused_qkv_a_proj_with_mqa)
module_list.append(layer.self_attn.q_b_proj)
else:
module_list.append(layer.self_attn.kv_a_proj_with_mqa)
module_list.append(layer.self_attn.q_proj)
for module in module_list:
requant_weight_ue8m0_inplace(
module.weight, module.weight_scale_inv, weight_block_size
)
if layer_id in moe_layers or is_nextn:
shared_experts = getattr(layer.mlp, "shared_experts", None)
if shared_experts is not None:
for module in [
shared_experts.gate_up_proj,
shared_experts.down_proj,
]:
requant_weight_ue8m0_inplace(
module.weight, module.weight_scale_inv, weight_block_size
)
experts = layer.mlp.experts
# TODO: move this logic to Fp8MoEMethod.process_weights_after_loading
if isinstance(experts, DeepEPMoE):
for w in [
(experts.w13_weight, experts.w13_weight_scale_inv),
(experts.w2_weight, experts.w2_weight_scale_inv),
]:
requant_weight_ue8m0_inplace(w[0], w[1], weight_block_size)
else:
mlp = layer.mlp
assert isinstance(mlp, DeepseekV2MLP)
for module in [
mlp.gate_up_proj,
mlp.down_proj,
]:
requant_weight_ue8m0_inplace(
module.weight, module.weight_scale_inv, weight_block_size
)
# TODO can move weight_requant_ue8m0 and transform_scale_ue8m0 into Fp8LinearMethod.process_weights_after_loading
def _transform_scale_ue8m0(self, is_nextn=False):
num_hidden_layers = 1 if is_nextn else self.config.num_hidden_layers
for layer_id in range(num_hidden_layers):
if is_nextn:
layer = self.model.decoder
else:
layer = self.model.layers[layer_id]
module_list = []
if self.config.q_lora_rank is not None:
module_list.append(layer.self_attn.q_b_proj)
for module in module_list:
transform_scale_ue8m0_inplace(
module.weight_scale_inv, mn=module.weight.shape[-2]
)
# TODO avoid code dup (currently combine from weight_requant_ue8m0 and transform_scale_ue8m0)
def _transform_scale_nextn_moe_ue8m0(self):

2
python/sglang/test/test_block_fp8_deep_gemm_blackwell.py
View File

@@ -222,7 +222,7 @@ class TestDeepGemmBlackwell(CustomTestCase):
with torch.inference_mode():
ref_out = native_w8a8_block_fp8_matmul(
A_q, B_q, A_s, B_s, block_size, out_dtype
A_qu[0], B_qu[0], A_qu[1], B_qu[1], block_size, out_dtype
)
out = torch.empty_like(ref_out)
fp8_gemm_nt(A_qu, B_qu, out)