[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>

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()

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):

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: