A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 15:00 UTC
[Attention][Spec Decode] Allow independent drafter attention backend selection (#39930)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Apr 27, 2026 at 22:38 UTC
fd74c90d9c3b5c35308f1f0ab308469235fa5277
6 files changed
+136
-11
tests/kernels/attention/test_attention_selector.py
+87
@@ -427,3 +427,90 @@ def test_per_head_quant_scales_backend_selection(
427
use_per_head_quant_scales=True,
428
)
429
assert backend_name in str(exc_info.value)
430
+
431
+
432
+@pytest.mark.parametrize(
433
+ "backend_name,use_non_causal,should_succeed",
434
+ [
435
+ ("FLASH_ATTN", True, True), # FlashAttn supports non-causal
436
+ ("FLASH_ATTN", False, True), # FlashAttn also works with causal
437
+ ("FLASHINFER", True, False), # FlashInfer does not support non-causal
438
+ ("FLASHINFER", False, True), # FlashInfer works with causal
439
+ ],
440
+)
441
+def test_non_causal_backend_selection(
442
+ backend_name: str, use_non_causal: bool, should_succeed: bool
443
+):
444
+ """Test that use_non_causal on AttentionConfig controls backend filtering.
445
+
446
+ DFlashProposer sets use_non_causal=True on the draft model's
447
+ AttentionConfig so only non-causal-capable backends are selected.
448
+ The target model keeps use_non_causal=False (default) and can use
449
+ any backend.
450
+ """
451
+ _cached_get_attn_backend.cache_clear()
452
+
453
+ attention_config = AttentionConfig(
454
+ backend=AttentionBackendEnum[backend_name],
455
+ use_non_causal=use_non_causal,
456
+ )
457
+ cache_config = CacheConfig(block_size=16)
458
+ vllm_config = VllmConfig(
459
+ attention_config=attention_config, cache_config=cache_config
460
+ )
461
+
462
+ if CudaPlatform is None:
463
+ pytest.skip("CudaPlatform not available")
464
+ with (
465
+ set_current_vllm_config(vllm_config),
466
+ patch("vllm.platforms.current_platform", CudaPlatform()),
467
+ ):
468
+ if should_succeed:
469
+ backend = get_attn_backend(
470
+ head_size=128,
471
+ dtype=torch.float16,
472
+ kv_cache_dtype=None,
473
+ )
474
+ assert backend.get_name() == backend_name
475
+ else:
476
+ with pytest.raises(ValueError) as exc_info:
477
+ get_attn_backend(
478
+ head_size=128,
479
+ dtype=torch.float16,
480
+ kv_cache_dtype=None,
481
+ )
482
+ assert "non-causal" in str(exc_info.value).lower()
483
+
484
+
485
+def test_non_causal_autoselect_backend():
486
+ """Test that when backend=None with use_non_causal=True, auto-selection
487
+ picks a compatible backend.
488
+
489
+ This simulates the DFlash scenario where the user doesn't specify
490
+ --attention-backend or --speculative-config.attention_backend.
491
+ The drafter inherits backend=None and auto-selects a backend that
492
+ supports non-causal attention.
493
+ """
494
+ _cached_get_attn_backend.cache_clear()
495
+
496
+ attention_config = AttentionConfig(
497
+ backend=None,
498
+ use_non_causal=True,
499
+ )
500
+ cache_config = CacheConfig(block_size=16)
501
+ vllm_config = VllmConfig(
502
+ attention_config=attention_config, cache_config=cache_config
503
+ )
504
+
505
+ if CudaPlatform is None:
506
+ pytest.skip("CudaPlatform not available")
507
+ with (
508
+ set_current_vllm_config(vllm_config),
509
+ patch("vllm.platforms.current_platform", CudaPlatform()),
510
+ ):
511
+ backend = get_attn_backend(
512
+ head_size=128,
513
+ dtype=torch.float16,
514
+ kv_cache_dtype=None,
515
+ )
516
+ assert backend.supports_non_causal()
vllm/config/attention.py
+3
@@ -54,6 +54,9 @@ class AttentionConfig:
54
use_fp4_indexer_cache: bool = False
55
"""If set, use fp4 indexer cache for dsv32 family model (not support yet)"""
56
57
+ use_non_causal: bool = False
58
+ """Whether to use non-causal (bidirectional) attention."""
59
+
60
def compute_hash(self) -> str:
61
"""
62
Provide a hash that uniquely identifies all the configs
vllm/config/speculative.py
+15
-1
@@ -5,7 +5,7 @@ import ast
5
import copy
6
from typing import TYPE_CHECKING, Any, Literal, get_args
7
8
-from pydantic import Field, SkipValidation, model_validator
8
+from pydantic import Field, SkipValidation, field_validator, model_validator
9
from typing_extensions import Self
10
11
from vllm.config import LoadConfig
@@ -17,6 +17,7 @@ from vllm.logger import init_logger
17
from vllm.transformers_utils.config import get_hf_text_config
18
from vllm.utils.hashing import safe_hash
19
from vllm.utils.import_utils import LazyLoader, has_arctic_inference
20
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
21
22
if TYPE_CHECKING:
23
from transformers import PretrainedConfig
@@ -106,6 +107,10 @@ class SpeculativeConfig:
107
inherits the target model's `--moe-backend` setting. Useful when the
108
drafter and generator require different MoE kernels (e.g. quantized
109
generator with unquantized drafter)."""
110
+ attention_backend: AttentionBackendEnum | None = None
111
+ """Attention backend to use for the draft model. When `None`, the backend is
112
+ automatically selected. Useful when the drafter requires a different attention
113
+ backend (e.g. DFlash needs a non-causal-capable backend like FLASH_ATTN)."""
114
max_model_len: int | None = Field(default=None, ge=1)
115
"""The maximum model length of the draft model. Used when testing the
116
ability to skip speculation for some sequences."""
@@ -911,6 +916,15 @@ class SpeculativeConfig:
916
917
return draft_parallel_config
918
919
+ @field_validator("attention_backend", mode="before")
920
+ @classmethod
921
+ def _parse_attention_backend(cls, value: Any) -> Any:
922
+ if isinstance(value, str):
923
+ if value.lower() == "auto":
924
+ return None
925
+ return AttentionBackendEnum[value.upper()]
926
+ return value
927
+
928
@model_validator(mode="after")
929
def _verify_args(self) -> Self:
930
if self.tensor_parallel_size is not None:
vllm/v1/attention/selector.py
+1
-6
@@ -81,11 +81,6 @@ def get_attn_backend(
81
else:
82
block_size = None
83
84
- speculative_config = vllm_config.speculative_config
85
- use_non_causal = (
86
- speculative_config is not None and speculative_config.method == "dflash"
87
- )
88
-
84
attn_selector_config = AttentionSelectorConfig(
85
head_size=head_size,
86
dtype=dtype,
@@ -97,7 +92,7 @@ def get_attn_backend(
92
use_mm_prefix=use_mm_prefix,
93
use_per_head_quant_scales=use_per_head_quant_scales,
94
attn_type=attn_type or AttentionType.DECODER,
100
- use_non_causal=use_non_causal,
95
+ use_non_causal=vllm_config.attention_config.use_non_causal,
96
use_batch_invariant=envs.VLLM_BATCH_INVARIANT,
97
)
98
vllm/v1/spec_decode/dflash.py
+12
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from dataclasses import replace
5
from typing import Any
6
7
import torch
@@ -67,6 +68,17 @@ class DFlashProposer(SpecDecodeBaseProposer):
68
# For DFlash we use the input embeddings to embed the mask token
69
self.parallel_drafting_hidden_state_tensor = None
70
71
+ @override
72
+ def _create_draft_vllm_config(self) -> VllmConfig:
73
+ base = super()._create_draft_vllm_config()
74
+ return replace(
75
+ base,
76
+ attention_config=replace(
77
+ base.attention_config,
78
+ use_non_causal=True,
79
+ ),
80
+ )
81
+
82
@override
83
def _raise_if_multimodal(self):
84
# Override to allow multimodal inputs since DFlash supports Qwen3.5 models
vllm/v1/spec_decode/llm_base_proposer.py
+18
-4
@@ -1276,15 +1276,29 @@ class SpecDecodeBaseProposer:
1276
Subclasses may override to apply additional config changes.
1277
"""
1278
spec_cfg = self.speculative_config
1279
+ base = self.vllm_config
1280
+
1281
if spec_cfg.moe_backend is not None:
1280
- return replace(
1281
- self.vllm_config,
1282
+ base = replace(
1283
+ base,
1284
kernel_config=replace(
1283
- self.vllm_config.kernel_config,
1285
+ base.kernel_config,
1286
moe_backend=spec_cfg.moe_backend,
1287
),
1288
)
1287
- return self.vllm_config
1289
+
1290
+ # Note (matt): Never inherit the attention backend from base, because there are
1291
+ # many opportunities for incompatibility, so we always independently autoselect
1292
+ # unless explicitly specified in the speculative config.
1293
+ base = replace(
1294
+ base,
1295
+ attention_config=replace(
1296
+ base.attention_config,
1297
+ backend=spec_cfg.attention_backend,
1298
+ ),
1299
+ )
1300
+
1301
+ return base
1302
1303
def _get_model(self) -> nn.Module:
1304
"""