[SM70] FP8-resident MTP experts for AWQ and NVFP4 targets - #595
Conversation
Quantize unquantized draft expert shards with per-row scales and reuse the existing SM70 FP8 weight-only MoE path. Preserve the target, shared heads, and standard rejection sampler. GPU and sampling acceptance remain pending. Assisted-by: Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com>
V100 probes show unpadded W2 K=160 is numerically invalid. Pad gate/up halves and down input consistently to 256; preserve zero contributions and reuse supported native dimensions. Assisted-by: Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com>
Four GPU cases passed against explicit routed-expert reconstruction at M=1,2,8,64. Assisted-by: Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com>
Reject unsupported draft architectures and synthetic acceptance for the opt-in FP8 expert feature. Validate hidden-dimension alignment before weight allocation. Assisted-by: Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com>
Document supported configuration, padding, full-model memory and output checks, and the limits of the quality evidence. Assisted-by: Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com>
Backport ModelOpt block-FP8 dispatch and MTP metadata remapping from vllm-project/vllm#55513 (60ad959b6f1a5c8f602edbd608c8decbc0788c50). Add scale-preserving TP padding and native FP8 allocation for independent AWQ/NVFP4 targets. Keep full-model expanded validation explicitly in progress. Assisted-by: Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com>
ModelOpt mixed checkpoints leave missing quantized_layers entries unquantized even without an explicit exclusion. Allow the online MTP-only FP8 option for that case while rejecting unsupported quantized formats. Assisted-by: Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com>
Record checkpoint-native paired memory and complete-response checks for both target formats, plus NVFP4 online conversion. Distinguish same-source controls from different-draft integration checks. Assisted-by: Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f436265275
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| warn users when they mistakenly provide the wrong argument.""" | ||
|
|
||
| # Draft model configuration | ||
| mtp_expert_quantization: Literal["fp8"] | None = None |
There was a problem hiding this comment.
Add mtp_expert_quantization to SpeculativeConfig.compute_hash(). For the same Qwen4Exp model, enabling this field replaces the ordinary expert implementation with padded SM70 FP8 parameters and a different kernel, yet the configuration currently hashes identically to the FP16 path. When compilation caches or AOT artifacts are shared between those configurations, vLLM can therefore select an artifact generated for the wrong expert graph or layout.
Useful? React with 👍 / 👎.
| hf_config = getattr(self.draft_model_config, "hf_config", None) | ||
| if self.method != "mtp" or getattr(hf_config, "architectures", []) != [ | ||
| "Qwen4ExpMTP" | ||
| ]: |
There was a problem hiding this comment.
Gate this option on the supported platform as well as the architecture. On ROCm, vllm.models.qwen4_exp.__getattr__ selects amd.mtp, whose _make_draft_vllm_config never reads mtp_expert_quantization; this validator nevertheless accepts the configuration, so an explicit "fp8" request silently leaves the MTP experts on their fallback quantization instead of providing the advertised FP8 residency. The NVIDIA implementation rejects non-SM70 use, so ROCm should likewise fail closed rather than ignore the option.
Useful? React with 👍 / 👎.
Include the online expert quantization option in the compilation hash and validate CUDA SM70 before dispatch can select an implementation that ignores the option. Validation: 52 CPU tests passed; pre-commit passed for both changed files. Assisted-by: Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com>
Hash the unset MTP expert option as well, so neither FP16 nor FP8 can reuse an old artifact whose key omitted the expert implementation. Unrelated speculative methods keep their existing hash. Validation: 52 CPU tests passed; pre-commit passed. Assisted-by: Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com>
|
@codex Review |
|
Codex Review: Didn't find any major issues. Keep it up! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
2 participants
Problem and implementation
Qwen4Exp MTP experts can consume substantial memory even when the target is AWQ or NVFP4. Keep the draft experts resident in FP8 while preserving the target model and standard rejection verification. Reduced proposal acceptance is allowed; the unchanged target distribution must be preserved.
Support both input routes:
"mtp_expert_quantization": "fp8".modelfield for a separate draft checkpoint; ModelOpt mixed checkpoints and regular block-FP8 metadata are recognized.This backports vllm-project/vllm#55513, merged as
60ad959b6f1a5c8f602edbd608c8decbc0788c50: ModelOptFP8_PB_WO/FP8_BLOCK_SCALESdispatch and AMD/NVIDIA MTP layer-metadata remapping. The new work is SM70 resident execution, TP padding that preserves original FP8 bytes and block-scale offsets, and the independent online expert option. It does not depend on #553, whose SM70 fallback retains FP16 expert weights. vllm-project/vllm#55498 overlaps the already-merged generic loading fix.On the tested TP4 shape, logical width 160 needs physical width 256. Native checkpoint loading preserves offsets 0/32/64/96 inside original scale blocks, without dequantizing and requantizing. The SM70 kernel uses FP16 scales; overflowing/underflowing scales are rejected. Online conversion is restricted to unquantized draft experts. Explicit opt-in is rejected on non-SM70 platforms, including ROCm. FP16 and online FP8 MTP use distinct compilation keys, and neither reuses the legacy key that omitted this option. SM70/FP16, Qwen4ExpMTP, ordinary TP without PP/EP, and standard rejection sampling are required.
Validation
Both loading routes completed full-model tests with AWQ and NVFP4 targets.
.venv/bin/python -m pytest tests/models/qwen4_exp/test_mtp_fp8_experts.py tests/models/qwen4_exp/test_mtp_fp8_checkpoint.py -q: 52 passed. Includes compilation-key separation (including rejection of legacy ambiguous keys), unsupported-platform rejection, both ModelOpt aliases, exclusions, metadata remapping, TP1/2/4/8 byte/scale preservation, incomplete checkpoints, the real allocator/TP loader with AWQ/NVFP4/mixed configuration names, and online conversion of unquantized experts omitted from mixed metadata..venv/bin/python -m pytest tests/models/qwen4_exp/test_mtp_fp8_experts_gpu.py -q: 20 passed, covering online and checkpoint-native methods, four TP4 offsets, M=1/2/8/64, reference reconstruction and exact CUDA Graph replay.All paired full-model arms use TP4/MTP3/C2, FP16 execution/KV, CUDA Graphs and fixed 5-GiB KV/rank. Memory values are idle snapshots after both greedy and stochastic campaigns, identical across all four ranks. Whole-GPU reductions include runtime overhead and must not be attributed entirely to the 225-MiB expert-payload reduction. Actual test prompts reach approximately 7700 tokens; configured 262144 context is not a full-context validation.
The native runtime uses extensions built from
752f86495f, with changed Python modules overlaid from this branch based onfe67339ddf; it is not a clean native rebuild. The first checkpoint-native AWQ run preceded two additional input guards; the 20 GPU tests and NVFP4 runs include them. Main target weights, shared-head handling and rejection sampling are not modified. Tests are bounded correctness/integration evidence, not a universal bitwise or benchmark-quality guarantee. Standard rejection sampling preserves the unchanged target distribution using actual proposal probabilities. No throughput or loading-peak improvement is claimed.See design and validation details. This submits the implementation validated in Leonccaa#13 for upstream review. No production deployment is included.
AI assistance: Codex assisted with implementation, tests and documentation under the submitter's direction.