From 32aef4438822f90310d2a709b71192f30da2fea8 Mon Sep 17 00:00:00 2001 From: Itay Etelis <92247226+Etelis@users.noreply.github.com> Date: Tue, 14 Jul 2026 16:07:26 +0300 Subject: [PATCH] [Bugfix] Include inline per-token-head scales in offloaded page transfer width (#48411) Signed-off-by: Itay Etelis Signed-off-by: Itay Etelis Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> Co-authored-by: Itay Etelis Co-authored-by: Itay Etelis --- tests/v1/core/test_kv_cache_utils.py | 31 +++++++++++++++++++ .../kv_connector/v1/offloading/worker.py | 2 +- vllm/v1/kv_cache_interface.py | 14 ++++++--- 3 files changed, 41 insertions(+), 6 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index e7520b1ee9..6c4779ddda 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -47,6 +47,7 @@ from vllm.v1.kv_cache_interface import ( KVCacheSpec, KVCacheSpecKind, KVCacheTensor, + KVQuantMode, MambaSpec, MLAAttentionSpec, SinkFullAttentionSpec, @@ -119,6 +120,7 @@ def new_kv_cache_spec( sliding_window=None, attention_chunk_size=None, indexes_kv_by_block_stride=False, + kv_quant_mode=KVQuantMode.NONE, ): return FullAttentionSpec( block_size=block_size, @@ -129,6 +131,7 @@ def new_kv_cache_spec( sliding_window=sliding_window, attention_chunk_size=attention_chunk_size, indexes_kv_by_block_stride=indexes_kv_by_block_stride, + kv_quant_mode=kv_quant_mode, ) @@ -2463,6 +2466,34 @@ def test_unify_kv_cache_page_size_padding_requires_backend_support(): kv_cache_utils.unify_kv_cache_spec_page_size(specs) +def test_unpadded_page_size_without_quant_matches_real_page(): + # Without quantization the offload transfer width is just the raw page. + spec = new_kv_cache_spec() + assert spec.unpadded_page_size_bytes == spec.real_page_size_bytes + assert spec.page_size_bytes == spec.unpadded_page_size_bytes + + +def test_unpadded_page_size_includes_per_token_head_scales(): + # Per-token-head quant carries inline fp32 scales that are carved from the + # raw KV allocation, so they must be budgeted into the offload width. + spec = new_kv_cache_spec( + dtype=torch.uint8, kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD + ) + scales = 2 * spec.block_size * spec.num_kv_heads * 4 + assert spec.unpadded_page_size_bytes == spec.real_page_size_bytes + scales + assert spec.page_size_bytes == spec.unpadded_page_size_bytes + + +def test_page_size_padded_wins(): + # An explicit padded page size takes precedence over the unpadded size. + spec = new_kv_cache_spec( + dtype=torch.uint8, + kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD, + page_size_padded=65536, + ) + assert spec.page_size_bytes == 65536 + + def test_unify_hybrid_kv_cache_specs(): # 1. has_full_attention and has_sliding_window before_spec_1 = new_kv_cache_spec() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index 29914e7388..60e0beb684 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -106,7 +106,7 @@ class OffloadingConnectorWorker: ) page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes unpadded_page_size_bytes[layer_name] = ( - layer_kv_cache_spec.real_page_size_bytes + layer_kv_cache_spec.unpadded_page_size_bytes ) elif isinstance(layer_kv_cache_spec, MambaSpec): diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index b0409b16de..d7f195e35b 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -182,19 +182,23 @@ class AttentionSpec(KVCacheSpec): indexes_kv_by_block_stride: bool = False @property - def page_size_bytes(self) -> int: - real_page_size = self.real_page_size_bytes + def unpadded_page_size_bytes(self) -> int: + unpadded = self.real_page_size_bytes # Per-token-head scales are stored in separate tensors managed # by the attention backend, but the memory is carved from the # raw KV cache allocation so it must be budgeted here. if self.kv_quant_mode.is_per_token_head: - real_page_size += ( + unpadded += ( 2 * self.block_size * self.num_kv_heads * get_dtype_size(torch.float32) ) + return unpadded + + @property + def page_size_bytes(self) -> int: if self.page_size_padded is not None: - assert self.page_size_padded >= real_page_size + assert self.page_size_padded >= self.unpadded_page_size_bytes return self.page_size_padded - return real_page_size + return self.unpadded_page_size_bytes @property def real_page_size_bytes(self) -> int: