[Bugfix] Use storage_block_size in KV cache reshape for compressed specs (DeepSeek V4) (#43988)
Signed-off-by: zixi-qi <zixi@inferact.ai> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.8
parent
4aaba00f92
commit
d07ad0693b
@@ -199,7 +199,12 @@ def _reshape_kv_cache(
|
||||
|
||||
if isinstance(kv_cache_spec, AttentionSpec):
|
||||
has_attn = True
|
||||
num_blocks_per_kv_block = kv_cache_spec.block_size // kernel_block_size
|
||||
# Use storage_block_size: it equals block_size for uncompressed
|
||||
# specs but is smaller for compressed ones (DeepSeek V4), which
|
||||
# store block_size tokens in block_size // compress_ratio slots.
|
||||
num_blocks_per_kv_block = (
|
||||
kv_cache_spec.storage_block_size // kernel_block_size
|
||||
)
|
||||
kernel_num_blocks = num_blocks * num_blocks_per_kv_block
|
||||
kv_cache_shape = group.backend.get_kv_cache_shape(
|
||||
kernel_num_blocks,
|
||||
|
||||
Reference in New Issue
Block a user