[Bugfix][Spec Decode] Honour the draft's attention_backend on Model Runner V2 - #54826
Conversation
…unner V2 The V1 proposer sets the draft's attention backend from the speculative config and deliberately never inherits the target's: # Note (matt): Never inherit the attention backend from base, because there # are many opportunities for incompatibility, so we always independently # autoselect unless explicitly specified in the speculative config. The V2 base speculator returned the target's VllmConfig unchanged from attn_vllm_config, which init_attn_backend then uses to build the draft's attention groups. So an explicit attention_backend in --speculative-config was silently dropped for every speculator family that does not override the property itself. dflash overrides it (for causality) and dspark resolves it in its own loader, but eagle, mtp, gemma4, autoregressive and multi_module_mtp all inherited the target's backend instead. Apply the override in the base property so every family picks it up. dflash already chains through super(), so it keeps its own behaviour on top, and dspark is unaffected because it resolves the backend before this point. Unset stays unchanged: the target's config is returned as before. Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com>
mergify
Bot
added
speculative-decoding
mrv2
…draft is built The previous approach overrode attention_config in the base speculator's attn_vllm_config property. That cannot work: init_attn_backend resolves the backend from each constructed layer via get_attn_backend(), and only uses vllm_config to look the layers up, so a config mutated after construction is ignored. Measured on an MTP draft asking for TRITON_ATTN: the property ran and the draft still resolved FlashInferBackend. Apply the override in load_eagle_model instead, before get_model(), alongside the existing kv_cache_dtype override, so the draft's layers are constructed with the requested backend. Assign it unconditionally, mirroring the V1 proposer: a None backend clears the target's so the draft autoselects independently rather than inheriting, which V1 does deliberately because draft and target attention shapes differ and not every backend serves both. Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
Hi @njhill @TheEpicDolphin, could you take a look? Currently we can't specify the attention backend of the drafter using runner v2 |
TheEpicDolphin
left a comment
TheEpicDolphin
left a comment
There was a problem hiding this comment.
Sorry for the delay. Thanks for the fix!
TheEpicDolphin
added
the
ready
|
/ci run |
|
✅ @stecasta, CI is now available for this PR.
|
|
✅ Triggered Buildkite CI #87129 for commit |
|
@stecasta the |
Assigning unconditionally cleared the target's backend when unset, letting the draft autoselect. Target and draft must share one KV cache layout, and the declared sets do not all intersect: FLEX_ATTENTION and TurboQuant accept only LBNHC, while b12x and SM100 FlashInfer exclude it. On ROCm this left no common layout and failed engine init in test_with_eagle3_spec_decoding. Restrict the override to the explicitly set case, which is the reported bug, and cover the unset case with a test. Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com>
|
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 (2)
Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review. 📝 SummarySummary by CodeRabbit
Walkthrough
ChangesEagle attention backend
Estimated code review effort: 2 (Simple) | ~10 minutes Merge Risk: ⚪ Minimal · up to Draft models now honor an explicitly configured attention backend while retaining existing unset behavior. Coverage confirms the override and configuration preservation, 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 |
|
Thanks @TheEpicDolphin! I had applied the override unconditionally, mirroring the V1 proposer, which clears the target's backend when attention_backend is unset and lets the draft autoselect. On ROCm the target is pinned to FLEX_ATTENTION ({LBNHC}) and the draft resolved to ROCM_ATTN ({LHBNC, LBHNC}), leaving no shared KV cache layout: ValueError: No KV cache layout satisfies every supported set: [['LBNHC'], ['LHBNC', 'LBHNC']] I do not think inheriting is the right long term answer though. The draft should autoselect subject to the target's layout constraint rather than the two being intersected after both are built. I opened #55312 for that, with a note on where the constraint would need to be threaded. Also in flight and related: #54788 (draft moe_backend dropped on V2, same class of bug) do you mind a quick review also there? :) |
|
/ci run |
|
✅ Triggered Buildkite CI #87232 for commit |
|
Ok it's passing now @TheEpicDolphin |
…unner V2 (vllm-project#54826) Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com> Signed-off-by: Jyotirmoy Roy <jyotirmoyroy649@gmail.com>
3 participants
Purpose
--speculative-config '{"attention_backend": ...}'is silently ignored on Model Runner V2.The three draft overrides have deliberately different contracts, per
SpeculativeConfig:moe_backendkv_cache_dtypeattention_backendThe V1 proposer implements this in
_create_draft_vllm_configby assigning the draft's backendunconditionally, so a
Noneerases the target's rather than inheriting it:V2 implements neither half: an explicit backend is dropped, and an unset one inherits. It affects
every speculator family that does not resolve the backend itself, which is
eagle,mtp,gemma4,autoregressive,multi_module_mtpanddflash2. Onlydflash(which overridesattn_vllm_config) anddspark(which resolves it indspark/utils.py) were unaffected.MTP is the common case and published recipes use it, e.g.
--speculative-config '{"method":"mtp","attention_backend":"TRITON_ATTN", ...}'. This becameuser-visible when MRV2 was made the default in #53183.
Fix
Apply the override in
load_eagle_model, beforeget_model(), alongside the existingkv_cache_dtypeoverride.Placement is the crux.
init_attn_backendresolves the backend from each already-constructed layer:and uses
vllm_configonly to look those layers up. So the backend is fixed at construction time andany config mutated afterwards is ignored. The assignment is unconditional, mirroring V1, so the unset
case clears the target's backend instead of inheriting it.
Validation
Unit,
tests/v1/spec_decode/test_draft_attention_backend_override.py, which captures the configload_eagle_modelactually hands toget_model:Both behavioural tests fail when the fix is reverted, so they are not vacuous.
Hardware, Qwen3.6-35B-A3B-NVFP4 with MTP-3 on an RTX PRO 6000 Blackwell Max-Q, V2, default MoE
backend, real routing, one node:
attention_backend: TRITON_ATTNFlashInferBackend(request dropped)attention_backend: TRITON_ATTNTritonAttentionBackendattention_backendunsetNote the throughput drops 9.9% in the middle row, and that is the point: the draft now runs the
backend the user asked for. Previously the request was discarded and the draft silently ran
FlashInfer. A no-op change could not move throughput at all, so this is also the proof the override
takes effect. The third row is the safety check on the unconditional assignment: clearing the backend
when unset does not break the default MTP path.
Why this is not a duplicate
I ran the checks in AGENTS.md. Nothing open touches the V2 draft attention-config path. #53450 pins
backends for components that autoselect, which is a different mechanism and config surface.
Model evaluation
Not applicable. This changes which attention backend the draft uses. Under real rejection sampling
the target verifies every drafted token, so the draft's kernel choice cannot alter output; it can only
move acceptance length. Affected configurations today run a backend the user did not request, so there
is no prior baseline to regress against.
AI assistance
AI assistance was used to author this change. I have reviewed every changed line, ran the tests and
the hardware validation above myself, and can defend the design.