spec-decode: honor the speculative-config attention backend for the DFlash draft

[voipmonitor]

Summary

The V2 DFlash draft loader (load_dflash_model) rewrites the draft's attention
config to set its causal/non-causal mode, but drops
speculative_config.attention_backend. The draft then auto-selects an
attention backend, which on Blackwell prefers FlashInfer regardless of what
the speculative config requests (e.g. TRITON_ATTN).

Symptoms (DFlash draft beside an MLA target, e.g. Kimi-K2.7-Code)

• Under DCP the engine fails to start:
  FlashInfer non-causal prefill is not supported with DCP yet.
• Without DCP the spec-decode cudagraph is downgraded FULL → PIECEWISE:
  CUDAGraphMode.FULL_AND_PIECEWISE is not supported with spec-decode for attention backend FlashInferBackend.
• The user's speculative_config.attention_backend is silently ignored.

Root cause / fix

load_dflash_model does replace(attention_config, use_non_causal=...) but
never sets backend=. The V1 llm_base_proposer already threads the
spec-config backend through (with the comment "never inherit the attention
backend from base ... unless explicitly specified in the speculative config").
This restores that one line for the V2 DFlash path. When the spec config leaves
attention_backend unset, behaviour is unchanged (auto-select).

Testing

Kimi-K2.7-Code TP8 + Kimi-K2.6 DFlash draft, V2 model runner, kv-fp8,
--speculative-config '{"method":"dflash", ..., "attention_backend":"TRITON_ATTN"}'.
Smoke + throughput at 0k context, concurrency 1:

Config	Before (auto/FlashInfer)	After (this PR, TRITON_ATTN)
DCP8 + DFlash	fails to start	boots, FULL draft graph, 113.6 tok/s (vs 81.5 target-only, +39%), coherent
DCP1 + DFlash	125 tok/s, PIECEWISE only	140 tok/s, coherent

Both configurations confirmed running on the V2 model runner with coherent
output.

Why this is not duplicating an existing PR

gh pr list --repo local-inference-lab/vllm --state open — the only open PR (#8,
codex/ds4-step35-mtp-local-argmax-spec-step) targets MTP local-argmax, an
unrelated area. No open PR touches the DFlash draft attention backend.

AI assistance

This change was drafted with AI assistance (Claude, Anthropic) and reviewed,
built, and validated end-to-end by the submitter on a Blackwell rig.

Summary by CodeRabbit

• Bug Fixes
  • Fixed speculative decoding with DFlash backend to correctly configure attention backend and causal mode settings for draft models.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 47c8285f-dd09-4244-907f-b0506d48f68e

📥 Commits

Reviewing files that changed from the base of the PR and between ad689cc and ba3c751.

📒 Files selected for processing (1)

• vllm/v1/worker/gpu/spec_decode/dflash/utils.py

---

📝 Walkthrough

Walkthrough

In load_dflash_model, the draft model's attention_config override is expanded to explicitly set backend to speculative_config.attention_backend in addition to the existing use_non_causal=not causal field. Comments are added explaining the backend selection behavior.

Changes

DFlash draft model attention config override

Layer / File(s)	Summary
Draft model attention_config backend field vllm/v1/worker/gpu/spec_decode/dflash/utils.py	The replace(...) call for the draft model's attention_config now includes backend=speculative_config.attention_backend alongside the existing use_non_causal=not causal, with comments describing backend selection.

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~3 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and specifically identifies the main change: ensuring the speculative-config's attention backend setting is honored in the DFlash draft loader.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/dflash-draft-attention-backend

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}