[Bugfix] Fix FlashInfer non-causal draft attention (DFlash/DSpark) on Blackwell (#48167)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Michael Goin committed
Jul 15, 2026 at 15:44 UTC
ecf4aa5ce2ccd4069f12318ca9d3fcef7c9f6257
12 files changed
+149
-47
tests/v1/spec_decode/test_dflash_causality.py
new
+55
@@ -0,0 +1,55 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Config-only resolution of DFlash draft attention causality.
4
+
5
+``dflash_has_any_non_causal`` decides pre-build whether the draft needs a
6
+non-causal-capable backend, so its branch table (explicit override, SWA-derived
7
+per-layer causality, and the no-``layer_types`` fallback) is worth pinning.
8
+"""
9
+
10
+from types import SimpleNamespace
11
+
12
+import pytest
13
+
14
+from vllm.model_executor.models.qwen3_dflash import (
15
+ _dflash_layer_causal,
16
+ dflash_has_any_non_causal,
17
+)
18
+
19
+
20
+def _config(num_hidden_layers, layer_types=None, causal_override=None):
21
+ dflash_config = None if causal_override is None else {"causal": causal_override}
22
+ return SimpleNamespace(
23
+ num_hidden_layers=num_hidden_layers,
24
+ layer_types=layer_types,
25
+ dflash_config=dflash_config,
26
+ )
27
+
28
+
29
+@pytest.mark.parametrize(
30
+ "config,expected",
31
+ [
32
+ # Override forces causality on every layer, ignoring layer_types.
33
+ (_config(2, layer_types=["full_attention"] * 2, causal_override=True), False),
34
+ # Override forces non-causal on every layer.
35
+ (
36
+ _config(2, layer_types=["sliding_attention"] * 2, causal_override=False),
37
+ True,
38
+ ),
39
+ # SWA-derived: full-attention layers are non-causal.
40
+ (_config(2, layer_types=["sliding_attention", "full_attention"]), True),
41
+ # SWA-derived: all-sliding is fully causal.
42
+ (_config(2, layer_types=["sliding_attention", "sliding_attention"]), False),
43
+ # No layer_types -> non-causal fallback.
44
+ (_config(2, layer_types=None), True),
45
+ (_config(2, layer_types=[]), True),
46
+ ],
47
+)
48
+def test_dflash_has_any_non_causal(config, expected):
49
+ assert dflash_has_any_non_causal(config) is expected
50
+
51
+
52
+def test_dflash_layer_causal_is_per_layer():
53
+ config = _config(2, layer_types=["sliding_attention", "full_attention"])
54
+ assert _dflash_layer_causal(config, 0) is True
55
+ assert _dflash_layer_causal(config, 1) is False
tools/pre_commit/generate_attention_backend_docs.py
+3
-1
@@ -1480,7 +1480,9 @@ def _get_backends_from_return(stmts: list) -> list[str]:
1480
1481
1482
def _is_sm100_check(test: ast.expr) -> bool:
1483
- """Check if test is `something.major == 10`."""
1483
+ """Check if test is `something.major == 10`, possibly inside an `and`."""
1484
+ if isinstance(test, ast.BoolOp) and isinstance(test.op, ast.And):
1485
+ return any(_is_sm100_check(value) for value in test.values)
1486
return (
1487
isinstance(test, ast.Compare)
1488
and isinstance(test.left, ast.Attribute)
vllm/model_executor/models/qwen3_dflash.py
+24
-12
@@ -52,6 +52,26 @@ from .utils import (
52
logger = init_logger(__name__)
53
54
55
+_SLIDING_ATTENTION = "sliding_attention"
56
+
57
+
58
+def _dflash_layer_causal(config: Qwen3Config, layer_idx: int) -> bool:
59
+ """``dflash_config.causal`` overrides all layers; else only SWA layers causal."""
60
+ override = (getattr(config, "dflash_config", None) or {}).get("causal")
61
+ if override is not None:
62
+ return override
63
+ layer_types = getattr(config, "layer_types", None)
64
+ return bool(layer_types) and layer_types[layer_idx] == _SLIDING_ATTENTION
65
+
66
+
67
+def dflash_has_any_non_causal(config: Qwen3Config) -> bool:
68
+ """Whether the draft needs a non-causal-capable backend, resolved from config
69
+ (config mirror of the model's ``get_draft_attn_causal``, usable pre-build)."""
70
+ return not all(
71
+ _dflash_layer_causal(config, i) for i in range(config.num_hidden_layers)
72
+ )
73
+
74
+
75
def _resolve_layer_attention(
76
config: Qwen3Config, layer_idx: int
77
) -> tuple[int | None, bool]:
@@ -79,12 +99,10 @@ def _resolve_layer_attention(
99
dflash_config = getattr(config, "dflash_config", None) or {}
100
layer_types = getattr(config, "layer_types", None)
101
use_swa = dflash_config.get("use_swa", False)
82
- config_causal = dflash_config.get("causal", None)
102
84
- SLIDING_ATTENTION = "sliding_attention"
103
any_sliding = False
104
if layer_types is not None:
87
- num_sliding = sum(lt == SLIDING_ATTENTION for lt in layer_types)
105
+ num_sliding = sum(lt == _SLIDING_ATTENTION for lt in layer_types)
106
any_sliding = num_sliding > 0
107
# Mixed sliding/full attention needs multiple KV groups (V2 runner only).
108
if (
@@ -97,16 +115,11 @@ def _resolve_layer_attention(
115
"VLLM_USE_V2_MODEL_RUNNER=1."
116
)
117
100
- default_causal = False
118
+ # ``use_swa`` forces SWA on every layer, even an all-full ``layer_types``.
119
if layer_types is None or (use_swa and not any_sliding):
102
- # An absent ``layer_types`` (or the all-"full_attention" one that may
103
- # be synthesized when the checkpoint omits it) must not override
104
- # ``dflash_config.use_swa``, which forces SWA on every layer.
120
is_sliding = use_swa
121
else:
107
- is_sliding = layer_types[layer_idx] == SLIDING_ATTENTION
108
- # Full-attention layers default non-causal; SWA layers default causal.
109
- default_causal = is_sliding
122
+ is_sliding = layer_types[layer_idx] == _SLIDING_ATTENTION
123
124
sliding_window = None
125
if is_sliding:
@@ -119,8 +132,7 @@ def _resolve_layer_attention(
132
"dflash_config.swa_window_size or the top-level sliding_window."
133
)
134
122
- causal = config_causal if config_causal is not None else default_causal
123
- return sliding_window, causal
135
+ return sliding_window, _dflash_layer_causal(config, layer_idx)
136
137
138
class DFlashQwen3Attention(nn.Module):
vllm/platforms/cuda.py
+6
-1
@@ -85,6 +85,7 @@ def _get_backend_priorities(
85
device_capability: DeviceCapability,
86
num_heads: int | None = None,
87
kv_cache_dtype: CacheDType | None = None,
88
+ use_non_causal: bool = False,
89
) -> list[AttentionBackendEnum]:
90
"""Get backend priorities with lazy import to avoid circular dependency."""
91
from vllm.utils.torch_utils import is_quantized_kv_cache
@@ -141,7 +142,10 @@ def _get_backend_priorities(
142
AttentionBackendEnum.FLASHMLA_SPARSE,
143
]
144
else:
144
- if device_capability.major == 10:
145
+ # SM100f defaults to FlashInfer for TRTLLM causal attention, but its non-causal
146
+ # cutlass path (used for dflash attention) is known to have problems.
147
+ # So prefer FlashAttention when non-causal on SM100f.
148
+ if device_capability.major == 10 and not use_non_causal:
149
return [
150
AttentionBackendEnum.FLASHINFER,
151
AttentionBackendEnum.FLASH_ATTN,
@@ -368,6 +372,7 @@ class CudaPlatformBase(Platform):
372
device_capability,
373
num_heads,
374
attn_selector_config.kv_cache_dtype,
375
+ attn_selector_config.use_non_causal,
376
)
377
for priority, backend in enumerate(backend_priorities):
378
try:
vllm/v1/attention/backends/flashinfer.py
+2
-1
@@ -927,7 +927,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
927
has_trtllm_support = False
928
break
929
930
- if has_trtllm_support:
930
+ # trtllm-gen only supports causal attention.
931
+ if has_trtllm_support and not vllm_config.attention_config.use_non_causal:
932
return AttentionCGSupport.UNIFORM_BATCH
933
else:
934
return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
vllm/v1/spec_decode/dflash.py
+5
-1
@@ -70,7 +70,11 @@ class DFlashProposer(SpecDecodeBaseProposer):
70
# For DFlash we use the input embeddings to embed the mask token
71
self.parallel_drafting_hidden_state_tensor = None
72
73
- self.dflash_causal = self.dflash_config.get("causal", False)
73
+ from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal
74
+
75
+ self.dflash_causal = not dflash_has_any_non_causal(
76
+ self.draft_model_config.hf_config
77
+ )
78
79
@override
80
def _create_draft_vllm_config(self) -> VllmConfig:
vllm/v1/worker/gpu/model_runner.py
+4
-3
@@ -472,9 +472,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
472
decode_query_len=self.decode_query_len,
473
lora_capture_cases=self.lora_capture_cases,
474
)
475
- if self.speculator is not None:
476
- self.speculator.init_cudagraph_manager(cudagraph_mode)
477
-
475
check_attention_cp_compatibility(self.vllm_config)
476
if isinstance(self.speculator, DraftModelSpeculator):
477
# HACK(woosuk)
@@ -485,6 +482,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
482
self.input_buffers,
483
self.attn_groups,
484
)
485
+ if self.speculator is not None:
486
+ # After set_attn, so the speculator can size its cudagraph mode
487
+ # to its own attention support.
488
+ self.speculator.init_cudagraph_manager(cudagraph_mode)
489
490
self.kv_caches: list[torch.Tensor] = []
491
kv_caches_dict = init_kv_cache(
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+35
-10
@@ -6,11 +6,12 @@ from typing import Any
6
import torch
7
import torch.nn as nn
8
9
-from vllm.config import VllmConfig
9
+from vllm.config import VllmConfig, replace
10
from vllm.config.compilation import CUDAGraphMode
11
from vllm.forward_context import BatchDescriptor, set_forward_context
12
from vllm.logger import init_logger
13
from vllm.triton_utils import tl, triton
14
+from vllm.v1.attention.backend import AttentionCGSupport
15
from vllm.v1.attention.backends.utils import PAD_SLOT_ID
16
from vllm.v1.kv_cache_interface import KVCacheConfig
17
from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
@@ -19,10 +20,7 @@ from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
20
from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
21
from vllm.v1.worker.gpu.model_states.interface import ModelState
22
from vllm.v1.worker.gpu.spec_decode.dflash.cudagraph import DFlashCudaGraphManager
22
-from vllm.v1.worker.gpu.spec_decode.dflash.utils import (
23
- get_dflash_causal,
24
- load_dflash_model,
25
-)
23
+from vllm.v1.worker.gpu.spec_decode.dflash.utils import load_dflash_model
24
from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
25
from vllm.v1.worker.gpu.spec_decode.utils import get_parallel_drafting_token_id
26
from vllm.v1.worker.utils import AttentionGroup
@@ -50,7 +48,11 @@ class DFlashSpeculator(DraftModelSpeculator):
48
self.draft_model_config.hf_config
49
)
50
53
- self.dflash_causal = get_dflash_causal(self.draft_model_config)
51
+ from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal
52
+
53
+ self.requires_non_causal = dflash_has_any_non_causal(
54
+ self.draft_model_config.hf_config
55
+ )
56
57
# Whether the anchor query position is itself a prediction. DFlash default uses
58
# the anchor as the bonus token (only mask tokens predict); DSpark samples from
@@ -84,9 +86,32 @@ class DFlashSpeculator(DraftModelSpeculator):
86
self.query_cudagraph_manager: DFlashCudaGraphManager | None = None
87
self.draft_kv_cache_group_id: int = -1
88
89
+ @property
90
+ def attn_vllm_config(self) -> VllmConfig:
91
+ # The draft's attention differs from the target's in causality.
92
+ return replace(
93
+ self.vllm_config,
94
+ attention_config=replace(
95
+ self.vllm_config.attention_config,
96
+ use_non_causal=self.requires_non_causal,
97
+ ),
98
+ )
99
+
100
def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
88
- # PIECEWISE cudagraphs are not supported for dflash
89
- if cudagraph_mode.decode_mode() == CUDAGraphMode.FULL:
101
+ wants_full = cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
102
+ supports_full = (
103
+ self.attn_cg_support.min_cg_support.value
104
+ >= AttentionCGSupport.UNIFORM_BATCH.value
105
+ )
106
+ if wants_full and not supports_full:
107
+ logger.warning(
108
+ "%s draft attention (%s) does not support full CUDA graphs; "
109
+ "running the draft eagerly.",
110
+ self._speculator_name,
111
+ self.attn_cg_support.min_cg_attn_backend,
112
+ )
113
+ # PIECEWISE cudagraphs are not supported for dflash.
114
+ if wants_full and supports_full:
115
cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY
116
else:
117
cudagraph_mode = CUDAGraphMode.NONE
@@ -158,8 +183,8 @@ class DFlashSpeculator(DraftModelSpeculator):
183
# of the kv-cache group its cache belongs to. Models that share a single group
184
# leave this as None and share one context slot mapping.
185
self._layer_group_idx: list[int] | None = None
161
- # Per-KV-group causal, falling back to the scalar dflash_causal.
162
- self._group_causal: dict[int, bool] | bool = self.dflash_causal
186
+ # Per-KV-group causal, falling back to whether the drafter is all-causal.
187
+ self._group_causal: dict[int, bool] | bool = not self.requires_non_causal
188
if hasattr(self.model, "get_draft_kv_cache_layer_names"):
189
layer_names = self.model.get_draft_kv_cache_layer_names()
190
name_to_gid = {
vllm/v1/worker/gpu/spec_decode/dflash/utils.py
+5
-11
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import torch.nn as nn
4
5
-from vllm.config import ModelConfig, VllmConfig, replace
5
+from vllm.config import VllmConfig, replace
6
from vllm.distributed.parallel_state import get_pp_group
7
from vllm.model_executor.model_loader import get_model
8
from vllm.v1.worker.gpu.spec_decode.eagle.utils import (
@@ -11,26 +11,20 @@ from vllm.v1.worker.gpu.spec_decode.eagle.utils import (
11
)
12
13
14
-def get_dflash_causal(draft_model_config: ModelConfig) -> bool:
15
- """Whether the DFlash draft uses causal (vs non-causal) attention."""
16
- dflash_config = getattr(draft_model_config.hf_config, "dflash_config", None) or {}
17
- return dflash_config.get("causal", False)
18
-
19
-
14
def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module:
15
from vllm.compilation.backends import set_model_tag
16
+ from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal
17
18
speculative_config = vllm_config.speculative_config
19
assert speculative_config is not None
20
draft_model_config = speculative_config.draft_model_config
26
- # Modify the attention config so that we select an attention backend that matches
27
- # the causal/non-causal mode of the dflash model.
28
- causal = get_dflash_causal(draft_model_config)
21
+ # Select an attention backend that supports the drafter's attention: mixing
22
+ # a non-causal layer onto a causal-only backend would fail.
23
draft_vllm_config = replace(
24
vllm_config,
25
attention_config=replace(
26
vllm_config.attention_config,
33
- use_non_causal=not causal,
27
+ use_non_causal=dflash_has_any_non_causal(draft_model_config.hf_config),
28
backend=speculative_config.attention_backend,
29
),
30
)
vllm/v1/worker/gpu/spec_decode/dspark/speculator.py
-2
@@ -59,8 +59,6 @@ class DSparkSpeculator(DFlashSpeculator):
59
self.max_num_tokens, draft_hidden, dtype=self.dtype, device=device
60
)
61
62
- self.dflash_causal = False
63
-
62
self._step_cols = torch.arange(
63
self.num_speculative_steps, dtype=torch.int32, device=device
64
)
vllm/v1/worker/gpu/spec_decode/dspark/utils.py
+2
-3
@@ -18,14 +18,13 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
18
draft_model_config = speculative_config.draft_model_config
19
20
from vllm.compilation.backends import set_model_tag
21
+ from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal
22
22
- # DSpark uses non-causal attention.
23
- causal = False
23
draft_vllm_config = replace(
24
vllm_config,
25
attention_config=replace(
26
vllm_config.attention_config,
28
- use_non_causal=not causal,
27
+ use_non_causal=dflash_has_any_non_causal(draft_model_config.hf_config),
28
backend=speculative_config.attention_backend,
29
),
30
)
vllm/v1/worker/gpu/spec_decode/speculator.py
+8
-2
@@ -175,6 +175,12 @@ class DraftModelSpeculator(BaseSpeculator):
175
num_unpadded_tokens,
176
)
177
178
+ @property
179
+ def attn_vllm_config(self) -> VllmConfig:
180
+ """Config for the draft's attention metadata builders. Overridden by
181
+ speculators whose attention mode differs from the target's."""
182
+ return self.vllm_config
183
+
184
def set_attn(
185
self,
186
model_state: ModelState,
@@ -185,9 +191,9 @@ class DraftModelSpeculator(BaseSpeculator):
191
) -> None:
192
self.model_state = model_state
193
self.kv_cache_config = kv_cache_config
188
- self.attn_groups, _, _ = init_attn_backend(
194
+ self.attn_groups, self.attn_cg_support, _ = init_attn_backend(
195
kv_cache_config,
190
- self.vllm_config,
196
+ self.attn_vllm_config,
197
self.device,
198
active_layer_names=self.draft_attn_layer_names,
199
)