diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index 330fa6c21..efa58209b 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -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 diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index 4746fc8cc..c54a77308 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -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. diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index 7a5553f4b..4374052ff 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -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) ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 997f6f389..561fcb88d 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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"