[Quantization][ModelOpt] W4A16 NVFP4 fused MoE + --override-activation-dtype flag by juhi10071998 · Pull Request #42428 · vllm-project/vllm

juhi10071998 · 2026-05-12T14:24:55Z

Summary

Adds ModelOpt NVFP4 W4A16 support for the fused-MoE path and an engine-level --override-activation-dtype flag that lets a user load a W4A4 NVFP4 checkpoint via the W4A16 path without editing the checkpoint's JSON.

Follow-up to #41769 (phase-1, merged), which added ModelOptNvFp4W4A16LinearMethod for dense Linear. This PR completes W4A16 NVFP4 support across the model: MoE + a user-facing override.

Tested: 7 unit tests + end-to-end smoke on nvidia/Qwen3.6-35B-A3B-NVFP4 (Qwen3.5-MoE-VL hybrid), with the flag (Marlin path) and without (FlashInferCutlass W4A4 path, regression check). See Test plan.

What's in this PR

Three commits, each scoped to one logical change:

1. `[Quantization][ModelOpt] Add W4A16 NVFP4 support to fused MoE`

Extends ModelOptNvFp4FusedMoE.__init__ to honor W4A16_NVFP4 on-disk checkpoints. When quant_config.quant_method == "W4A16_NVFP4", the MoE class passes activation_key=None to select_nvfp4_moe_backend.
Every W4A4 backend's _supports_quant_scheme requires (kNvfp4Static, kNvfp4Dynamic) exactly, so they all reject themselves; Marlin's _supports_quant_scheme (fused_marlin_moe.py:607) only checks weight_key, so it accepts. Marlin's MoE prep already nulls activation scales in convert_to_nvfp4_moe_kernel_format (oracle/nvfp4.py:362-363) and routes through nvfp4_w4a16_moe_quant_config (oracle/nvfp4.py:433) — no other change needed.

2. `[Quantization][ModelOpt] Add --override-activation-dtype CLI flag` (4 files, +80/−8)

Adds an engine-level dtype override that mirrors the existing --override-attention-dtype precedent (vllm/config/model.py:315). When set to bfloat16 or float16, a W4A4 NVFP4 checkpoint is routed through the W4A16 LinearMethod and Marlin MoE backend regardless of the on-disk quant_algo.

Pipeline:

ModelConfig.override_activation_dtype: Literal["auto", "bfloat16", "float16"] (inline at field site so get_kwargs(ModelConfig) auto-derives argparse choices=).
weight_utils._maybe_inject_engine_overrides() injects a __override_activation_dtype__ private key into hf_quant_config at the two ModelOpt code paths. Inject-and-leave, matching the total_num_heads precedent at compressed_tensors.py:250-251.
ModelOptQuantConfigBase.from_config reads the key once, propagates as a kwarg to every subclass's _from_config. ModelOptNvFp4Config pulls it via kwargs.get(...); the other three ModelOpt Configs absorb-and-discard via their existing **kwargs: Any.
ModelOptNvFp4FusedMoE's use_a16 extends to OR on quant_config.override_activation_dtype in ("bfloat16", "float16").

3. `[Quantization][ModelOpt] Add unit tests` (+99 LOC in `tests/quantization/test_modelopt.py`)

Two parametrized tests, 7 effective cases:

Why a new flag rather than `--hf-overrides`?

--hf-overrides is the natural alternative for "patch the on-disk quant config." It has three paths to the quant parser and none are usable here:

Path A (--hf-overrides '{"quantization_config": {...}}') — ModelConfig._apply_dict_overrides at vllm/config/model.py:454 does wholesale replace on plain dict attributes. To change one field, the user has to type the entire quantization_config block verbatim.
Path B (hf_overrides.quantization_config_file) and Path C (hf_overrides.quantization_config_dict_json) at weight_utils.py:338,354 are both gated by if hf_quant_config is None — unreachable when the ckpt already carries an on-disk quant config.

--override-activation-dtype is one token, semantically specific, and mirrors the existing single-subsystem override pattern. Full design + alternatives discussion in commit-2's added comments + the PR diff.

Test plan

Done

Unit tests — 7/7 pass on CPU, no GPU / no ckpt needed

pytest tests/quantization/test_modelopt.py::test_modelopt_nvfp4_moe_dispatches_to_marlin_when_w4a16 \
       tests/quantization/test_modelopt.py::test_modelopt_nvfp4_override_activation_dtype_pipeline -v
# 7 passed in 2.07s

End-to-end smoke — `nvidia/Qwen3.6-35B-A3B-NVFP4`

Qwen3.5-MoE-VL hybrid (gated-delta linear attention + conv1d + experts, ~35B). One GPU, eager mode.

With override (LLM(override_activation_dtype="bfloat16", ...)) — exercises this PR's new path:

WARNING modelopt.py:1036  Detected ModelOpt NVFP4 checkpoint (quant_algo=NVFP4)
INFO    modelopt.py:1060  ModelOpt NVFP4 W4A16 override active (override_activation_dtype=bfloat16):
                          loading a NVFP4 checkpoint via the W4A16 LinearMethod and Marlin MoE backend.
INFO    nvfp4.py:282      Using 'MARLIN' NvFp4 MoE backend out of potential backends:
                          ['FLASHINFER_TRTLLM', 'FLASHINFER_CUTEDSL', 'FLASHINFER_CUTEDSL_BATCHED',
                           'FLASHINFER_CUTLASS', 'VLLM_CUTLASS', 'MARLIN', 'EMULATION'].
Model loading took 20.54 GiB memory and 34.06 seconds
...
=== prompt:     'The capital of France is'
=== completion: ' Paris.\n\n<think>\n\n</think>\n\nThat is correct. Paris is the capital and'
=== prompt:     'Q: What is 2 + 2? A:'
=== completion: '\n\n<think>\n\n</think>\n\n4'

The three log signals together prove the full pipeline wired up: on-disk quant_algo=NVFP4 (no JSON edit), override INFO line fires, MoE oracle picks Marlin (only backend that accepts activation_key=None).

Without the flag — regression check (same ckpt, plain LLM(...)):

WARNING modelopt.py:1023  Detected ModelOpt NVFP4 checkpoint (quant_algo=NVFP4)
INFO    nvfp4.py:282      Using 'FLASHINFER_CUTLASS' NvFp4 MoE backend ...
=== prompt:     'The capital of France is'
=== completion: ' Paris.\n\n<think>\n\n</think>\n\nThat is correct. Paris is the capital and'

W4A4 path stays intact; the override INFO log does not fire; FlashInferCutlass NVFP4 MoE backend selected; coherent generation. Confirms zero behavioral drift when the flag is unset.

Duplicate-work check

Builds on [Quantization] Add ModelOpt NVFP4 W4A16 (4-bit weights, fp16/bf16 activations) support #41769 (merged, phase-1 dense Linear W4A16).
[Quantization] Rework quantization_config to use QuantKey and allow for activation override #41566 (open) is not a duplicate: it targets MXFP4/MXFP8 on gpt-oss via a QuantSpec(weight, activation) rework of OnlineQuantScheme; this PR targets NVFP4 on ModelOpt via a new ModelConfig.override_activation_dtype field. Different quant method, different mechanism, no file overlap.
No other open PRs match modelopt nvfp4 moe w4a16 or override-activation-dtype.

Extends ModelOptNvFp4FusedMoE to honor W4A16_NVFP4 checkpoints. When the on-disk quant_algo is W4A16_NVFP4, the MoE class passes activation_key=None to select_nvfp4_moe_backend. W4A4 backends reject the scheme (their _supports_quant_scheme requires (kNvfp4Static, kNvfp4Dynamic) exactly); Marlin survives (it only checks weight_key). Marlin's MoE prep already nulls activation scales in convert_to_nvfp4_moe_kernel_format and routes through nvfp4_w4a16_moe_quant_config — no other change needed. Follow-up to PR vllm-project#41769 (dense Linear W4A16). Signed-off-by: Juhi Mittal <juhim@nvidia.com> Co-authored-by: Claude

Engine-level override that loads a W4A4 NVFP4 ckpt via the W4A16 path without editing the ckpt's JSON. Mirrors --override-attention-dtype in shape (single-subsystem dtype override on ModelConfig). - ModelConfig.override_activation_dtype: Literal["auto","bfloat16","float16"] - weight_utils._maybe_inject_engine_overrides() injects a private dict key into hf_quant_config (inject-and-leave, matches total_num_heads precedent). - ModelOptQuantConfigBase.from_config reads it once, propagates as a kwarg to every subclass's _from_config. - ModelOptNvFp4Config honors it (forces W4A16 LinearMethod); the other three ModelOpt Configs absorb-and-discard via **kwargs. - ModelOptNvFp4FusedMoE extends use_a16 to OR on the override. Signed-off-by: Juhi Mittal <juhim@nvidia.com> Co-authored-by: Claude

…e flag Two new parametrized tests appended to tests/quantization/test_modelopt.py: - test_modelopt_nvfp4_moe_dispatches_to_marlin_when_w4a16 (3 cases) — ModelOptNvFp4FusedMoE must pass activation_key=None to select_nvfp4_moe_backend for both routes that should trigger W4A16 (on-disk quant_algo W4A16_NVFP4, or --override-activation-dtype set). Mocks the oracle to capture call args; verifies use_a16 and activation_key without needing a GPU or real ckpt. - test_modelopt_nvfp4_override_activation_dtype_pipeline (4 cases) — End-to-end dict -> ModelOptNvFp4Config.from_config -> correct LinearMethodCls. Exercises base read of __override_activation_dtype__, kwarg propagation, and the priority rule (override > on-disk algo). Both run in ~2s on CPU; all 7 cases pass locally. Signed-off-by: Juhi Mittal <juhim@nvidia.com> Co-authored-by: Claude

gemini-code-assist

Code Review

This pull request introduces the --override-activation-dtype engine-level flag, allowing users to force W4A16 execution (via Marlin) for ModelOpt NVFP4 checkpoints even if they were originally serialized as W4A4. The changes include updates to the engine arguments, model configuration, and the ModelOpt quantization logic to propagate and honor this override. However, a critical issue was identified where this override is not correctly propagated in mixed-precision configurations (ModelOptMixedPrecisionConfig), which would cause the flag to be silently ignored for those models.

gemini-code-assist · 2026-05-12T14:28:13Z

+        # Engine-level overrides flow in via private dict keys injected
+        # by weight_utils._maybe_inject_engine_overrides. Read them here
+        # once and propagate to every subclass's _from_config via kwargs.
+        # Subclasses that honor a knob name the kwarg in their signature
+        # (e.g. ModelOptNvFp4Config); others absorb-and-discard via
+        # **kwargs. The keys stay in original_config (no pop) — matches
+        # the total_num_heads precedent in compressed_tensors.
+        override_activation_dtype = config.get("__override_activation_dtype__", "auto")
+
        return cls._from_config(
            quant_method=quant_method,
            kv_cache_quant_method=kv_cache_quant_method,
            exclude_modules=exclude_modules,
            group_size=group_size,
            original_config=config,
+            override_activation_dtype=override_activation_dtype,
        )


The logic for propagating the override_activation_dtype is incomplete for mixed-precision models. The comment on lines 362-364 mentions that some subclasses will "absorb-and-discard" the kwarg, but ModelOptMixedPrecisionConfig should not do this. It contains an nvfp4_config and needs to propagate the override to it.

Currently, ModelOptMixedPrecisionConfig._from_config does not pass the override_activation_dtype to the ModelOptNvFp4Config it creates. This will cause the --override-activation-dtype flag to be silently ignored for mixed-precision models that include NVFP4 layers, which is a critical bug.

To fix this, ModelOptMixedPrecisionConfig._from_config should be updated to pass override_activation_dtype to its ModelOptNvFp4Config instance.

There appears to be a related pre-existing issue where ModelOptMixedPrecisionConfig.get_quant_method hardcodes ModelOptNvFp4LinearMethod instead of using the dynamically selected self.nvfp4_config.LinearMethodCls. This would also prevent the override from working correctly for linear layers in mixed-precision models.

claude

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

mgoin · 2026-05-12T14:39:52Z

Actually I would say #41566 is a duplicate since it is my general approach on activation overrides. It is just starting with MXFP4/MXFP8 since we have kernels for gpt-oss

juhi10071998 · 2026-05-12T17:21:15Z

Actually I would say #41566 is a duplicate since it is my general approach on activation overrides. It is just starting with MXFP4/MXFP8 since we have kernels for gpt-oss

Thanks @mgoin , I will review and evaluate if we can use your approach for our ModelOpt path to ingest nvfp4 ckpts (and route to w4a16).
For now I have marked this PR as draft and created a new PR here #42440 to route to Marlin FusedMoE kernel based on the on-disk quant_algo field. (orthogonal to the override). Could you review it when you have a chance, it is a small PR.

cc- @pavanimajety

mergify · 2026-05-15T23:53:38Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @juhi10071998.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

juhi10071998 added 3 commits May 12, 2026 00:34

gemini-code-assist Bot reviewed May 12, 2026

View reviewed changes

juhi10071998 marked this pull request as ready for review May 12, 2026 14:31

juhi10071998 requested review from 22quinn, ProExpertProg, WoosukKwon, hmellor, houseroad, mgoin, pavanimajety, robertgshaw2-redhat, tlrmchlsmth, yewentao256, youkaichao and zyongye as code owners May 12, 2026 14:31

claude Bot reviewed May 12, 2026

View reviewed changes

juhi10071998 mentioned this pull request May 12, 2026

[Quantization][ModelOpt] Add W4A16 NVFP4 support to fused MoE #42440

Closed

juhi10071998 mentioned this pull request May 12, 2026

[Quantization] Rework quantization_config to use QuantKey and allow for activation override #41566

Merged

4 tasks

juhi10071998 marked this pull request as draft May 15, 2026 23:50

mergify Bot added the needs-rebase label May 15, 2026

MerkyorLynn mentioned this pull request Jun 5, 2026

Document ModelOpt W4A16 NVFP4 Marlin path #44672

Open

mergify Bot added the quantization label Jul 23, 2026

This was referenced Jul 24, 2026

[Core][Kernel] Support dynamic activation scaling for ModelOpt NVFP4 #49715

Open

[Bugfix] Detect ModelOpt W4A16 NVFP4 configs vllm-project/vllm-omni#5378

Closed

stecasta mentioned this pull request Aug 31, 2026

[Quant][NVFP4][MoE] Declare expert weight layout so W4A16/W4A4 can be selected dynamically #54617

Draft

Navigation Menu

Uh oh!

[Quantization][ModelOpt] W4A16 NVFP4 fused MoE + --override-activation-dtype flag - #42428

[Quantization][ModelOpt] W4A16 NVFP4 fused MoE + --override-activation-dtype flag#42428
juhi10071998 wants to merge 3 commits into
vllm-project:mainfrom
juhi10071998:w4a16nvfp4_fused_moe

juhi10071998 commented May 12, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot May 12, 2026

Uh oh!

claude Bot left a comment

Uh oh!

mgoin commented May 12, 2026

Uh oh!

juhi10071998 commented May 12, 2026 •

edited

Loading

Uh oh!

mergify Bot commented May 15, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

juhi10071998 commented May 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What's in this PR

1. [Quantization][ModelOpt] Add W4A16 NVFP4 support to fused MoE

2. [Quantization][ModelOpt] Add --override-activation-dtype CLI flag (4 files, +80/−8)

3. [Quantization][ModelOpt] Add unit tests (+99 LOC in tests/quantization/test_modelopt.py)

Why a new flag rather than --hf-overrides?

Test plan

Done

Unit tests — 7/7 pass on CPU, no GPU / no ckpt needed

End-to-end smoke — nvidia/Qwen3.6-35B-A3B-NVFP4

Duplicate-work check

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot May 12, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

mgoin commented May 12, 2026

Uh oh!

juhi10071998 commented May 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mergify Bot commented May 15, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

juhi10071998 commented May 12, 2026 •

edited

Loading

1. `[Quantization][ModelOpt] Add W4A16 NVFP4 support to fused MoE`

2. `[Quantization][ModelOpt] Add --override-activation-dtype CLI flag` (4 files, +80/−8)

3. `[Quantization][ModelOpt] Add unit tests` (+99 LOC in `tests/quantization/test_modelopt.py`)

Why a new flag rather than `--hf-overrides`?

End-to-end smoke — `nvidia/Qwen3.6-35B-A3B-NVFP4`

juhi10071998 commented May 12, 2026 •

edited

Loading