specdec_bench: forward kv_cache_dtype to the vLLM engine - #2496
yeyu-nvidia wants to merge 1 commit into
Conversation
|
Navigate logical layers of code changes, visualize relationships, and explore their blast radius. No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
Included review availability: Your plan provides up to 12 included reviews per hour; 10 remain after this review. 📝 WalkthroughWalkthroughThe vLLM model wrapper documents ChangesvLLM argument forwarding
Priority: ⬇️ Low Estimated code review effort: 1 (Trivial) | ~2 minutes Change: Bug fix Merge Risk: ⚪ Minimal · up to The wrapper now forwards explicitly configured KV-cache dtype values while preserving default behavior when unset, with no merge-blocking risk identified. 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
@ChenhanYu tagging you as the main author of The part worth a second opinion is the blast radius rather than the diff: an explicit Context for why it matters: without this, |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
|
/claude review Scope note for the review: the diff is one line, so the interesting questions are behavioural rather than syntactic.
|
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
|
||
| # Forwarded from ``--runtime_params`` ``engine_args.<key>``; extend as needed. | ||
| PASSTHROUGH_ENGINE_ARGS = ( | ||
| "kv_cache_dtype", |
There was a problem hiding this comment.
[SUGGESTION] The forwarding itself is correct, but this key differs from its tuple-mates in blast radius and that is invisible here.
Issue: every other entry in PASSTHROUGH_ENGINE_ARGS is a mamba_* knob scoped to the target model's own layers. kv_cache_dtype lands on CacheConfig and therefore applies to every attention layer the engine builds — including the EAGLE3/MTP/draft-model layers. The tuple gives a reader no hint that one of these keys crosses the target/draft boundary.
Why it matters: the failure mode is specific to spec-dec, and it moves exactly the metric this harness measures. If the target carries FP8 KV scales but the draft checkpoint does not, setting kv_cache_dtype: fp8 quantizes the draft's KV cache too, and vLLM falls back to a default scale of 1.0 for layers with no loaded k_scale/v_scale. The draft then drafts from a degraded cache, acceptance length drops, and — as with the bug this PR fixes — nothing in the logs says why. A benchmark whose output is AL should make that condition visible rather than leave it to the operator's memory.
Suggested fix: keep the passthrough as-is (rejecting the combination would be wrong — see the summary comment), and add a one-line note that this key is engine-global, so the next reader does not assume mamba_*-style target-only scoping:
# Forwarded from ``--runtime_params`` ``engine_args.<key>``; extend as needed.
# Note: ``kv_cache_dtype`` is engine-global — it applies to the draft's
# attention layers as well as the target's; vLLM exposes no per-model knob.
PASSTHROUGH_ENGINE_ARGS = (
"kv_cache_dtype",The README caveat you offered in the PR description is worth adding too, but the inline note is the one that reaches whoever extends this tuple next.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2496 +/- ##
=======================================
Coverage 71.17% 71.18%
=======================================
Files 603 603
Lines 66785 66785
=======================================
+ Hits 47537 47542 +5
+ Misses 19248 19243 -5
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Claude review summary Scope applied: per the scope note, I reviewed the one-line diff for behaviour rather than syntax, and answered the three questions directly. Files changed: 1 ( Findings — CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1 The mechanism is correct. 1. Global dtype for target and draft — acceptable; do not reject the combination Rejecting would be the wrong call, for two reasons. First, the wrapper cannot know the draft KV precision without reaching into the draft checkpoint quant config — and for the MTP families where the draft layer ships inside the target (Qwen 3.5 etc., the What is worth surfacing is the narrower case where it actually bites — a draft with no 2. Adding the key does not conflict with a checkpoint that declares
That means the two are not in an override relationship — an explicit 3. EAGLE3/MTP paths reading KV dtype independently — nothing in-repo Grepped One cross-engine note, non-blocking: Risk assessment Low. Opt-in, additive, no default change, and the new value is recorded in 🤖 Generated with Claude Code |
``AsyncEngineArgs`` never received ``kv_cache_dtype``, so
``--runtime_params '{"engine_args": {"kv_cache_dtype": "fp8"}}'`` was accepted
and silently dropped. Benchmarking a checkpoint with a quantized KV cache then
ran BF16 KV against FP8 weights with nothing in the logs to say so, which shifts
acceptance length without any visible error.
The checkpoint declaration and the engine flag are two independent switches and
both have to be on: ``kv_cache_quant_algo: FP8`` makes vLLM attach a KV-cache
quant method and load ``k_scale``/``v_scale``, but those scales are consumed
only when ``CacheConfig.cache_dtype`` is an fp8 variant. ``kv_cache_dtype:
"auto"`` resolves to the model dtype, not to the checkpoint's declared KV algo,
so before this change the scales were loaded and then ignored.
Add the key to PASSTHROUGH_ENGINE_ARGS. It is forwarded only when the user sets
it, so the default is unchanged. Note inline that this key is engine-global --
it lands on CacheConfig and applies to the draft's attention layers as well as
the target's, unlike the target-scoped ``mamba_*`` keys it sits beside.
Signed-off-by: Ye Yu <yeyu@nvidia.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
f628fb8 to
9a0bad0
Compare
|
Thanks — both points taken, and pushed as The Inline SUGGESTION applied more or less verbatim, plus the # Note: ``kv_cache_dtype`` is engine-global -- it lands on ``CacheConfig`` and so
# applies to the draft's attention layers as well as the target's; vLLM exposes
# no per-model knob. Unlike the ``mamba_*`` keys, which are scoped to the target.The unscaled-draft failure case and the reasoning for documenting rather than rejecting it are now in the PR description too. On deriving the default from @ChenhanYu the diff is now 4 lines (1 functional + 3 comment) — re-request when you have a moment. |
This branch has not been deployed
2 participants
What does this PR do?
Type of change: Bug fix
Adds
kv_cache_dtypetoPASSTHROUGH_ENGINE_ARGSin the specdec_bench vLLM model wrapper.AsyncEngineArgsnever receivedkv_cache_dtype, so setting it through--runtime_paramswas accepted and silently dropped. Benchmarking a checkpoint with a quantized KV cache therefore ran BF16 KV against FP8/NVFP4 weights, with nothing in the logs to indicate it. That shifts measured acceptance length with no visible error — the run looks successful and the number is simply wrong.The checkpoint declaration and the engine flag are two independent switches, and both have to be on.
kv_cache_quant_algo: FP8makes vLLM attach a KV-cache quant method to the attention layers and loadk_scale/v_scale; those scales are consumed only whenCacheConfig.cache_dtypeis an fp8 variant.kv_cache_dtype: "auto"resolves to the model dtype — per vLLM's own docstring, "If 'auto', will use model data type" — not to the checkpoint's declared KV algo. So before this change the scales were loaded and then ignored.Found while measuring AL for NVFP4/FP8 quantized Nemotron 3.5 Super with an MTP draft, where the target checkpoint carries an FP8 KV cache.
Usage
python3 run.py \ --model_dir <quantized-checkpoint-with-fp8-kv> \ --engine VLLM \ --runtime_params '{"engine_args": {"kv_cache_dtype": "fp8"}}'Before this change the flag was dropped and the engine recorded
kv_cache_dtype: auto. After it,get_serving_config()records the effective value per run, which makes the setting checkable rather than assertable.Testing
Verified on Nemotron 3.5 Super 120B-A12B with an MTP draft on 4x GB200: the engine-recorded
kv_cache_dtypenow reflects the requested value instead ofauto, and the FP8-KV vs BF16-KV arms produce distinguishable acceptance lengths where previously both arms silently ran the same configuration.No new unit test: the change adds one string to an existing passthrough tuple whose forwarding mechanism is already exercised by the surrounding
mamba_*keys.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
kv_cache_dtypeis engine-global: it lands onCacheConfigand so applies to the draft's attention layers as well as the target's, unlike the target-scopedmamba_*keys it now sits beside. vLLM exposes no per-model knob. This is called out in an inline comment on the tuple, since that is what reaches whoever extends it next.The case where that bites: a target carrying FP8 KV scales plus a draft checkpoint that does not. Setting
kv_cache_dtype: fp8quantizes the draft's KV cache too, and layers with no loadedk_scale/v_scalefall back to a default scale of 1.0 — the draft then drafts from a degraded cache and AL drops silently. Rejecting the combination would be wrong (for in-target MTP drafts there is no separate checkpoint to inspect, and a uniformly-FP8 target plus in-target MTP layer is a perfectly normal config), so this is documented rather than enforced.Cross-engine note, out of scope here:
models/sglang.pyand the TRT-LLM wrapper have no equivalent forwarding, so the same--runtime_paramskey stays silently dropped on those engines.🤖 Generated with Claude Code
Summary by CodeRabbit