From dbc1f4f8fa8d41dd0371283f9a1bbb8b209333da Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 23 Oct 2026 12:06:09 +0800 Subject: [PATCH] [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. --- .../test_hybrid_dcp_block_alignment.py | 31 +++++++- .../v1/core/test_sm70_hybrid_page_rounding.py | 38 ++++++++++ vllm/platforms/interface.py | 70 ++++++++++++------- vllm/v1/core/kv_cache_utils.py | 24 ++++++- 4 files changed, 137 insertions(+), 26 deletions(-) diff --git a/tests/platforms/test_hybrid_dcp_block_alignment.py b/tests/platforms/test_hybrid_dcp_block_alignment.py index 13fca729f..3f3f8e356 100644 --- a/tests/platforms/test_hybrid_dcp_block_alignment.py +++ b/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)) diff --git a/tests/v1/core/test_sm70_hybrid_page_rounding.py b/tests/v1/core/test_sm70_hybrid_page_rounding.py index d6a82e19c..6d036799e 100644 --- a/tests/v1/core/test_sm70_hybrid_page_rounding.py +++ b/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 diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index a69c5f696..ab854eeba 100644 --- a/vllm/platforms/interface.py +++ b/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: diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 60667a02c..797c1353b 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/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