[Core] Size compact hybrid pages by factor search (#1252)

A cache type whose scale plane makes its physical page larger than the
target page could previously only be compacted by halving the block size
repeatedly, and the search stopped at the first integer ratio. When the
ratio is non-integral the compaction skipped the type, leaving the
unification to pad it, which wastes a whole page per block.

Search the divisors of the shared logical grid instead and pick the
largest block size whose physical page still fits the target page, so a
type keeps its logical block when possible and only shrinks when the
padding would be larger than the page itself.

tests/platforms/test_hybrid_dcp_block_alignment.py

@@ -9,7 +9,7 @@ import torch
 from vllm.model_executor.models import ModelRegistry
 from vllm.platforms.interface import Platform
 from vllm.v1.attention.backend import MultipleOf
-from vllm.v1.kv_cache_interface import MambaSpec
+from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec, get_kv_quant_mode
 
 # FP16, one KV head of size 64: 2 * 64 * 2 bytes of K and V per token.
 ATTN_BYTES_PER_TOKEN = 256
@@ -84,3 +84,32 @@ def test_dcp_share_keeps_the_block_alignment(monkeypatch, mamba_tokens, dcp) ->
     # No larger than the next aligned size: existing geometries do not move.
     assert block_size - KERNEL_BLOCK * dcp < mamba_tokens
     assert config.cache_config.mamba_block_size == block_size
+
+
+@pytest.mark.parametrize("mamba_tokens", [20, 21, 41])
+@pytest.mark.parametrize("cache_dtype", ["auto", "fp8_e4m3", "int8"])
+def test_hybrid_budget_uses_the_whole_physical_page(
+    monkeypatch, mamba_tokens, cache_dtype
+):
+    config, model_cls = _vllm_config(mamba_tokens, 1)
+    config.cache_config.cache_dtype = cache_dtype
+    monkeypatch.setattr(
+        ModelRegistry, "resolve_model_cls", lambda *args, **kwargs: (model_cls, "fake")
+    )
+    Platform._align_hybrid_block_size(config, _LooseBackend)
+    cache = config.cache_config
+    dtype = torch.float16 if cache_dtype == "auto" else torch.uint8
+
+    def page(block):
+        return FullAttentionSpec(
+            block_size=block,
+            num_kv_heads=1,
+            head_size=64,
+            dtype=dtype,
+            kv_quant_mode=get_kv_quant_mode(cache_dtype),
+        ).page_size_bytes
+
+    state_bytes = mamba_tokens * ATTN_BYTES_PER_TOKEN
+    assert page(cache.block_size) >= state_bytes
+    assert page(cache.block_size - KERNEL_BLOCK) < state_bytes
+    assert cache.mamba_page_size_padded in (None, page(cache.block_size))

tests/v1/core/test_sm70_hybrid_page_rounding.py

@@ -121,3 +121,41 @@ def test_non_integral_hybrid_pages_preserve_logical_blocks(heads):
     assert {spec.page_size_bytes for spec in unified.values()} == {page}
     assert specs == original
     assert kv_cache_utils.unify_kv_cache_spec_page_size(unified) == unified
+
+
+@pytest.mark.parametrize("cache_dtype", ["fp8_e4m3", "int8", "int8_per_token_head"])
+@pytest.mark.parametrize("heads", [1, 2])
+def test_compact_draft_preserves_prefix_grid_with_scale_overhead(cache_dtype, heads):
+    target = FullAttentionSpec(
+        block_size=2048,
+        num_kv_heads=heads,
+        head_size=256,
+        dtype=torch.uint8 if cache_dtype.startswith("fp8") else torch.int8,
+        kv_quant_mode=get_kv_quant_mode(cache_dtype),
+    )
+    draft = SlidingWindowSpec(
+        block_size=2048,
+        num_kv_heads=2 * heads,
+        head_size=128,
+        dtype=torch.float16,
+        sliding_window=2048,
+    )
+    state = MambaSpec(
+        block_size=2048,
+        shapes=((10, 2560 * heads), (12 * heads, 128, 128)),
+        dtypes=(torch.float16, torch.float32),
+        num_speculative_blocks=7,
+        mamba_cache_mode="align",
+        page_size_padded=target.page_size_bytes,
+    )
+    specs = {"target": target, "draft": draft, "state": state}
+    config = SimpleNamespace(
+        attention_config=SimpleNamespace(backend=AttentionBackendEnum.FLASH_ATTN_V100)
+    )
+    unified = kv_cache_utils.unify_kv_cache_spec_page_size(specs, config)
+    assert unified["target"] == target
+    assert unified["state"] == state
+    assert unified["draft"].block_size == 1024
+    assert {s.page_size_bytes for s in unified.values()} == {target.page_size_bytes}
+    assert specs["draft"] == draft
+    assert kv_cache_utils.unify_kv_cache_spec_page_size(unified, config) == unified

vllm/platforms/interface.py

@@ -577,6 +577,7 @@ class Platform:
         from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
         from vllm.v1.attention.backend import MultipleOf
         from vllm.v1.kv_cache_interface import (
+            AttentionSpec,
             FullAttentionSpec,
             MambaSpec,
             MLAAttentionSpec,
@@ -594,16 +595,20 @@ class Platform:
 
         kv_quant_mode = get_kv_quant_mode(cache_config.cache_dtype)
 
-        # Compute attention page size for 1 token
+        # A one-token page is only an estimate for selecting a block. Padding
+        # and separate scale regions need not grow linearly with token count.
+        attention_page_specs: list[AttentionSpec]
         if model_config.use_mla:
-            attn_page_size_1_token = MLAAttentionSpec(
-                block_size=1,
-                num_kv_heads=model_config.get_num_kv_heads(parallel_config),
-                head_size=model_config.get_head_size(),
-                dtype=kv_cache_dtype,
-                cache_dtype_str=cache_config.cache_dtype,
-                kv_quant_mode=kv_quant_mode,
-            ).page_size_bytes
+            attention_page_specs = [
+                MLAAttentionSpec(
+                    block_size=1,
+                    num_kv_heads=model_config.get_num_kv_heads(parallel_config),
+                    head_size=model_config.get_head_size(),
+                    dtype=kv_cache_dtype,
+                    cache_dtype_str=cache_config.cache_dtype,
+                    kv_quant_mode=kv_quant_mode,
+                )
+            ]
         elif cache_config.cache_dtype.startswith("turboquant_"):
             # TQ has a packed K|V layout; the standard FullAttentionSpec
             # formula over-sizes it and trips unify_kv_cache_spec_page_size
@@ -618,7 +623,7 @@ class Platform:
             tq_cfg = TurboQuantConfig.from_cache_dtype(
                 cache_config.cache_dtype, model_config.get_head_size()
             )
-            tq_page = TQFullAttentionSpec(
+            tq_spec = TQFullAttentionSpec(
                 block_size=1,
                 num_kv_heads=model_config.get_num_kv_heads(parallel_config),
                 head_size=model_config.get_head_size(),
@@ -626,28 +631,39 @@ class Platform:
                 dtype=kv_cache_dtype,
                 kv_quant_mode=kv_quant_mode,
                 tq_slot_size=tq_cfg.slot_size_aligned,
-            ).page_size_bytes
+            )
+            attention_page_specs = [tq_spec]
             if cache_config.kv_cache_dtype_skip_layers:
-                skip_page = FullAttentionSpec(
+                skip_spec = FullAttentionSpec(
                     block_size=1,
                     num_kv_heads=model_config.get_num_kv_heads(parallel_config),
                     head_size=model_config.get_head_size(),
                     dtype=model_config.dtype,
-                ).page_size_bytes
+                )
                 # lcm, not max: skip_page is often not a multiple of
                 # tq_page, so max would leave per-layer page sizes
                 # un-unifiable downstream.
-                attn_page_size_1_token = lcm(tq_page, skip_page)
-            else:
-                attn_page_size_1_token = tq_page
+                attention_page_specs.append(skip_spec)
         else:
-            attn_page_size_1_token = FullAttentionSpec(
-                block_size=1,
-                num_kv_heads=model_config.get_num_kv_heads(parallel_config),
-                head_size=model_config.get_head_size(),
-                dtype=kv_cache_dtype,
-                kv_quant_mode=kv_quant_mode,
-            ).page_size_bytes
+            attention_page_specs = [
+                FullAttentionSpec(
+                    block_size=1,
+                    num_kv_heads=model_config.get_num_kv_heads(parallel_config),
+                    head_size=model_config.get_head_size(),
+                    dtype=kv_cache_dtype,
+                    kv_quant_mode=kv_quant_mode,
+                )
+            ]
+
+        def attention_page_bytes(block_tokens: int) -> int:
+            return lcm(
+                *(
+                    spec.copy_with_new_block_size(block_tokens).page_size_bytes
+                    for spec in attention_page_specs
+                )
+            )
+
+        attn_page_size_1_token = attention_page_bytes(1)
 
         # Compute mamba page size
         model_cls, _ = ModelRegistry.resolve_model_cls(
@@ -706,6 +722,7 @@ class Platform:
             chunk_size = lcm(base_chunk_size, kernel_block_alignment_size)
             attn_block_size = chunk_size * cdiv(attn_tokens_per_mamba_state, chunk_size)
             cache_config.mamba_block_size = attn_block_size
+            block_step = chunk_size
         else:
             # Without prefix caching, use minimum block size that satisfies
             # both backend alignment and mamba page size compatibility
@@ -716,6 +733,11 @@ class Platform:
             indexer_align = cls._get_indexer_block_alignment(vllm_config)
             if indexer_align:
                 attn_block_size = indexer_align * cdiv(attn_block_size, indexer_align)
+            block_step = lcm(kernel_block_alignment_size, indexer_align or 1)
+        while attention_page_bytes(attn_block_size) < mamba_page_size:
+            attn_block_size += block_step
+        if cache_config.mamba_cache_mode == "all":
+            cache_config.mamba_block_size = attn_block_size
 
         if cache_config.block_size < attn_block_size:
             cache_config.block_size = attn_block_size
@@ -742,7 +764,7 @@ class Platform:
                 cache_config.mamba_block_size = cache_config.block_size
 
         # Pad mamba page size to exactly match attention page size
-        attn_page_size = cache_config.block_size * attn_page_size_1_token
+        attn_page_size = attention_page_bytes(cache_config.block_size)
         assert attn_page_size >= mamba_page_size
 
         if attn_page_size == mamba_page_size:

vllm/v1/core/kv_cache_utils.py

@@ -1109,9 +1109,31 @@ def _compact_sm70_hybrid_sliding_window_pages(
             type(spec) is not SlidingWindowSpec
             or spec.page_size_padded is not None
             or spec.page_size_bytes <= target_page
-            or spec.page_size_bytes % target_page
         ):
             continue
+        if spec.page_size_bytes % target_page:
+            # Scale bytes can make the physical ratio non-integral. Preserve
+            # the existing prefix/checkpoint grid by using a divisor of the
+            # logical block and padding its physical page to the common size.
+            grid = math.gcd(
+                *(
+                    other.block_size
+                    for other in specs.values()
+                    if isinstance(other, (AttentionSpec, MambaSpec))
+                )
+            )
+            divisors: set[int] = set()
+            for factor in range(1, math.isqrt(grid) + 1):
+                if grid % factor == 0:
+                    divisors.update((factor, grid // factor))
+            for block_size in sorted(divisors, reverse=True):
+                if block_size >= spec.block_size or block_size % 16:
+                    continue
+                compact = replace(spec, block_size=block_size)
+                if compact.page_size_bytes <= target_page:
+                    result[name] = replace(compact, page_size_padded=target_page)
+                    break
+            continue
         ratio = spec.page_size_bytes // target_page
         if spec.block_size % (16 * ratio):
             continue