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

B300 NextN fp4: decide keep-vs-drop from the checkpoint, not the MoE runner backend

Browse Source 
The runner-backend heuristic is fundamentally the wrong signal: it reads AUTO in the EAGLE draft
worker (verified live), and an fp4-capable runner does not imply an fp4 MTP — DeepSeek-V3-0324-FP4
runs --moe-runner-backend flashinfer_trtllm yet ships a bf16 MTP (CI test_deepseek_v3_fp4_mtp_small).
is_layer_excluded is also disqualified: R1/V3 leave the MTP unquantized-but-absent-from-`ignore`.

Complete fix (opus-designed): inspect the checkpoint's model.safetensors.index.json and keep fp4 iff
the NextN/MTP experts (model.layers.{num_hidden_layers}.mlp.experts.*weight_scale*) ship fp4 scales —
the ground truth. GLM-5.2-NVFP4 has them (verified: 1536 layer-78 scale keys) -> keep; R1/V3 don't
-> drop. Unreadable index -> drop (PR #7376 default, bf16 MTP fallback, never crashes). Removed the
backend heuristic + debug log + unused get_moe_runner_backend import. Resolves both the 6144-vs-3072
draft crash AND the flood of "model.decoder.mlp.experts.*_scale not found in params_dict" (same root
cause: dropping fp4 left the unquant MTP with no scale params for the checkpoint's fp4 scales).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-21 16:48:05 +00:00
parent a3d0c2394b
commit 1d96197d3d
1 changed files with 102 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -45,7 +45,6 @@ from sglang.srt.layers.dp_attention import (
)
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.moe import get_moe_runner_backend
from sglang.srt.layers.quantization import Fp8Config
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.vocab_parallel_embedding import (
@@ -79,44 +78,116 @@ def _log_eagle_accept_cp_draft_hidden_debug(key: str, message: str, *args):
)
def _load_safetensors_index_weight_map(
model_dir: str, revision: Optional[str]
) -> Optional[dict]:
"""Return the safetensors index ``weight_map``, or None if it can't be read.
Reads a local snapshot dir first (the common case — server-arg model paths are
resolved to local dirs at init); falls back to fetching just the index file
from the HF hub for repo-id paths.
"""
import glob
import json
import os
if os.path.isdir(model_dir):
for index_path in glob.glob(
os.path.join(model_dir, "*.safetensors.index.json")
):
try:
with open(index_path) as f:
return json.load(f).get("weight_map", {})
except Exception as e: # noqa: BLE001
logger.warning("Failed to read %s: %s", index_path, e)
return None
try:
import huggingface_hub
from huggingface_hub import hf_hub_download
from sglang.srt.model_loader.weight_utils import (
download_safetensors_index_file_from_hf,
)
index_file = "model.safetensors.index.json"
download_safetensors_index_file_from_hf(model_dir, index_file, None, revision)
index_path = hf_hub_download(
repo_id=model_dir,
filename=index_file,
revision=revision,
local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
)
with open(index_path) as f:
return json.load(f).get("weight_map", {})
except Exception as e: # noqa: BLE001
logger.debug("Could not fetch safetensors index for %s: %s", model_dir, e)
return None
def _nextn_moe_is_fp4_quantized_in_checkpoint(
config: PretrainedConfig,
) -> Optional[bool]:
"""True/False if the checkpoint's NextN/MTP MoE experts ship fp4 scales, else None.
The MTP/NextN decoder lives at HF prefix ``model.layers.{num_hidden_layers}``
(the same index the loader uses to remap onto ``model.decoder``). fp4-quantized
MoE experts always carry ``weight_scale`` / ``weight_scale_2`` tensors; an
unquantized (bf16) MTP MoE (DeepSeek-R1-FP4 / V3-0324-FP4) does not.
"""
nextn_layer_id = getattr(config, "num_hidden_layers", None)
if nextn_layer_id is None:
return None
server_args = get_global_server_args()
if server_args.speculative_draft_model_path:
model_dir = server_args.speculative_draft_model_path
revision = server_args.speculative_draft_model_revision
else:
model_dir = server_args.model_path
revision = server_args.revision
if not model_dir:
return None
weight_map = _load_safetensors_index_weight_map(model_dir, revision)
if weight_map is None:
return None
expert_prefix = f"model.layers.{nextn_layer_id}.mlp.experts."
return any(
name.startswith(expert_prefix) and "weight_scale" in name
for name in weight_map
)
def _should_drop_nextn_modelopt_fp4_quant_config(
quant_config: Optional[QuantizationConfig],
config: PretrainedConfig,
) -> bool:
"""Whether to drop the modelopt_fp4 quant config for the NextN/MTP module.
DeepSeek-R1-FP4 / V3-0324-FP4 ship an *unquantized* MTP module (no fp4 scales
for the MTP MoE, and their hf_quant_config does NOT list the MTP layer in
`ignore`); for those we must drop fp4 so the whole MTP loads bf16 (PR #7376).
GLM-5.2-NVFP4 instead ships fp4 routed experts on the MTP layer and runs on a
flashinfer fp4 MoE backend that consumes them — there we keep fp4 and let the
exclude-list make attn / indexer / shared_experts / gate bf16.
DeepSeek-R1-FP4 / V3-0324-FP4 ship an *unquantized* MTP MoE (no fp4 scales in
the checkpoint) and do NOT list the MTP layer in `ignore`; for those we drop
fp4 so the whole MTP loads bf16 (PR #7376). GLM-5.2-NVFP4 ships fp4 routed
experts on the MTP layer; there we keep fp4 and let the exclude-list make
attn / indexer / shared_experts / gate bf16.
The decision is strictly per-checkpoint. The MoE runner backend is the WRONG
signal: it reads AUTO in the EAGLE draft worker, and an fp4-capable runner does
not imply an fp4 MTP (e.g. DeepSeek-V3-0324-FP4 runs flashinfer_trtllm yet ships
a bf16 MTP). So we inspect the checkpoint's safetensors index for the MTP MoE's
weight_scale tensors — the ground truth.
"""
if quant_config is None or quant_config.get_name() != "modelopt_fp4":
return False
# The NextN/MTP draft is built by the EAGLE draft worker, where the effective
# MoE runner may be the *speculative* backend (which is AUTO when
# --speculative-moe-runner-backend is unset), not the main runner. Keep fp4 if
# EITHER the main or the speculative runner is a flashinfer fp4 backend.
from sglang.srt.layers.moe.utils import get_speculative_moe_runner_backend
def _is_fp4_capable(b) -> bool:
return (
b.is_flashinfer_trtllm()
or b.is_flashinfer_trtllm_routed()
or b.is_flashinfer_cutedsl()
or b.is_flashinfer_cutlass()
is_fp4 = _nextn_moe_is_fp4_quantized_in_checkpoint(config)
if is_fp4 is None:
# Couldn't read the index. Preserve the historical default (PR #7376): drop
# fp4 so the MTP loads bf16 — safe for R1/V3, and degrades a (would-be) fp4
# MTP to bf16 rather than crashing on an unreadable checkpoint.
logger.warning(
"Could not determine NextN MoE fp4 quantization from the checkpoint "
"index; dropping modelopt_fp4 for the NextN module (bf16 MTP fallback)."
)
main_backend = get_moe_runner_backend()
spec_backend = get_speculative_moe_runner_backend()
keep_fp4 = _is_fp4_capable(main_backend) or _is_fp4_capable(spec_backend)
logger.warning(
"[NextN-fp4] main_runner=%s spec_runner=%s keep_fp4=%s",
main_backend,
spec_backend,
keep_fp4,
)
return not keep_fp4
return True
return not is_fp4
class DeepseekModelNextN(nn.Module):
@@ -137,7 +208,7 @@ class DeepseekModelNextN(nn.Module):
else:
moe_quant_config_override = None
if _should_drop_nextn_modelopt_fp4_quant_config(quant_config):
if _should_drop_nextn_modelopt_fp4_quant_config(quant_config, config):
logger.warning(
"Overriding DeepseekV3ForCausalLMNextN quant config for modelopt_fp4 Deepseek model."
)