iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[quantization] Add per-token NVFP4 quantization MoE support for ReLU2 (#56740)

Browse Source 
Signed-off-by: Xuanteng Huang <xuantengh@nvidia.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>
This commit is contained in:
Xuanteng Huang
2026-10-07 22:01:09 +00:00
committed by GitHub
co-authored by Claude Fable 5.1 OpenAI Codex
parent b803edbc9f
commit f38c6778ce
18 changed files with 450 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-NVFP4-per-token.yaml
+14
View File
@@ -0,0 +1,14 @@
model_name: "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4"
accuracy_threshold: 0.93
num_questions: 1319
num_fewshot: 5
startup_max_wait_seconds: 1200
server_args: >-
--enforce-eager
--max-model-len 4096
--tensor-parallel-size 2
--enable-expert-parallel
--mamba-backend flashinfer
--moe-backend flashinfer_trtllm
--quantization-config.moe.activation nvfp4_per_token
--speculative-config '{"method":"mtp","num_speculative_tokens":5}'
tests/evals/gsm8k/configs/models-blackwell-ep.txt
+1
View File
@@ -1,3 +1,4 @@
Qwen3-Next-80B-A3B-NVFP4-EP2.yaml
Qwen3-Next-FP8-EP2.yaml
Nemotron-3-Super-120B-A12B-NVFP4.yaml
Nemotron-3-Super-120B-A12B-NVFP4-per-token.yaml
tests/kernels/moe/test_trtllm_nvfp4_moe.py
+112 -18
View File
@@ -16,8 +16,7 @@ import torch
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from tests.kernels.moe.utils import check_deferred_moe_finalize, make_test_quant_config
from tests.kernels.quantization.nvfp4_utils import (
FLOAT4_E2M1_MAX,
FLOAT8_E4M3_MAX,
convert_swizzled_to_linear,
dequantize_nvfp4_to_dtype,
)
from tests.kernels.utils import torch_moe
@@ -39,6 +38,13 @@ from vllm.model_executor.layers.fused_moe.config import (
from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
TrtLlmNvFp4ExpertsModular,
)
from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
prepare_static_weights_for_trtllm_fp4_moe,
reorder_w1w3_to_w3w1,
)
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
NVFP4_PER_TOKEN_BASE_GLOBAL_SCALE,
)
from vllm.platforms import current_platform
from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
from vllm.utils.math_utils import next_power_of_2
@@ -96,10 +102,14 @@ SITU = op_registry[_SITU_OP_NAME]
def _make_trtllm_fp4_moe_kernel(
moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig
moe_config: FusedMoEConfig,
quant_config: FusedMoEQuantConfig,
per_token_activation: bool = False,
) -> mk.FusedMoEKernel:
experts = TrtLlmNvFp4ExpertsModular(
moe_config=moe_config, quant_config=quant_config
moe_config=moe_config,
quant_config=quant_config,
per_token_activation=per_token_activation,
)
# Mimic the production weight-loader path so per-expert tensors that
# are normally precomputed in process_weights_after_loading (g1_scale_c
@@ -126,16 +136,37 @@ def _make_trtllm_fp4_moe_kernel(
ACTIVATION_CASES = [
pytest.param(MoEActivation.SILU, MoEActivation.SILU, None, id="silu"),
pytest.param(MoEActivation.SILU, SILU_WITH_CLAMP, _SWIGLU_LIMIT, id="silu_clamp"),
pytest.param(MoEActivation.SITU, SITU, None, id="situ"),
pytest.param(MoEActivation.SILU, MoEActivation.SILU, None, False, id="silu"),
pytest.param(
MoEActivation.SILU,
MoEActivation.SILU,
None,
True,
id="silu_per_token",
),
pytest.param(
MoEActivation.SILU,
SILU_WITH_CLAMP,
_SWIGLU_LIMIT,
False,
id="silu_clamp",
),
pytest.param(MoEActivation.SITU, SITU, None, False, id="situ"),
pytest.param(
MoEActivation.RELU2_NO_MUL,
MoEActivation.RELU2_NO_MUL,
None,
False,
id="relu2_no_mul",
),
pytest.param(MoEActivation.GELU, MoEActivation.GELU, None, id="gelu"),
pytest.param(
MoEActivation.RELU2_NO_MUL,
MoEActivation.RELU2_NO_MUL,
None,
True,
id="relu2_no_mul_per_token",
),
pytest.param(MoEActivation.GELU, MoEActivation.GELU, None, False, id="gelu"),
]
@@ -143,7 +174,10 @@ ACTIVATION_CASES = [
@pytest.mark.parametrize("e", [128])
@pytest.mark.parametrize("topk", [8])
@pytest.mark.parametrize("dtype", [torch.bfloat16])
@pytest.mark.parametrize("activation,torch_activation,swiglu_limit", ACTIVATION_CASES)
@pytest.mark.parametrize(
"activation,torch_activation,swiglu_limit,per_token_activation",
ACTIVATION_CASES,
)
@torch.inference_mode()
def test_trtllm_fp4_moe_no_graph(
m: int,
@@ -155,6 +189,7 @@ def test_trtllm_fp4_moe_no_graph(
activation: MoEActivation,
torch_activation: MoEActivation | type[SiluAndMulWithClamp] | type[SituAndMul],
swiglu_limit: float | None,
per_token_activation: bool,
workspace_init,
):
# FlashInfer's trtllm_batched_gemm_runner has no precompiled tile
@@ -202,6 +237,42 @@ def test_trtllm_fp4_moe_no_graph(
# clamp/output-scale coupling in the FlashInfer kernel wrapper.
quant_config.g1_alphas.fill_(_LARGE_OUTPUT1_SCALE)
w1_q_linear = w1_q
w2_q_linear = w2_q
w1_scale_reference = quant_config.w1_scale
w2_scale_reference = quant_config.w2_scale
assert w1_scale_reference is not None
assert w2_scale_reference is not None
w1_scale_linear = torch.stack(
[
convert_swizzled_to_linear(scale, w1_q.size(1), k, quant_blocksize)
for scale in w1_scale_reference
]
)
w2_scale_linear = torch.stack(
[
convert_swizzled_to_linear(scale, w2_q.size(1), n, quant_blocksize)
for scale in w2_scale_reference
]
)
w1_q_for_kernel = w1_q
if is_gated_act:
w1_q_for_kernel, w1_scale_linear = reorder_w1w3_to_w3w1(
w1_q.clone(), w1_scale_linear
)
w1_q, w1_scale, w2_q, w2_scale = prepare_static_weights_for_trtllm_fp4_moe(
w1_q_for_kernel,
w2_q,
w1_scale_linear,
w2_scale_linear,
hidden_size=k,
intermediate_size=n,
num_experts=e,
is_gated_activation=is_gated_act,
)
quant_config._w1.scale = w1_scale
quant_config._w2.scale = w2_scale
score = torch.randn((m, e), device="cuda", dtype=dtype)
topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False)
@@ -224,7 +295,11 @@ def test_trtllm_fp4_moe_no_graph(
),
)
trtllm_experts = _make_trtllm_fp4_moe_kernel(moe_config, quant_config)
trtllm_experts = _make_trtllm_fp4_moe_kernel(
moe_config,
quant_config,
per_token_activation=per_token_activation,
)
if activation == MoEActivation.SITU:
torch.testing.assert_close(
trtllm_experts.fused_experts.g1_scale_c, quant_config.a2_gscale
@@ -245,10 +320,20 @@ def test_trtllm_fp4_moe_no_graph(
# Reference: round-trip activations and weights through FP4
# quant/dequant so the comparison isolates kernel/activation behavior
# from quantization error.
a_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / a.abs().max()).to(
torch.float32
)
a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a, a_global_scale)
if per_token_activation:
from flashinfer import SfLayout, nvfp4_quantize
a_fp4, a_scale_interleaved, per_token_scale = nvfp4_quantize(
a,
NVFP4_PER_TOKEN_BASE_GLOBAL_SCALE,
sfLayout=SfLayout.layout_linear,
per_token_activation=True,
)
a_global_scale = 1.0 / per_token_scale.unsqueeze(1)
else:
assert quant_config.a1_gscale is not None
a_global_scale = quant_config.a1_gscale[0]
a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a, a_global_scale)
a_in_dtype = dequantize_nvfp4_to_dtype(
a_fp4,
a_scale_interleaved,
@@ -256,6 +341,7 @@ def test_trtllm_fp4_moe_no_graph(
dtype=a.dtype,
device=a.device,
block_size=quant_blocksize,
is_sf_linear_layout=per_token_activation,
)
w1_d = torch.empty(
@@ -264,16 +350,16 @@ def test_trtllm_fp4_moe_no_graph(
w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype)
for idx in range(e):
w1_d[idx] = dequantize_nvfp4_to_dtype(
w1_q[idx],
quant_config.w1_scale[idx],
w1_q_linear[idx],
w1_scale_reference[idx],
(1 / quant_config.g1_alphas[idx]),
dtype=dtype,
device=w1_q.device,
block_size=quant_blocksize,
)
w2_d[idx] = dequantize_nvfp4_to_dtype(
w2_q[idx],
quant_config.w2_scale[idx],
w2_q_linear[idx],
w2_scale_reference[idx],
(1 / quant_config.g2_alphas[idx]),
dtype=dtype,
device=w2_q.device,
@@ -284,6 +370,13 @@ def test_trtllm_fp4_moe_no_graph(
a_in_dtype, w1_d, w2_d, score, topk, activation=torch_activation
)
if swiglu_limit is None:
cosine = torch.nn.functional.cosine_similarity(
torch_output.float().flatten(),
trtllm_output.float().flatten(),
dim=0,
)
assert cosine > 0.99
torch.testing.assert_close(torch_output, trtllm_output, atol=2e-1, rtol=2e-1)
@@ -416,5 +509,6 @@ if __name__ == "__main__":
MoEActivation.GELU,
MoEActivation.GELU,
None,
False,
None,
)
tests/model_executor/model_loader/test_weight_cache.py
+46
View File
@@ -356,6 +356,52 @@ def test_weight_cache_key_distinguishes_dp_ranks():
assert key.mismatched_fields(replace(key, pp_rank=0)) == ["pp_rank"]
@pytest.mark.parametrize("hf_quant_config", [None, {"quant_algo": "NVFP4"}])
def test_weight_cache_key_distinguishes_nvfp4_activation_override(
tmp_path, hf_quant_config
):
"""Cache keys isolate activation modes and reject ambiguous legacy hashes."""
import json
from dataclasses import replace
from types import SimpleNamespace
import torch
from vllm.config.quantization import QuantizationConfigArgs
from vllm.model_executor.model_loader.weight_cache.protocol import WeightCacheKey
from vllm.utils.hashing import safe_hash
model_config = SimpleNamespace(
model=str(tmp_path),
dtype=torch.bfloat16,
quantization="modelopt_fp4",
quantization_config=None,
revision=None,
hf_config=SimpleNamespace(
architectures=["NemotronHForCausalLM"],
quantization_config=hf_quant_config,
),
)
static_key = WeightCacheKey.from_model_config(model_config, tp_size=1, tp_rank=0)
model_config.quantization_config = QuantizationConfigArgs(
moe={"activation": "nvfp4_per_token"}
)
per_token_key = WeightCacheKey.from_model_config(model_config, tp_size=1, tp_rank=0)
assert static_key.mismatched_fields(per_token_key) == ["quant_config_hash"]
legacy_hash = (
""
if hf_quant_config is None
else safe_hash(
json.dumps(hf_quant_config, sort_keys=True).encode(),
usedforsecurity=False,
).hexdigest()
)
assert static_key.mismatched_fields(
replace(static_key, quant_config_hash=legacy_hash)
) == ["quant_config_hash"]
def test_ipc_loader_copy_mode_reports_no_external_weight_memory():
"""Copy mode clones the weights into the engine, so nothing is external
(vllm_config is never touched)."""
tests/quantization/test_blackwell_moe.py
+13
View File
@@ -297,3 +297,16 @@ def test_nemotron_fp4_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
hf_overrides=HF_OVERRIDE_TEXT,
extra_args=["--moe-backend=flashinfer_trtllm"],
)
def test_nemotron_nvfp4_per_token_moe_flashinfer_trtllm(
monkeypatch: pytest.MonkeyPatch,
):
can_initialize(
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4",
hf_overrides=HF_OVERRIDE_TEXT,
extra_args=[
"--moe-backend=flashinfer_trtllm",
"--quantization-config.moe.activation=nvfp4_per_token",
],
)
tests/quantization/test_modelopt.py
+76 -1
View File
@@ -6,6 +6,7 @@ Run `pytest tests/quantization/test_modelopt.py`.
"""
import os
from contextlib import nullcontext
from typing import Any, NoReturn
from unittest.mock import MagicMock, Mock, patch
@@ -17,7 +18,9 @@ from tests.quantization.utils import (
load_model_without_vllm_runner,
)
from vllm.config import VllmConfig, set_current_vllm_config
from vllm.config.load import LoadConfig
from vllm.config.model import ModelConfig
from vllm.config.quantization import QuantizationConfigArgs
from vllm.forward_context import set_forward_context
from vllm.model_executor.kernels.linear import (
FlashInferCuteDslNvFp4W4A16LinearKernel,
@@ -40,12 +43,14 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
kMxfp8Dynamic,
kMxfp8Static,
kNvfp4Dynamic,
kNvfp4DynamicToken,
kNvfp4Static,
)
from vllm.model_executor.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from vllm.model_executor.model_loader.weight_utils import get_quant_config
from vllm.platforms import current_platform
@@ -113,11 +118,13 @@ def _mixed_precision_config(quantized_layers: dict) -> ModelOptMixedPrecisionCon
)
def test_modelopt_nvfp4_quantizes_parallel_lm_head():
@pytest.mark.parametrize("moe_activation", [None, "nvfp4_per_token"])
def test_modelopt_nvfp4_quantizes_parallel_lm_head(moe_activation):
config = ModelOptNvFp4Config(
is_checkpoint_nvfp4_serialized=True,
kv_cache_quant_algo=None,
exclude_modules=[],
quantization_args=QuantizationConfigArgs(moe={"activation": moe_activation}),
)
method = config.get_quant_method(_mock_lm_head(), prefix="lm_head")
@@ -790,6 +797,74 @@ def test_modelopt_nvfp4_moe_dispatches_to_marlin_when_w4a16(
assert kwargs["activation_key"] is kNvfp4Dynamic
@pytest.mark.parametrize("quantization", ["modelopt_fp4", "modelopt_mixed"])
@pytest.mark.parametrize("per_token", [False, True])
def test_modelopt_nvfp4_moe_activation_override(quantization, per_token):
"""Runtime overrides reach NVFP4 configs without changing HF metadata."""
hf_config = {
"quant_method": "modelopt",
"quant_algo": "MIXED_PRECISION"
if quantization == "modelopt_mixed"
else "NVFP4",
"quantized_layers": {"experts": {"quant_algo": "NVFP4"}},
}
args = (
QuantizationConfigArgs(moe={"activation": "nvfp4_per_token"})
if per_token
else None
)
config = get_quant_config(
Mock(
quantization=quantization,
quantization_config=args,
hf_config=Mock(quantization_config=hf_config),
),
LoadConfig(),
)
if isinstance(config, ModelOptMixedPrecisionConfig):
config = config.nvfp4_config
assert isinstance(config, ModelOptNvFp4Config)
assert config.moe_activation_override == (kNvfp4DynamicToken if per_token else None)
assert "_online_quantization_args" not in hf_config
@pytest.mark.parametrize(
"quant_method,activation,backend,error",
[
("NVFP4", kNvfp4DynamicToken, "FLASHINFER_TRTLLM", None),
("W4A16_NVFP4", kNvfp4DynamicToken, "FLASHINFER_TRTLLM", "W4A4 checkpoint"),
(
"NVFP4",
kMxfp8Dynamic,
"FLASHINFER_TRTLLM",
"Unsupported.*activation override",
),
("NVFP4", kNvfp4DynamicToken, "MARLIN", "requires the FlashInfer TRTLLM"),
],
)
def test_modelopt_nvfp4_moe_validates_activation_override(
quant_method, activation, backend, error
):
from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import NvFp4MoeBackend
from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4FusedMoE
config = ModelOptNvFp4Config(
quant_method=quant_method,
quantization_args=QuantizationConfigArgs(moe={"activation": activation}),
)
expected = pytest.raises(ValueError, match=error) if error else nullcontext()
with (
patch(
"vllm.model_executor.layers.quantization.modelopt.select_nvfp4_moe_backend",
return_value=(NvFp4MoeBackend[backend], Mock()),
) as select_backend,
expected,
):
method = ModelOptNvFp4FusedMoE(config, Mock())
assert method.per_token_activation
assert select_backend.call_args.kwargs["activation_key"] == activation
@pytest.mark.parametrize(
"per_layer_algo, expected_weight, expected_activation",
[
tests/quantization/test_online.py
+25 -7
View File
@@ -103,6 +103,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
amax_for_tp_weight_quant,
kMxfp8Dynamic,
kMxfp8Static,
kNvfp4DynamicToken,
kNvfp4Static,
weight_amax,
)
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -211,15 +213,26 @@ def test_qianfan_online_fp8_keeps_vision_layers_unquantized(
assert args.ignore == []
@pytest.mark.parametrize(
("is_act_and_mul", "expected_hidden_alignment"),
[(True, 256), (False, 512)],
)
def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed(
monkeypatch,
monkeypatch, is_act_and_mul, expected_hidden_alignment
) -> None:
method = object.__new__(Nvfp4OnlineMoEMethod)
method.moe = SimpleNamespace(is_act_and_mul=True)
method.nvfp4_backend = object()
method.experts_cls = object
method.moe_quant_config = None
method.moe_kernel = None
monkeypatch.setattr(current_platform, "is_device_capability_family", lambda _: True)
select_backend = Mock(return_value=(object(), object))
monkeypatch.setattr(
"vllm.model_executor.layers.quantization.online.nvfp4.select_nvfp4_moe_backend",
select_backend,
)
moe_config = SimpleNamespace(is_act_and_mul=is_act_and_mul)
method = Nvfp4OnlineMoEMethod(moe=moe_config)
select_backend.assert_called_once_with(
config=moe_config,
weight_key=kNvfp4Static,
activation_key=kNvfp4DynamicToken,
)
layer = Mock()
converted_weights = [
@@ -274,7 +287,12 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed(
assert method.moe_kernel is kernel
assert convert_weights.call_count == 3
assert all(
call.kwargs["trtllm_hidden_alignment"] == expected_hidden_alignment
for call in convert_weights.call_args_list
)
make_kernel.assert_called_once()
assert make_kernel.call_args.kwargs["per_token_activation"] is True
get_quant_config.assert_called_once()
assert process_weights.call_count == 3
tests/quantization/test_quantization_config_args.py
+37 -14
View File
@@ -26,6 +26,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
kInt4Static32,
kInt8StaticChannelSym,
kMxfp8Dynamic,
kNvfp4DynamicToken,
kNvfp4Static,
)
# ---- QuantSpec ------------------------------------------------------------
@@ -68,13 +70,17 @@ def test_args_linear_string_resolves_via_quant_key_names():
assert args.moe is None
def test_args_moe_string_resolves_via_online_shorthand():
@pytest.mark.parametrize(
"shorthand,weight",
[("fp8_per_block", kFp8Static128BlockSym), ("nvfp4_per_token", kNvfp4Static)],
)
def test_args_moe_string_resolves_via_online_shorthand(shorthand, weight):
# An online-shorthand name pulls the matching slot from _ONLINE_SHORTHANDS
# (so `linear: "fp8_per_block"` and `moe: "fp8_per_block"` produce the
# same per-layer-kind spec the `--quantization fp8_per_block` shorthand
# would).
args = quant_config_args(moe="fp8_per_block")
assert args.moe == QuantSpec(weight=kFp8Static128BlockSym)
args = quant_config_args(moe=shorthand)
assert args.moe == QuantSpec(weight=weight)
def test_args_string_shorthand_missing_slot_raises():
@@ -84,9 +90,14 @@ def test_args_string_shorthand_missing_slot_raises():
quant_config_args(linear="int8_per_channel_weight_only")
def test_args_accepts_dict_form():
args = quant_config_args(moe={"activation": "mxfp8"})
assert args.moe == QuantSpec(weight=None, activation=kMxfp8Dynamic)
@pytest.mark.parametrize(
"activation,expected",
[("mxfp8", kMxfp8Dynamic), ("nvfp4_per_token", kNvfp4DynamicToken)],
)
def test_args_accepts_dict_form(activation, expected):
args = quant_config_args(moe={"activation": activation})
assert args.moe == QuantSpec(weight=None, activation=expected)
assert args.linear is None
def test_targets_reject_non_string_keys():
@@ -110,22 +121,34 @@ def test_resolve_colliding_shorthand_is_deferred(quantization: str):
assert resolve_quantization_config(quantization, None) is None
def test_resolve_int8_shorthand_leaves_linear_unset():
# int8_per_channel_weight_only is MoE-only; linear stays None so that
# OnlineQuantizationConfig leaves Linear layers in full precision.
args = resolve_quantization_config("int8_per_channel_weight_only", None)
@pytest.mark.parametrize(
"shorthand,weight",
[
("int8_per_channel_weight_only", kInt8StaticChannelSym),
("nvfp4_per_token", kNvfp4Static),
],
)
def test_resolve_moe_only_shorthand_leaves_linear_unset(shorthand, weight):
args = resolve_quantization_config(shorthand, None)
assert args is not None
assert args.linear is None
assert args.moe == QuantSpec(weight=kInt8StaticChannelSym)
assert args.moe == QuantSpec(weight=weight)
def test_resolve_quantization_config_only():
@pytest.mark.parametrize("quantization", [None, "modelopt_fp4"])
@pytest.mark.parametrize(
"activation,expected",
[("mxfp8", kMxfp8Dynamic), ("nvfp4_per_token", kNvfp4DynamicToken)],
)
def test_resolve_activation_override(quantization, activation, expected):
# When only `quantization_config` is given (e.g. for an already-quantized
# checkpoint that needs an activation override), it's returned as-is.
args = resolve_quantization_config(None, {"moe": {"activation": "mxfp8"}})
args = resolve_quantization_config(
quantization, {"moe": {"activation": activation}}
)
assert args is not None
assert args.linear is None
assert args.moe == QuantSpec(weight=None, activation=kMxfp8Dynamic)
assert args.moe == QuantSpec(weight=None, activation=expected)
def test_resolve_merges_explicit_over_shorthand():
vllm/config/quantization.py
+2
View File
@@ -27,6 +27,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
kMxfp4Dynamic,
kMxfp4Static,
kMxfp8Dynamic,
kNvfp4DynamicToken,
kNvfp4Static,
)
@@ -40,6 +41,7 @@ QUANT_KEY_NAMES: dict[str, QuantKey] = {
"fp8_per_block_dynamic": kFp8Dynamic128Sym,
"mxfp8": kMxfp8Dynamic,
"mxfp4": kMxfp4Dynamic,
"nvfp4_per_token": kNvfp4DynamicToken,
"int8_per_channel_static": kInt8StaticChannelSym,
# Load-time MXFP4-to-int4 for Kimi-K3 on gfx942.
"int4_per_group_32": kInt4Static32,
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+2
View File
@@ -334,6 +334,7 @@ def convert_to_nvfp4_moe_kernel_format(
a2_scale: torch.Tensor | None,
is_act_and_mul: bool,
use_a16: bool = False,
trtllm_hidden_alignment: int = 256,
) -> tuple[
torch.Tensor,
torch.Tensor,
@@ -419,6 +420,7 @@ def convert_to_nvfp4_moe_kernel_format(
w2_scale_2=w2_scale_2,
a2_scale=a2_scale,
is_act_and_mul=is_act_and_mul,
trtllm_hidden_alignment=trtllm_hidden_alignment,
)
elif nvfp4_backend == NvFp4MoeBackend.HUMMING:
from vllm.model_executor.layers.quantization.utils.humming import (
vllm/model_executor/layers/quantization/modelopt.py
+49 -8
View File
@@ -12,7 +12,7 @@ from torch.nn.parameter import Parameter
import vllm.envs as envs
from vllm.config import get_current_vllm_config_or_none
from vllm.config.quantization import QuantSpec
from vllm.config.quantization import QuantizationConfigArgs, QuantSpec
from vllm.logger import init_logger
from vllm.model_executor.kernels.linear import (
init_fp8_linear_kernel,
@@ -63,6 +63,9 @@ from vllm.model_executor.layers.quantization.base_config import (
)
from vllm.model_executor.layers.quantization.fp8 import Fp8Config, Fp8MoEMethod
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
trtllm_nvfp4_hidden_alignment,
)
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
FP8_SCALE_SENTINEL,
process_fp8_input_tensor_strategy_moe,
@@ -90,6 +93,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
kMxfp8Dynamic,
kMxfp8Static,
kNvfp4Dynamic,
kNvfp4DynamicToken,
kNvfp4Static,
)
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
@@ -724,12 +728,18 @@ class ModelOptNvFp4Config(ModelOptQuantConfigBase):
kv_cache_quant_algo: str | None = None,
exclude_modules: list[str] | None = None,
group_size: int = 16,
*,
quantization_args: QuantizationConfigArgs | None = None,
) -> None:
if exclude_modules is None:
exclude_modules = []
super().__init__(exclude_modules)
self.quant_method = quant_method
self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized
moe_spec = quantization_args.moe if quantization_args is not None else None
self.moe_activation_override: QuantKey | None = (
moe_spec.activation if moe_spec is not None else None
)
if is_checkpoint_nvfp4_serialized:
self.group_size = group_size
self.kv_cache_quant_algo = kv_cache_quant_algo
@@ -815,6 +825,7 @@ class ModelOptNvFp4Config(ModelOptQuantConfigBase):
kv_cache_quant_method,
exclude_modules,
group_size,
quantization_args=original_config.get("_online_quantization_args"),
)
@@ -835,18 +846,36 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
) -> None:
super().__init__(moe_config)
self.quant_config = quant_config
# W4A16 mode fires for W4A16_NVFP4 on-disk checkpoints. With
# activation_key=None every W4A4 backend's _supports_quant_scheme
# rejects itself (they all require (kNvfp4Static, kNvfp4Dynamic)
# exactly); only Marlin survives. Marlin's MoE path drops
# activation scales in convert_to_nvfp4_moe_kernel_format, so no
# other change is needed.
self.use_a16 = quant_config.quant_method == "W4A16_NVFP4"
activation_key = None if self.use_a16 else kNvfp4Dynamic
if quant_config.moe_activation_override is not None:
if self.use_a16:
raise ValueError(
"NVFP4 activation overrides require a W4A4 checkpoint."
)
activation_key = quant_config.moe_activation_override
if activation_key not in (kNvfp4Dynamic, kNvfp4DynamicToken):
raise ValueError(
"Unsupported ModelOpt NVFP4 MoE activation override: "
f"{activation_key}. Use quantization_config.moe.activation="
"'nvfp4_per_token' or omit the override."
)
self.per_token_activation = activation_key == kNvfp4DynamicToken
self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
config=self.moe,
weight_key=kNvfp4Static,
activation_key=None if self.use_a16 else kNvfp4Dynamic,
activation_key=activation_key,
)
if (
self.per_token_activation
and self.nvfp4_backend != NvFp4MoeBackend.FLASHINFER_TRTLLM
):
raise ValueError(
"Per-token NVFP4 activation for pre-quantized weights "
"requires the FlashInfer TRTLLM MoE backend."
)
self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(
self.nvfp4_backend
@@ -984,6 +1013,10 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
self._build_moe_kernel(layer)
return
if self.per_token_activation:
layer.w13_input_scale.data.fill_(1.0)
layer.w2_input_scale.data.fill_(1.0)
# Use a single gscale for w13.
if self.moe.is_act_and_mul and not torch.allclose(
layer.w13_weight_scale_2[:, 0], layer.w13_weight_scale_2[:, 1]
@@ -1016,6 +1049,10 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
a2_scale=layer.w2_input_scale,
is_act_and_mul=self.moe.is_act_and_mul,
use_a16=self.use_a16,
trtllm_hidden_alignment=trtllm_nvfp4_hidden_alignment(
per_token_activation=self.per_token_activation,
is_act_and_mul=self.moe.is_act_and_mul,
),
)
replace_parameter(layer, "w13_weight", w13)
@@ -1039,6 +1076,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
experts_cls=self.experts_cls,
backend=self.nvfp4_backend,
routing_tables=layer._expert_routing_tables(),
per_token_activation=self.per_token_activation,
)
self.moe_kernel.fused_experts.process_weights_after_loading(layer)
@@ -1588,6 +1626,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
group_size: int | None,
**kwargs: Any,
) -> "ModelOptMixedPrecisionConfig":
quantization_args = original_config.get("_online_quantization_args")
if "quantization" in original_config:
quantized_layers = original_config["quantization"].get(
"quantized_layers", {}
@@ -1626,6 +1665,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
kv_cache_quant_algo=kv_cache_quant_method,
exclude_modules=[],
group_size=group_size,
quantization_args=quantization_args,
)
# Sibling config for layers that declare quant_algo: "W4A16_NVFP4".
# get_quant_method resolves this sub-config to the (kNvfp4Static, None)
@@ -1638,6 +1678,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
kv_cache_quant_algo=kv_cache_quant_method,
exclude_modules=[],
group_size=group_size,
quantization_args=quantization_args,
)
mxfp8_config = ModelOptMxFp8Config(
vllm/model_executor/layers/quantization/online/nvfp4.py
+15 -2
View File
@@ -19,6 +19,9 @@ from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
from vllm.model_executor.layers.quantization.online.moe_base import (
OnlineMoEMethodBase,
)
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
trtllm_nvfp4_hidden_alignment,
)
from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
FLOAT4_E2M1_MAX,
)
@@ -87,6 +90,8 @@ class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase):
(SM100) only.
"""
activation_quant_key = kNvfp4DynamicToken
def __init__(
self,
*,
@@ -100,9 +105,13 @@ class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase):
self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
config=self.moe,
weight_key=kNvfp4Static,
activation_key=kNvfp4DynamicToken,
activation_key=self.activation_quant_key,
)
@property
def per_token_activation(self) -> bool:
return self.activation_quant_key == kNvfp4DynamicToken
def process_weights_after_loading(self, layer: Module) -> None:
if getattr(layer, "_already_called_process_weights_after_loading", False):
return
@@ -156,6 +165,10 @@ class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase):
w2_scale_2=layer.w2_weight_scale_2,
a2_scale=layer.w2_input_scale,
is_act_and_mul=self.moe.is_act_and_mul,
trtllm_hidden_alignment=trtllm_nvfp4_hidden_alignment(
per_token_activation=self.per_token_activation,
is_act_and_mul=self.moe.is_act_and_mul,
),
)
replace_parameter(layer, "w13_weight", w13)
@@ -176,7 +189,7 @@ class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase):
experts_cls=self.experts_cls,
backend=self.nvfp4_backend,
routing_tables=layer._expert_routing_tables(),
per_token_activation=True,
per_token_activation=self.per_token_activation,
)
else:
# Reload creates new scale tensors; derived kernel scales must use
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+6 -1
View File
@@ -344,6 +344,7 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass(
w2_scale_2: torch.Tensor,
a2_scale: torch.Tensor,
is_act_and_mul: bool,
trtllm_hidden_alignment: int = 256,
) -> tuple[
torch.Tensor,
torch.Tensor,
@@ -398,7 +399,11 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass(
# Shuffle weights and scales for FI TRTLLM NVFP4 MoE kernels.
if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:
w13, w13_scale, w2, w2_scale, padded_hidden = align_fp4_moe_hidden_dim_for_fi(
w13, w13_scale, w2, w2_scale
w13,
w13_scale,
w2,
w2_scale,
min_alignment=trtllm_hidden_alignment,
)
if layer.moe_config.hidden_dim_unpadded is None:
layer.moe_config.hidden_dim_unpadded = layer.moe_config.hidden_dim
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+14
View File
@@ -280,6 +280,20 @@ def align_fp4_moe_weights_for_fi(
return padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_intermediate
def trtllm_nvfp4_hidden_alignment(
per_token_activation: bool, is_act_and_mul: bool
) -> int:
"""Hidden-dim alignment for the TRTLLM-Gen NVFP4 MoE cubins.
The per-token (dynamic activation scale) non-gated variant has far fewer
tile configs for 256-aligned K and its heuristic default tactic fails with
"No valid config found" (e.g. Nemotron's 2688 -> 2816), so pad to 512.
"""
if per_token_activation and not is_act_and_mul:
return 512
return 256
def align_fp4_moe_hidden_dim_for_fi(
w13: torch.Tensor,
w13_scale: torch.Tensor,
vllm/model_executor/model_loader/weight_cache/protocol.py
+14 -5
View File
@@ -28,6 +28,7 @@ from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
import vllm.version
from vllm.config import ModelConfig
from vllm.config.utils import normalize_value
from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
from vllm.model_executor.model_loader.weight_cache.utils import (
format_socket_role_suffix,
@@ -227,12 +228,17 @@ def verify_peer_is_owner(conn: socket.socket) -> None:
raise PermissionError(f"Rejecting weight cache connection from uid {peer_uid}")
def _hash_quant_config(quant_config: Any) -> str:
if quant_config is None:
return ""
def _hash_quant_config(quant_config: Any, runtime_quant_config: Any = None) -> str:
if hasattr(quant_config, "to_dict"):
quant_config = quant_config.to_dict()
payload = json.dumps(quant_config, sort_keys=True, default=str)
payload = json.dumps(
{
"checkpoint": quant_config,
"runtime": normalize_value(runtime_quant_config),
},
sort_keys=True,
default=str,
)
return safe_hash(payload.encode(), usedforsecurity=False).hexdigest()
@@ -333,7 +339,10 @@ class WeightCacheKey:
pp_rank=pp_rank,
dtype=str(model_config.dtype),
quantization=model_config.quantization,
quant_config_hash=_hash_quant_config(quant_config),
quant_config_hash=_hash_quant_config(
quant_config=quant_config,
runtime_quant_config=model_config.quantization_config,
),
revision=model_config.revision,
vllm_version=vllm.version.__version__,
is_draft=is_draft,
vllm/model_executor/model_loader/weight_utils.py
+14 -5
View File
@@ -200,6 +200,9 @@ def get_quant_config(
QuantizationConfigArgs,
resolve_quantization_config,
)
from vllm.model_executor.layers.quantization.modelopt import (
ModelOptQuantConfigBase,
)
from vllm.model_executor.layers.quantization.online.base import (
OnlineQuantizationConfig,
)
@@ -220,6 +223,14 @@ def get_quant_config(
)
return checkpoint_config
def make_checkpoint_config(config: dict[str, Any]) -> QuantizationConfig:
if issubclass(quant_cls, ModelOptQuantConfigBase):
assert online_args is None or isinstance(
online_args, QuantizationConfigArgs
)
config = {**config, "_online_quantization_args": online_args}
return maybe_compose_online_quantization(quant_cls.from_config(config))
# Read the quantization config from the HF model config, if available.
hf_quant_config = getattr(model_config.hf_config, "quantization_config", None)
# some vision model may keep quantization_config in their text_config
@@ -260,9 +271,7 @@ def get_quant_config(
):
pass # fall through to file-based loading below
else:
return maybe_compose_online_quantization(
quant_cls.from_config(hf_quant_config)
)
return make_checkpoint_config(hf_quant_config)
# if hf_quant_config is None, we will try to get config from
# hf_overrides
@@ -376,14 +385,14 @@ def get_quant_config(
if model_config.quantization in ("modelopt", "modelopt_mixed"):
if config.get("producer", {}).get("name") == "modelopt":
return maybe_compose_online_quantization(quant_cls.from_config(config))
return make_checkpoint_config(config)
else:
raise ValueError(
f"Unsupported quantization config"
f" found for {model_config.quantization} in {f}."
)
return maybe_compose_online_quantization(quant_cls.from_config(config))
return make_checkpoint_config(config)
def get_sparse_attention_config(
vllm/models/deepseek_v4/quant_config.py
+5
View File
@@ -110,6 +110,11 @@ class DeepseekV4FP8Config(Fp8Config):
kv_cache_quant_algo=None,
exclude_modules=[],
group_size=16,
quantization_args=(
self.online_quantization_config.args
if self.online_quantization_config is not None
else None
),
)
return self._nvfp4_config
vllm/models/deepseek_v41/quant_config.py
+5
View File
@@ -114,6 +114,11 @@ class DeepseekV4FP8Config(Fp8Config):
kv_cache_quant_algo=None,
exclude_modules=[],
group_size=16,
quantization_args=(
self.online_quantization_config.args
if self.online_quantization_config is not None
else None
),
)
return self._nvfp4_config