Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
Manny7717
left a comment
Manny7717
left a comment
There was a problem hiding this comment.
Verified locally on head 779e4183 (worktree + standalone probes; no DFlash hardware here, so mechanism-level verification).
Root cause confirmed by execution: kv_cache_layout is field(default=None, init=False) in CacheConfig (vllm/config/cache.py:89) and load_dflash_model builds the draft config via dataclasses.replace(vllm_config, ...) — which drops init=False fields. Executed probe: replace(cc, cache_dtype='auto') yields a draft whose kv_cache_layout is None even after the target resolved LHBNC. The engine's later set_kv_cache_layout collective RPC (core.py:294 → executor collective_rpc → worker_base.py:112) only ever touches the target config, so the draft's first forward dies in get_resolved_kv_cache_layout with exactly the reported "KV cache layout has not been resolved yet" (reproduced the raise).
Fix mechanics verified by probes:
- Mirror step:
draft_cache.kv_cache_layout = vllm_config.cache_config.kv_cache_layoutcarries an already-resolved value; guarded bydraft_cache is not vllm_config.cache_configso the passthrough (no kv_cache_dtype override) case stays untouched. - Registration:
_draft_cache_configssibling list on the target config is consumed by the worker'sset_kv_cache_layout, which now propagates to every registered draft copy — verified the full chain (mirror → register → record → both configs end at LHBNC). - Fail-closed preserved:
record_kv_cache_layoutstill raises on a conflicting layout (already resolved to LHBNC; cannot change it to LBHNC) — the propagation can't silently flip an inconsistent draft. get_resolved_kv_cache_layoutraise on a stale (None-layout) draft still reproduces the pre-fix failure mode, confirming the propagation is what fixes it.
No regressions: dflash unit suites (test_dflash_causality, test_dflash_prepare_inputs, test_dflash_lookahead): head 2 failed / 14 passed / 16 errors vs base 73723b7 identical 2f/14p/16e — the failures/errors are CUDA-less torch.accelerator fixture noise, byte-identical on both. ruff (pinned 0.14.0, CI version) clean on both changed files. Author ran the real hardware path (2x DGX Spark, TP=2, Qwen3.8-27B + DFlash2 drafter) before/after.
Non-blocking nit: no automated test added — a config-level unit test pinning that load_dflash_model's replace path carries kv_cache_layout (or that set_kv_cache_layout reaches registered drafts) would guard this against future dataclass reordering; the author's e2e evidence covers the live path today.
|
Added a test, thanks for the nit. |
4d0fbbd to
f8dc114
Compare
AndreasKaratzas
left a comment
AndreasKaratzas
left a comment
There was a problem hiding this comment.
LGTM -- probably a second pair of eyes would be good for the worker base change
| @@ -0,0 +1,58 @@ | |||
| # SPDX-License-Identifier: Apache-2.0 | |||
There was a problem hiding this comment.
Not sure if this test file is needed though.
AndreasKaratzas
added
ready
|
✅ @ptorsten, CI is now available for this PR.
|
|
Hi @ptorsten, the pre-commit checks have failed. Please run: uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-filesThen, commit the changes and push to your branch. For future commits, |
FlashInferImpl captured cache_config from the current vllm config at construction and read kv_cache_layout from it lazily, since the engine core resolves the layout after model load (vllm-project#51718) and records it on the worker's CacheConfig. A draft model with a kv_cache_dtype override is built under a replace()d config whose CacheConfig is a copy, so its impls never see the resolution and the first draft forward, during memory profiling, raises "KV cache layout has not been resolved yet". Put the resolved layout on FlashInferMetadata. The builder reads it from the worker's config at build time and every impl-side read happens in forward() with the metadata in hand, so the impl needs no config of its own. This covers the DFlash, EAGLE and DSpark drafters alike, which all derive their config the same way. Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Patrik Torstensson <patrik.torstensson@gmail.com>
b0a9782 to
44aac86
Compare
ptorsten
changed the title
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (1)
Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review. 📝 SummarySummary by CodeRabbit
WalkthroughFlashInfer now resolves the KV cache layout in ChangesFlashInfer KV cache layout
Estimated code review effort: 2 (Simple) | ~10 minutes Merge Risk: ⚪ Minimal · up to This change carries the resolved KV-cache layout through FlashInfer metadata so draft-model attention uses the correct layout during execution. The updated paths consistently use the propagated value, with no current merge-blocking risk identified. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
I opened #55384 as a draft alternative for design comparison, not as a second fix intended to merge alongside this PR. Instead of carrying the resolved layout through FlashInfer metadata, it makes CacheConfig copies share an explicit ResolvedKVCacheLayout state object. That addresses the same late-resolution issue generically across DFlash, EAGLE, DSpark, and other draft configs that override cache dtype, at the cost of a broader CacheConfig API change. Patrik is credited as a co-author. |
This branch has not been deployed
4 participants
Purpose
FlashInferImplreadskv_cache_layoutfrom acache_configcaptured at construction. #51718 resolves the layout after model load and records it on the worker'sCacheConfig. A drafter with akv_cache_dtypeoverride (DFlash, EAGLE, DSpark alike) is built under areplace()d config, so its impls hold a copy that never sees the resolution, and the first draft forward during profiling raises "KV cache layout has not been resolved yet".Carry the layout on
FlashInferMetadatainstead. The builder already reads it from the worker's config to plan the wrappers, and every impl-side read is inforward(), so the impl needs no config of its own. FlashInfer is the only backend whose impl read the layout this way.Replaces the first revision, which propagated the layout into draft config copies from
worker_base.py. No open PR touches drafter layout resolution.Test
pytest tests/v1/attention/test_attention_backends.py -k "flashinfer or causal_backend_correctness"- 110 passed on GB10 (sm_121).2x DGX Spark (GB10), TP=2, FLASHINFER, target Qwen3.8-27B-NVFP4 with KV
auto, DFlash2 drafter z-lab/Qwen3.8-27B-DFlash2 withkv_cache_dtype: fp8,--linear-backend flashinfer_cutlass, 8 speculative tokens.mainat 8f816a3 fails in memory profiling with the error above fromFlashInferImpl.kv_cache_layout. With this commit it reaches READY (34 piecewise + 2 full target graphs, 12 DFlash2 full graphs) and serves the 12-scenario tool-calling suite 12/12 at 154 tok/s mean decode, 72% draft acceptance. A W4A16 drafter does not load onmainfor an unrelated reason (_build_context_kv_buffersreadsqkv_proj.weightafter Marlin repacks it), hence the bf16 drafter.AI assistance was used; every line reviewed and the tests above run by me.