[Bugfix] Include inline per-token-head scales in offloaded page transfer width (#48411)
Signed-off-by: Itay Etelis <itay.etelis@ibm.com> Signed-off-by: Itay Etelis <Itay.etelis@gmail.com> Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> Co-authored-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Itay Etelis <Itay.etelis@gmail.com>
Itay Etelis committed
Jul 14, 2026 at 16:07 UTC
32aef4438822f90310d2a709b71192f30da2fea8
3 files changed
+41
-6
tests/v1/core/test_kv_cache_utils.py
+31
@@ -47,6 +47,7 @@ from vllm.v1.kv_cache_interface import (
47
KVCacheSpec,
48
KVCacheSpecKind,
49
KVCacheTensor,
50
+ KVQuantMode,
51
MambaSpec,
52
MLAAttentionSpec,
53
SinkFullAttentionSpec,
@@ -119,6 +120,7 @@ def new_kv_cache_spec(
120
sliding_window=None,
121
attention_chunk_size=None,
122
indexes_kv_by_block_stride=False,
123
+ kv_quant_mode=KVQuantMode.NONE,
124
):
125
return FullAttentionSpec(
126
block_size=block_size,
@@ -129,6 +131,7 @@ def new_kv_cache_spec(
131
sliding_window=sliding_window,
132
attention_chunk_size=attention_chunk_size,
133
indexes_kv_by_block_stride=indexes_kv_by_block_stride,
134
+ kv_quant_mode=kv_quant_mode,
135
)
136
137
@@ -2463,6 +2466,34 @@ def test_unify_kv_cache_page_size_padding_requires_backend_support():
2466
kv_cache_utils.unify_kv_cache_spec_page_size(specs)
2467
2468
2469
+def test_unpadded_page_size_without_quant_matches_real_page():
2470
+ # Without quantization the offload transfer width is just the raw page.
2471
+ spec = new_kv_cache_spec()
2472
+ assert spec.unpadded_page_size_bytes == spec.real_page_size_bytes
2473
+ assert spec.page_size_bytes == spec.unpadded_page_size_bytes
2474
+
2475
+
2476
+def test_unpadded_page_size_includes_per_token_head_scales():
2477
+ # Per-token-head quant carries inline fp32 scales that are carved from the
2478
+ # raw KV allocation, so they must be budgeted into the offload width.
2479
+ spec = new_kv_cache_spec(
2480
+ dtype=torch.uint8, kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD
2481
+ )
2482
+ scales = 2 * spec.block_size * spec.num_kv_heads * 4
2483
+ assert spec.unpadded_page_size_bytes == spec.real_page_size_bytes + scales
2484
+ assert spec.page_size_bytes == spec.unpadded_page_size_bytes
2485
+
2486
+
2487
+def test_page_size_padded_wins():
2488
+ # An explicit padded page size takes precedence over the unpadded size.
2489
+ spec = new_kv_cache_spec(
2490
+ dtype=torch.uint8,
2491
+ kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD,
2492
+ page_size_padded=65536,
2493
+ )
2494
+ assert spec.page_size_bytes == 65536
2495
+
2496
+
2497
def test_unify_hybrid_kv_cache_specs():
2498
# 1. has_full_attention and has_sliding_window
2499
before_spec_1 = new_kv_cache_spec()
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+1
-1
@@ -106,7 +106,7 @@ class OffloadingConnectorWorker:
106
)
107
page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes
108
unpadded_page_size_bytes[layer_name] = (
109
- layer_kv_cache_spec.real_page_size_bytes
109
+ layer_kv_cache_spec.unpadded_page_size_bytes
110
)
111
112
elif isinstance(layer_kv_cache_spec, MambaSpec):
vllm/v1/kv_cache_interface.py
+9
-5
@@ -182,19 +182,23 @@ class AttentionSpec(KVCacheSpec):
182
indexes_kv_by_block_stride: bool = False
183
184
@property
185
- def page_size_bytes(self) -> int:
186
- real_page_size = self.real_page_size_bytes
185
+ def unpadded_page_size_bytes(self) -> int:
186
+ unpadded = self.real_page_size_bytes
187
# Per-token-head scales are stored in separate tensors managed
188
# by the attention backend, but the memory is carved from the
189
# raw KV cache allocation so it must be budgeted here.
190
if self.kv_quant_mode.is_per_token_head:
191
- real_page_size += (
191
+ unpadded += (
192
2 * self.block_size * self.num_kv_heads * get_dtype_size(torch.float32)
193
)
194
+ return unpadded
195
+
196
+ @property
197
+ def page_size_bytes(self) -> int:
198
if self.page_size_padded is not None:
195
- assert self.page_size_padded >= real_page_size
199
+ assert self.page_size_padded >= self.unpadded_page_size_bytes
200
return self.page_size_padded
197
- return real_page_size
201
+ return self.unpadded_page_size_bytes
202
203
@property
204
def real_page_size_bytes(self) -> int: