From c4a69351b465f6a19166e045663dcb5c7dd69fa1 Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Thu, 11 Jun 2026 14:44:07 +0800 Subject: [PATCH 1/7] Support Qwen3.6 ModelOpt mixed NVFP4 --- python/sglang/srt/configs/model_config.py | 12 +- python/sglang/srt/layers/logits_processor.py | 12 + .../srt/layers/quantization/modelopt_quant.py | 224 ++++++++++++++++-- .../srt/layers/vocab_parallel_embedding.py | 8 +- .../sglang/srt/model_loader/weight_utils.py | 21 +- python/sglang/srt/models/qwen3_5_mtp.py | 33 ++- .../sglang/srt/speculative/eagle_worker_v2.py | 10 + 7 files changed, 287 insertions(+), 33 deletions(-) diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index db705eae96f0..915b30192188 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -1031,10 +1031,14 @@ def _parse_modelopt_quant_config(self, quant_config_dict: dict) -> Optional[dict quant_algo = json_quant_configs.get("quant_algo", None) if quant_algo == "MIXED_PRECISION": - architectures = getattr(self.hf_config, "architectures", []) or [] - if getattr(self.hf_config, "model_type", None) == "nemotron_h" or any( - arch.startswith("NemotronH") for arch in architectures - ): + quantized_layers = json_quant_configs.get("quantized_layers", {}) + has_modelopt_nvfp4_layers = any( + str(layer_info.get("quant_algo", "")).upper() + in ("NVFP4", "W4A16_NVFP4") + for layer_info in quantized_layers.values() + if isinstance(layer_info, dict) + ) + if has_modelopt_nvfp4_layers: return {"quant_method": "modelopt_mixed", "quant_algo": quant_algo} return {"quant_method": "w4afp8", "quant_algo": quant_algo} elif quant_algo and ("FP4" in quant_algo or "NVFP4" in quant_algo): diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 15a3d6aac9b1..57dc55a2f4e7 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -880,9 +880,21 @@ def _compute_lm_head( lm_head: VocabParallelEmbedding, embedding_bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: + quant_method = getattr(lm_head, "quant_method", None) if hasattr(lm_head, "set_lora") and hasattr(lm_head, "apply_lora"): # This is a LoRA-wrapped module, use its forward method logits = lm_head(hidden_states) + elif ( + quant_method is not None + and callable(getattr(quant_method, "apply", None)) + and type(quant_method).__name__ + not in ( + "UnquantizedEmbeddingMethod", + "UnquantizedLinearMethod", + "PackWeightMethod", + ) + ): + logits = quant_method.apply(lm_head, hidden_states, embedding_bias) elif hasattr(lm_head, "weight"): # Normal linear layer if self.use_fp32_lm_head: diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index d362b764eb69..6da55790d247 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -5,7 +5,7 @@ import logging from enum import IntEnum -from typing import TYPE_CHECKING, Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import regex as re import torch @@ -69,6 +69,7 @@ is_sm120_supported, next_power_of_2, round_up, + set_weight_attrs, ) from sglang.srt.utils.custom_op import register_custom_op from sglang.srt.utils.patch_torch import register_fake_if_exists @@ -299,8 +300,9 @@ def _get_quant_method( ) -> Optional[QuantizeMethodBase]: from sglang.srt.layers.linear import LinearBase from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead - if isinstance(layer, LinearBase): + if isinstance(layer, (LinearBase, ParallelLMHead)): if is_layer_skipped( prefix, self.exclude_modules, self.packed_modules_mapping ) or self.is_layer_excluded(prefix): @@ -451,6 +453,8 @@ def from_config(cls, config: Dict[str, Any]) -> ModelOptFp8Config: and kv_cache_scheme.get("num_bits") == 8 ): kv_cache_quant_method = "FP8" + else: + kv_cache_quant_method = config.get("kv_cache_quant_algo") # Map 'ignore' field to 'exclude_modules' exclude_modules = config.get("ignore") @@ -553,17 +557,16 @@ def create_weights( if self.quant_config.is_checkpoint_fp8_serialized: # Register weight and input scales for scale_name in ["weight_scale", "input_scale"]: - layer.register_parameter( - scale_name, - PerTensorScaleParameter( - data=torch.full( - (len(output_partition_sizes),), - torch.finfo(torch.float32).min, - dtype=torch.float32, - ), - weight_loader=weight_loader, + scale = PerTensorScaleParameter( + data=torch.full( + (len(output_partition_sizes),), + torch.finfo(torch.float32).min, + dtype=torch.float32, ), + weight_loader=weight_loader, ) + set_weight_attrs(scale, {"needs_scalar_to_array": True}) + layer.register_parameter(scale_name, scale) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: """Requantizes weights after loading using the maximum scale.""" @@ -614,11 +617,13 @@ def __init__( quantized_layers: Dict[str, Dict[str, Any]], fp8_config: ModelOptFp8Config, nvfp4_config: "ModelOptFp4Config", + w4a16_nvfp4_config: "ModelOptFp4Config", ) -> None: super().__init__(kv_cache_quant_algo, exclude_modules, packed_modules_mapping) self.quantized_layers = quantized_layers self.fp8_config = fp8_config self.nvfp4_config = nvfp4_config + self.w4a16_nvfp4_config = w4a16_nvfp4_config @classmethod def override_quantization_method(cls, hf_quant_config, user_quant): @@ -662,6 +667,8 @@ def from_config(cls, config: Dict[str, Any]) -> "ModelOptMixedPrecisionConfig": kv_cache_quant_algo = "NVFP4" else: kv_cache_quant_algo = "auto" + else: + kv_cache_quant_algo = config.get("kv_cache_quant_algo") exclude_modules = config.get("ignore") quantized_layers = config.get("quantized_layers", {}) else: @@ -682,7 +689,10 @@ def from_config(cls, config: Dict[str, Any]) -> "ModelOptMixedPrecisionConfig": group_size = None for layer_info in quantized_layers.values(): - if layer_info.get("quant_algo", "").upper() == "NVFP4": + if layer_info.get("quant_algo", "").upper() in ( + "NVFP4", + "W4A16_NVFP4", + ): group_size = layer_info.get("group_size", 16) break if group_size is None: @@ -702,6 +712,14 @@ def from_config(cls, config: Dict[str, Any]) -> "ModelOptMixedPrecisionConfig": packed_modules_mapping=packed_modules_mapping, group_size=group_size, ) + w4a16_nvfp4_config = ModelOptFp4Config( + is_checkpoint_nvfp4_serialized=True, + kv_cache_quant_algo=kv_cache_quant_algo, + exclude_modules=[], + packed_modules_mapping=packed_modules_mapping, + group_size=group_size, + use_per_token_activation=False, + ) return cls( kv_cache_quant_algo=kv_cache_quant_algo, @@ -710,6 +728,7 @@ def from_config(cls, config: Dict[str, Any]) -> "ModelOptMixedPrecisionConfig": quantized_layers=quantized_layers, fp8_config=fp8_config, nvfp4_config=nvfp4_config, + w4a16_nvfp4_config=w4a16_nvfp4_config, ) def apply_weight_name_mapper(self, hf_to_sglang_mapper: "WeightsMapper"): @@ -720,17 +739,21 @@ def apply_weight_name_mapper(self, hf_to_sglang_mapper: "WeightsMapper"): ) def _resolve_quant_algo(self, prefix: str) -> Optional[str]: - if prefix in self.quantized_layers: - return self.quantized_layers[prefix]["quant_algo"].upper() + for candidate in self._quantized_layer_prefix_candidates(prefix): + if candidate in self.quantized_layers: + return self.quantized_layers[candidate]["quant_algo"].upper() proj_name = prefix.rsplit(".", 1)[-1] if self.packed_modules_mapping and proj_name in self.packed_modules_mapping: algos = set() base = prefix.rsplit(".", 1)[0] - for shard_name in self.packed_modules_mapping[proj_name]: - shard_prefix = f"{base}.{shard_name}" - if shard_prefix in self.quantized_layers: - algos.add(self.quantized_layers[shard_prefix]["quant_algo"].upper()) + for base_candidate in self._quantized_layer_prefix_candidates(base): + for shard_name in self.packed_modules_mapping[proj_name]: + shard_prefix = f"{base_candidate}.{shard_name}" + if shard_prefix in self.quantized_layers: + algos.add( + self.quantized_layers[shard_prefix]["quant_algo"].upper() + ) if len(algos) == 1: return algos.pop() if len(algos) > 1: @@ -739,22 +762,42 @@ def _resolve_quant_algo(self, prefix: str) -> Optional[str]: "All shards must use the same quantization." ) - prefix_dot = prefix + "." - for key, info in self.quantized_layers.items(): - if key.startswith(prefix_dot): - return info["quant_algo"].upper() + for candidate in self._quantized_layer_prefix_candidates(prefix): + prefix_dot = candidate + "." + for key, info in self.quantized_layers.items(): + if key.startswith(prefix_dot): + return info["quant_algo"].upper() return None + @staticmethod + def _quantized_layer_prefix_candidates(prefix: str) -> Tuple[str, ...]: + candidates = [prefix] + + if prefix.endswith(".lm_head"): + candidates.append("lm_head") + + if prefix.startswith("language_model.model."): + candidates.append( + "model.language_model." + prefix[len("language_model.model.") :] + ) + elif prefix.startswith("model.language_model."): + candidates.append( + "language_model.model." + prefix[len("model.language_model.") :] + ) + + return tuple(dict.fromkeys(candidates)) + def get_quant_method( self, layer: torch.nn.Module, prefix: str ) -> Optional[QuantizeMethodBase]: from sglang.srt.layers.linear import LinearBase from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead quant_algo = self._resolve_quant_algo(prefix) - if isinstance(layer, LinearBase): + if isinstance(layer, (LinearBase, ParallelLMHead)): if is_layer_skipped( prefix, self.exclude_modules, self.packed_modules_mapping ) or self.is_layer_excluded(prefix): @@ -763,6 +806,8 @@ def get_quant_method( return ModelOptFp8LinearMethod(self.fp8_config) if quant_algo == "NVFP4": return ModelOptFp4LinearMethod(self.nvfp4_config) + if quant_algo == "W4A16_NVFP4": + return ModelOptFp4W4A16LinearMethod(self.w4a16_nvfp4_config) return UnquantizedLinearMethod() if self.kv_cache_quant_algo and isinstance(layer, RadixAttention): @@ -775,6 +820,8 @@ def get_quant_method( return ModelOptFp8MoEMethod(self.fp8_config) if quant_algo == "NVFP4": return ModelOptNvFp4FusedMoEMethod(self.nvfp4_config) + if quant_algo == "W4A16_NVFP4": + return ModelOptNvFp4FusedMoEMethod(self.w4a16_nvfp4_config) return None return None @@ -1279,7 +1326,7 @@ def from_config(cls, config: Dict[str, Any]) -> ModelOptFp4Config: else: kv_cache_quant_algo = "auto" else: - kv_cache_quant_algo = "auto" + kv_cache_quant_algo = config.get("kv_cache_quant_algo") or "auto" group_size = config.get("group_size") # If group_size is not at top level, try to extract from config_groups @@ -1415,13 +1462,14 @@ def create_weights( data=torch.empty(len(output_partition_sizes), dtype=torch.float32), weight_loader=weight_loader, ) - + set_weight_attrs(input_scale, {"needs_scalar_to_array": True}) layer.register_parameter("input_scale", input_scale) weight_scale_2 = PerTensorScaleParameter( data=torch.empty(len(output_partition_sizes), dtype=torch.float32), weight_loader=weight_loader, ) + set_weight_attrs(weight_scale_2, {"needs_scalar_to_array": True}) layer.register_parameter("weight_scale_2", weight_scale_2) weight_scale = ModelWeightParameter( @@ -1461,6 +1509,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: ) copy_or_rebind_param(layer, "input_global_scale", input_scale_2) copy_or_rebind_param(layer, "weight_global_scale", weight_scale_2) + layer.quant_config = self.quant_config prepare_nvfp4_layer_for_marlin(layer) layer.weights_padding_cols = 0 return @@ -1679,6 +1728,131 @@ def apply( return out.view(*output_shape) +class ModelOptFp4W4A16LinearMethod(LinearMethodBase): + """Linear method for ModelOpt W4A16_NVFP4 checkpoints. + + Loads packed NVFP4 weights with fp16/bf16 activations. ModelOpt may still + provide input_scale tensors for fused loader compatibility; they are + consumed during loading and discarded before runtime. + """ + + def __init__(self, quant_config: ModelOptFp4Config): + self.quant_config = quant_config + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + del input_size, output_size + if not self.quant_config.is_checkpoint_nvfp4_serialized: + raise ValueError( + "W4A16_NVFP4 quantization was selected, " + "dynamic quantization is not supported." + ) + + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.params_dtype = params_dtype + layer.quant_config = self.quant_config + + if input_size_per_partition % 16 != 0: + raise ValueError( + "Unsupported model when input feature size is not a multiple of 16" + ) + + weight = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition // 2, + dtype=torch.uint8, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + weight_scale_2 = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + set_weight_attrs(weight_scale_2, {"needs_scalar_to_array": True}) + layer.register_parameter("weight_scale_2", weight_scale_2) + + weight_scale = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition // self.quant_config.group_size, + dtype=torch.float8_e4m3fn, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight_scale", weight_scale) + + # Some ModelOpt checkpoints may still include input_scale entries in + # fused-loader paths. W4A16 does not use them, but registering the + # placeholder lets the generic loader consume those tensors harmlessly. + input_scale = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + set_weight_attrs(input_scale, {"needs_scalar_to_array": True}) + layer.register_parameter("input_scale", input_scale) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + if hasattr(layer, "input_scale"): + del layer.input_scale + + if torch.unique(layer.weight_scale_2).numel() != 1: + logger.warning( + "In W4A16_NVFP4 linear, weight_scale_2 differs across fused " + "parallel layers. Accuracy may be degraded." + ) + + copy_or_rebind_param( + layer, + "weight_global_scale", + layer.weight_scale_2.max().to(torch.float32), + ) + del layer.weight_scale_2 + + if self.quant_config.group_size != 16: + raise ValueError( + f"W4A16_NVFP4 Marlin requires group_size=16, got {self.quant_config.group_size}." + ) + layer.quant_config = self.quant_config + prepare_nvfp4_layer_for_marlin(layer) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return apply_fp4_marlin_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + weight_global_scale=layer.weight_global_scale, + workspace=layer.workspace, + size_n=layer.output_size_per_partition, + size_k=layer.input_size_per_partition, + bias=bias, + ) + + class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): """ MoE Method for FP4 Quantization with Blockscales and PerTensorScales diff --git a/python/sglang/srt/layers/vocab_parallel_embedding.py b/python/sglang/srt/layers/vocab_parallel_embedding.py index cbba8cd0637c..128bd668f1dd 100644 --- a/python/sglang/srt/layers/vocab_parallel_embedding.py +++ b/python/sglang/srt/layers/vocab_parallel_embedding.py @@ -300,7 +300,7 @@ def __init__( # If we are making an embedding layer, then our quantization linear # method must implement the embedding operation. If we are another # layer type like ParallelLMHead, this is not important. - is_embedding_layer = type(self.__class__) is VocabParallelEmbedding + is_embedding_layer = type(self) is VocabParallelEmbedding quant_method_implements_embedding = method_has_implemented_embedding( type(quant_method) ) @@ -462,6 +462,12 @@ def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): # If parameter does not have output dim, then it should # be copied onto all gpus (e.g. g_idx for act_order gptq). if output_dim is None: + if ( + loaded_weight.ndim == 0 + and param.data.ndim == 1 + and param.data.numel() == 1 + ): + loaded_weight = loaded_weight.reshape(1) assert param.data.shape == loaded_weight.shape param.data.copy_(loaded_weight) return diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index e53a3193b001..30a51ce69442 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -76,6 +76,7 @@ RUNAI_STREAMER_TENSOR_ATTR = "_sglang_runai_streamer_tensor" + # Matches routed-expert weight keys in both HF-style layouts # (``...mlp.experts..{gate,up,down}_proj.weight``) and DeepSeek V4 # layouts (``...ffn.experts..w{1,2,3}.weight``). ``shared_experts`` is @@ -260,8 +261,24 @@ def get_quant_config( if hf_quant_config is not None: if not isinstance(hf_quant_config, dict): hf_quant_config = hf_quant_config.to_dict() - hf_quant_config["packed_modules_mapping"] = packed_modules_mapping - return quant_cls.from_config(hf_quant_config) + + # For modelopt_mixed, config.json's quantization_config may not + # contain all runtime metadata. Fall through to the file-based + # hf_quant_config.json path when the per-layer map or KV-cache + # quantization metadata is missing. + modelopt_mixed_config_incomplete = ( + model_config.quantization == "modelopt_mixed" + and ( + "quantized_layers" not in hf_quant_config + or ( + "kv_cache_quant_algo" not in hf_quant_config + and "kv_cache_scheme" not in hf_quant_config + ) + ) + ) + if not modelopt_mixed_config_incomplete: + hf_quant_config["packed_modules_mapping"] = packed_modules_mapping + return quant_cls.from_config(hf_quant_config) # In case of bitsandbytes/QLoRA, get quant config from the adapter model. if model_config.quantization == "bitsandbytes": diff --git a/python/sglang/srt/models/qwen3_5_mtp.py b/python/sglang/srt/models/qwen3_5_mtp.py index 6b39a590ccd5..4335ec09448c 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/python/sglang/srt/models/qwen3_5_mtp.py @@ -58,7 +58,10 @@ def __init__( config = copy.deepcopy(config) # The MTP model is unquantized in the nvfp4 checkpoint. - if quant_config and quant_config.get_name() == "modelopt_fp4": + if quant_config and quant_config.get_name() in ( + "modelopt_fp4", + "modelopt_mixed", + ): quant_config = None if ( is_npu() @@ -133,6 +136,34 @@ def set_embed_and_head(self, embed, head): torch.cuda.empty_cache() torch.cuda.synchronize() + def set_lm_head_from_target(self, target_lm_head): + if self.config.tie_word_embeddings: + return + + for param_name, param in target_lm_head.named_parameters(recurse=False): + if hasattr(self.lm_head, param_name): + delattr(self.lm_head, param_name) + self.lm_head.register_parameter(param_name, param) + + for buffer_name, buffer in target_lm_head.named_buffers(recurse=False): + if hasattr(self.lm_head, buffer_name): + delattr(self.lm_head, buffer_name) + self.lm_head.register_buffer(buffer_name, buffer) + + for attr_name in ( + "quant_method", + "quant_config", + "logical_widths", + "input_size_per_partition", + "output_size_per_partition", + "params_dtype", + "weights_padding_cols", + "workspace", + "_accepts_prequantized_fp4", + ): + if hasattr(target_lm_head, attr_name): + setattr(self.lm_head, attr_name, getattr(target_lm_head, attr_name)) + @torch.no_grad() def forward( self, diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 732a364bc6c1..0812f1e061a3 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -275,6 +275,14 @@ def init_token_map(self): def init_lm_head(self): embed, head = self.target_worker.model_runner.model.get_embed_and_head() + target_lm_head = getattr(self.target_worker.model_runner.model, "lm_head", None) + + def maybe_share_lm_head_quant_attrs(): + if target_lm_head is not None and hasattr( + self.draft_runner.model, "set_lm_head_from_target" + ): + self.draft_runner.model.set_lm_head_from_target(target_lm_head) + if self.speculative_algorithm.is_eagle3(): # most cases EAGLE3 models don't share lm_head # but some models (e.g. nvidia/gpt-oss-120b-Eagle3) shares @@ -283,6 +291,7 @@ def init_lm_head(self): and self.draft_runner.model.load_lm_head_from_target ): self.draft_runner.model.set_embed_and_head(embed, head) + maybe_share_lm_head_quant_attrs() else: self.draft_runner.model.set_embed(embed) @@ -300,6 +309,7 @@ def init_lm_head(self): # Share the embedding and lm_head self.draft_runner.model.set_embed_and_head(embed, head) + maybe_share_lm_head_quant_attrs() def init_attention_backend(self): # Create multi-step attn backends and cuda graph runners From 64b320242c6d484b41f9b0954aa40abd19ae219c Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Thu, 11 Jun 2026 17:20:27 +0800 Subject: [PATCH 2/7] Address review comments --- python/sglang/srt/configs/model_config.py | 2 +- python/sglang/srt/layers/logits_processor.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 915b30192188..e780c5b28860 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -1031,7 +1031,7 @@ def _parse_modelopt_quant_config(self, quant_config_dict: dict) -> Optional[dict quant_algo = json_quant_configs.get("quant_algo", None) if quant_algo == "MIXED_PRECISION": - quantized_layers = json_quant_configs.get("quantized_layers", {}) + quantized_layers = json_quant_configs.get("quantized_layers") or {} has_modelopt_nvfp4_layers = any( str(layer_info.get("quant_algo", "")).upper() in ("NVFP4", "W4A16_NVFP4") diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 57dc55a2f4e7..8fd07cd7712b 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -886,6 +886,7 @@ def _compute_lm_head( logits = lm_head(hidden_states) elif ( quant_method is not None + and hasattr(lm_head, "weight") and callable(getattr(quant_method, "apply", None)) and type(quant_method).__name__ not in ( From a6784503fc206f96b7713335cb0616411a008410 Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Thu, 11 Jun 2026 23:38:40 +0800 Subject: [PATCH 3/7] Update ModelOpt mixed precision routing tests --- .../unit/model_loader/test_modelopt_loader.py | 46 ++++++++++++++++++- 1 file changed, 44 insertions(+), 2 deletions(-) diff --git a/test/registered/unit/model_loader/test_modelopt_loader.py b/test/registered/unit/model_loader/test_modelopt_loader.py index 35ea5609b992..6a5d438aa2a3 100644 --- a/test/registered/unit/model_loader/test_modelopt_loader.py +++ b/test/registered/unit/model_loader/test_modelopt_loader.py @@ -626,14 +626,56 @@ def test_non_modelopt_quant_method_unchanged(self): class TestModelOptMixedPrecisionConfig(CustomTestCase): - def test_nemotron_mixed_precision_uses_modelopt_mixed(self): + def test_nemotron_mixed_precision_with_nvfp4_layers_uses_modelopt_mixed(self): model_config = ModelConfig.__new__(ModelConfig) model_config.hf_config = MagicMock() model_config.hf_config.model_type = "nemotron_h" model_config.hf_config.architectures = ["NemotronHForCausalLM"] result = model_config._parse_modelopt_quant_config( - {"quantization": {"quant_algo": "MIXED_PRECISION"}} + { + "quantization": { + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "backbone.layers.0.mixer.in_proj": {"quant_algo": "FP8"}, + "backbone.layers.0.mixer.out_proj": {"quant_algo": "FP8"}, + "backbone.layers.1.mixer.experts.0.up_proj": { + "quant_algo": "NVFP4", + "group_size": 16, + }, + "backbone.layers.1.mixer.experts.0.down_proj": { + "quant_algo": "NVFP4", + "group_size": 16, + }, + }, + } + } + ) + + self.assertEqual(result["quant_method"], "modelopt_mixed") + + def test_qwen_mixed_precision_with_w4a16_nvfp4_layers_uses_modelopt_mixed(self): + model_config = ModelConfig.__new__(ModelConfig) + model_config.hf_config = MagicMock() + model_config.hf_config.model_type = "qwen3_5_moe" + model_config.hf_config.architectures = ["Qwen3_5MoeForConditionalGeneration"] + + result = model_config._parse_modelopt_quant_config( + { + "quantization": { + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "lm_head": {"quant_algo": "W4A16_NVFP4", "group_size": 16}, + "model.language_model.layers.0.mlp.shared_expert.up_proj": { + "quant_algo": "W4A16_NVFP4", + "group_size": 16, + }, + "model.language_model.layers.0.linear_attn.in_proj_qkv": { + "quant_algo": "FP8" + }, + }, + } + } ) self.assertEqual(result["quant_method"], "modelopt_mixed") From ce53c6e23dfcf8a6ca2675f8e61fc6fcc27ff3e5 Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Thu, 11 Jun 2026 23:49:18 +0800 Subject: [PATCH 4/7] Guard quantized LM head apply by weight layout --- python/sglang/srt/layers/logits_processor.py | 42 +++++++++++++++----- 1 file changed, 31 insertions(+), 11 deletions(-) diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 8fd07cd7712b..80b5cff05bba 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -66,6 +66,36 @@ _is_npu = is_npu() _is_cpu = is_cpu() +_UNQUANTIZED_LM_HEAD_METHODS = { + "UnquantizedEmbeddingMethod", + "UnquantizedLinearMethod", + "PackWeightMethod", +} + + +def should_apply_lm_head_quant_method(lm_head, quant_method) -> bool: + if ( + quant_method is None + or not hasattr(lm_head, "weight") + or not callable(getattr(quant_method, "apply", None)) + ): + return False + + method_name = type(quant_method).__name__ + if method_name in _UNQUANTIZED_LM_HEAD_METHODS: + return False + + # Some draft models share an unquantized target lm_head tensor while still + # carrying the draft model's stale ModelOpt quant_method. Only use the + # ModelOpt lm_head kernel when the actual weight layout matches it. + if method_name in ("ModelOptFp4LinearMethod", "ModelOptFp4W4A16LinearMethod"): + return lm_head.weight.dtype == torch.uint8 + if method_name == "ModelOptFp8LinearMethod": + return lm_head.weight.dtype == torch.float8_e4m3fn + + return True + + # When set, LogitsProcessor.forward returns an empty output and skips the # LM head + tensor-parallel all-gather. FlashInfer autotune only profiles # attention/MoE/GEMM kernels, so the LM-head all-gather is wasted work -- @@ -884,17 +914,7 @@ def _compute_lm_head( if hasattr(lm_head, "set_lora") and hasattr(lm_head, "apply_lora"): # This is a LoRA-wrapped module, use its forward method logits = lm_head(hidden_states) - elif ( - quant_method is not None - and hasattr(lm_head, "weight") - and callable(getattr(quant_method, "apply", None)) - and type(quant_method).__name__ - not in ( - "UnquantizedEmbeddingMethod", - "UnquantizedLinearMethod", - "PackWeightMethod", - ) - ): + elif should_apply_lm_head_quant_method(lm_head, quant_method): logits = quant_method.apply(lm_head, hidden_states, embedding_bias) elif hasattr(lm_head, "weight"): # Normal linear layer From 1972460a285d395158a9ab4b1bac9f9c0932ba04 Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Tue, 16 Jun 2026 07:20:40 +0800 Subject: [PATCH 5/7] Address review comments --- python/sglang/srt/layers/logits_processor.py | 2 +- .../srt/layers/quantization/modelopt_quant.py | 76 +++++++++++-------- python/sglang/srt/models/qwen3_5_mtp.py | 24 +----- .../sglang/srt/speculative/eagle_worker_v2.py | 13 ++-- .../unit/model_loader/test_modelopt_loader.py | 2 +- 5 files changed, 57 insertions(+), 60 deletions(-) diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 80b5cff05bba..5e0db598d265 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -88,7 +88,7 @@ def should_apply_lm_head_quant_method(lm_head, quant_method) -> bool: # Some draft models share an unquantized target lm_head tensor while still # carrying the draft model's stale ModelOpt quant_method. Only use the # ModelOpt lm_head kernel when the actual weight layout matches it. - if method_name in ("ModelOptFp4LinearMethod", "ModelOptFp4W4A16LinearMethod"): + if method_name in ("ModelOptFp4LinearMethod", "ModelOptNvFp4A16LinearMethod"): return lm_head.weight.dtype == torch.uint8 if method_name == "ModelOptFp8LinearMethod": return lm_head.weight.dtype == torch.float8_e4m3fn diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index 32a2323f536d..d008ca8d456b 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -82,6 +82,25 @@ ) from sglang.srt.models.utils import WeightsMapper + +def _make_per_tensor_scale_parameter( + shape, + weight_loader, + *, + fill_value: Optional[float] = None, + needs_scalar_to_array: bool = False, +) -> PerTensorScaleParameter: + data = ( + torch.empty(shape, dtype=torch.float32) + if fill_value is None + else torch.full(shape, fill_value, dtype=torch.float32) + ) + scale = PerTensorScaleParameter(data=data, weight_loader=weight_loader) + if needs_scalar_to_array: + set_weight_attrs(scale, {"needs_scalar_to_array": True}) + return scale + + try: from flashinfer import mm_fp4 as flashinfer_fp4_gemm from flashinfer import reorder_rows_for_gated_act_gemm, shuffle_matrix_sf_a @@ -557,15 +576,12 @@ def create_weights( if self.quant_config.is_checkpoint_fp8_serialized: # Register weight and input scales for scale_name in ["weight_scale", "input_scale"]: - scale = PerTensorScaleParameter( - data=torch.full( - (len(output_partition_sizes),), - torch.finfo(torch.float32).min, - dtype=torch.float32, - ), + scale = _make_per_tensor_scale_parameter( + (len(output_partition_sizes),), weight_loader=weight_loader, + fill_value=torch.finfo(torch.float32).min, + needs_scalar_to_array=True, ) - set_weight_attrs(scale, {"needs_scalar_to_array": True}) layer.register_parameter(scale_name, scale) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: @@ -617,13 +633,13 @@ def __init__( quantized_layers: Dict[str, Dict[str, Any]], fp8_config: ModelOptFp8Config, nvfp4_config: ModelOptFp4Config, - w4a16_nvfp4_config: ModelOptFp4Config, + nvfp4a16_config: ModelOptFp4Config, ) -> None: super().__init__(kv_cache_quant_algo, exclude_modules, packed_modules_mapping) self.quantized_layers = quantized_layers self.fp8_config = fp8_config self.nvfp4_config = nvfp4_config - self.w4a16_nvfp4_config = w4a16_nvfp4_config + self.nvfp4a16_config = nvfp4a16_config @classmethod def override_quantization_method(cls, hf_quant_config, user_quant): @@ -712,7 +728,7 @@ def from_config(cls, config: Dict[str, Any]) -> ModelOptMixedPrecisionConfig: packed_modules_mapping=packed_modules_mapping, group_size=group_size, ) - w4a16_nvfp4_config = ModelOptFp4Config( + nvfp4a16_config = ModelOptFp4Config( is_checkpoint_nvfp4_serialized=True, kv_cache_quant_algo=kv_cache_quant_algo, exclude_modules=[], @@ -728,7 +744,7 @@ def from_config(cls, config: Dict[str, Any]) -> ModelOptMixedPrecisionConfig: quantized_layers=quantized_layers, fp8_config=fp8_config, nvfp4_config=nvfp4_config, - w4a16_nvfp4_config=w4a16_nvfp4_config, + nvfp4a16_config=nvfp4a16_config, ) def apply_weight_name_mapper(self, hf_to_sglang_mapper: WeightsMapper): @@ -807,7 +823,7 @@ def get_quant_method( if quant_algo == "NVFP4": return ModelOptFp4LinearMethod(self.nvfp4_config) if quant_algo == "W4A16_NVFP4": - return ModelOptFp4W4A16LinearMethod(self.w4a16_nvfp4_config) + return ModelOptNvFp4A16LinearMethod(self.nvfp4a16_config) return UnquantizedLinearMethod() if self.kv_cache_quant_algo and isinstance(layer, RadixAttention): @@ -821,7 +837,7 @@ def get_quant_method( if quant_algo == "NVFP4": return ModelOptNvFp4FusedMoEMethod(self.nvfp4_config) if quant_algo == "W4A16_NVFP4": - return ModelOptNvFp4FusedMoEMethod(self.w4a16_nvfp4_config) + return ModelOptNvFp4FusedMoEMethod(self.nvfp4a16_config) return None return None @@ -1458,18 +1474,18 @@ def create_weights( ) layer.register_parameter("weight", weight) - input_scale = PerTensorScaleParameter( - data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + input_scale = _make_per_tensor_scale_parameter( + (len(output_partition_sizes),), weight_loader=weight_loader, + needs_scalar_to_array=True, ) - set_weight_attrs(input_scale, {"needs_scalar_to_array": True}) layer.register_parameter("input_scale", input_scale) - weight_scale_2 = PerTensorScaleParameter( - data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_scale_2 = _make_per_tensor_scale_parameter( + (len(output_partition_sizes),), weight_loader=weight_loader, + needs_scalar_to_array=True, ) - set_weight_attrs(weight_scale_2, {"needs_scalar_to_array": True}) layer.register_parameter("weight_scale_2", weight_scale_2) weight_scale = ModelWeightParameter( @@ -1728,8 +1744,8 @@ def apply( return out.view(*output_shape) -class ModelOptFp4W4A16LinearMethod(LinearMethodBase): - """Linear method for ModelOpt W4A16_NVFP4 checkpoints. +class ModelOptNvFp4A16LinearMethod(LinearMethodBase): + """Linear method for ModelOpt NVFP4A16 checkpoints. Loads packed NVFP4 weights with fp16/bf16 activations. ModelOpt may still provide input_scale tensors for fused loader compatibility; they are @@ -1782,11 +1798,11 @@ def create_weights( ) layer.register_parameter("weight", weight) - weight_scale_2 = PerTensorScaleParameter( - data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_scale_2 = _make_per_tensor_scale_parameter( + (len(output_partition_sizes),), weight_loader=weight_loader, + needs_scalar_to_array=True, ) - set_weight_attrs(weight_scale_2, {"needs_scalar_to_array": True}) layer.register_parameter("weight_scale_2", weight_scale_2) weight_scale = ModelWeightParameter( @@ -1802,13 +1818,13 @@ def create_weights( layer.register_parameter("weight_scale", weight_scale) # Some ModelOpt checkpoints may still include input_scale entries in - # fused-loader paths. W4A16 does not use them, but registering the + # fused-loader paths. NVFP4A16 does not use them, but registering the # placeholder lets the generic loader consume those tensors harmlessly. - input_scale = PerTensorScaleParameter( - data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + input_scale = _make_per_tensor_scale_parameter( + (len(output_partition_sizes),), weight_loader=weight_loader, + needs_scalar_to_array=True, ) - set_weight_attrs(input_scale, {"needs_scalar_to_array": True}) layer.register_parameter("input_scale", input_scale) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: @@ -1817,7 +1833,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if torch.unique(layer.weight_scale_2).numel() != 1: logger.warning( - "In W4A16_NVFP4 linear, weight_scale_2 differs across fused " + "In NVFP4A16 linear, weight_scale_2 differs across fused " "parallel layers. Accuracy may be degraded." ) @@ -1830,7 +1846,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if self.quant_config.group_size != 16: raise ValueError( - f"W4A16_NVFP4 Marlin requires group_size=16, got {self.quant_config.group_size}." + f"NVFP4A16 Marlin requires group_size=16, got {self.quant_config.group_size}." ) layer.quant_config = self.quant_config prepare_nvfp4_layer_for_marlin(layer) diff --git a/python/sglang/srt/models/qwen3_5_mtp.py b/python/sglang/srt/models/qwen3_5_mtp.py index d390671ee1ce..6888bf19dda9 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/python/sglang/srt/models/qwen3_5_mtp.py @@ -140,29 +140,7 @@ def set_lm_head_from_target(self, target_lm_head): if self.config.tie_word_embeddings: return - for param_name, param in target_lm_head.named_parameters(recurse=False): - if hasattr(self.lm_head, param_name): - delattr(self.lm_head, param_name) - self.lm_head.register_parameter(param_name, param) - - for buffer_name, buffer in target_lm_head.named_buffers(recurse=False): - if hasattr(self.lm_head, buffer_name): - delattr(self.lm_head, buffer_name) - self.lm_head.register_buffer(buffer_name, buffer) - - for attr_name in ( - "quant_method", - "quant_config", - "logical_widths", - "input_size_per_partition", - "output_size_per_partition", - "params_dtype", - "weights_padding_cols", - "workspace", - "_accepts_prequantized_fp4", - ): - if hasattr(target_lm_head, attr_name): - setattr(self.lm_head, attr_name, getattr(target_lm_head, attr_name)) + self.lm_head = target_lm_head @torch.no_grad() def forward( diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index e5a90d3fa347..33b2a54c3ff6 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -278,9 +278,12 @@ def init_lm_head(self): embed, head = self.target_worker.model_runner.model.get_embed_and_head() target_lm_head = getattr(self.target_worker.model_runner.model, "lm_head", None) - def maybe_share_lm_head_quant_attrs(): - if target_lm_head is not None and hasattr( - self.draft_runner.model, "set_lm_head_from_target" + def maybe_share_target_lm_head(): + if ( + target_lm_head is not None + and self.hot_token_id is None + and getattr(self.draft_runner.model, "hot_token_id", None) is None + and hasattr(self.draft_runner.model, "set_lm_head_from_target") ): self.draft_runner.model.set_lm_head_from_target(target_lm_head) @@ -292,7 +295,7 @@ def maybe_share_lm_head_quant_attrs(): and self.draft_runner.model.load_lm_head_from_target ): self.draft_runner.model.set_embed_and_head(embed, head) - maybe_share_lm_head_quant_attrs() + maybe_share_target_lm_head() else: self.draft_runner.model.set_embed(embed) @@ -310,7 +313,7 @@ def maybe_share_lm_head_quant_attrs(): # Share the embedding and lm_head self.draft_runner.model.set_embed_and_head(embed, head) - maybe_share_lm_head_quant_attrs() + maybe_share_target_lm_head() def init_attention_backend(self): # Create multi-step attn backends and cuda graph runners diff --git a/test/registered/unit/model_loader/test_modelopt_loader.py b/test/registered/unit/model_loader/test_modelopt_loader.py index 6a5d438aa2a3..80b737cd6190 100644 --- a/test/registered/unit/model_loader/test_modelopt_loader.py +++ b/test/registered/unit/model_loader/test_modelopt_loader.py @@ -654,7 +654,7 @@ def test_nemotron_mixed_precision_with_nvfp4_layers_uses_modelopt_mixed(self): self.assertEqual(result["quant_method"], "modelopt_mixed") - def test_qwen_mixed_precision_with_w4a16_nvfp4_layers_uses_modelopt_mixed(self): + def test_qwen_mixed_precision_with_nvfp4a16_layers_uses_modelopt_mixed(self): model_config = ModelConfig.__new__(ModelConfig) model_config.hf_config = MagicMock() model_config.hf_config.model_type = "qwen3_5_moe" From a6e43bc8060fa14873d778464d6bbc3afd6ca5bc Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Tue, 16 Jun 2026 07:24:38 +0800 Subject: [PATCH 6/7] Force-disable NVFP4A16 per-token activation --- python/sglang/srt/layers/quantization/modelopt_quant.py | 5 +++-- .../registered/unit/model_loader/test_modelopt_loader.py | 9 +++++++++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index a5f4ff24cf7f..9e5ec3fda46a 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -1252,8 +1252,9 @@ def __init__( ) self.group_size = group_size self.use_per_token_activation = ( - use_per_token_activation - or envs.SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION.get() + envs.SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION.get() + if use_per_token_activation is None + else use_per_token_activation ) @classmethod diff --git a/test/registered/unit/model_loader/test_modelopt_loader.py b/test/registered/unit/model_loader/test_modelopt_loader.py index 80b737cd6190..4d25aac202b0 100644 --- a/test/registered/unit/model_loader/test_modelopt_loader.py +++ b/test/registered/unit/model_loader/test_modelopt_loader.py @@ -688,6 +688,15 @@ def test_mixed_precision_override_does_not_hijack_w4afp8(self): ) ) + @patch( + "sglang.srt.layers.quantization.modelopt_quant.envs.SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION.get", + return_value=True, + ) + def test_explicit_nvfp4_per_token_activation_false_overrides_env(self, _): + config = ModelOptFp4Config(use_per_token_activation=False) + + self.assertFalse(config.use_per_token_activation) + def test_mixed_precision_uses_nvfp4_min_capability(self): self.assertEqual( ModelOptMixedPrecisionConfig.get_min_capability(), From 6f6a7c1caaf406c1b26fb520e067bddc7f766da2 Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Tue, 16 Jun 2026 08:11:34 +0800 Subject: [PATCH 7/7] Fix ModelOpt FP4 LM head guard --- python/sglang/srt/layers/logits_processor.py | 46 +++++++++++++++++-- .../unit/model_loader/test_modelopt_loader.py | 46 +++++++++++++++++++ 2 files changed, 88 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 5e0db598d265..5df782bd788f 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -73,6 +73,10 @@ } +def _has_lm_head_runtime_attrs(lm_head, attr_names: Tuple[str, ...]) -> bool: + return all(hasattr(lm_head, attr_name) for attr_name in attr_names) + + def should_apply_lm_head_quant_method(lm_head, quant_method) -> bool: if ( quant_method is None @@ -87,11 +91,45 @@ def should_apply_lm_head_quant_method(lm_head, quant_method) -> bool: # Some draft models share an unquantized target lm_head tensor while still # carrying the draft model's stale ModelOpt quant_method. Only use the - # ModelOpt lm_head kernel when the actual weight layout matches it. - if method_name in ("ModelOptFp4LinearMethod", "ModelOptNvFp4A16LinearMethod"): - return lm_head.weight.dtype == torch.uint8 + # ModelOpt lm_head kernel when the runtime quantization state matches it. + if method_name == "ModelOptFp4LinearMethod": + if lm_head.weight.dtype == torch.int32 and _has_lm_head_runtime_attrs( + lm_head, + ( + "weight_scale", + "weight_global_scale", + "workspace", + "input_size_per_partition", + "output_size_per_partition", + ), + ): + return True + return lm_head.weight.dtype == torch.uint8 and _has_lm_head_runtime_attrs( + lm_head, + ( + "weight_scale_interleaved", + "alpha", + "input_scale_inv", + "input_size_per_partition", + "output_size_per_partition", + ), + ) + if method_name == "ModelOptNvFp4A16LinearMethod": + return lm_head.weight.dtype == torch.int32 and _has_lm_head_runtime_attrs( + lm_head, + ( + "weight_scale", + "weight_global_scale", + "workspace", + "input_size_per_partition", + "output_size_per_partition", + ), + ) if method_name == "ModelOptFp8LinearMethod": - return lm_head.weight.dtype == torch.float8_e4m3fn + return ( + lm_head.weight.dtype == torch.float8_e4m3fn + and _has_lm_head_runtime_attrs(lm_head, ("weight_scale", "input_scale")) + ) return True diff --git a/test/registered/unit/model_loader/test_modelopt_loader.py b/test/registered/unit/model_loader/test_modelopt_loader.py index 4d25aac202b0..56bdf684f0bf 100644 --- a/test/registered/unit/model_loader/test_modelopt_loader.py +++ b/test/registered/unit/model_loader/test_modelopt_loader.py @@ -8,15 +8,19 @@ import unittest from unittest.mock import MagicMock, patch +import torch import torch.nn as nn from sglang.srt.configs.device_config import DeviceConfig from sglang.srt.configs.load_config import LoadConfig from sglang.srt.configs.model_config import ModelConfig +from sglang.srt.layers.logits_processor import should_apply_lm_head_quant_method from sglang.srt.layers.modelopt_utils import QUANT_CFG_CHOICES from sglang.srt.layers.quantization.modelopt_quant import ( ModelOptFp4Config, + ModelOptFp4LinearMethod, ModelOptMixedPrecisionConfig, + ModelOptNvFp4A16LinearMethod, ) from sglang.srt.model_loader.loader import ModelOptModelLoader from sglang.srt.models.utils import WeightsMapper @@ -697,6 +701,48 @@ def test_explicit_nvfp4_per_token_activation_false_overrides_env(self, _): self.assertFalse(config.use_per_token_activation) + def test_lm_head_guard_accepts_modelopt_fp4_marlin_runtime_state(self): + lm_head = nn.Module() + lm_head.weight = nn.Parameter( + torch.empty(128, 496640, dtype=torch.int32), requires_grad=False + ) + lm_head.weight_scale = nn.Parameter(torch.empty(1)) + lm_head.weight_global_scale = nn.Parameter(torch.empty(1)) + lm_head.workspace = torch.empty(1) + lm_head.input_size_per_partition = 2048 + lm_head.output_size_per_partition = 128000 + + self.assertTrue( + should_apply_lm_head_quant_method( + lm_head, ModelOptNvFp4A16LinearMethod(ModelOptFp4Config()) + ) + ) + + def test_lm_head_guard_rejects_stale_modelopt_fp4_method_on_dense_head(self): + lm_head = nn.Module() + lm_head.weight = nn.Parameter(torch.empty(128000, 2048)) + + self.assertFalse( + should_apply_lm_head_quant_method( + lm_head, ModelOptFp4LinearMethod(ModelOptFp4Config()) + ) + ) + + def test_lm_head_guard_rejects_stale_modelopt_fp4_attrs_on_dense_head(self): + lm_head = nn.Module() + lm_head.weight = nn.Parameter(torch.empty(128000, 2048)) + lm_head.weight_scale = nn.Parameter(torch.empty(1)) + lm_head.weight_global_scale = nn.Parameter(torch.empty(1)) + lm_head.workspace = torch.empty(1) + lm_head.input_size_per_partition = 2048 + lm_head.output_size_per_partition = 128000 + + self.assertFalse( + should_apply_lm_head_quant_method( + lm_head, ModelOptNvFp4A16LinearMethod(ModelOptFp4Config()) + ) + ) + def test_mixed_precision_uses_nvfp4_min_capability(self): self.assertEqual( ModelOptMixedPrecisionConfig.get_min_capability(),