[Bugfix] Use storage_block_size in KV cache reshape for compressed specs (DeepSeek V4)

[zixi-qi]

Summary

_reshape_kv_cache (V2 model runner) overruns the KV-cache buffer during
initialization for compressed attention specs — specifically the
DeepSeek-V4 fp8 Lightning-Indexer cache. All workers crash at startup:

RuntimeError: setStorage: sizes [3944832, 2, 584], strides [1728, 584, 1], ...
requiring a storage size of 6816669136 are out of bounds for storage of size 53255232
  at vllm/v1/worker/gpu/attn_utils.py  _reshape_kv_cache (page_size_padded branch)

Reproducible with VLLM_USE_V2_MODEL_RUNNER=1, a DeepSeek-V4 model,
--block-size 256, --kv-cache-dtype fp8. The overrun factor is exactly
compress_ratio (128 = block_size 256 / kernel_block_size 2).

Root cause

For a compressed MLAAttentionSpec (DeepSeek-V4: compress_ratio = 128), a
256-token logical block is stored as storage_block_size = block_size // compress_ratio = 2 slots, and kernel_block_size is set to
storage_block_size. But the kernel-block count per logical block was computed
from block_size:

num_blocks_per_kv_block = block_size // kernel_block_size   # 256 // 2 = 128  (== compress_ratio!)
kernel_num_blocks       = num_blocks * num_blocks_per_kv_block

block_size // kernel_block_size equals compress_ratio for a compressed spec,
so kernel_num_blocks (hence kv_cache_shape[0]) is compress_ratio× too
large. The page_size_padded branch then builds a strided view whose block
dimension steps one full padded page (stride[0] = page_size_bytes // dtype_size) while iterating compress_ratio× too many blocks — walking past
the allocation by exactly that factor. The branch's own comment notes it
"assumes kv_cache_shape[0] == num_blocks", which this violates.

The buffer holds num_blocks padded pages (one per logical block); a compressed
logical block is exactly one kernel block (storage_block_size slots) plus
padding, not compress_ratio of them.

Introduced by #38831, which changed kv_cache_shape[0] from num_blocks to
kernel_num_blocks to support hybrid kernel block sizes. Uncompressed specs
were unaffected because there storage_block_size == block_size.

Fix

Compute the count from storage_block_size (what the buffer is actually laid
out in) rather than block_size:

num_blocks_per_kv_block = storage_block_size // kernel_block_size

• Compressed specs (DeepSeek-V4): kernel_block_size == storage_block_size
  ⇒ 1. kv_cache_shape[0] == num_blocks, restoring the padded-stride
  invariant and the pre-[ModelRunnerV2][Hybrid model] Support kernel block size in hybrid model #38831 shape.
• Uncompressed specs: KVCacheSpec.storage_block_size returns block_size,
  so the expression is byte-for-byte unchanged — no behavior change for any
  non-compressed model.

Test plan / results

• Arithmetic check: the buggy strided view requires 6,816,669,136 B
  (matches the reported error) against a 53,255,232 B allocation; the fixed
  view requires 53,254,672 B — fits.
• End-to-end (DeepSeek-V4-Flash, 4×GB200, VLLM_USE_V2_MODEL_RUNNER=1,
  --block-size 256, --kv-cache-dtype fp8): without the fix the server
  crashes at KV-cache init (the trace above); with the fix it initializes the
  fp8 indexer cache and reaches Application startup complete.
• GSM8K (5-shot, lm_eval): flexible-extract = 0.9567 against the running
  DeepSeek-V4-Flash server, confirming correct end-to-end execution.

Not a duplicate

Searched open PRs for _reshape_kv_cache, storage_block_size,
num_blocks_per_kv_block, page_size_padded, and DeepSeek-V4 KV cache. The
nearest, #43607, only optimizes block_table.py block mapping and does not
touch _reshape_kv_cache; no open PR addresses this overrun.

Note

AI assistance was used for the root-cause analysis and the patch; the change has
been reviewed and is being submitted by the human author. Opened as a draft for
maintainer review.

[njhill]

Thanks @zixi-qi

cc @MengqingCao

[MengqingCao]

Sorry for missing the deepseek-v4 scenario, and thanks for this fix!
@njhill Do we still have no test case on deepseek-v4 with MRV2 now? If so, I'd like to add a test case in #43607.