[Model Runner V2][DFlash] Enable dflash attention backend selection (#46770)

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
2026-06-25 19:29:25 -07:00
parent 3daea7ceb9
commit 5314665bad
1 changed files with 3 additions and 1 deletions
@@ -26,7 +26,9 @@ 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,
+            backend=speculative_config.attention_backend,
        ),
    )
    with set_model_tag("dflash_head"):