From ba3c7514c38738d48af53a34a34c3b95f9596243 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sun, 14 Jun 2026 01:15:11 +0000 Subject: [PATCH] spec-decode: honor the speculative-config attention backend for the DFlash draft load_dflash_model rewrites the draft's attention config to set its causal/non-causal mode but does not propagate speculative_config.attention_backend. The draft therefore always auto-selects an attention backend, which on Blackwell prefers FlashInfer regardless of what the speculative config requests (e.g. TRITON_ATTN). With a DFlash draft beside an MLA target (Kimi-K2.7-Code): - under DCP the engine fails to start with "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. Thread the speculative-config backend into the draft's attention config, matching what llm_base_proposer already does ("never inherit the attention backend from base ... unless explicitly specified in the speculative config"). When the spec config leaves it unset the behaviour is unchanged (auto-select). Validated on Kimi-K2.7-Code TP8 + Kimi-K2.6 DFlash draft, V2 runner, attention_backend=TRITON_ATTN, kv-fp8: - DCP8: previously failed to start; now boots, FULL draft graph captured, 0k C1 decode 113.6 tok/s vs 81.5 target-only (+39%), output coherent. - DCP1: 0k C1 decode 140 tok/s vs 125 with the auto-selected FlashInfer. Co-Authored-By: Claude Fable 5 --- vllm/v1/worker/gpu/spec_decode/dflash/utils.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py index 1c371a65caca..da46b90cb149 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py @@ -26,7 +26,13 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo draft_vllm_config = replace( vllm_config, attention_config=replace( - vllm_config.attention_config, use_non_causal=not causal + vllm_config.attention_config, + use_non_causal=not causal, + # Honor the speculative-config attention backend for the draft + # (matches llm_base_proposer): otherwise auto-select picks + # FlashInfer, which downgrades the spec-decode cudagraph to + # PIECEWISE and cannot do non-causal prefill under DCP. + backend=speculative_config.attention_backend, ), ) with set_model_tag("dflash_head"):