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

Refactor fp8 nextn layer for DeepSeek nvfp4 checkpoint (#15353)

Browse Source
This commit is contained in:
Baizhou Zhang
2025-12-28 11:57:09 +08:00
committed by GitHub
parent 8e43980ebb
commit 656f4d69a1
6 changed files with 115 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
python/sglang/srt/environ.py
View File

@@ -253,6 +253,7 @@ class Envs:
SGLANG_MOE_NVFP4_DISPATCH = EnvBool(False)
SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN = EnvBool(False)
SGLANG_PER_TOKEN_GROUP_QUANT_8BIT_V2 = EnvBool(False)
SGLANG_NVFP4_CKPT_FP8_NEXTN_MOE = EnvBool(False)
# Flashinfer
SGLANG_IS_FLASHINFER_AVAILABLE = EnvBool(True)

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

@@ -667,6 +667,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
),
requires_grad=False,
)
# w13_weight and w2_weight are always requanted together
w13_weight_scale.format_ue8m0 = False
w2_weight_scale.format_ue8m0 = False
layer.register_parameter("w13_weight_scale_inv", w13_weight_scale)
layer.register_parameter("w2_weight_scale_inv", w2_weight_scale)
assert self.quant_config.activation_scheme == "dynamic"
@@ -814,6 +817,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
),
)
and get_moe_runner_backend().is_deep_gemm()
and not layer.w13_weight_scale_inv.format_ue8m0
):
assert isinstance(
layer, DeepEPMoE
@@ -825,6 +829,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
requant_weight_ue8m0_inplace(
layer.w2_weight, layer.w2_weight_scale_inv, weight_block_size
)
layer.w13_weight_scale_inv.format_ue8m0 = True
layer.w2_weight_scale_inv.format_ue8m0 = True
return
# If checkpoint is fp16 or bfloat16, quantize in place.

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

@@ -71,12 +71,12 @@ class DeepseekModelNextN(nn.Module):
super().__init__()
if enable_nextn_moe_bf16_cast_to_fp8(quant_config):
# refer to real DeepSeek V3 quant config
moe_quant_config = Fp8Config(
moe_quant_config_override = Fp8Config(
is_checkpoint_fp8_serialized=True,
weight_block_size=[128, 128],
)
else:
moe_quant_config = None
moe_quant_config_override = None
if quant_config is not None and quant_config.get_name() == "modelopt_fp4":
logger.warning(
@@ -115,7 +115,7 @@ class DeepseekModelNextN(nn.Module):
config,
0,
quant_config=quant_config,
moe_quant_config=moe_quant_config,
moe_quant_config_override=moe_quant_config_override,
is_nextn=True,
prefix=add_prefix(layer_name, prefix),
alt_stream=self.alt_stream,

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

@@ -94,7 +94,7 @@ from sglang.srt.layers.moe import (
get_moe_runner_backend,
should_use_flashinfer_cutlass_moe_fp4_allgather,
)
from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE, get_moe_impl_class
from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
from sglang.srt.layers.moe.kt_ep_wrapper import KTEPWrapperMethod
from sglang.srt.layers.moe.token_dispatcher.base import (
@@ -119,7 +119,6 @@ from sglang.srt.layers.quantization.fp8_utils import (
inverse_transform_scale_ue8m0,
normalize_e4m3fn_to_e4m3fnuz,
quant_weight_ue8m0,
transform_scale_ue8m0_inplace,
)
from sglang.srt.layers.quantization.int8_utils import (
block_dequant as int8_block_dequant,
@@ -232,11 +231,16 @@ _is_cublas_ge_129 = is_nvidia_cublas_cu12_version_ge_12_9()
logger = logging.getLogger(__name__)
# Optional quantization for DeepSeek nvfp4 checkpoint
NVFP4_CKPT_FP8_ATTN_QUANT_MODULES = ["q_b_proj"]
def enable_nextn_moe_bf16_cast_to_fp8(quant_config):
return (
quant_config is not None
envs.SGLANG_NVFP4_CKPT_FP8_NEXTN_MOE.get()
and quant_config is not None
and quant_config.get_name() == "modelopt_fp4"
and get_moe_a2a_backend().is_deepep()
and get_moe_runner_backend().is_deep_gemm()
)
@@ -2742,7 +2746,7 @@ class DeepseekV2DecoderLayer(nn.Module):
config: PretrainedConfig,
layer_id: int,
quant_config: Optional[QuantizationConfig] = None,
moe_quant_config: Optional[QuantizationConfig] = None,
moe_quant_config_override: Optional[QuantizationConfig] = None,
is_nextn: bool = False,
prefix: str = "",
alt_stream: Optional[torch.cuda.Stream] = None,
@@ -2795,7 +2799,7 @@ class DeepseekV2DecoderLayer(nn.Module):
if self.is_layer_sparse:
self.mlp = DeepseekV2MoE(
config=config,
quant_config=moe_quant_config or quant_config,
quant_config=moe_quant_config_override or quant_config,
prefix=add_prefix("mlp", prefix),
layer_id=self.layer_id,
alt_stream=alt_stream,
@@ -3621,47 +3625,6 @@ 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 is_nextn and enable_nextn_moe_bf16_cast_to_fp8(self.quant_config):
self._transform_scale_nextn_moe_ue8m0()
# TODO avoid code dup (currently combine from weight_requant_ue8m0 and transform_scale_ue8m0)
def _transform_scale_nextn_moe_ue8m0(self):
layer = self.model.decoder
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,
]:
transform_scale_ue8m0_inplace(
module.weight_scale_inv, mn=module.weight.shape[-2]
)
experts = layer.mlp.experts
w13_weight_fp8 = (
experts.w13_weight,
(
experts.w13_weight_scale_inv
if hasattr(experts, "w13_weight_scale_inv")
else experts.w13_weight_scale
),
)
w2_weight_fp8 = (
experts.w2_weight,
(
experts.w2_weight_scale_inv
if hasattr(experts, "w2_weight_scale_inv")
else experts.w2_weight_scale
),
)
if isinstance(experts, DeepEPMoE):
for w in [
w13_weight_fp8,
w2_weight_fp8,
]:
transform_scale_ue8m0_inplace(w[1], mn=w[0].shape[-2])
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]], is_nextn=False):
if is_nextn:
@@ -3677,12 +3640,9 @@ class DeepseekV2ForCausalLM(nn.Module):
else:
raise ValueError("num_nextn_predict_layers is not in the config")
if envs.SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN.get():
weights = self._quant_attn_to_fp8_ue8m0(weights, is_nextn=is_nextn)
if is_nextn and enable_nextn_moe_bf16_cast_to_fp8(self.quant_config):
weights = self._quant_nextn_moe_to_fp8_ue8m0(
weights, nextn_layer_id=nextn_layer_id
)
weights = self._maybe_quant_weights_to_fp8_ue8m0(
weights, NVFP4_CKPT_FP8_ATTN_QUANT_MODULES, is_nextn
)
stacked_params_mapping = [
# (param_name, shard_name, shard_id)
@@ -3945,62 +3905,6 @@ class DeepseekV2ForCausalLM(nn.Module):
self.post_load_weights(is_nextn=is_nextn, weight_names=weight_names)
def _quant_attn_to_fp8_ue8m0(self, weights, is_nextn):
weights_dict = dict(weights)
# temporarily only support DeepSeek V3/R1
weight_block_size = [128, 128]
for layer_id in tqdm.trange(
self.config.num_hidden_layers + int(is_nextn),
desc="quant attn to fp8 ue8m0",
):
for stem in [
# may put tensors like `o_proj` here for DeepSeek FP4 ckpt v1
"q_b_proj",
]:
partial_name = f"model.layers.{layer_id}.self_attn.{stem}"
original_weight = weights_dict[f"{partial_name}.weight"]
out_w, out_s = quant_weight_ue8m0(
original_weight, weight_block_size=weight_block_size
)
weights_dict[f"{partial_name}.weight"] = out_w
weights_dict[f"{partial_name}.weight_scale_inv"] = out_s
return list(weights_dict.items())
# TODO avoid code dup
def _quant_nextn_moe_to_fp8_ue8m0(self, weights, nextn_layer_id: int):
weights_dict = dict(weights)
# temporarily only support DeepSeek V3/R1
weight_block_size = [128, 128]
for layer_id in [nextn_layer_id]:
for expert_sub_name in [
"shared_experts",
*[
f"experts.{expert_id}"
for expert_id in range(self.config.n_routed_experts)
],
]:
for stem in [
"gate_proj",
"up_proj",
"down_proj",
]:
partial_name = (
f"model.layers.{layer_id}.mlp.{expert_sub_name}.{stem}"
)
original_weight = weights_dict[f"{partial_name}.weight"]
out_w, out_s = quant_weight_ue8m0(
original_weight, weight_block_size=weight_block_size
)
weights_dict[f"{partial_name}.weight"] = out_w
weights_dict[f"{partial_name}.weight_scale_inv"] = out_s
return list(weights_dict.items())
def get_embed_and_head(self):
return self.model.embed_tokens.weight, self.lm_head.weight
@@ -4034,6 +3938,76 @@ class DeepseekV2ForCausalLM(nn.Module):
# of the (i-1)th layer as aux hidden state
self.model.layers_to_capture = [val + 1 for val in layer_ids]
# Mark the ue8m0 flag of nextn moe weights as True to avoid requantization
def _mark_nextn_moe_weights_as_ue8m0(self):
experts = self.model.decoder.mlp.experts
w13_scale = (
experts.w13_weight_scale_inv
if hasattr(experts, "w13_weight_scale_inv")
else experts.w13_weight_scale
)
w2_scale = (
experts.w2_weight_scale_inv
if hasattr(experts, "w2_weight_scale_inv")
else experts.w2_weight_scale
)
w13_scale.format_ue8m0 = True
w2_scale.format_ue8m0 = True
def _maybe_quant_weights_to_fp8_ue8m0(
self, weights, attn_quant_modules, is_nextn=False
):
# Quantize some weights to fp8 ue8m0 for DeepSeek nvfp4 checkpoint
partial_names = []
nextn_layer_id = (
0 if self.config.num_hidden_layers == 1 else self.config.num_hidden_layers
)
weights_dict = dict(weights)
weight_block_size = [128, 128]
if envs.SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN.get():
layer_ids = (
list(range(self.config.num_hidden_layers))
if not is_nextn
else [nextn_layer_id]
)
for layer_id in layer_ids:
for stem in attn_quant_modules:
partial_names.append(f"model.layers.{layer_id}.self_attn.{stem}")
if is_nextn and enable_nextn_moe_bf16_cast_to_fp8(self.quant_config):
for expert_sub_name in [
"shared_experts",
*[
f"experts.{expert_id}"
for expert_id in range(self.config.n_routed_experts)
],
]:
for stem in [
"gate_proj",
"up_proj",
"down_proj",
]:
partial_names.append(
f"model.layers.{nextn_layer_id}.mlp.{expert_sub_name}.{stem}"
)
for partial_name in tqdm.tqdm(
partial_names,
desc="quant weights to fp8 ue8m0",
):
original_weight = weights_dict[f"{partial_name}.weight"]
out_w, out_s = quant_weight_ue8m0(
original_weight, weight_block_size=weight_block_size
)
weights_dict[f"{partial_name}.weight"] = out_w
weights_dict[f"{partial_name}.weight_scale_inv"] = out_s
if is_nextn and enable_nextn_moe_bf16_cast_to_fp8(self.quant_config):
self._mark_nextn_moe_weights_as_ue8m0()
return list(weights_dict.items())
AttentionBackendRegistry.register("ascend", handle_attention_ascend)
AttentionBackendRegistry.register("flashinfer", handle_attention_flashinfer)

21
python/sglang/srt/server_args.py
View File

@@ -1138,6 +1138,27 @@ class ServerArgs:
"Use flashinfer_trtllm as MoE runner backend on sm100 for DeepseekV3ForCausalLM"
)
if (
self.quantization == "modelopt_fp4"
and self.speculative_algorithm == "EAGLE"
and (
self.speculative_moe_runner_backend is None
or self.speculative_moe_a2a_backend is None
)
):
if envs.SGLANG_NVFP4_CKPT_FP8_NEXTN_MOE.get():
self.speculative_moe_runner_backend = "deep_gemm"
self.speculative_moe_a2a_backend = "deepep"
logger.info(
"Use deep_gemm moe runner and deepep a2a backend for bf16 nextn layer in deepseek fp4 checkpoint."
)
else:
self.speculative_moe_runner_backend = "triton"
self.speculative_moe_a2a_backend = "none"
logger.info(
"Use triton fused moe by default for bf16 nextn layer in deepseek fp4 checkpoint."
)
elif model_arch in ["GptOssForCausalLM"]:
# Set attention backend for GPT-OSS
if self.is_attention_backend_not_set():