[Attention] Allow selecting a different attention backend per KV-cache group (#48012)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Nicolò Lucchesi committed
Jul 17, 2026 at 21:19 UTC
5784507da458656ef8248119590822445d4c67f2
4 files changed
+258
-3
tests/v1/attention/test_backend_per_kind.py
new
+67
@@ -0,0 +1,67 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for per-KV-group attention backend selection (backend_per_kind)."""
4
+
5
+import pytest
6
+
7
+from vllm.config.attention import AttentionConfig
8
+from vllm.v1.attention.backend import AttentionType
9
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
10
+from vllm.v1.attention.selector import get_attn_spec_kind
11
+from vllm.v1.kv_cache_interface import KVCacheSpecKind
12
+
13
+
14
+@pytest.mark.parametrize(
15
+ "signals,expected",
16
+ [
17
+ (dict(use_mla=False, has_sliding_window=False), "full"),
18
+ (dict(use_mla=True, has_sliding_window=False), "mla"),
19
+ (dict(use_mla=True, has_sliding_window=True), "sw_mla"),
20
+ (dict(use_mla=False, has_sliding_window=True), "sw"),
21
+ ],
22
+)
23
+def test_get_attn_spec_kind_decoder(signals, expected):
24
+ kind_by_name = {
25
+ "full": KVCacheSpecKind.FULL_ATTENTION,
26
+ "mla": KVCacheSpecKind.MLA_ATTENTION,
27
+ "sw_mla": KVCacheSpecKind.SLIDING_WINDOW_MLA,
28
+ "sw": KVCacheSpecKind.SLIDING_WINDOW,
29
+ }
30
+ kind = get_attn_spec_kind(attn_type=AttentionType.DECODER, **signals)
31
+ assert kind is kind_by_name[expected]
32
+
33
+
34
+@pytest.mark.parametrize(
35
+ "attn_type,expected",
36
+ [
37
+ (AttentionType.ENCODER_ONLY, KVCacheSpecKind.ENCODER_ONLY_ATTENTION),
38
+ (AttentionType.ENCODER_DECODER, KVCacheSpecKind.CROSS_ATTENTION),
39
+ ],
40
+)
41
+def test_get_attn_spec_kind_attn_type(attn_type, expected):
42
+ kind = get_attn_spec_kind(
43
+ use_mla=False,
44
+ has_sliding_window=False,
45
+ attn_type=attn_type,
46
+ )
47
+ assert kind is expected
48
+
49
+
50
+def test_backend_per_kind_parses_strings():
51
+ cfg = AttentionConfig(
52
+ backend_per_kind={
53
+ "mla_attention": "FLASHINFER_MLA",
54
+ "sliding_window_mla": "triton_mla", # case-insensitive
55
+ }
56
+ )
57
+ assert cfg.backend_per_kind["mla_attention"] is AttentionBackendEnum.FLASHINFER_MLA
58
+ assert cfg.backend_per_kind["sliding_window_mla"] is AttentionBackendEnum.TRITON_MLA
59
+
60
+
61
+def test_backend_per_kind_rejects_unknown_kind():
62
+ with pytest.raises(ValueError, match="Unknown KV cache group kind"):
63
+ AttentionConfig(backend_per_kind={"not_a_kind": "TRITON_MLA"})
64
+
65
+
66
+def test_backend_per_kind_defaults_empty():
67
+ assert AttentionConfig().backend_per_kind == {}
tests/v1/e2e/general/test_attention_backend_per_kind.py
new
+94
@@ -0,0 +1,94 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""End-to-end checks that ``backend_per_kind`` selects the requested attention
4
+backend for each KV-cache group at runtime.
5
+
6
+Uses ``google/gemma-3-1b-it``, which interleaves full-attention and
7
+sliding-window layers, so the model produces separate ``full_attention`` and
8
+``sliding_window`` KV-cache groups.
9
+"""
10
+
11
+import pytest
12
+
13
+from vllm import LLM
14
+from vllm.config.attention import AttentionConfig
15
+from vllm.platforms import current_platform
16
+
17
+MODEL = "google/gemma-3-1b-it"
18
+
19
+
20
+def _collect_group_backends(worker) -> list[tuple[str, str]]:
21
+ """Runs on the worker: returns (spec_kind, backend_name) per attn group."""
22
+ from vllm.v1.kv_cache_interface import get_kv_cache_spec_kind
23
+
24
+ out: list[tuple[str, str]] = []
25
+ for kv_group in worker.model_runner.attn_groups:
26
+ for attn_group in kv_group:
27
+ kind = get_kv_cache_spec_kind(attn_group.kv_cache_spec)
28
+ out.append((kind.value, attn_group.backend.get_name()))
29
+ return out
30
+
31
+
32
+@pytest.mark.skipif(
33
+ not current_platform.is_cuda(), reason="backend names are CUDA-specific"
34
+)
35
+@pytest.mark.parametrize(
36
+ "backend_per_kind",
37
+ [
38
+ {"full_attention": "FLASH_ATTN", "sliding_window": "TRITON_ATTN"},
39
+ # Swapped, to prove the mapping is causal rather than the default.
40
+ {"full_attention": "TRITON_ATTN", "sliding_window": "FLASH_ATTN"},
41
+ ],
42
+)
43
+def test_backend_per_kind_splits_groups(backend_per_kind, monkeypatch):
44
+ # collective_rpc ships the callable to the EngineCore subprocess; the
45
+ # secure msgpack encoder can't serialize functions, so opt into the
46
+ # pickle fallback (same pattern as test_pooling_chunked_prefill).
47
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
48
+ llm = LLM(
49
+ model=MODEL,
50
+ attention_config=AttentionConfig(backend_per_kind=backend_per_kind),
51
+ enforce_eager=True,
52
+ max_model_len=2048,
53
+ gpu_memory_utilization=0.4,
54
+ )
55
+
56
+ group_backends = llm.llm_engine.collective_rpc(_collect_group_backends)[0]
57
+ kinds = {kind for kind, _ in group_backends}
58
+
59
+ # gemma3 must actually split into both kinds for this test to be meaningful.
60
+ assert "full_attention" in kinds
61
+ assert "sliding_window" in kinds
62
+
63
+ for kind, backend_name in group_backends:
64
+ if kind in backend_per_kind:
65
+ assert backend_name == backend_per_kind[kind], (
66
+ f"{kind} group used {backend_name}, expected {backend_per_kind[kind]}"
67
+ )
68
+
69
+
70
+@pytest.mark.skipif(
71
+ not current_platform.is_cuda(), reason="backend names are CUDA-specific"
72
+)
73
+def test_backend_per_kind_overrides_global_backend(monkeypatch):
74
+ """A per-kind entry wins over the global ``backend`` for its kind; other
75
+ kinds fall back to the global backend."""
76
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
77
+ llm = LLM(
78
+ model=MODEL,
79
+ attention_config=AttentionConfig(
80
+ backend="FLASH_ATTN",
81
+ backend_per_kind={"sliding_window": "TRITON_ATTN"},
82
+ ),
83
+ enforce_eager=True,
84
+ max_model_len=2048,
85
+ gpu_memory_utilization=0.4,
86
+ )
87
+
88
+ group_backends = llm.llm_engine.collective_rpc(_collect_group_backends)[0]
89
+
90
+ for kind, backend_name in group_backends:
91
+ if kind == "sliding_window":
92
+ assert backend_name == "TRITON_ATTN"
93
+ elif kind == "full_attention":
94
+ assert backend_name == "FLASH_ATTN"
vllm/config/attention.py
+38
@@ -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 field
5
from typing import Any, Literal
6
7
from pydantic import field_validator
@@ -19,6 +20,17 @@ class AttentionConfig:
20
backend: AttentionBackendEnum | None = None
21
"""Attention backend to use. Use "auto" or None for automatic selection."""
22
23
+ backend_per_kind: dict[str, AttentionBackendEnum] = field(default_factory=dict)
24
+ """Per-KV-cache-group attention backend overrides, keyed by
25
+ `KVCacheSpecKind` (e.g. `{"mla_attention": "FLASHINFER_MLA",
26
+ "sliding_window_mla": "TRITON_MLA"}`). This lets a model that splits its
27
+ layers across multiple KV-cache groups (e.g. interleaved full and
28
+ sliding-window attention) use a different backend per group.
29
+
30
+ An entry overrides `backend` for layers of the matching kind; kinds not
31
+ listed fall back to `backend` (or automatic selection). A selected backend
32
+ that is invalid for that kind raises at startup."""
33
+
34
flash_attn_version: Literal[2, 3, 4] | None = None
35
"""Force vllm to use a specific flash-attention version (2, 3, or 4).
36
Only valid when using the flash-attention backend."""
@@ -120,3 +132,29 @@ class AttentionConfig:
132
if isinstance(value, str):
133
return MLAPrefillBackendEnum[value.upper()]
134
return value
135
+
136
+ @field_validator("backend_per_kind", mode="before")
137
+ @classmethod
138
+ def validate_backend_per_kind_before(cls, value: Any) -> Any:
139
+ """Parse the `backend_per_kind` map from strings.
140
+
141
+ Keys must be valid `KVCacheSpecKind` values; values are parsed like
142
+ `backend` (enum name, case-insensitive).
143
+ """
144
+ from vllm.v1.kv_cache_interface import KVCacheSpecKind
145
+
146
+ if not isinstance(value, dict):
147
+ return value
148
+ valid_kinds = {kind.value for kind in KVCacheSpecKind}
149
+ parsed: dict[str, AttentionBackendEnum] = {}
150
+ for kind, backend in value.items():
151
+ if kind not in valid_kinds:
152
+ raise ValueError(
153
+ f"Unknown KV cache group kind '{kind}' in "
154
+ f"backend_per_kind. Valid kinds are: "
155
+ f"{', '.join(sorted(valid_kinds))}."
156
+ )
157
+ if isinstance(backend, str):
158
+ backend = AttentionBackendEnum[backend.upper()]
159
+ parsed[kind] = backend
160
+ return parsed
vllm/v1/attention/selector.py
+59
-3
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from functools import cache
5
-from typing import NamedTuple, cast, get_args
5
+from typing import TYPE_CHECKING, NamedTuple, cast, get_args
6
7
import torch
8
@@ -15,6 +15,9 @@ from vllm.v1.attention.backends.registry import (
15
MambaAttentionBackendEnum,
16
)
17
18
+if TYPE_CHECKING:
19
+ from vllm.v1.kv_cache_interface import KVCacheSpecKind
20
+
21
logger = init_logger(__name__)
22
23
@@ -53,6 +56,46 @@ class AttentionSelectorConfig(NamedTuple):
56
)
57
58
59
+def get_attn_spec_kind(
60
+ use_mla: bool,
61
+ has_sliding_window: bool,
62
+ attn_type: str,
63
+) -> "KVCacheSpecKind":
64
+ """Derive the KV-cache group kind a layer belongs to from its signals.
65
+
66
+ Mirrors ``get_kv_cache_spec_kind`` (which derives the kind from the
67
+ produced ``KVCacheSpec``) so users can target groups by kind when
68
+ setting ``AttentionConfig.backend_per_kind``.
69
+
70
+ ``SINK_FULL_ATTENTION`` is intentionally not derived here: it is produced
71
+ only by the ``StaticSinkAttention`` layer, whereas a plain ``Attention``
72
+ layer with attention sinks (e.g. gpt-oss) still yields a
73
+ ``FullAttentionSpec``/``SlidingWindowSpec``. Sinks therefore do not change
74
+ the kind.
75
+
76
+ Args:
77
+ use_mla: Whether the layer uses multi-head latent attention.
78
+ has_sliding_window: Whether the layer applies a sliding window.
79
+ attn_type: The layer's ``AttentionType``.
80
+
81
+ Returns:
82
+ The ``KVCacheSpecKind`` the layer maps to.
83
+ """
84
+ from vllm.v1.kv_cache_interface import KVCacheSpecKind
85
+
86
+ if attn_type == AttentionType.ENCODER_ONLY:
87
+ return KVCacheSpecKind.ENCODER_ONLY_ATTENTION
88
+ if attn_type == AttentionType.ENCODER_DECODER:
89
+ return KVCacheSpecKind.CROSS_ATTENTION
90
+ if use_mla:
91
+ if has_sliding_window:
92
+ return KVCacheSpecKind.SLIDING_WINDOW_MLA
93
+ return KVCacheSpecKind.MLA_ATTENTION
94
+ if has_sliding_window:
95
+ return KVCacheSpecKind.SLIDING_WINDOW
96
+ return KVCacheSpecKind.FULL_ATTENTION
97
+
98
+
99
def get_attn_backend(
100
head_size: int,
101
dtype: torch.dtype,
@@ -91,6 +134,7 @@ def get_attn_backend(
134
kv_transfer_config is not None and kv_transfer_config.is_kv_transfer_instance
135
)
136
137
+ attn_type = attn_type or AttentionType.DECODER
138
attn_selector_config = AttentionSelectorConfig(
139
head_size=head_size,
140
dtype=dtype,
@@ -101,15 +145,27 @@ def get_attn_backend(
145
use_sparse=use_sparse,
146
use_mm_prefix=use_mm_prefix,
147
use_per_head_quant_scales=use_per_head_quant_scales,
104
- attn_type=attn_type or AttentionType.DECODER,
148
+ attn_type=attn_type,
149
has_sliding_window=has_sliding_window,
150
use_non_causal=vllm_config.attention_config.use_non_causal,
151
use_batch_invariant=envs.VLLM_BATCH_INVARIANT,
152
use_kv_connector=use_kv_connector,
153
)
154
155
+ # A per-KV-group override (keyed by KVCacheSpecKind) takes precedence over
156
+ # the global backend; kinds not present in the map fall back to it.
157
+ attention_config = vllm_config.attention_config
158
+ backend = attention_config.backend
159
+ if attention_config.backend_per_kind:
160
+ kind = get_attn_spec_kind(
161
+ use_mla=use_mla,
162
+ has_sliding_window=has_sliding_window,
163
+ attn_type=attn_type,
164
+ )
165
+ backend = attention_config.backend_per_kind.get(kind.value, backend)
166
+
167
return _cached_get_attn_backend(
112
- backend=vllm_config.attention_config.backend,
168
+ backend=backend,
169
attn_selector_config=attn_selector_config,
170
num_heads=num_heads,
171
)