[Quantization] Rework quantization_config to use QuantKey and allow for activation override#41566
Conversation
|
Documentation preview: https://vllm--41566.org.readthedocs.build/en/41566/ |
There was a problem hiding this comment.
Code Review
This pull request refactors the online quantization configuration system, replacing the OnlineQuantScheme enum with a more flexible QuantizationConfigArgs structure. It renames linear_scheme_override and moe_scheme_override to linear and moe, respectively, and introduces QuantSpec to allow independent control over weight and activation quantization. Additionally, the PR migrates activation selection logic (e.g., for MXFP8) from environment variables to the explicit quantization_config and enables user-defined overrides to coexist with checkpoint-based quantization. I have no feedback to provide.
| # Upgrade BF16-act variants to MXFP8-act variants when the user | ||
| # requests MXFP8 activations via quantization_config. | ||
| if activation_qk == kMxfp8Dynamic: | ||
| requested_backend = { | ||
| Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16: ( | ||
| Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8 | ||
| ), | ||
| Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16: ( | ||
| Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8 | ||
| ), | ||
| }.get(requested_backend, requested_backend) | ||
| backend_act = _backend_activation_key(requested_backend) | ||
| if activation_qk is not None and activation_qk != backend_act: | ||
| raise ValueError( | ||
| f"moe_backend={runner_backend!r} runs with activation=" | ||
| f"{backend_act}, but activation={activation_qk} was requested" | ||
| ) |
There was a problem hiding this comment.
can we leave it to is_supported_config to find the matching config? as long as user overriden activation_key is passed, it should find the right backend.
supported, reason = k_cls.is_supported_config(
k_cls, config, kMxfp4Static, activation_key, activation_format
)There was a problem hiding this comment.
Mostly addressed. The runner_backend path now drops the manual upgrade dict and just iterates the candidate list from map_mxfp4_backend, leaving variant selection to is_supported_config. The one wrinkle: _supports_quant_scheme on TrtLlmMxfp4ExpertsMonolithic accepts both (kMxfp4Static, None) and (kMxfp4Static, kMxfp8Dynamic) for the same kernel class, so when both FLASHINFER_TRTLLM_MXFP4_BF16 and FLASHINFER_TRTLLM_MXFP4_MXFP8 are candidates, is_supported_config accepts both and the first listed one wins. To match the user's requested activation deterministically, I added a thin _activation_matches(backend) = _backend_activation_key(backend) == requested_activation_key filter at the oracle level (only when an override is set). Happy to push that filter down into the kernels in a follow-up if you'd prefer.
There was a problem hiding this comment.
accepts both (kMxfp4Static, None) and (kMxfp4Static, kMxfp8Dynamic) for the same kernel class
I ran into a similar situation in #41436 for aiter backend, where w4a4 and w4a16 are both supported. There to resolve it I create a separate shim for w4a4. Either solution works for me, happy to hear your thoughts!
Replaces the OnlineQuantScheme enum and global_scheme/linear_scheme_override/moe_scheme_override fields with a QuantSpec(weight, activation) per layer kind, addressable by name from a hand-curated QUANT_KEY_NAMES table. Routes quantization_config.moe.activation through the MXFP4 MoE oracle so users can opt into MXFP8 activations on already-quantized checkpoints (gpt-oss) without setting env vars. The oracle combines the override with any caller-supplied activation_key from vllm-project#39136 and raises on conflict. map_mxfp4_backend now returns a list of candidate backends per vendor name; both the runner_backend path and the priority list filter by activation_key match. Eval configs and test_blackwell_moe migrate to --quantization-config.moe.activation mxfp8. weight_utils stops rejecting quantization_config alongside a checkpoint quant config. --quantization-config is registered as a CLI arg. Adds parsing UX unit tests and expands docs/features/quantization/online.md. Signed-off-by: mgoin <mgoin64@gmail.com>
bdb8a4c to
5be803c
Compare
is_supported_config alone cannot distinguish backend variants that share an experts class (TRTLLM_BF16 and TRTLLM_MXFP8 both use TrtLlmMxfp4ExpertsMonolithic), so iterating the candidate list with the override propagated through act_key would silently pick the BF16 backend running MXFP8 activations. When no activation is requested and the candidate list contains a BF16 alternative, drop the non-BF16 entries; otherwise keep them (so flashinfer_trtllm_afp8 -> [TRTLLM_MXFP8] still picks MXFP8). When an activation is requested, keep only candidates whose intrinsic activation matches. Signed-off-by: mgoin <mgoin64@gmail.com>
Per Luka: gpt-oss model_kwargs setup belongs in tests/compile/fusions_e2e/models.py keyed on is_blackwell(), not as a special case in conftest. Trim verbose comments and docstrings added in this PR; let docs/features/quantization/online.md be the canonical reference instead of repeating format names inline where they will go stale. Signed-off-by: mgoin <mgoin64@gmail.com>
Drops the duplicated _make_log_backend / _make_log_unsupported / _return_or_raise nested closures from select_mxfp4_moe_backend and select_deepseek_v4_mxfp4_moe_backend. _return_or_raise gains a scope parameter (defaulting to logger.info_once default of "local") so the deepseek path keeps its existing scope="local" logging. _filter_by_activation also moves to module scope and takes requested_activation_key as a parameter instead of capturing it. Signed-off-by: mgoin <mgoin64@gmail.com>
|
This pull request has merge conflicts that must be resolved before it can be |
…ework # Conflicts: # vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
| "mxfp8": kMxfp8Dynamic, | ||
| "mxfp4": kMxfp4Dynamic, | ||
| "int8_per_channel_static": kInt8StaticChannelSym, | ||
| } |
There was a problem hiding this comment.
hi @mgoin - for our usecase we wanted to ingest the nvfp4 ckpt (weights and activation in nvfp4- input_scales in ckpt) as the W4A16_NVFP4 and use Marlin backend (i.e leave the activation in bf16 and ignore input_scale).
Is it possible to add a key for bfloat16/ float16 in this so we can leverage this to set the activation key in fusedMoE ModelOpt class
CT ckpts config have the input_dtype which sets the use_a16 but prevents using pure nvfp4 ckpt as-is w/o changing the config manually.
We took a parallel path in our open draft #42428 and introduced the override_activation_dtype flag which gets passed to ModelOptNVFP4Config from BaseConfig and eventually to the [ModelOptNvFp4FusedMoE class] for changing activation_key here
Appreciate your insights if we should have the override flag and your QuantSpec surface co-exist or just add the bfloat16/ float16 keys here.
Should we call the activation_override as modelopt_override_activation_dtype instead of override_activation_dtype?
There was a problem hiding this comment.
If you want to use the marlin backend with an existing NVFP4 W4A4 model, you can just use --linear-backend marlin or --moe-backend marlin. Is that sufficient?
We can consider allowing overriding the checkpoint activation type with the path added in this PR too
There was a problem hiding this comment.
thanks Michael, we may want to have a finer control than the --moe-backed marlin in my understanding.
When you say added in this PR- does it refer to the current PR you have, or introduce the new override_activation_dtype standalone flag from cli?
…ework # Conflicts: # vllm/model_executor/layers/fused_moe/oracle/mxfp4.py # vllm/model_executor/layers/quantization/online/base.py
…or activation override (vllm-project#41566)
…or activation override (vllm-project#41566)
…or activation override (vllm-project#41566)
…or activation override (vllm-project#41566)
…or activation override (vllm-project#41566) Signed-off-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
…or activation override (vllm-project#41566)
Purpose
Replaces the
OnlineQuantSchemeenum and global_scheme/linear_scheme_override/moe_scheme_override fields with aQuantSpec(weight, activation)per layer kind, addressable by name from aQUANT_KEY_NAMEStable. The CLI shorthands (fp8_per_tensor,fp8_per_block,mxfp8,int8_per_channel_weight_only) keep working and now desugar into the new structure leveraging QuantKey pairs directly.Routes
--quantization_config.moe.activationthrough the MXFP4 MoE oracle so users can opt into MXFP8 activations on already-quantized checkpoints (gpt-oss) without setting env vars. When combined with--moe-backend flashinfer_cutlassorflashinfer_trtllm, the oracle upgrades the BF16 backend variant to its MXFP8 counterpart.Also migrates the gpt-oss eval configs and test_blackwell_moe MXFP8 tests off VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8* env vars and onto the new flag to test this
Adds tests/quantization/test_quantization_config_args.py for the parsing UX
Test Plan
Test Result
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.