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 2, 17:01 UTC
[ModelRunnerV2][Hybrid model] Support kernel block size in hybrid model (#38831)
Signed-off-by: MengqingCao <cmq0113@163.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Signed-off-by: Mengqing Cao <cmq0113@163.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Mengqing Cao committed
May 28, 2026 at 08:55 UTC
1223732dda9d0616969397dd78697be1d40941b6
5 files changed
+82
-75
vllm/config/vllm.py
+6
-2
@@ -1946,8 +1946,12 @@ class VllmConfig:
1946
model_config = self.model_config
1947
speculative_config = self.speculative_config
1948
1949
- if model_config is not None and model_config.has_inner_state:
1950
- unsupported.append("hybrid/mamba models")
1949
+ if (
1950
+ model_config is not None
1951
+ and model_config.has_inner_state
1952
+ and self.cache_config.mamba_cache_mode == "align"
1953
+ ):
1954
+ unsupported.append("hybrid/mamba models with align cache mode")
1955
1956
if self.parallel_config.prefill_context_parallel_size > 1:
1957
unsupported.append("prefill context parallelism")
vllm/v1/worker/gpu/attn_utils.py
+70
-63
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from collections.abc import Sequence
3
+from collections.abc import Iterable, Sequence
4
from dataclasses import dataclass
5
from typing import Any, cast
6
@@ -11,7 +11,6 @@ from vllm.model_executor.layers.attention import Attention
11
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
12
from vllm.utils.torch_utils import get_dtype_size
13
from vllm.v1.attention.backend import (
14
- AttentionBackend,
14
AttentionCGSupport,
15
CommonAttentionMetadata,
16
)
@@ -65,13 +64,8 @@ def init_attn_backend(
64
vllm_config: VllmConfig,
65
device: torch.device,
66
active_layer_names: set[str] | None = None,
68
-) -> tuple[
69
- dict[str, type[AttentionBackend]],
70
- list[list[AttentionGroup]],
71
- AttentionCGSupportInfo,
72
- list[int],
73
-]:
74
- attn_backends: dict[str, type[AttentionBackend]] = {}
67
+) -> tuple[list[list[AttentionGroup]], AttentionCGSupportInfo, list[int]]:
68
+ # Phase 1: discover attention groups for each kv cache group.
69
attn_groups: list[list[AttentionGroup]] = []
70
71
# Add KV-sharing layers to their target's kv cache group so they are
@@ -96,7 +90,6 @@ def init_attn_backend(
90
91
for layer_name in layer_names:
92
attn_backend = attn_layers[layer_name].get_attn_backend()
99
- attn_backends[layer_name] = attn_backend
93
94
layer_kv_cache_spec: KVCacheSpec = kv_cache_group_spec.kv_cache_spec
95
if isinstance(layer_kv_cache_spec, UniformTypeKVCacheSpecs):
@@ -122,7 +115,6 @@ def init_attn_backend(
115
min_cg_support = AttentionCGSupport.ALWAYS
116
min_cg_attn_backend = None
117
for kv_cache_group_id, groups in enumerate(attn_groups):
125
- kv_cache_group_spec = kv_cache_config.kv_cache_groups[kv_cache_group_id]
118
kernel_block_size = None
119
if kv_cache_group_id < len(kernel_block_sizes):
120
kernel_block_size = kernel_block_sizes[kv_cache_group_id]
@@ -143,21 +135,16 @@ def init_attn_backend(
135
# Check cudagraph support for the attention backend
136
cg_support = builder.get_cudagraph_support(
137
vllm_config,
146
- cast(AttentionSpec, kv_cache_group_spec.kv_cache_spec),
138
+ cast(AttentionSpec, group.kv_cache_spec),
139
)
140
if cg_support.value < min_cg_support.value:
141
min_cg_support = cg_support
142
min_cg_attn_backend = group.backend.__name__
143
152
- return (
153
- attn_backends,
154
- attn_groups,
155
- AttentionCGSupportInfo(
156
- min_cg_support=min_cg_support,
157
- min_cg_attn_backend=min_cg_attn_backend,
158
- ),
159
- kernel_block_sizes,
144
+ attn_cg_support_info = AttentionCGSupportInfo(
145
+ min_cg_support=min_cg_support, min_cg_attn_backend=min_cg_attn_backend
146
)
147
+ return attn_groups, attn_cg_support_info, kernel_block_sizes
148
149
150
def _allocate_kv_cache(
@@ -180,25 +167,31 @@ def _allocate_kv_cache(
167
168
169
def _reshape_kv_cache(
183
- kv_cache_config: KVCacheConfig,
170
+ attn_groups: Sequence[AttentionGroup],
171
kv_cache_raw_tensors: dict[str, torch.Tensor],
185
- attn_backends: dict[str, type[AttentionBackend]],
172
cache_dtype: str,
173
kernel_block_sizes: list[int],
174
shared_kv_cache_layers: dict[str, str],
175
) -> dict[str, Any]:
176
kv_caches: dict[str, Any] = {}
177
has_attn, has_mamba = False, False
192
- for kv_cache_group_id, kv_cache_group_spec in enumerate(
193
- kv_cache_config.kv_cache_groups
194
- ):
195
- for layer_name in kv_cache_group_spec.layer_names:
178
+
179
+ for group in attn_groups:
180
+ if group.kv_cache_group_id >= len(kernel_block_sizes):
181
+ continue
182
+
183
+ kv_cache_spec = group.kv_cache_spec
184
+ if kv_cache_spec.storage_block_size != kv_cache_spec.block_size:
185
+ # use storage_block_size as the kernel block size for groups
186
+ # that apply a compression on block size (eg. DeepSeek V4).
187
+ kernel_block_size = kv_cache_spec.storage_block_size
188
+ else:
189
+ kernel_block_size = kernel_block_sizes[group.kv_cache_group_id]
190
+
191
+ for layer_name in group.layer_names:
192
if layer_name in shared_kv_cache_layers:
193
# Shared layer — tensor will be aliased to its target later.
194
continue
199
- kv_cache_spec = kv_cache_group_spec.kv_cache_spec
200
- if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
201
- kv_cache_spec = kv_cache_spec.kv_cache_specs[layer_name]
195
196
kv_raw_tensor = kv_cache_raw_tensors[layer_name]
197
assert kv_raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0
@@ -206,22 +199,11 @@ def _reshape_kv_cache(
199
200
if isinstance(kv_cache_spec, AttentionSpec):
201
has_attn = True
209
- attn_backend = attn_backends[layer_name]
210
-
211
- if kv_cache_group_id < len(kernel_block_sizes):
212
- kernel_block_size = kernel_block_sizes[kv_cache_group_id]
213
- num_blocks *= kv_cache_spec.block_size // kernel_block_size
214
- else:
215
- kernel_block_size = kv_cache_spec.block_size
216
-
217
- if kv_cache_spec.storage_block_size != kv_cache_spec.block_size:
218
- shape_block_size = kv_cache_spec.storage_block_size
219
- else:
220
- shape_block_size = kernel_block_size
221
-
222
- kv_cache_shape = attn_backend.get_kv_cache_shape(
223
- num_blocks,
224
- shape_block_size,
202
+ num_blocks_per_kv_block = kv_cache_spec.block_size // kernel_block_size
203
+ kernel_num_blocks = num_blocks * num_blocks_per_kv_block
204
+ kv_cache_shape = group.backend.get_kv_cache_shape(
205
+ kernel_num_blocks,
206
+ kernel_block_size,
207
kv_cache_spec.num_kv_heads,
208
kv_cache_spec.head_size,
209
cache_dtype_str=cache_dtype,
@@ -229,7 +211,7 @@ def _reshape_kv_cache(
211
212
# FIXME(woosuk): Add kv_cache_stride_order to all attention backends.
213
try:
232
- kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
214
+ kv_cache_stride_order = group.backend.get_kv_cache_stride_order()
215
assert len(kv_cache_stride_order) == len(kv_cache_shape)
216
except (AttributeError, NotImplementedError):
217
kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
@@ -290,7 +272,12 @@ def _reshape_kv_cache(
272
)
273
274
if has_attn and has_mamba:
293
- _update_hybrid_attention_layout(kv_caches, kv_cache_config)
275
+ _update_hybrid_attention_layout(
276
+ attn_groups=attn_groups,
277
+ kv_caches=kv_caches,
278
+ kernel_block_sizes=kernel_block_sizes,
279
+ cache_dtype=cache_dtype,
280
+ )
281
282
# Map any sharing layers to their target layer's KV cache.
283
for layer_name, target_layer_name in shared_kv_cache_layers.items():
@@ -300,18 +287,38 @@ def _reshape_kv_cache(
287
288
289
def _update_hybrid_attention_layout(
303
- kv_caches: dict[str, Any], kv_cache_config: KVCacheConfig
290
+ attn_groups: Iterable[AttentionGroup],
291
+ kv_caches: dict[str, Any],
292
+ kernel_block_sizes: list[int],
293
+ cache_dtype: str,
294
) -> None:
305
- for kv_cache_group_spec in kv_cache_config.kv_cache_groups:
306
- for layer_name in kv_cache_group_spec.layer_names:
295
+ for group in attn_groups:
296
+ if group.kv_cache_group_id >= len(kernel_block_sizes):
297
+ continue
298
+
299
+ kv_cache_spec = group.kv_cache_spec
300
+ if not isinstance(kv_cache_spec, AttentionSpec):
301
+ continue
302
+ block_dim = group.backend.get_kv_cache_block_dim(
303
+ kernel_block_sizes[group.kv_cache_group_id],
304
+ kv_cache_spec.num_kv_heads,
305
+ kv_cache_spec.head_size,
306
+ cache_dtype_str=cache_dtype,
307
+ )
308
+ # if the first dim of the kvcache's layout is already num_blocks, continue
309
+ if block_dim == 0:
310
+ continue
311
+
312
+ assert block_dim == 1, (
313
+ "Expected the dim `num_blocks` at the second dim when updating"
314
+ " the kvcache's layout of full attention layer"
315
+ )
316
+
317
+ for layer_name in group.layer_names:
318
if layer_name not in kv_caches:
319
# Shared layer — will be aliased to its target after this pass.
320
continue
310
- kv_cache_spec = kv_cache_group_spec.kv_cache_spec
311
- if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
312
- kv_cache_spec = kv_cache_spec.kv_cache_specs[layer_name]
313
- if not isinstance(kv_cache_spec, AttentionSpec):
314
- continue
321
+
322
kv_cache = kv_caches[layer_name]
323
if kv_cache.shape[0] == 2:
324
assert kv_cache.shape[1] != 2, (
@@ -329,10 +336,10 @@ def _update_hybrid_attention_layout(
336
337
338
def init_kv_cache(
332
- runner_kv_caches: list[torch.Tensor],
339
+ runner_kv_caches: list[torch.Tensor | list[torch.Tensor]],
340
forward_context: dict[str, Any],
341
kv_cache_config: KVCacheConfig,
335
- attn_backends: dict[str, type[AttentionBackend]],
342
+ attn_groups: list[list[AttentionGroup]],
343
device: torch.device,
344
cache_dtype: str,
345
kernel_block_sizes: list[int],
@@ -342,13 +349,13 @@ def init_kv_cache(
349
kv_cache_raw_tensors = _allocate_kv_cache(
350
kv_cache_config, shared_kv_cache_layers, device
351
)
352
+ flattened_attn_groups = list(group for groups in attn_groups for group in groups)
353
kv_caches = _reshape_kv_cache(
346
- kv_cache_config,
347
- kv_cache_raw_tensors,
348
- attn_backends,
349
- cache_dtype,
350
- kernel_block_sizes,
351
- shared_kv_cache_layers,
354
+ attn_groups=flattened_attn_groups,
355
+ kv_cache_raw_tensors=kv_cache_raw_tensors,
356
+ kernel_block_sizes=kernel_block_sizes,
357
+ cache_dtype=cache_dtype,
358
+ shared_kv_cache_layers=shared_kv_cache_layers,
359
)
360
bind_kv_cache(kv_caches, forward_context, runner_kv_caches)
361
return kv_caches
vllm/v1/worker/gpu/block_table.py
+2
-5
@@ -44,9 +44,7 @@ class BlockTables:
44
for i in range(self.num_kv_cache_groups):
45
max_num_blocks = max_num_blocks_per_group[i] * self.blocks_per_kv_block[i]
46
block_table = StagedWriteTensor(
47
- (self.max_num_reqs, max_num_blocks),
48
- dtype=torch.int32,
49
- device=device,
47
+ (self.max_num_reqs, max_num_blocks), dtype=torch.int32, device=device
48
)
49
self.block_tables.append(block_table)
50
@@ -132,7 +130,6 @@ class BlockTables:
130
self.num_blocks.gpu,
131
self.num_blocks.gpu.stride(0),
132
num_reqs,
135
- self.input_block_tables[0].shape[1], # max_num_blocks
133
BLOCK_SIZE=1024, # type: ignore
134
)
135
return tuple(bt[:num_reqs_padded] for bt in self.input_block_tables)
@@ -192,7 +189,6 @@ def _gather_block_tables_kernel(
189
num_blocks_ptr, # [num_kv_cache_groups, max_num_reqs]
190
num_blocks_stride,
191
num_reqs, # actual number of requests (for padding)
195
- max_num_blocks, # stride for zeroing padded rows
192
BLOCK_SIZE: tl.constexpr,
193
):
194
# kv cache group id
@@ -200,6 +196,7 @@ def _gather_block_tables_kernel(
196
batch_idx = tl.program_id(1)
197
198
stride = tl.load(block_table_strides + group_id)
199
+ max_num_blocks = stride # stride equals max_num_blocks for this group.
200
dst_block_table_ptr = _load_ptr(dst_block_table_ptrs + group_id, tl.int32)
201
dst_row_ptr = dst_block_table_ptr + batch_idx * stride
202
vllm/v1/worker/gpu/model_runner.py
+3
-4
@@ -393,10 +393,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
393
) + spec.num_speculative_blocks
394
max_num_blocks_per_group.append(max_num_blocks)
395
396
- (self.attn_backends, self.attn_groups, attn_cg_support, kernel_block_sizes) = (
397
- init_attn_backend(self.kv_cache_config, self.vllm_config, self.device)
396
+ self.attn_groups, attn_cg_support, kernel_block_sizes = init_attn_backend(
397
+ self.kv_cache_config, self.vllm_config, self.device
398
)
399
-
399
self.block_tables = BlockTables(
400
block_sizes=block_sizes,
401
max_num_reqs=self.max_num_reqs,
@@ -440,7 +439,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
439
self.kv_caches,
440
self.compilation_config.static_forward_context,
441
self.kv_cache_config,
443
- self.attn_backends,
442
+ self.attn_groups,
443
self.device,
444
self.cache_config.cache_dtype,
445
kernel_block_sizes,
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+1
-1
@@ -173,7 +173,7 @@ class EagleSpeculator:
173
) -> None:
174
self.model_state = model_state
175
self.kv_cache_config = kv_cache_config
176
- _, self.attn_groups, _, _ = init_attn_backend(
176
+ self.attn_groups, _, _ = init_attn_backend(
177
kv_cache_config,
178
self.vllm_config,
179
self.device,