[Bugfix] Fix handling 5D KV cache in kv_postprocess_layout_on_receive (#47791)
Signed-off-by: Daniel Socek <daniel.socek@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Daniel Socek committed
Jul 26, 2026 at 10:22 UTC
7154856f3dcb1d3fdd5a136f7d2c5987f22244f5
1 file changed
+5
-1
vllm/distributed/kv_transfer/kv_connector/utils.py
+5
-1
@@ -258,8 +258,12 @@ def kv_postprocess_layout_on_receive(cache, indices):
258
This method corrects layout mismatches from direct memory copies by
259
permuting the tensor dimensions.
260
261
+ 4D cache:
262
- **Source Layout:** `[num_blocks, n_kv_head, block_size, head_dim]`
263
- **Target Layout:** `[num_blocks, block_size, n_kv_head, head_dim]`
264
+ 5D cache:
265
+ - **Source Layout:** `[num_blocks, kv_dim, n_kv_head, block_size, head_dim]`
266
+ - **Target Layout:** `[num_blocks, kv_dim, block_size, n_kv_head, head_dim]`
267
268
Implementation:
269
- x = blocks_to_update.reshape(src_shape) # view local kv with sender layout
@@ -270,7 +274,7 @@ def kv_postprocess_layout_on_receive(cache, indices):
274
blocks_to_update = cache.index_select(0, indices)
275
target_shape = list(blocks_to_update.shape)
276
target_shape[0] = -1
273
- inv_order = [0, 2, 1, 3]
277
+ inv_order = [0, 1, 3, 2, 4] if blocks_to_update.ndim == 5 else [0, 2, 1, 3]
278
src_shape = tuple(target_shape[i] for i in inv_order)
279
blocks_to_update = cache.index_select(0, indices)
280
permuted_blocks = blocks_to_update.reshape(src_shape).permute(*inv_order)