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 5, 18:30 UTC
[kv_offload+HMA] Fix num_blocks with different per-layer page sizes and improve assert message (#38554)
Signed-off-by: Kfir Toledo <kfir.toledo@ibm.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
Kfir Toledo committed
Mar 31, 2026 at 09:00 UTC
6cc7abdc66b98ef66141dd89f94f5f265e427f23
2 files changed
+8
-4
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+2
-3
@@ -83,6 +83,8 @@ class OffloadingConnectorWorker:
83
if layer_name in layers
84
}
85
86
+ num_blocks = self.spec.kv_cache_config.num_blocks
87
+
88
# layer_name -> list of matching KV cache tensors
89
# such that each tensor starts with the num_blocks dimension.
90
# FlashAttention layers which use the (2, num_blocks, ...) layout
@@ -132,7 +134,6 @@ class OffloadingConnectorWorker:
134
num_blocks_logical_dim
135
)
136
if num_blocks_physical_dim == 0:
135
- num_blocks = layer_kv_cache.shape[num_blocks_logical_dim]
137
storage = layer_kv_cache.untyped_storage()
138
page = layer_kv_cache_spec.page_size_bytes
139
tensors_per_block[layer_name] = (
@@ -154,7 +155,6 @@ class OffloadingConnectorWorker:
155
assert num_blocks_physical_dim == 1
156
157
# unbind the tensor to separate K and V tensors
157
- num_blocks = layer_kv_cache.shape[num_blocks_logical_dim]
158
half_page_size = layer_kv_cache_spec.page_size_bytes // 2
159
storage = layer_kv_cache.untyped_storage()
160
raw = (
@@ -181,7 +181,6 @@ class OffloadingConnectorWorker:
181
assert len(state_tensors) > 0
182
first_state_tensor = state_tensors[0]
183
assert first_state_tensor.storage_offset() == 0
184
- num_blocks = first_state_tensor.shape[0]
184
tensor = (
185
torch.tensor(
186
[],
vllm/v1/kv_offload/spec.py
+6
-1
@@ -93,7 +93,12 @@ class OffloadingSpec(ABC):
93
)
94
95
for block_size in self.gpu_block_size:
96
- assert block_size % self.hash_block_size == 0
96
+ assert block_size % self.hash_block_size == 0, (
97
+ f"gpu_block_size={block_size} not divisible by "
98
+ f"hash_block_size={self.hash_block_size}. "
99
+ f"Hybrid models (e.g. Mamba+Attention) need "
100
+ f"--enable-prefix-caching to align block sizes."
101
+ )
102
103
# offloaded_block_size / gpu_block_size
104
self.block_size_factor: int = 1