[Model] Use standard ModelOpt config for Inkling NVFP4 (#48990)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jul 17, 2026 at 21:13 UTC
02c01f442b8bfd0affa1f2ceedda01f2e1384303
7 files changed
+56
-110
tests/config/test_model_arch_config.py
+17
@@ -131,6 +131,23 @@ def test_head_size_falls_back_when_head_dim_is_zero():
131
assert convertor.get_head_size() == 128
132
133
134
+def test_legacy_modelopt_config_without_producer_is_normalized():
135
+ quantization_config = {
136
+ "quantization": {
137
+ "quant_algo": "NVFP4",
138
+ "group_size": 16,
139
+ "kv_cache_quant_algo": None,
140
+ "exclude_modules": [],
141
+ "modelopt_quant_config": {"quant_cfg": {}},
142
+ }
143
+ }
144
+ hf_config = PretrainedConfig(quantization_config=quantization_config)
145
+
146
+ convertor = ModelArchConfigConvertorBase(hf_config, hf_config)
147
+
148
+ assert convertor.get_quantization_config()["quant_method"] == "modelopt_fp4"
149
+
150
+
151
@pytest.mark.parametrize("model", BASE_MODELS_TO_TEST)
152
def test_base_model_arch_config(model: str):
153
"""Test model architecture config for base models."""
tests/models/inkling/test_moe_weight_layout.py
+25
@@ -7,6 +7,7 @@ import pytest
7
import torch
8
9
from vllm.lora.utils import get_supported_lora_modules
10
+from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4Config
11
from vllm.models.inkling.nvidia import moe
12
from vllm.models.inkling.nvidia.model import _TmlForCausalLMBase
13
from vllm.platforms import current_platform
@@ -87,6 +88,30 @@ def test_custom_embedding_is_not_a_lora_target() -> None:
88
assert "lm_head" in supported
89
90
91
+def test_inkling_mapper_maps_modelopt_exclusions() -> None:
92
+ quant_config = ModelOptNvFp4Config.from_config(
93
+ {
94
+ "quantization": {
95
+ "quant_algo": "NVFP4",
96
+ "group_size": 16,
97
+ "kv_cache_quant_algo": None,
98
+ "exclude_modules": [
99
+ "model.llm.layers.2.mlp.experts",
100
+ "model.llm.layers.2.mlp.shared_experts",
101
+ ],
102
+ }
103
+ }
104
+ )
105
+
106
+ quant_config.apply_vllm_mapper(
107
+ _TmlForCausalLMBase.hf_to_vllm_mapper.get_unstacked_mapper()
108
+ )
109
+
110
+ assert quant_config.is_layer_excluded("model.layers.2.mlp.experts")
111
+ assert quant_config.is_layer_excluded("model.layers.2.mlp.shared_experts")
112
+ assert not quant_config.is_layer_excluded("model.layers.3.mlp.experts")
113
+
114
+
115
@pytest.mark.parametrize(("projection", "amax"), [("w13", 4.375), ("w2", 2960.0)])
116
def test_moe_loads_calibrated_input_scale(projection: str, amax: float) -> None:
117
experts = SimpleNamespace(
vllm/models/inkling/nvfp4.py
deleted
-64
@@ -1,64 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""NVFP4 (ModelOpt) support for the Inkling mixture-of-experts.
4
-
5
-Only the routed MoE experts are quantized in the Inkling checkpoint;
6
-attention, the dense MLP, and the shared "sink" experts stay bf16 (they are
7
-in the checkpoint ``exclude_modules``). The routed experts are served by
8
-vLLM's standard ModelOpt NVFP4 fused-MoE stack (see ``moe.py``); this module
9
-keeps the checkpoint detection.
10
-"""
11
-
12
-from __future__ import annotations
13
-
14
-FLOAT8_E4M3_MAX = 448.0
15
-FLOAT4_E2M1_MAX = 6.0
16
-
17
-
18
-class InklingNvfp4Config:
19
- """Lightweight NVFP4 descriptor parsed from the checkpoint quant config.
20
-
21
- Holds the (mapped) ``exclude_modules`` so the model can decide, per MoE
22
- layer and per expert group, whether the weights are NVFP4 or plain bf16.
23
- """
24
-
25
- def __init__(self, group_size: int, exclude_modules: list[str]) -> None:
26
- self.group_size = group_size
27
- self.exclude_modules = set(exclude_modules)
28
-
29
- @staticmethod
30
- def _is_nvfp4(quant_cfg: dict) -> bool:
31
- wq = quant_cfg["modelopt_quant_config"]["quant_cfg"]["*weight_quantizer"]
32
- return tuple(wq["num_bits"]) == (2, 1) and tuple(
33
- wq["block_sizes"].get("scale_bits", [])
34
- ) == (4, 3)
35
-
36
- @classmethod
37
- def from_hf_config(cls, hf_config) -> InklingNvfp4Config | None:
38
- quant_cfg = getattr(hf_config, "quantization_config", None)
39
- text_config = getattr(hf_config, "text_config", None)
40
- if quant_cfg is None and text_config is not None:
41
- quant_cfg = getattr(text_config, "quantization_config", None)
42
- if quant_cfg is None:
43
- return None
44
- # ModelOpt <=0.29 nests everything under "quantization".
45
- if "quantization" in quant_cfg:
46
- quant_cfg = quant_cfg["quantization"]
47
- if not cls._is_nvfp4(quant_cfg):
48
- return None
49
- group_size = quant_cfg.get("group_size", 16)
50
- if group_size != 16:
51
- raise ValueError("Inkling NVFP4 only supports group size 16")
52
- exclude = list(quant_cfg.get("exclude_modules", []) or [])
53
- return cls(group_size=group_size, exclude_modules=exclude)
54
-
55
- def experts_quantized(self, layer_id: int) -> bool:
56
- """Whether the routed experts of ``layer_id`` are NVFP4 (vs excluded)."""
57
- return f"model.llm.layers.{layer_id}.mlp.experts" not in self.exclude_modules
58
-
59
- def shared_experts_quantized(self, layer_id: int) -> bool:
60
- """Whether the shared sink experts of ``layer_id`` are NVFP4."""
61
- return (
62
- f"model.llm.layers.{layer_id}.mlp.shared_experts"
63
- not in self.exclude_modules
64
- )
vllm/models/inkling/nvidia/model.py
+2
-15
@@ -47,7 +47,6 @@ from vllm.multimodal import MULTIMODAL_REGISTRY
47
from vllm.sequence import IntermediateTensors
48
49
from ..configs import InklingMMConfig, InklingModelConfig
50
-from ..nvfp4 import InklingNvfp4Config
50
from .attention import InklingAttention, compute_log_scaling_tau
51
from .layernorm import InklingRMSNorm
52
from .logits_processor import InklingLogitsProcessor
@@ -123,7 +122,6 @@ class InklingDecoderLayer(nn.Module):
122
is_local: bool,
123
quant_config: QuantizationConfig | None,
124
prefix: str,
126
- nvfp4_config: InklingNvfp4Config | None = None,
125
force_dense_mlp: bool = False,
126
) -> None:
127
super().__init__()
@@ -172,14 +170,10 @@ class InklingDecoderLayer(nn.Module):
170
prefix=f"{prefix}.mlp",
171
)
172
else:
175
- # InklingMoE decides per layer (from the checkpoint exclude list)
176
- # whether the routed experts are NVFP4 or bf16; the shared sink
177
- # experts are always bf16.
173
self.mlp = InklingMoE(
174
config,
180
- layer_id,
175
prefix=f"{prefix}.mlp",
182
- nvfp4_config=nvfp4_config,
176
+ quant_config=quant_config,
177
)
178
179
# Short convolution on the attention-output and MLP-output residual
@@ -261,7 +255,6 @@ class InklingModel(nn.Module):
255
config: InklingModelConfig,
256
quant_config: QuantizationConfig | None,
257
prefix: str,
264
- nvfp4_config: InklingNvfp4Config | None = None,
258
) -> None:
259
super().__init__()
260
self.config = config
@@ -278,7 +271,7 @@ class InklingModel(nn.Module):
271
def get_layer(prefix: str) -> InklingDecoderLayer:
272
idx = _layer_id(prefix + ".") or int(prefix.split(".")[-1])
273
return InklingDecoderLayer(
281
- config, idx, idx in local_ids, quant_config, prefix, nvfp4_config
274
+ config, idx, idx in local_ids, quant_config, prefix
275
)
276
277
self.start_layer, self.end_layer, self.layers = make_layers(
@@ -408,11 +401,6 @@ class _TmlForCausalLMBase(nn.Module, SupportsPP, SupportsLoRA):
401
) -> None:
402
quant_config = vllm_config.quant_config
403
self.config = text_config
411
- # NVFP4 experts are detected directly from the checkpoint quant config;
412
- # only the MoE experts are quantized (attention/dense MLP stay bf16).
413
- self.nvfp4_config = InklingNvfp4Config.from_hf_config(
414
- vllm_config.model_config.hf_config
415
- )
404
# Read by the MRV2 runner to publish per-request short-conv metadata.
405
# Short convolution is intrinsic to Inkling, so this is always set.
406
self.uses_sconv = True
@@ -420,7 +408,6 @@ class _TmlForCausalLMBase(nn.Module, SupportsPP, SupportsLoRA):
408
config=text_config,
409
quant_config=quant_config,
410
prefix=maybe_prefix(prefix, "model"),
423
- nvfp4_config=self.nvfp4_config,
411
)
412
initialize_lamport_rs_conv(
413
text_config.hidden_size,
vllm/models/inkling/nvidia/moe.py
+5
-28
@@ -43,20 +43,19 @@ from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
43
from vllm.utils.torch_utils import aux_stream
44
45
from ..configs import InklingModelConfig
46
-from ..nvfp4 import FLOAT4_E2M1_MAX, FLOAT8_E4M3_MAX
46
47
if TYPE_CHECKING:
48
from vllm.model_executor.layers.fused_moe.routed_experts import (
49
RoutedExperts,
50
)
52
-
53
- from ..nvfp4 import InklingNvfp4Config
51
+ from vllm.model_executor.layers.quantization import QuantizationConfig
52
53
# ---------------------------------------------------------------------------
54
# Gate / expert selection
55
# ---------------------------------------------------------------------------
56
57
_INKLING_LL_BF16_MAX_TOKENS = 64
58
+_NVFP4_INPUT_SCALE_DENOMINATOR = torch.finfo(torch.float8_e4m3fn).max * 6.0
59
60
61
def _linear_with_fp32_out(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
@@ -412,10 +411,9 @@ class InklingMoE(nn.Module):
411
def __init__(
412
self,
413
config: InklingModelConfig,
415
- layer_id: int,
414
*,
415
prefix: str = "",
418
- nvfp4_config: InklingNvfp4Config | None = None,
416
+ quant_config: QuantizationConfig | None = None,
417
) -> None:
418
super().__init__()
419
# Overfit to the served checkpoint: sigmoid gate renormalized after
@@ -436,22 +434,6 @@ class InklingMoE(nn.Module):
434
use_gate_bias=config.use_gate_bias,
435
)
436
439
- moe_quant_config = None
440
- if nvfp4_config is not None and nvfp4_config.experts_quantized(layer_id):
441
- from vllm.model_executor.layers.quantization.modelopt import (
442
- ModelOptNvFp4Config,
443
- )
444
-
445
- # The Inkling checkpoint is ModelOpt NVFP4; exclusion is decided per
446
- # layer right here, so no exclude list is needed.
447
- moe_quant_config = ModelOptNvFp4Config(
448
- quant_method="NVFP4",
449
- is_checkpoint_nvfp4_serialized=True,
450
- kv_cache_quant_algo=None,
451
- exclude_modules=[],
452
- group_size=nvfp4_config.group_size,
453
- )
454
-
437
# TRTLLM MoE kernels assume equal, contiguous per-rank expert slabs
438
# (local_expert_offset = ep_rank * local_num_experts), so pad the
439
# expert count to a multiple of the EP size. A no-op for the usual
@@ -464,7 +446,7 @@ class InklingMoE(nn.Module):
446
hidden_size=config.hidden_size,
447
intermediate_size=config.intermediate_size,
448
renormalize=False,
467
- quant_config=moe_quant_config,
449
+ quant_config=quant_config,
450
prefix=f"{prefix}.experts",
451
custom_routing_function=self._select_routed,
452
router_logits_dtype=torch.float32,
@@ -476,11 +458,6 @@ class InklingMoE(nn.Module):
458
self.experts.moe_config.skip_final_all_reduce = True
459
460
self._routed_sel: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None
479
- # The sinks are always bf16; fail loudly on a checkpoint that
480
- # quantizes them instead of silently misloading.
481
- assert nvfp4_config is None or not nvfp4_config.shared_experts_quantized(
482
- layer_id
483
- ), f"layer {layer_id}: NVFP4 shared experts are not supported"
461
462
sink_experts_cls = (
463
InklingSinkExpertsLinear
@@ -589,7 +566,7 @@ class InklingMoE(nn.Module):
566
f"bad {projection} input_amax: {amax}"
567
)
568
input_scale = getattr(experts, f"{projection}_input_scale")
592
- input_scale.data.fill_(amax / (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX))
569
+ input_scale.data.fill_(amax / _NVFP4_INPUT_SCALE_DENOMINATOR)
570
return [f"experts.routed_experts.{projection}_input_scale"]
571
572
param = getattr(experts, key)
vllm/models/inkling/nvidia/mtp.py
-1
@@ -73,7 +73,6 @@ class InklingMTPDepthLayer(nn.Module):
73
is_local=is_local,
74
quant_config=None,
75
prefix=f"{prefix}.transformer_block",
76
- nvfp4_config=None,
76
force_dense_mlp=True,
77
)
78
vllm/transformers_utils/model_arch_config_convertor.py
+7
-2
@@ -215,8 +215,13 @@ class ModelArchConfigConvertorBase:
215
else:
216
# Set quant_method for ModelOpt models.
217
producer_name = quant_cfg.get("producer", {}).get("name")
218
- if producer_name == "modelopt":
219
- quant_algo = quant_cfg.get("quantization", {}).get("quant_algo")
218
+ modelopt_quant_cfg = quant_cfg.get("quantization", {})
219
+ is_legacy_modelopt = (
220
+ isinstance(modelopt_quant_cfg, dict)
221
+ and "modelopt_quant_config" in modelopt_quant_cfg
222
+ )
223
+ if producer_name == "modelopt" or is_legacy_modelopt:
224
+ quant_algo = modelopt_quant_cfg.get("quant_algo")
225
if quant_algo is not None:
226
quant_algo_upper = str(quant_algo).upper()
227
if quant_algo_upper in {