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, 16:30 UTC
[BugFix][kv_offload]: Fix kernel block size detection (#35125)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
Or Ozeri committed
Feb 26, 2026 at 18:29 UTC
f2ad952f40a98e0bb7f89763c51a73124ccc20a6
1 file changed
+14
-10
vllm/v1/kv_offload/worker/cpu_gpu.py
+14
-10
@@ -259,16 +259,20 @@ class CpuGpuOffloadingHandlers:
259
assert gpu_shape[0] == 2
260
split_k_and_v = True
261
262
- try:
263
- kv_cache_stride_order = attn_backend.get_kv_cache_stride_order(
264
- include_num_layers_dimension=has_layers_dim
265
- )
266
- assert len(kv_cache_stride_order) == len(gpu_shape)
267
- except (AttributeError, NotImplementedError):
268
- kv_cache_stride_order = tuple(range(len(gpu_shape)))
269
-
270
- # permute test_shape according to stride_order
271
- test_shape = tuple(test_shape[i] for i in kv_cache_stride_order)
262
+ if has_layers_dim:
263
+ # in the cross layers case, the registered kv cache tensor
264
+ # shape matches the physical layout, whereas test_shape
265
+ # is the logical layout.
266
+ # To match them, we need to permute test_shape
267
+ try:
268
+ kv_cache_stride_order = attn_backend.get_kv_cache_stride_order(
269
+ include_num_layers_dimension=has_layers_dim
270
+ )
271
+ assert len(kv_cache_stride_order) == len(gpu_shape)
272
+ except (AttributeError, NotImplementedError):
273
+ kv_cache_stride_order = tuple(range(len(gpu_shape)))
274
+
275
+ test_shape = tuple(test_shape[i] for i in kv_cache_stride_order)
276
277
# find block_size (16) dimension index
278
block_size_idx = test_shape.index(16)