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 7, 22:45 UTC
[Model Runner V2] Support attention group (#35036)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Feb 21, 2026 at 16:42 UTC
b71fbd06e215a1a09600220c947a1bb2d5494de9
5 files changed
+87
-55
vllm/v1/worker/gpu/attn_utils.py
+55
-30
@@ -7,17 +7,14 @@ import torch
7
8
from vllm.config import VllmConfig, get_layers_from_vllm_config
9
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
10
-from vllm.v1.attention.backend import (
11
- AttentionBackend,
12
- AttentionMetadataBuilder,
13
- CommonAttentionMetadata,
14
-)
10
+from vllm.v1.attention.backend import AttentionBackend, CommonAttentionMetadata
11
from vllm.v1.kv_cache_interface import (
12
AttentionSpec,
13
KVCacheConfig,
14
KVCacheSpec,
15
+ UniformTypeKVCacheSpecs,
16
)
20
-from vllm.v1.worker.utils import bind_kv_cache
17
+from vllm.v1.worker.utils import AttentionGroup, bind_kv_cache
18
19
20
def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
@@ -35,29 +32,56 @@ def init_attn_backend(
32
kv_cache_config: KVCacheConfig, vllm_config: VllmConfig, device: torch.device
33
):
34
attn_backends: dict[str, type[AttentionBackend]] = {}
38
- attn_metadata_builders: list[AttentionMetadataBuilder] = []
39
- flashinfer_workspace: torch.Tensor | None = None
40
- for kv_cache_group_spec in kv_cache_config.kv_cache_groups:
35
+ attn_groups: list[list[AttentionGroup]] = []
36
+ attn_backend_workspace: torch.Tensor | None = None
37
+ for kv_cache_group_id, kv_cache_group_spec in enumerate(
38
+ kv_cache_config.kv_cache_groups
39
+ ):
40
layer_names = kv_cache_group_spec.layer_names
42
- any_layer_name = next(iter(layer_names))
41
42
layer_type = cast(type[Any], AttentionLayerBase)
43
attn_layers = get_layers_from_vllm_config(vllm_config, layer_type, layer_names)
46
- attn_backend = attn_layers[any_layer_name].get_attn_backend()
44
+
45
+ group_map: dict[tuple[tuple[str, str], KVCacheSpec], AttentionGroup] = {}
46
+ group_order: list[tuple[tuple[str, str], KVCacheSpec]] = []
47
+
48
for layer_name in layer_names:
49
+ attn_backend = attn_layers[layer_name].get_attn_backend()
50
attn_backends[layer_name] = attn_backend
51
50
- attn_metadata_builder = attn_backend.get_builder_cls()(
51
- kv_cache_group_spec.kv_cache_spec, layer_names, vllm_config, device
52
- )
53
- attn_metadata_builders.append(attn_metadata_builder) # type: ignore
54
-
55
- if attn_backend.get_name() == "FLASHINFER":
56
- if flashinfer_workspace is None:
57
- flashinfer_workspace = attn_metadata_builder._get_workspace_buffer()
52
+ layer_kv_cache_spec: KVCacheSpec = kv_cache_group_spec.kv_cache_spec
53
+ if isinstance(layer_kv_cache_spec, UniformTypeKVCacheSpecs):
54
+ layer_kv_cache_spec = layer_kv_cache_spec.kv_cache_specs[layer_name]
55
+
56
+ key = (attn_backend.full_cls_name(), layer_kv_cache_spec)
57
+ if key not in group_map:
58
+ group_map[key] = AttentionGroup(
59
+ attn_backend,
60
+ [layer_name],
61
+ layer_kv_cache_spec,
62
+ kv_cache_group_id,
63
+ )
64
+ group_order.append(key)
65
else:
59
- attn_metadata_builder.set_workspace_buffer(flashinfer_workspace)
60
- return attn_backends, attn_metadata_builders
66
+ group_map[key].layer_names.append(layer_name)
67
+
68
+ groups = [group_map[key] for key in group_order]
69
+ for group in groups:
70
+ group.create_metadata_builders(
71
+ vllm_config=vllm_config,
72
+ device=device,
73
+ kernel_block_size=None,
74
+ num_metadata_builders=1,
75
+ )
76
+ builder = group.get_metadata_builder(0)
77
+ if attn_backend_workspace is None:
78
+ if hasattr(builder, "_get_workspace_buffer"):
79
+ attn_backend_workspace = builder._get_workspace_buffer()
80
+ else:
81
+ if hasattr(builder, "set_workspace_buffer"):
82
+ builder.set_workspace_buffer(attn_backend_workspace)
83
+ attn_groups.append(groups)
84
+ return attn_backends, attn_groups
85
86
87
def _allocate_kv_cache(kv_cache_config: KVCacheConfig, device: torch.device):
@@ -144,7 +168,7 @@ def build_slot_mappings_by_layer(
168
169
170
def build_attn_metadata(
147
- attn_metadata_builders: list[AttentionMetadataBuilder],
171
+ attn_groups: list[list[AttentionGroup]],
172
num_reqs: int,
173
num_tokens: int,
174
query_start_loc_gpu: torch.Tensor,
@@ -162,8 +186,8 @@ def build_attn_metadata(
186
dcp_local_seq_lens = dcp_local_seq_lens[:num_reqs]
187
188
attn_metadata: dict[str, Any] = {}
165
- kv_cache_groups = kv_cache_config.kv_cache_groups
166
- for i, kv_cache_spec in enumerate(kv_cache_groups):
189
+ num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)
190
+ for i in range(num_kv_cache_groups):
191
block_table = block_tables[i]
192
slot_mapping = slot_mappings[i]
193
@@ -181,10 +205,11 @@ def build_attn_metadata(
205
dcp_local_seq_lens=dcp_local_seq_lens,
206
)
207
184
- attn_metadata_builder = attn_metadata_builders[i]
185
- metadata = attn_metadata_builder.build(
186
- common_prefix_len=0, common_attn_metadata=common_attn_metadata
187
- )
188
- for layer_name in kv_cache_spec.layer_names:
189
- attn_metadata[layer_name] = metadata
208
+ for attn_group in attn_groups[i]:
209
+ attn_metadata_builder = attn_group.get_metadata_builder(0)
210
+ metadata = attn_metadata_builder.build(
211
+ common_prefix_len=0, common_attn_metadata=common_attn_metadata
212
+ )
213
+ for layer_name in attn_group.layer_names:
214
+ attn_metadata[layer_name] = metadata
215
return attn_metadata
vllm/v1/worker/gpu/cudagraph_utils.py
+17
-7
@@ -13,7 +13,6 @@ from vllm.config.compilation import CUDAGraphMode
13
from vllm.distributed.parallel_state import graph_capture, is_global_first_rank
14
from vllm.forward_context import BatchDescriptor, set_forward_context
15
from vllm.utils.math_utils import cdiv
16
-from vllm.v1.attention.backend import AttentionMetadataBuilder
16
from vllm.v1.kv_cache_interface import KVCacheConfig
17
from vllm.v1.worker.gpu.attn_utils import (
18
build_attn_metadata,
@@ -22,6 +21,7 @@ from vllm.v1.worker.gpu.attn_utils import (
21
from vllm.v1.worker.gpu.block_table import BlockTables
22
from vllm.v1.worker.gpu.dp_utils import make_num_tokens_across_dp
23
from vllm.v1.worker.gpu.input_batch import InputBuffers
24
+from vllm.v1.worker.utils import AttentionGroup
25
26
27
class CudaGraphManager:
@@ -83,7 +83,7 @@ class CudaGraphManager:
83
mrope_positions: torch.Tensor | None,
84
inputs_embeds: torch.Tensor | None,
85
block_tables: BlockTables,
86
- attn_metadata_builders: list[AttentionMetadataBuilder],
86
+ attn_groups: list[list[AttentionGroup]],
87
kv_cache_config: KVCacheConfig,
88
has_lora: bool = False,
89
uniform_decode: bool = False,
@@ -116,7 +116,7 @@ class CudaGraphManager:
116
num_tokens,
117
input_buffers,
118
block_tables,
119
- attn_metadata_builders,
119
+ attn_groups,
120
self.max_model_len,
121
kv_cache_config,
122
uniform_decode_query_len=(
@@ -232,7 +232,7 @@ class CudaGraphManager:
232
mrope_positions: torch.Tensor | None,
233
inputs_embeds: torch.Tensor | None,
234
block_tables: BlockTables,
235
- attn_metadata_builders: list[AttentionMetadataBuilder],
235
+ attn_groups: list[list[AttentionGroup]],
236
kv_cache_config: KVCacheConfig,
237
has_lora: bool = False,
238
) -> None:
@@ -244,7 +244,7 @@ class CudaGraphManager:
244
mrope_positions=mrope_positions,
245
inputs_embeds=inputs_embeds,
246
block_tables=block_tables,
247
- attn_metadata_builders=attn_metadata_builders,
247
+ attn_groups=attn_groups,
248
kv_cache_config=kv_cache_config,
249
has_lora=has_lora,
250
)
@@ -286,6 +286,16 @@ class CudaGraphManager:
286
cudagraph_mode = self.cudagraph_mode.decode_mode()
287
else:
288
cudagraph_mode = self.cudagraph_mode.mixed_mode()
289
+
290
+ if (
291
+ cudagraph_mode == CUDAGraphMode.FULL
292
+ and cudagraph_size is not None
293
+ and cudagraph_size not in self.graphs
294
+ ):
295
+ # If graph wasn't captured yet, fall back to eager.
296
+ # This might happen when the dummy run is called before capture.
297
+ cudagraph_mode = CUDAGraphMode.NONE
298
+ cudagraph_size = None
299
return cudagraph_mode, cudagraph_size
300
301
def run_fullgraph(self, num_tokens: int) -> torch.Tensor:
@@ -354,7 +364,7 @@ def prepare_inputs_to_capture(
364
num_tokens: int,
365
input_buffers: InputBuffers,
366
block_tables: BlockTables,
357
- attn_metadata_builders: list[AttentionMetadataBuilder],
367
+ attn_groups: list[list[AttentionGroup]],
368
max_model_len: int,
369
kv_cache_config: KVCacheConfig,
370
uniform_decode_query_len: int = 0,
@@ -386,7 +396,7 @@ def prepare_inputs_to_capture(
396
)
397
398
attn_metadata = build_attn_metadata(
389
- attn_metadata_builders=attn_metadata_builders,
399
+ attn_groups=attn_groups,
400
num_reqs=num_reqs,
401
num_tokens=num_tokens,
402
query_start_loc_gpu=query_start_loc,
vllm/v1/worker/gpu/model_runner.py
+5
-8
@@ -283,7 +283,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
283
cp_interleave=self.cp_interleave,
284
)
285
286
- self.attn_backends, self.attn_metadata_builders = init_attn_backend(
286
+ self.attn_backends, self.attn_groups = init_attn_backend(
287
self.kv_cache_config, self.vllm_config, self.device
288
)
289
check_attention_cp_compatibility(self.vllm_config)
@@ -291,7 +291,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
291
# HACK(woosuk)
292
self.speculator.set_attn(
293
self.kv_cache_config,
294
- self.attn_metadata_builders,
294
+ self.attn_groups,
295
self.block_tables,
296
)
297
@@ -305,9 +305,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
305
)
306
self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)
307
308
- # Attention groups are not supported.
309
- self.attn_groups = [] # type: ignore
310
-
308
def prepare_dummy_attn_metadata(self, input_batch: InputBatch) -> None:
309
block_tables = self.block_tables.get_dummy_block_tables(input_batch.num_reqs)
310
slot_mappings = self.block_tables.get_dummy_slot_mappings(
@@ -317,7 +314,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
314
slot_mappings, self.kv_cache_config
315
)
316
attn_metadata = build_attn_metadata(
320
- attn_metadata_builders=self.attn_metadata_builders,
317
+ attn_groups=self.attn_groups,
318
num_reqs=input_batch.num_reqs,
319
num_tokens=input_batch.num_tokens,
320
query_start_loc_gpu=input_batch.query_start_loc,
@@ -477,7 +474,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
474
mrope_positions=mrope_positions,
475
inputs_embeds=inputs_embeds,
476
block_tables=self.block_tables,
480
- attn_metadata_builders=self.attn_metadata_builders,
477
+ attn_groups=self.attn_groups,
478
kv_cache_config=self.kv_cache_config,
479
has_lora=self.lora_config is not None,
480
)
@@ -712,7 +709,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
709
710
# Layer name -> attention metadata.
711
attn_metadata = build_attn_metadata(
715
- attn_metadata_builders=self.attn_metadata_builders,
712
+ attn_groups=self.attn_groups,
713
num_reqs=num_reqs,
714
num_tokens=num_tokens,
715
query_start_loc_gpu=query_start_loc,
vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py
+5
-5
@@ -7,7 +7,6 @@ import torch
7
8
from vllm.config import VllmConfig
9
from vllm.config.compilation import CUDAGraphMode
10
-from vllm.v1.attention.backend import AttentionMetadataBuilder
10
from vllm.v1.kv_cache_interface import KVCacheConfig
11
from vllm.v1.worker.gpu.block_table import BlockTables
12
from vllm.v1.worker.gpu.cudagraph_utils import (
@@ -17,6 +16,7 @@ from vllm.v1.worker.gpu.cudagraph_utils import (
16
)
17
from vllm.v1.worker.gpu.dp_utils import make_num_tokens_across_dp
18
from vllm.v1.worker.gpu.input_batch import InputBuffers
19
+from vllm.v1.worker.utils import AttentionGroup
20
21
22
class EagleCudaGraphManager:
@@ -60,7 +60,7 @@ class EagleCudaGraphManager:
60
generate_fn: Callable,
61
input_buffers: InputBuffers,
62
block_tables: BlockTables,
63
- attn_metadata_builders: list[AttentionMetadataBuilder],
63
+ attn_groups: list[list[AttentionGroup]],
64
kv_cache_config: KVCacheConfig,
65
) -> None:
66
assert capture_cg_mode in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL], (
@@ -77,7 +77,7 @@ class EagleCudaGraphManager:
77
num_tokens,
78
input_buffers,
79
block_tables,
80
- attn_metadata_builders,
80
+ attn_groups,
81
self.max_model_len,
82
kv_cache_config,
83
uniform_decode_query_len=1,
@@ -150,7 +150,7 @@ class EagleCudaGraphManager:
150
generate_fn: Callable,
151
input_buffers: InputBuffers,
152
block_tables: BlockTables,
153
- attn_metadata_builders: list[AttentionMetadataBuilder],
153
+ attn_groups: list[list[AttentionGroup]],
154
kv_cache_config: KVCacheConfig,
155
) -> None:
156
if self.cudagraph_mode == CUDAGraphMode.NONE:
@@ -165,7 +165,7 @@ class EagleCudaGraphManager:
165
generate_fn=generate_fn,
166
input_buffers=input_buffers,
167
block_tables=block_tables,
168
- attn_metadata_builders=attn_metadata_builders,
168
+ attn_groups=attn_groups,
169
kv_cache_config=kv_cache_config,
170
)
171
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+5
-5
@@ -10,7 +10,6 @@ from vllm.config.compilation import CUDAGraphMode
10
from vllm.forward_context import BatchDescriptor, set_forward_context
11
from vllm.logger import init_logger
12
from vllm.triton_utils import tl, triton
13
-from vllm.v1.attention.backend import AttentionMetadataBuilder
13
from vllm.v1.kv_cache_interface import KVCacheConfig
14
from vllm.v1.worker.gpu.attn_utils import (
15
build_attn_metadata,
@@ -21,6 +20,7 @@ from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
20
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
21
from vllm.v1.worker.gpu.spec_decode.eagle.cudagraph import EagleCudaGraphManager
22
from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
23
+from vllm.v1.worker.utils import AttentionGroup
24
25
logger = init_logger(__name__)
26
@@ -78,11 +78,11 @@ class EagleSpeculator:
78
def set_attn(
79
self,
80
kv_cache_config: KVCacheConfig,
81
- attn_metadata_builders: list[AttentionMetadataBuilder],
81
+ attn_groups: list[list[AttentionGroup]],
82
block_tables: BlockTables,
83
) -> None:
84
self.kv_cache_config = kv_cache_config
85
- self.attn_metadata_builders = attn_metadata_builders
85
+ self.attn_groups = attn_groups
86
self.block_tables = block_tables
87
88
@torch.inference_mode()
@@ -174,7 +174,7 @@ class EagleSpeculator:
174
self.generate_draft,
175
self.input_buffers,
176
self.block_tables,
177
- self.attn_metadata_builders,
177
+ self.attn_groups,
178
self.kv_cache_config,
179
)
180
@@ -298,7 +298,7 @@ class EagleSpeculator:
298
299
# FIXME(woosuk): This is UNSAFE!!
300
attn_metadata = build_attn_metadata(
301
- attn_metadata_builders=self.attn_metadata_builders,
301
+ attn_groups=self.attn_groups,
302
num_reqs=num_reqs,
303
num_tokens=num_reqs,
304
query_start_loc_gpu=query_start_loc,