From 66b6c684ab5a32dabcd2e69daaa60ff5d5198392 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Wed, 15 Jul 2026 12:26:02 +0200 Subject: [PATCH] [PD][Bugfix] Fix validation of cache shape for attn backends enforcing different `kernel_block_size` (#48125) Signed-off-by: NickLucche --- .../kv_transfer/kv_connector/v1/nixl/base_worker.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 4724396893..d48cb2a5ff 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -1169,7 +1169,12 @@ class NixlBaseConnectorWorker: "All non-MLA kv cache tensors must have the same size" ) - if cache.shape[0] != num_blocks: + # When there's a mismatch between kbs<>bs, we rely on HMA to ensure + # caches are either [NB, PS] or [NB*r, PS/r] where r is bs/kbs. + if ( + self._physical_blocks_per_logical_kv_block == 1 + and cache.shape[0] != num_blocks + ): raise AssertionError( "All kv cache tensors must have the same number of " f"blocks; layer={layer_name}, "