biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[kv_offload+HMA] Fix num_blocks with different per-layer page sizes and improve assert message (#38554)

Browse Source 
Signed-off-by: Kfir Toledo <kfir.toledo@ibm.com>
Co-authored-by: Or Ozeri <oro@il.ibm.com>
This commit is contained in:
Kfir Toledo
2026-03-31 09:00:40 +03:00
committed by GitHub
parent d53cb9cb8e
commit 6cc7abdc66
2 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
View File

@@ -83,6 +83,8 @@ class OffloadingConnectorWorker:
if layer_name in layers
}
num_blocks = self.spec.kv_cache_config.num_blocks
# layer_name -> list of matching KV cache tensors
# such that each tensor starts with the num_blocks dimension.
# FlashAttention layers which use the (2, num_blocks, ...) layout
@@ -132,7 +134,6 @@ class OffloadingConnectorWorker:
num_blocks_logical_dim
)
if num_blocks_physical_dim == 0:
num_blocks = layer_kv_cache.shape[num_blocks_logical_dim]
storage = layer_kv_cache.untyped_storage()
page = layer_kv_cache_spec.page_size_bytes
tensors_per_block[layer_name] = (
@@ -154,7 +155,6 @@ class OffloadingConnectorWorker:
assert num_blocks_physical_dim == 1
# unbind the tensor to separate K and V tensors
num_blocks = layer_kv_cache.shape[num_blocks_logical_dim]
half_page_size = layer_kv_cache_spec.page_size_bytes // 2
storage = layer_kv_cache.untyped_storage()
raw = (
@@ -181,7 +181,6 @@ class OffloadingConnectorWorker:
assert len(state_tensors) > 0
first_state_tensor = state_tensors[0]
assert first_state_tensor.storage_offset() == 0
num_blocks = first_state_tensor.shape[0]
tensor = (
torch.tensor(
[],

7
vllm/v1/kv_offload/spec.py
View File

@@ -93,7 +93,12 @@ class OffloadingSpec(ABC):
)
for block_size in self.gpu_block_size:
assert block_size % self.hash_block_size == 0
assert block_size % self.hash_block_size == 0, (
f"gpu_block_size={block_size} not divisible by "
f"hash_block_size={self.hash_block_size}. "
f"Hybrid models (e.g. Mamba+Attention) need "
f"--enable-prefix-caching to align block sizes."
)
# offloaded_block_size / gpu_block_size
self.block_size_factor: int = 1