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

B300 NextN fp4: run the EAGLE draft on the same trtllm MoE backend as main (elegant root fix)

Browse Source 
ROOT PROBLEM (opus-diagnosed): the NextN/MTP draft MoE is a normal DeepSeek-V3 fp4 MoE, identical in
kind to the main model's — but our fork forced the speculative MoE backend to "auto" (+ a blanket
assert forbidding trtllm-as-speculative), so the draft built under speculative_moe_backend_context →
MOE_RUNNER_BACKEND=AUTO → plain FusedMoE → the non-flashinfer process_weights 'else' branch that
doesn't reconcile EP scale sharding → w13_input_scale(global 256) × w13_weight_scale_2(local 64) crash.

The trtllm-as-speculative ban is over-broad: the Renormalize-only routing constraint is for the BF16
trtllm MoE; the FP4 trtllm kernel supports the DeepSeek routing the NextN carries (FlashInferFP4MoE.
forward_impl). So the elegant fix is to let the draft run the SAME flashinfer_trtllm path as main when
the NextN is fp4 — scoped via the checkpoint's safetensors index (ground truth: do the MTP experts
ship fp4 scales). DeepSeek-R1/V3-FP4 ship a bf16 MTP -> still drop to bf16 + keep the non-trtllm path.

- weight_utils: add shared nextn_moe_is_fp4_quantized_in_checkpoint + load_safetensors_index_weight_map.
- deepseek_nextn: delegate the keep/drop probe to the shared helper (one source of truth).
- server_args._handle_speculative_decoding: for an fp4 NextN, default the speculative backend to the
  main trtllm backend (instead of "auto"); narrow the assert to allow trtllm only for an fp4 MTP.
- modelopt_quant: REVERT the else-branch input-scale slice band-aid (de726b4e7e) — no longer hit; the
  fp4 draft now takes the trtllm .max() scalar branch, same as main.

Bug B (k_rope 256-vs-64 in the trtllm NSA *prefill* path under CP) is independent and still pending.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-21 17:32:01 +00:00
parent de726b4e7e
commit b3a4f4174f
4 changed files with 116 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
python/sglang/srt/layers/quantization/modelopt_quant.py
View File

@@ -2027,17 +2027,6 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
else:
w13_input_scale = layer.w13_input_scale.max(dim=-1).values.to(torch.float32)
w2_input_scale = layer.w2_input_scale
# w13_input_scale/w2_input_scale are loaded GLOBAL (num_experts) because
# create_weights tags them _sglang_require_global_experts, while
# w13_weight_scale_2 is EP-local (num_local_experts). Under EP>1 this
# non-flashinfer path must slice the global input scales down to this
# rank's local experts so the gemm1/gemm2 alpha products line up
# (mirrors the cutedsl branch above). No-op for moe_ep_size == 1.
if layer.moe_ep_size > 1 and w13_input_scale.shape[0] == layer.num_experts:
start = layer.moe_ep_rank * layer.num_local_experts
end = start + layer.num_local_experts
w13_input_scale = w13_input_scale[start:end]
w2_input_scale = w2_input_scale[start:end]
if self.quant_config.use_per_token_activation:
# FlashInfer computes activation scales dynamically per token, so

65
python/sglang/srt/model_loader/weight_utils.py
View File

@@ -555,6 +555,71 @@ def download_safetensors_index_file_from_hf(
logger.info("No %s found in local cache.", index_file)
def load_safetensors_index_weight_map(
model_dir: str, revision: Optional[str] = None
) -> Optional[dict]:
"""Return a checkpoint's safetensors index ``weight_map``, or None if unreadable.
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:
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(
model_dir: Optional[str],
revision: Optional[str],
num_hidden_layers: Optional[int],
) -> 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}``.
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.
Used to decide both whether to keep modelopt_fp4 for the NextN module and
whether the EAGLE draft should run the fp4-capable trtllm MoE backend.
"""
if not model_dir or num_hidden_layers is None:
return None
weight_map = load_safetensors_index_weight_map(model_dir, revision)
if weight_map is None:
return None
expert_prefix = f"model.layers.{num_hidden_layers}.mlp.experts."
return any(
name.startswith(expert_prefix) and "weight_scale" in name
for name in weight_map
)
# For models like Mistral-7B-v0.3, there are both sharded
# safetensors files and a consolidated safetensors file.
# Passing both of these to the weight loader functionality breaks.

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

@@ -78,53 +78,6 @@ 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]:
@@ -135,9 +88,10 @@ def _nextn_moe_is_fp4_quantized_in_checkpoint(
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
from sglang.srt.model_loader.weight_utils import (
nextn_moe_is_fp4_quantized_in_checkpoint,
)
server_args = get_global_server_args()
if server_args.speculative_draft_model_path:
model_dir = server_args.speculative_draft_model_path
@@ -145,15 +99,8 @@ def _nextn_moe_is_fp4_quantized_in_checkpoint(
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
return nextn_moe_is_fp4_quantized_in_checkpoint(
model_dir, revision, getattr(config, "num_hidden_layers", None)
)

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

@@ -3277,22 +3277,55 @@ class ServerArgs:
):
self.speculative_draft_model_revision = "main"
# FlashInfer trtllm moe bf16 only support RenormalizeNaive routing method and Deepseek routing method
# It is hard to tell the routing method in draft model, and the moe layer in draft model is not the bottleneck among
# end to end, so we just avoid using trtllm_moe for speculative decoding.
from sglang.srt.layers.moe.utils import MoeRunnerBackend
from sglang.srt.model_loader.weight_utils import (
nextn_moe_is_fp4_quantized_in_checkpoint,
)
# An fp4 NextN/MTP layer (e.g. GLM-5.2-NVFP4) is a genuine DeepSeek-V3 fp4 MoE,
# identical in kind to the main model's MoE — and the fp4 trtllm kernel supports
# the DeepSeek routing it carries. So the EAGLE draft can run the SAME
# flashinfer_trtllm path as the main model (consistent, well-tested, correct EP
# scale handling). The trtllm restriction below is real only for a *bf16* MTP
# (DeepSeek-R1/V3-FP4): the bf16 trtllm MoE requires RenormalizeNaive routing,
# and the draft routing is hard to guarantee there. We distinguish the two from
# the checkpoint's safetensors index (ground truth: do the MTP experts ship fp4
# scales), the same signal used to keep/drop modelopt_fp4 for the NextN module.
_main_is_trtllm = self.moe_runner_backend in [
"flashinfer_trtllm",
"flashinfer_trtllm_routed",
]
_spec_model_dir = self.speculative_draft_model_path or self.model_path
_spec_revision = (
self.speculative_draft_model_revision
if self.speculative_draft_model_path
else self.revision
)
_nextn_is_fp4 = _main_is_trtllm and (
nextn_moe_is_fp4_quantized_in_checkpoint(
_spec_model_dir,
_spec_revision,
getattr(self.get_model_config().hf_config, "num_hidden_layers", None),
)
is True
)
if self.speculative_moe_runner_backend is None:
self.speculative_moe_runner_backend = (
"auto"
if self.moe_runner_backend
in ["flashinfer_trtllm", "flashinfer_trtllm_routed"]
else self.moe_runner_backend
if _nextn_is_fp4:
# fp4 NextN -> draft uses the same fp4 trtllm backend as main.
self.speculative_moe_runner_backend = self.moe_runner_backend
elif _main_is_trtllm:
self.speculative_moe_runner_backend = "auto"
else:
self.speculative_moe_runner_backend = self.moe_runner_backend
elif MoeRunnerBackend(
self.speculative_moe_runner_backend
).is_flashinfer_trtllm():
assert _nextn_is_fp4, (
"Speculative MoE runner backend flashinfer_trtllm is only supported for "
"an fp4 NextN/MTP layer (the bf16 trtllm MoE requires RenormalizeNaive "
"routing); use triton or auto for a bf16 MTP."
)
else:
assert not MoeRunnerBackend(
self.speculative_moe_runner_backend
).is_flashinfer_trtllm(), "Currently speculative MoE runner backend doesn't support flashinfer_trtllm, please use triton or auto backend for speculative moe runner instead."
if self.speculative_algorithm == "NEXTN":
self.speculative_algorithm = "EAGLE"