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 3, 09:15 UTC
[Bugfix][MRV2] Fix KVCache tensor explicit `kernel_block_size` dim (#42766)
Signed-off-by: NickLucche <nlucches@redhat.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Nicolò Lucchesi committed
May 19, 2026 at 05:25 UTC
fba010dd74e2f94e4f7223b164ec9097d1b8a6af
6 files changed
+68
-30
tests/test_config.py
-7
@@ -148,13 +148,6 @@ def test_is_default_v2_model_runner_model(model_config, expected):
148
assert VllmConfig._is_default_v2_model_runner_model(config) is expected
149
150
151
-def test_use_v2_model_runner_defaults_to_v1_when_kv_connector_present():
152
- config = SimpleNamespace(kv_transfer_config=object())
153
- with patch.object(envs, "VLLM_USE_V2_MODEL_RUNNER", None):
154
- result = VllmConfig.use_v2_model_runner.fget(config)
155
- assert result is False
156
-
157
-
151
@pytest.mark.skip_global_cleanup
152
def test_with_hf_config_populates_missing_architectures_from_causal_lm_mapping(
153
monkeypatch,
vllm/config/vllm.py
-4
@@ -494,10 +494,6 @@ class VllmConfig:
494
if use_v2_model_runner is not None:
495
return use_v2_model_runner
496
497
- # KVCache layout changes are breaking, let's stick with v1 for now (see #42846)
498
- if self.kv_transfer_config is not None:
499
- return False
500
-
497
if not self._is_default_v2_model_runner_model():
498
return False
499
vllm/v1/worker/gpu/attn_utils.py
+49
-12
@@ -22,7 +22,11 @@ from vllm.v1.kv_cache_interface import (
22
UniformTypeKVCacheSpecs,
23
)
24
from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
25
-from vllm.v1.worker.utils import AttentionGroup, bind_kv_cache
25
+from vllm.v1.worker.utils import (
26
+ AttentionGroup,
27
+ bind_kv_cache,
28
+ prepare_kernel_block_sizes,
29
+)
30
31
32
@dataclass(frozen=True)
@@ -51,13 +55,12 @@ def init_attn_backend(
55
dict[str, type[AttentionBackend]],
56
list[list[AttentionGroup]],
57
AttentionCGSupportInfo,
58
+ list[int],
59
]:
60
attn_backends: dict[str, type[AttentionBackend]] = {}
61
attn_groups: list[list[AttentionGroup]] = []
57
- attn_backend_workspace: torch.Tensor | None = None
58
- # Find minimum cudagraph support across all attention backends
59
- min_cg_support = AttentionCGSupport.ALWAYS
60
- min_cg_attn_backend = None
62
+
63
+ # Phase 1: discover attention groups for each kv cache group.
64
for kv_cache_group_id, kv_cache_group_spec in enumerate(
65
kv_cache_config.kv_cache_groups
66
):
@@ -91,12 +94,26 @@ def init_attn_backend(
94
else:
95
group_map[key].layer_names.append(layer_name)
96
94
- groups = [group_map[key] for key in group_order]
97
+ attn_groups.append([group_map[key] for key in group_order])
98
+
99
+ # Phase 2: pick a kernel block size per kv cache group that is supported
100
+ # by all backends within that group.
101
+ kernel_block_sizes = prepare_kernel_block_sizes(kv_cache_config, attn_groups)
102
+
103
+ # Phase 3: create metadata builders and determine cudagraph support.
104
+ attn_backend_workspace: torch.Tensor | None = None
105
+ min_cg_support = AttentionCGSupport.ALWAYS
106
+ min_cg_attn_backend = None
107
+ for kv_cache_group_id, groups in enumerate(attn_groups):
108
+ kv_cache_group_spec = kv_cache_config.kv_cache_groups[kv_cache_group_id]
109
+ kernel_block_size = None
110
+ if kv_cache_group_id < len(kernel_block_sizes):
111
+ kernel_block_size = kernel_block_sizes[kv_cache_group_id]
112
for group in groups:
113
group.create_metadata_builders(
114
vllm_config=vllm_config,
115
device=device,
99
- kernel_block_size=None,
116
+ kernel_block_size=kernel_block_size,
117
num_metadata_builders=1,
118
)
119
builder = group.get_metadata_builder(0)
@@ -113,8 +130,7 @@ def init_attn_backend(
130
)
131
if cg_support.value < min_cg_support.value:
132
min_cg_support = cg_support
116
- min_cg_attn_backend = attn_backend.__name__
117
- attn_groups.append(groups)
133
+ min_cg_attn_backend = group.backend.__name__
134
135
return (
136
attn_backends,
@@ -123,6 +139,7 @@ def init_attn_backend(
139
min_cg_support=min_cg_support,
140
min_cg_attn_backend=min_cg_attn_backend,
141
),
142
+ kernel_block_sizes,
143
)
144
145
@@ -148,10 +165,13 @@ def _reshape_kv_cache(
165
kv_cache_raw_tensors: dict[str, torch.Tensor],
166
attn_backends: dict[str, type[AttentionBackend]],
167
cache_dtype: str,
168
+ kernel_block_sizes: list[int],
169
) -> dict[str, Any]:
170
kv_caches: dict[str, Any] = {}
171
has_attn, has_mamba = False, False
154
- for kv_cache_group_spec in kv_cache_config.kv_cache_groups:
172
+ for kv_cache_group_id, kv_cache_group_spec in enumerate(
173
+ kv_cache_config.kv_cache_groups
174
+ ):
175
for layer_name in kv_cache_group_spec.layer_names:
176
kv_cache_spec = kv_cache_group_spec.kv_cache_spec
177
if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
@@ -164,9 +184,21 @@ def _reshape_kv_cache(
184
if isinstance(kv_cache_spec, AttentionSpec):
185
has_attn = True
186
attn_backend = attn_backends[layer_name]
187
+
188
+ if kv_cache_group_id < len(kernel_block_sizes):
189
+ kernel_block_size = kernel_block_sizes[kv_cache_group_id]
190
+ num_blocks *= kv_cache_spec.block_size // kernel_block_size
191
+ else:
192
+ kernel_block_size = kv_cache_spec.block_size
193
+
194
+ if kv_cache_spec.storage_block_size != kv_cache_spec.block_size:
195
+ shape_block_size = kv_cache_spec.storage_block_size
196
+ else:
197
+ shape_block_size = kernel_block_size
198
+
199
kv_cache_shape = attn_backend.get_kv_cache_shape(
200
num_blocks,
169
- kv_cache_spec.storage_block_size,
201
+ shape_block_size,
202
kv_cache_spec.num_kv_heads,
203
kv_cache_spec.head_size,
204
cache_dtype_str=cache_dtype,
@@ -275,10 +307,15 @@ def init_kv_cache(
307
attn_backends: dict[str, type[AttentionBackend]],
308
device: torch.device,
309
cache_dtype: str,
310
+ kernel_block_sizes: list[int],
311
) -> dict[str, Any]:
312
kv_cache_raw_tensors = _allocate_kv_cache(kv_cache_config, device)
313
kv_caches = _reshape_kv_cache(
281
- kv_cache_config, kv_cache_raw_tensors, attn_backends, cache_dtype
314
+ kv_cache_config,
315
+ kv_cache_raw_tensors,
316
+ attn_backends,
317
+ cache_dtype,
318
+ kernel_block_sizes,
319
)
320
bind_kv_cache(kv_caches, forward_context, runner_kv_caches)
321
return kv_caches
vllm/v1/worker/gpu/block_table.py
+12
-2
@@ -17,11 +17,13 @@ class BlockTables:
17
max_num_batched_tokens: int,
18
max_num_blocks_per_group: list[int],
19
device: torch.device,
20
+ kernel_block_sizes: list[int],
21
cp_size: int = 1,
22
cp_rank: int = 0,
23
cp_interleave: int = 1,
24
):
25
self.block_sizes = block_sizes
26
+ self.kernel_block_sizes = kernel_block_sizes
27
self.max_num_reqs = max_num_reqs
28
self.max_num_batched_tokens = max_num_batched_tokens
29
self.device = device
@@ -32,10 +34,15 @@ class BlockTables:
34
35
self.num_kv_cache_groups = len(self.block_sizes)
36
assert len(max_num_blocks_per_group) == self.num_kv_cache_groups
37
+
38
+ self.blocks_per_kv_block = [
39
+ bs // kbs for bs, kbs in zip(block_sizes, kernel_block_sizes)
40
+ ]
41
+
42
# num_kv_cache_groups x [max_num_reqs, max_num_blocks]
43
self.block_tables: list[StagedWriteTensor] = []
44
for i in range(self.num_kv_cache_groups):
38
- max_num_blocks = max_num_blocks_per_group[i]
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,
@@ -84,7 +91,7 @@ class BlockTables:
91
device=self.device,
92
)
93
self.block_sizes_tensor = torch.tensor(
87
- self.block_sizes, dtype=torch.int32, device=self.device
94
+ self.kernel_block_sizes, dtype=torch.int32, device=self.device
95
)
96
self.input_block_table_ptrs = self._make_ptr_tensor(self.input_block_tables)
97
@@ -97,6 +104,9 @@ class BlockTables:
104
for i in range(self.num_kv_cache_groups):
105
start = self.num_blocks.np[i, req_index] if not overwrite else 0
106
block_ids = new_block_ids[i]
107
+ bpk = self.blocks_per_kv_block[i]
108
+ if bpk > 1:
109
+ block_ids = [b * bpk + k for b in block_ids for k in range(bpk)]
110
self.block_tables[i].stage_write(req_index, start, block_ids)
111
self.num_blocks.np[i, req_index] = start + len(block_ids)
112
vllm/v1/worker/gpu/model_runner.py
+6
-4
@@ -392,20 +392,21 @@ class GPUModelRunner(LoRAModelRunnerMixin):
392
) + spec.num_speculative_blocks
393
max_num_blocks_per_group.append(max_num_blocks)
394
395
+ (self.attn_backends, self.attn_groups, attn_cg_support, kernel_block_sizes) = (
396
+ init_attn_backend(self.kv_cache_config, self.vllm_config, self.device)
397
+ )
398
+
399
self.block_tables = BlockTables(
400
block_sizes=block_sizes,
401
max_num_reqs=self.max_num_reqs,
402
max_num_batched_tokens=self.max_num_tokens,
403
max_num_blocks_per_group=max_num_blocks_per_group,
404
device=self.device,
405
+ kernel_block_sizes=kernel_block_sizes,
406
cp_size=self.dcp_size,
407
cp_rank=self.dcp_rank,
408
cp_interleave=self.cp_interleave,
409
)
405
-
406
- self.attn_backends, self.attn_groups, attn_cg_support = init_attn_backend(
407
- self.kv_cache_config, self.vllm_config, self.device
408
- )
410
initialize_mamba_ssu_backend(
411
self.vllm_config.mamba_config, self.kv_cache_config
412
)
@@ -443,6 +444,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
444
self.attn_backends,
445
self.device,
446
self.cache_config.cache_dtype,
447
+ kernel_block_sizes,
448
)
449
self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)
450
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,