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