diff --git a/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-NVFP4-per-token.yaml b/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-NVFP4-per-token.yaml new file mode 100644 index 0000000000..62587ae10a --- /dev/null +++ b/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-NVFP4-per-token.yaml @@ -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}' diff --git a/tests/evals/gsm8k/configs/models-blackwell-ep.txt b/tests/evals/gsm8k/configs/models-blackwell-ep.txt index c2a1905198..34090d7205 100644 --- a/tests/evals/gsm8k/configs/models-blackwell-ep.txt +++ b/tests/evals/gsm8k/configs/models-blackwell-ep.txt @@ -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 diff --git a/tests/kernels/moe/test_trtllm_nvfp4_moe.py b/tests/kernels/moe/test_trtllm_nvfp4_moe.py index 292e4549bb..4ee0955051 100644 --- a/tests/kernels/moe/test_trtllm_nvfp4_moe.py +++ b/tests/kernels/moe/test_trtllm_nvfp4_moe.py @@ -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, ) diff --git a/tests/model_executor/model_loader/test_weight_cache.py b/tests/model_executor/model_loader/test_weight_cache.py index 731c732625..7960a23791 100644 --- a/tests/model_executor/model_loader/test_weight_cache.py +++ b/tests/model_executor/model_loader/test_weight_cache.py @@ -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).""" diff --git a/tests/quantization/test_blackwell_moe.py b/tests/quantization/test_blackwell_moe.py index da70491bbc..3a4e2b704b 100644 --- a/tests/quantization/test_blackwell_moe.py +++ b/tests/quantization/test_blackwell_moe.py @@ -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", + ], + ) diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 7b1131a00e..33c798cba0 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -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", [ diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index 3f27757b90..c6b9894009 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -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 diff --git a/tests/quantization/test_quantization_config_args.py b/tests/quantization/test_quantization_config_args.py index 1d6ced0d48..7e2e21a3ac 100644 --- a/tests/quantization/test_quantization_config_args.py +++ b/tests/quantization/test_quantization_config_args.py @@ -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(): diff --git a/vllm/config/quantization.py b/vllm/config/quantization.py index d950084726..face2c3a7b 100644 --- a/vllm/config/quantization.py +++ b/vllm/config/quantization.py @@ -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, diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 330e7b6aa1..80aa9122f4 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -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 ( diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index e33226cad3..98af93c30f 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -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( diff --git a/vllm/model_executor/layers/quantization/online/nvfp4.py b/vllm/model_executor/layers/quantization/online/nvfp4.py index efc97705c5..6c780c2c46 100644 --- a/vllm/model_executor/layers/quantization/online/nvfp4.py +++ b/vllm/model_executor/layers/quantization/online/nvfp4.py @@ -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 diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index 24b95ef89a..c5a299eea6 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -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 diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index be28de9b66..5a6ee76428 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -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, diff --git a/vllm/model_executor/model_loader/weight_cache/protocol.py b/vllm/model_executor/model_loader/weight_cache/protocol.py index 72053805cd..535c4de962 100644 --- a/vllm/model_executor/model_loader/weight_cache/protocol.py +++ b/vllm/model_executor/model_loader/weight_cache/protocol.py @@ -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, diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index c6a72c4c33..ac1da8d698 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -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( diff --git a/vllm/models/deepseek_v4/quant_config.py b/vllm/models/deepseek_v4/quant_config.py index be8846fc79..f3aaf72d9b 100644 --- a/vllm/models/deepseek_v4/quant_config.py +++ b/vllm/models/deepseek_v4/quant_config.py @@ -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 diff --git a/vllm/models/deepseek_v41/quant_config.py b/vllm/models/deepseek_v41/quant_config.py index bed60ae758..384e948879 100644 --- a/vllm/models/deepseek_v41/quant_config.py +++ b/vllm/models/deepseek_v41/quant_config.py @@ -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