[PD][Bugfix] Fix validation of cache shape for attn backends enforcing different `kernel_block_size` (#48125)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Nicolò Lucchesi committed
Jul 15, 2026 at 12:26 UTC
66b6c684ab5a32dabcd2e69daaa60ff5d5198392
1 file changed
+6
-1
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+6
-1
@@ -1169,7 +1169,12 @@ class NixlBaseConnectorWorker:
1169
"All non-MLA kv cache tensors must have the same size"
1170
)
1171
1172
- if cache.shape[0] != num_blocks:
1172
+ # When there's a mismatch between kbs<>bs, we rely on HMA to ensure
1173
+ # caches are either [NB, PS] or [NB*r, PS/r] where r is bs/kbs.
1174
+ if (
1175
+ self._physical_blocks_per_logical_kv_block == 1
1176
+ and cache.shape[0] != num_blocks
1177
+ ):
1178
raise AssertionError(
1179
"All kv cache tensors must have the same number of "
1180
f"blocks; layer={layer_name}, "