diff --git a/PROTOCOL.md b/PROTOCOL.md index 98057d8baff5..2cefd20e546a 100644 --- a/PROTOCOL.md +++ b/PROTOCOL.md @@ -101,6 +101,40 @@ per N layers in each block (`get_tensor_slots`) and gives every fused layer a view of that same page. Without it the summed page is charged once per layer and usable capacity drops N-fold. +### Separate state pool for hybrid models — `vllm/v1/core/`, `vllm/v1/worker/` + +A hybrid model's O(1) recurrent layers (Mamba/GDN) hold a fixed number of +blocks per running request, never one per token. Upstream puts them in the +attention pool: every recurrent layer gets a tensor of `num_blocks` state +pages and its page is unified with the attention page, which spends most of +the KV budget on state no request can use (Qwen3.5-0.8B at 10 GB: 5,728 +tokens shared vs 355,696 split; turbo-attn ATTRIBUTION §7.3). + +The seam gives those groups their own pool when it is safe to +(`_mamba_pool_eligible`: Mamba cache mode `none`/`align`, no HiSparse, no KV +connector, no context parallelism, no hidden-state layers): + +- **Planning** (`kv_cache_utils.py`): the recurrent pages are not unified + with attention (attention layers are unified among themselves); the state + pool is sized `max_num_seqs * sum(per-request blocks per group) + 1` and, + with a sampler-warmup reserve (`VLLM_SAMPLER_RESERVE_MIB`), comes off the + budget before attention is planned, so auto-fit, the override and the + admission check all see only the attention pool. The two pools are two + regions of the one backing allocation, each laid out exactly as upstream + lays out a shared pool (`_layout_kv_cache_tensors`), so groups within a + pool still overlay each other and the pools never alias. + `KVCacheConfig.mamba_pool_num_blocks` / `mamba_pool_offset`, + `KVCacheTensor.num_blocks`. +- **Scheduler** (`kv_cache_coordinator.py`, `kv_cache_manager.py`): a second + `BlockPool` for the recurrent groups; admission checks each pool against its + own free blocks (watermark and reservations count main-pool blocks); + events and usage cover both. Frees route by `KVCacheBlock.pool`, which + upstream already does. +- **Worker** (`worker/utils.py`, `worker/gpu/*`): the two pools' block ids + overlap, so CoW copies carry state-pool ids shifted past `num_blocks` and + apply only to the caches in that pool's region; profiling dummy contexts + and warmup number each pool's blocks separately. + --- ## Sync workflow when moving to a new upstream release @@ -113,7 +147,7 @@ and usable capacity drops N-fold. extension point that covers a hook, move the capability into the plugin and drop the hook. 3. `pytest tests/v1/core/test_kv_seam_invariants.py - tests/v1/core/test_aggregated_layer_count.py + tests/v1/core/test_aggregated_layer_count.py tests/v1/core/test_mamba_pool.py tests/v1/attention/test_mla_wrapper_selection.py tests/v1/worker/test_gpu_worker.py -k hook tests/engine/test_arg_utils.py -k plugin_kv_cache_dtype` (CPU). diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 1f599139e604..97aa949609f8 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -1081,9 +1081,7 @@ def test_plugin_kv_cache_dtype_auto_selects_turbo_attn_backend(): ) engine_args = EngineArgs.from_cli_args(args) vllm_config = engine_args.create_engine_config() - assert vllm_config.attention_config.backend == ( - AttentionBackendEnum.TURBO_ATTN - ) + assert vllm_config.attention_config.backend == (AttentionBackendEnum.TURBO_ATTN) # 2. Plugin dtype + explicit --attention-backend → user choice honoured. args = parser.parse_args( @@ -1098,9 +1096,7 @@ def test_plugin_kv_cache_dtype_auto_selects_turbo_attn_backend(): ) engine_args = EngineArgs.from_cli_args(args) vllm_config = engine_args.create_engine_config() - assert vllm_config.attention_config.backend == ( - AttentionBackendEnum.FLASH_ATTN - ) + assert vllm_config.attention_config.backend == (AttentionBackendEnum.FLASH_ATTN) # 3. Builtin dtype "auto" → no auto-override (stays None). args = parser.parse_args( diff --git a/tests/v1/core/test_mamba_pool.py b/tests/v1/core/test_mamba_pool.py new file mode 100644 index 000000000000..0e8e6f900fc9 --- /dev/null +++ b/tests/v1/core/test_mamba_pool.py @@ -0,0 +1,361 @@ +# SPDX-License-Identifier: Apache-2.0 +"""A hybrid model's O(1) state layers draw from their own block pool. + +Mamba/GDN layers hold a fixed number of blocks per running request, never one +per token. Sharing the attention pool gives every recurrent layer a tensor of +``num_blocks`` state pages and unifies its page with the attention page, which +spends the budget on state no request can use (on Qwen3.5-0.8B at 10 GB that +was 5,728 tokens against 355,696 with separate pools). These tests pin the +split layout: the planner, the scheduler's pools, and the worker's copies. +""" + +from dataclasses import dataclass + +import pytest +import torch + +from vllm.config import DeviceConfig, ModelConfig, SchedulerConfig, VllmConfig +from vllm.sampling_params import SamplingParams +from vllm.v1.core.kv_cache_manager import KVCacheManager +from vllm.v1.core.kv_cache_utils import ( + KVCacheBlockCopy, + get_kv_cache_configs, + get_request_block_hasher, + init_none_hash, +) +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + MambaSpec, +) +from vllm.v1.request import Request +from vllm.v1.worker.utils import copy_kv_cache_blocks_inplace + +pytestmark = pytest.mark.cpu_test + +BLOCK_SIZE = 16 +MAX_MODEL_LEN = 1024 +MAX_NUM_SEQS = 8 +NUM_LAYERS = 24 # every 4th is attention: 6 attention + 18 GDN layers + + +@pytest.fixture(autouse=True) +def _fixed_sampler_reserve(monkeypatch): + monkeypatch.setenv("VLLM_SAMPLER_RESERVE_MIB", "64") + + +def _attention_spec(): + return FullAttentionSpec( + block_size=BLOCK_SIZE, num_kv_heads=2, head_size=256, dtype=torch.bfloat16 + ) + + +def _mamba_spec(): + # A GDN state: far larger than one attention page. + return MambaSpec( + block_size=BLOCK_SIZE, + shapes=((3, 8192), (16, 128, 128)), + dtypes=(torch.bfloat16, torch.bfloat16), + mamba_cache_mode="none", + ) + + +def _specs(attention_spec=None): + attention_spec = attention_spec or _attention_spec() + return { + f"layer_{i}": attention_spec if (i + 1) % 4 == 0 else _mamba_spec() + for i in range(NUM_LAYERS) + } + + +def _config(max_model_len=4096): + model_config = ModelConfig(max_model_len=max_model_len) + config = VllmConfig( + model_config=model_config, + device_config=DeviceConfig(device="cpu"), + scheduler_config=SchedulerConfig( + max_num_batched_tokens=max_model_len, + max_num_seqs=MAX_NUM_SEQS, + enable_chunked_prefill=True, + max_model_len=max_model_len, + is_encoder_decoder=model_config.is_encoder_decoder, + ), + ) + config.cache_config.kv_cache_layout = "LBNHC" + config.cache_config.mamba_cache_mode = "none" + config.scheduler_config.disable_hybrid_kv_cache_manager = False + return config + + +def _plan(specs, available=4 * 1024**3, vllm_config=None): + return get_kv_cache_configs(vllm_config or _config(), [specs], [available])[0] + + +def _mamba_groups(config): + return [ + g + for g in range(len(config.kv_cache_groups)) + if isinstance(config.kv_cache_groups[g].kv_cache_spec, MambaSpec) + ] + + +def test_state_groups_get_a_pool_sized_for_running_requests(): + config = _plan(_specs()) + mamba_groups = _mamba_groups(config) + + assert len(mamba_groups) == 3 + # One state block per running request per group ("none" mode), + null. + assert config.mamba_pool_num_blocks == MAX_NUM_SEQS * len(mamba_groups) + 1 + assert all(config.uses_mamba_pool(g) for g in mamba_groups) + for g in mamba_groups: + assert config.group_num_blocks(g) == config.mamba_pool_num_blocks + assert config.group_num_blocks(0 if 0 not in mamba_groups else 3) == ( + config.num_blocks + ) + + +def test_split_pool_buys_the_attention_capacity_the_shared_pool_wasted(): + split = _plan(_specs()) + shared_config = _config() + shared_config.cache_config.mamba_cache_mode = "all" + shared_specs = { + k: ( + MambaSpec( + block_size=v.block_size, + shapes=v.shapes, + dtypes=v.dtypes, + mamba_cache_mode="all", + ) + if isinstance(v, MambaSpec) + else v + ) + for k, v in _specs().items() + } + shared = _plan(shared_specs, vllm_config=shared_config) + + assert shared.mamba_pool_num_blocks is None + # Same bytes; 16.9x the attention blocks here with bf16 attention (more + # with a compressed attention page, whose block a state page dwarfs). + assert split.num_blocks > 15 * shared.num_blocks + + +def test_pools_never_alias_and_fit_the_allocation(): + config = _plan(_specs()) + (size,) = {t.size for t in config.kv_cache_tensors} + main_end = 0 + for t in config.kv_cache_tensors: + spec = next( + g.kv_cache_spec + for g in config.kv_cache_groups + if t.layers[0] in g.layer_names + ) + blocks = config.num_blocks_of(t) + end = t.offset + (len(t.layers) - 1) * t.layer_stride + ( + blocks - 1 + ) * t.block_stride + spec.page_size_bytes + assert end <= size + if t.num_blocks is None: + main_end = max(main_end, end) + assert t.offset < config.mamba_pool_offset + else: + assert t.num_blocks == config.mamba_pool_num_blocks + assert t.offset >= config.mamba_pool_offset + assert main_end <= config.mamba_pool_offset + + +@pytest.mark.parametrize( + "feature", + ["kv_transfer_config", "decode_context_parallel_size", "mamba_cache_mode_all"], +) +def test_shared_pool_kept_where_ids_cross_groups(feature): + """Features that address blocks across groups by one id space keep the + shared pool.""" + from vllm.v1.core.kv_cache_utils import _mamba_pool_eligible + + vllm_config = _config() + specs = list(_specs().values()) + assert _mamba_pool_eligible(vllm_config, specs) + if feature == "kv_transfer_config": + object.__setattr__(vllm_config, "kv_transfer_config", object()) + elif feature == "decode_context_parallel_size": + object.__setattr__(vllm_config.parallel_config, feature, 2) + else: + vllm_config.cache_config.mamba_cache_mode = "all" + assert not _mamba_pool_eligible(vllm_config, specs) + + +# ── Scheduler side ─────────────────────────────────────────────────────────── + + +def _manager_config(num_blocks=64, mamba_blocks=5): + attention = FullAttentionSpec( + block_size=BLOCK_SIZE, num_kv_heads=1, head_size=1, dtype=torch.float32 + ) + # "none" mode: the state block spans the whole sequence (upstream sets + # mamba_block_size = max_model_len), one block per request per group. + mamba = MambaSpec(block_size=MAX_MODEL_LEN, shapes=((1, 1),), dtypes=(torch.float32,)) + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["attn"], attention), + KVCacheGroupSpec(["gdn_a"], mamba), + KVCacheGroupSpec(["gdn_b"], mamba), + ], + mamba_pool_num_blocks=mamba_blocks, + mamba_pool_offset=0, + ) + + +def _request(request_id, num_tokens): + from vllm.utils.hashing import sha256 + + init_none_hash(sha256) + params = SamplingParams(max_tokens=4) + params.update_from_generation_config({}, eos_token_id=100) + return Request( + request_id=request_id, + prompt_token_ids=list(range(num_tokens)), + sampling_params=params, + pooling_params=None, + block_hasher=get_request_block_hasher(BLOCK_SIZE, sha256), + ) + + +def _manager(config): + return KVCacheManager( + config, + max_model_len=MAX_MODEL_LEN, + scheduler_block_size=MAX_MODEL_LEN, + hash_block_size=BLOCK_SIZE, + enable_caching=False, + ) + + +def test_state_groups_allocate_from_their_own_pool(): + manager = _manager(_manager_config()) + coordinator = manager.coordinator + attn_mgr, gdn_a, gdn_b = coordinator.single_type_managers + + assert attn_mgr.block_pool is coordinator.block_pool + assert gdn_a.block_pool is gdn_b.block_pool is coordinator.mamba_block_pool + assert coordinator.block_pool.num_gpu_blocks == 64 + assert coordinator.mamba_block_pool.num_gpu_blocks == 5 + + blocks = manager.allocate_slots(_request("r0", 3 * BLOCK_SIZE), 3 * BLOCK_SIZE) + assert blocks is not None + attn_blocks, a_blocks, b_blocks = blocks.blocks + assert len(attn_blocks) == 3 + assert all(b.pool is coordinator.block_pool for b in attn_blocks) + assert all(b.pool is coordinator.mamba_block_pool for b in a_blocks + b_blocks) + + +def test_admission_checks_each_pool_against_its_own_blocks(): + # 4 usable state blocks: two requests fill them (one per group each). + manager = _manager(_manager_config(mamba_blocks=5)) + assert manager.allocate_slots(_request("r0", BLOCK_SIZE), BLOCK_SIZE) + assert manager.allocate_slots(_request("r1", BLOCK_SIZE), BLOCK_SIZE) + assert manager.coordinator.block_pool.get_num_free_blocks() > 50 + # Plenty of attention blocks, no state blocks: refused, not over-allocated. + assert manager.allocate_slots(_request("r2", BLOCK_SIZE), BLOCK_SIZE) is None + + +def test_deferred_frees_return_blocks_to_their_own_pools(): + manager = _manager(_manager_config()) + coordinator = manager.coordinator + main_free = coordinator.block_pool.get_num_free_blocks() + mamba_free = coordinator.mamba_block_pool.get_num_free_blocks() + request = _request("r0", 2 * BLOCK_SIZE) + assert manager.allocate_slots(request, 2 * BLOCK_SIZE) + + # The scheduler's deferred path: pop every group's blocks, free them + # through the main pool in one call. + blocks = manager.pop_blocks_for_free(request) + manager.block_pool.free_blocks(reversed(blocks)) + + assert coordinator.block_pool.get_num_free_blocks() == main_free + assert coordinator.mamba_block_pool.get_num_free_blocks() == mamba_free + + +def test_usage_reports_the_fuller_pool(): + manager = _manager(_manager_config(mamba_blocks=5)) + manager.allocate_slots(_request("r0", BLOCK_SIZE), BLOCK_SIZE) + assert manager.usage == pytest.approx(2 / 4) + + +# ── Worker side ────────────────────────────────────────────────────────────── + + +def test_block_copies_reach_only_their_own_pool(monkeypatch): + import vllm.v1.worker.utils as worker_utils + + monkeypatch.setattr( + worker_utils, + "async_tensor_h2d", + lambda data, device=None, dtype=None: torch.tensor(data, device=device), + ) + num_main, num_mamba, page = 6, 3, 4 + buf = torch.zeros(num_main * page + num_mamba * page, dtype=torch.int8) + main = buf[: num_main * page].view(num_main, page) + mamba = buf[num_main * page :].view(num_mamba, page) + main.copy_(torch.arange(num_main).repeat_interleave(page).view(num_main, page)) + mamba.copy_( + (10 + torch.arange(num_mamba)).repeat_interleave(page).view(num_mamba, page) + ) + + copies = [KVCacheBlockCopy(1, 4), KVCacheBlockCopy(num_main + 2, num_main + 1)] + copy_kv_cache_blocks_inplace( + [main, mamba], num_main, copies, mamba_pool=(num_main * page, num_mamba) + ) + + assert main[:, 0].tolist() == [0, 1, 2, 3, 1, 5] + assert mamba[:, 0].tolist() == [10, 12, 12] + + +# ── Composite (fused) attention pages on the split layout ──────────────────── + + +@dataclass(frozen=True, kw_only=True) +class _FusedAttentionSpec(FullAttentionSpec): + """Six attention layers packed into one shared page (smart-mix).""" + + layer_slot_bytes: tuple[int, ...] = (544, 544, 544, 480, 544, 608) + + @property + def num_heads(self) -> int: + return 1 + + @property + def state_content_size_bytes(self) -> int: + return sum(self.layer_slot_bytes) + + @property + def aggregated_layer_count(self) -> int: + return len(self.layer_slot_bytes) + + @classmethod + def merge(cls, specs): + return specs[0] + + +def test_composite_attention_shares_one_dense_page_beside_the_state_pool(): + fused = _FusedAttentionSpec( + block_size=BLOCK_SIZE, num_kv_heads=2, head_size=256, dtype=torch.uint8 + ) + config = _plan(_specs(fused)) + attention = [ + t + for t in config.kv_cache_tensors + if t.num_blocks is None and t.layers[0] in {f"layer_{i}" for i in range(3, 24, 4)} + ] + + assert config.mamba_pool_num_blocks is not None + assert len(attention) == 6 + # Every fused layer views the same page region, laid out block-dense so + # the plugin can carve its layer-major regions out of it. + assert len({(t.offset, t.block_stride) for t in attention}) == 1 + assert attention[0].block_stride == fused.page_size_bytes + assert attention[0].offset + config.num_blocks * fused.page_size_bytes <= ( + config.mamba_pool_offset + ) diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 310d019ba158..1388203011f9 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -80,17 +80,17 @@ def register_cache_dtype(name: str, torch_dtype) -> None: STR_DTYPE_TO_TORCH_DTYPE. """ from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE + _PLUGIN_CACHE_DTYPES.add(name) STR_DTYPE_TO_TORCH_DTYPE[name] = torch_dtype def validate_cache_dtype(name: str) -> str: - """argparse ``type=`` callable. Validates against builtins + plugins.""" + """Argparse ``type=`` callable. Validates against builtins + plugins.""" allowed = _BUILTIN_CACHE_DTYPES | _PLUGIN_CACHE_DTYPES if name not in allowed: raise ValueError( - f"Unknown --kv-cache-dtype {name!r}. " - f"Allowed: {sorted(allowed)}" + f"Unknown --kv-cache-dtype {name!r}. Allowed: {sorted(allowed)}" ) return name diff --git a/vllm/envs.py b/vllm/envs.py index 64b5ca2c4eb5..9b5c23e6ecda 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -250,6 +250,7 @@ Literal["LBNHC", "LBHNC", "LHBNC", "NHD", "HND", "BLHNC", "BLNHC", "BHLNC"] | None ) = None + VLLM_SAMPLER_RESERVE_MIB: int = 0 VLLM_SSM_CONV_STATE_LAYOUT: Literal["SD", "DS"] | None = None VLLM_COMPUTE_NANS_IN_LOGITS: bool = False VLLM_RAISE_ON_LOGIT_NANS: bool = False @@ -1835,6 +1836,10 @@ def _resolve_rust_cli_path() -> str | None: None, ["LBNHC", "LBHNC", "LHBNC", "NHD", "HND", "BLHNC", "BLNHC", "BHLNC"], ), + # Sampler-warmup headroom (MiB) a hybrid model's split KV layout keeps free + # after allocating the KV cache. 0 = estimate it from the vocab size and + # max_num_seqs. + "VLLM_SAMPLER_RESERVE_MIB": lambda: int(os.getenv("VLLM_SAMPLER_RESERVE_MIB", "0")), # SSM conv state layout used for Mamba models. # - SD: (state_len, dim) — dim contiguous (default) # - DS: (dim, state_len) — TP-sharded dim on dim1, diff --git a/vllm/v1/attention/selector.py b/vllm/v1/attention/selector.py index 4a0cd410114b..259cc31d9f9a 100644 --- a/vllm/v1/attention/selector.py +++ b/vllm/v1/attention/selector.py @@ -121,6 +121,7 @@ def get_attn_backend( """Selects which attention backend to use and lazily imports it.""" if kv_cache_dtype is not None: from vllm.config.cache import validate_cache_dtype + # Raises ValueError with a clear message if the dtype isn't # builtin or plugin-registered. validate_cache_dtype(kv_cache_dtype) @@ -215,6 +216,7 @@ def _mla_wrapper_cls(backend) -> "type[AttentionBackend] | None": Returns: The backend class, or None if it is unimportable or not a wrapper. + """ try: cls = backend.get_class() @@ -242,6 +244,7 @@ def _mla_wrapper_for_dtype( Returns: The wrapper backend class, or None when no plugin claims the dtype. + """ if kv_cache_dtype is None or kv_cache_dtype == "auto": return None diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index daf0cd12c0e8..e6b25c6c21e1 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from abc import ABC, abstractmethod +from collections import defaultdict from collections.abc import Sequence from typing import NamedTuple @@ -100,6 +101,19 @@ def __init__( enable_kv_cache_events=enable_kv_cache_events, metrics_collector=metrics_collector, ) + # The O(1) state groups' own pool (KVCacheConfig.mamba_pool_num_blocks). + self.mamba_block_pool: BlockPool | None = None + if kv_cache_config.mamba_pool_num_blocks is not None: + self.mamba_block_pool = BlockPool( + num_gpu_blocks=kv_cache_config.mamba_pool_num_blocks, + enable_caching=enable_caching, + hash_block_size=hash_block_size, + enable_kv_cache_events=enable_kv_cache_events, + metrics_collector=metrics_collector, + ) + self.block_pools: tuple[BlockPool, ...] = tuple( + pool for pool in (self.block_pool, self.mamba_block_pool) if pool + ) # KV cache group indices that get the EAGLE last-block drop. self.eagle_group_ids: set[int] = { @@ -136,7 +150,9 @@ def __init__( kv_cache_spec=kv_cache_group.kv_cache_spec, max_in_flight_tokens=max_in_flight_tokens, max_model_len=max_model_len, - block_pool=self.block_pool, + block_pool=self.mamba_block_pool + if kv_cache_config.uses_mamba_pool(i) + else self.block_pool, role=kv_cache_group.role, enable_caching=enable_caching, kv_cache_group_id=i, @@ -206,14 +222,43 @@ def get_num_blocks_to_allocate( The number of blocks to allocate. """ - num_blocks_to_allocate = 0 + return sum( + self.get_num_blocks_to_allocate_per_pool( + request_id, + num_tokens, + new_computed_blocks, + num_encoder_tokens, + total_computed_tokens, + num_local_computed_tokens, + num_tokens_main_model, + apply_admission_cap=apply_admission_cap, + prefill_end=prefill_end, + ).values() + ) + + def get_num_blocks_to_allocate_per_pool( + self, + request_id: str, + num_tokens: int, + new_computed_blocks: tuple[Sequence[KVCacheBlock], ...], + num_encoder_tokens: int, + total_computed_tokens: int, + num_local_computed_tokens: int, + num_tokens_main_model: int, + apply_admission_cap: bool = False, + prefill_end: int = 0, + ) -> dict[BlockPool, int]: + """``get_num_blocks_to_allocate`` broken down by the pool each group + draws from, so each pool is checked against its own free blocks.""" + # Keyed by the manager's own pool: a HiSparse manager wraps its pool. + needed: defaultdict[BlockPool, int] = defaultdict(int) if self.retention_interval != 0: prefill_end = 0 for i, manager in enumerate(self.single_type_managers): if isinstance(manager, CrossAttentionManager): # For cross-attention, we issue a single static allocation # of blocks based on the number of encoder input tokens. - num_blocks_to_allocate += manager.get_num_blocks_to_allocate( + needed[manager.block_pool] += manager.get_num_blocks_to_allocate( request_id, num_encoder_tokens, [], @@ -223,7 +268,7 @@ def get_num_blocks_to_allocate( apply_admission_cap=apply_admission_cap, ) else: - num_blocks_to_allocate += manager.get_num_blocks_to_allocate( + needed[manager.block_pool] += manager.get_num_blocks_to_allocate( request_id, num_tokens, new_computed_blocks[i], @@ -233,7 +278,7 @@ def get_num_blocks_to_allocate( apply_admission_cap=apply_admission_cap, prefill_end=prefill_end, ) - return num_blocks_to_allocate + return dict(needed) def allocate_new_computed_blocks( self, diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 2d39a3865124..93c13ec2dc4f 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -9,6 +9,7 @@ from vllm.distributed.kv_events import BlockStored, KVCacheEvent from vllm.logger import init_logger from vllm.utils.math_utils import cdiv +from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_coordinator import ( HybridKVCacheCoordinator, get_kv_cache_coordinator, @@ -195,6 +196,8 @@ def __init__( manager.shared_prefix_checkpoint = True self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) self.block_pool = self.coordinator.block_pool + # The O(1) state groups' pool, when they have their own; None otherwise. + self.mamba_block_pool = self.coordinator.mamba_block_pool self.retained_hit_group_ids = tuple( manager.kv_cache_group_id for manager in self.coordinator.single_type_managers @@ -231,7 +234,26 @@ def usage(self) -> float: The KV cache usage (between 0.0 and 1.0). """ - return self.block_pool.get_usage() + return max(pool.get_usage() for pool in self.coordinator.block_pools) + + def _fits( + self, + needed: dict[BlockPool, int], + watermark_blocks: int, + reserved_blocks: int = 0, + ) -> bool: + """Whether every pool has the blocks ``needed`` from it. The watermark + and the in-flight reservation are counted in main-pool blocks.""" + if self.mamba_block_pool is None: + return sum(needed.values()) + watermark_blocks <= ( + self.block_pool.get_num_free_blocks() - reserved_blocks + ) + mamba_needed = needed.pop(self.mamba_block_pool, 0) + if mamba_needed > self.mamba_block_pool.get_num_free_blocks(): + return False + return sum(needed.values()) + watermark_blocks <= ( + self.block_pool.get_num_free_blocks() - reserved_blocks + ) def make_prefix_cache_stats(self) -> PrefixCacheStats | None: """Get (and reset) the prefix cache stats. @@ -520,7 +542,7 @@ def allocate_slots( # First check and fail if the full request sequence won't fit. full_num_tokens = min(request.num_tokens, self.max_model_len) - num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate( + needed = self.coordinator.get_num_blocks_to_allocate_per_pool( request_id=request.request_id, num_tokens=full_num_tokens, new_computed_blocks=new_computed_block_list, @@ -531,8 +553,7 @@ def allocate_slots( apply_admission_cap=True, prefill_end=prefill_end, ) - required_blocks = num_blocks_to_allocate + watermark_blocks - if required_blocks > self.block_pool.get_num_free_blocks(): + if not self._fits(needed, watermark_blocks): return None num_tokens_main_model = total_computed_tokens + num_new_tokens @@ -555,7 +576,7 @@ def allocate_slots( num_prompt_tokens=request.num_prompt_tokens, ) - num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate( + needed = self.coordinator.get_num_blocks_to_allocate_per_pool( request_id=request.request_id, num_tokens=num_tokens_need_slot, new_computed_blocks=new_computed_block_list, @@ -569,9 +590,7 @@ def allocate_slots( # Keep `reserved_blocks` free for other in-flight sequences, and an # additional watermark of headroom for waiting/preempted admissions. - available_blocks = self.block_pool.get_num_free_blocks() - reserved_blocks - required_blocks = num_blocks_to_allocate + watermark_blocks - if required_blocks > available_blocks: + if not self._fits(needed, watermark_blocks, reserved_blocks): # Cannot allocate new blocks return None @@ -726,7 +745,11 @@ def take_events(self) -> list[KVCacheEvent]: A list of KV cache events. """ - events = self.block_pool.take_events() + events = [ + event + for pool in self.coordinator.block_pools + for event in pool.take_events() + ] for event in events: if not isinstance(event, BlockStored): continue @@ -898,10 +921,16 @@ def take_kv_cache_block_copies( pending_copies: list[tuple[KVCacheBlock, KVCacheBlock]] = [] for mgr in self.coordinator.single_type_managers: pending_copies.extend(mgr.take_pending_cow_copies()) + # Mamba-pool block ids are shifted past the main pool's, so the worker + # can tell which pool's caches a copy belongs to (KVCacheConfig + # .mamba_pool_offset); the two pools' ids otherwise collide. + num_main = self.kv_cache_config.num_blocks copies = [ KVCacheBlockCopy( - src_block_id=source_block.block_id, - dst_block_id=cow_block.block_id, + src_block_id=source_block.block_id + + (num_main if source_block.pool is self.mamba_block_pool else 0), + dst_block_id=cow_block.block_id + + (num_main if cow_block.pool is self.mamba_block_pool else 0), ) for source_block, cow_block in pending_copies ] diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 801ddfeadd81..6031566da990 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1104,7 +1104,10 @@ def get_max_concurrency_for_kv_cache_config( """ num_blocks_per_request = 0 host_blocks_per_request = 0 - for group in kv_cache_config.kv_cache_groups: + for group_id, group in enumerate(kv_cache_config.kv_cache_groups): + if kv_cache_config.uses_mamba_pool(group_id): + # Sized for every running request; never the token limit. + continue required = cdiv( group.kv_cache_spec.max_memory_usage_bytes(vllm_config), group.kv_cache_spec.page_size_bytes, @@ -1686,6 +1689,247 @@ def validate_kv_cache_layout( ) +def _layout_kv_cache_tensors( + kv_cache_groups: list[KVCacheGroupSpec], + num_blocks: int, + bytes_per_block: int, + layout: "KVCacheLayout", + size: int, + base_offset: int = 0, + pool_num_blocks: int | None = None, +) -> list[KVCacheTensor]: + """Place ``kv_cache_groups`` in one region of the backing allocation. + + The region starts ``base_offset`` bytes in and holds ``num_blocks`` blocks + of ``bytes_per_block``; ``size`` is the whole backing allocation. + ``pool_num_blocks`` tags the tensors of a pool other than the main one. + """ + interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None + + # Groups alias from byte 0. Spec regions are laid out differently: + # + # block-outer (the same packing repeats for every block): + # group 0: | blk 0 [ A | B | pad ] | blk 1 [ A | B | pad ] | ... + # group 1: | blk 0 [ C | D ] | blk 1 [ C | D ] | ... + # |<--- bytes_per_block -->| + # + # layer-outer (only supported for uniform page sizes or single-group models): + # group 0: | A [ blk 0 | blk 1 | ... ] | B [ blk 0 | blk 1 | ... ] | + # group 1: | C [ blk 0 | blk 1 | ... ] | D [ blk 0 | blk 1 | ... ] | + + kv_cache_tensors = [] + for group in kv_cache_groups: + group_spec = group.kv_cache_spec + layers_by_spec: defaultdict[KVCacheSpec, list[str]] = defaultdict(list) + if isinstance(group_spec, UniformTypeKVCacheSpecs): + for layer_name, spec in group_spec.kv_cache_specs.items(): + layers_by_spec[spec].append(layer_name) + elif group.layer_names: + layers_by_spec[group_spec].extend(group.layer_names) + + byte_offset = 0 + for spec, layer_names in layers_by_spec.items(): + agg = spec.aggregated_layer_count + # Layers fused into one shared page take one page slot per fused + # set (see get_tensor_slots); every layer of a set views it. + num_slots = len(layer_names) // agg if agg > 1 else len(layer_names) + layer_stride, block_stride, _, _, _ = compute_layout_strides( + spec, + num_blocks, + num_slots, + layout, + fixed_strides=(None, interleaved_block_stride, None, None, None), + ) + offset = ( + byte_offset + * max(layer_stride, spec.page_size_bytes) + // spec.page_size_bytes + ) + if agg > 1: + for slot, slot_layers in enumerate( + get_tensor_slots(KVCacheGroupSpec(layer_names, spec)) + ): + kv_cache_tensors.extend( + KVCacheTensor( + size=size, + layers=[layer_name], + layer_stride=layer_stride, + block_stride=block_stride, + offset=base_offset + offset + slot * layer_stride, + num_blocks=pool_num_blocks, + ) + for layer_name in slot_layers + ) + else: + kv_cache_tensors.append( + KVCacheTensor( + size=size, + layers=layer_names, + layer_stride=layer_stride, + block_stride=block_stride, + offset=base_offset + offset, + num_blocks=pool_num_blocks, + ) + ) + byte_offset += num_slots * spec.page_size_bytes + + return kv_cache_tensors + + +def _mamba_pool_eligible( + vllm_config: VllmConfig, kv_cache_specs: Iterable[KVCacheSpec] +) -> bool: + """Whether the O(1) state layers get their own pool (see + ``KVCacheConfig.mamba_pool_num_blocks``). + + Needs both kinds of layer, a Mamba cache mode with a per-request bound + ("none" or "align"), and none of the features that address blocks + across groups by a single id space: HiSparse host pools, KV transfer + connectors, context parallelism, and hidden-state cache layers. + """ + specs = list(kv_cache_specs) + if not any(isinstance(s, MambaSpec) and s.mamba_cache_mode != "all" for s in specs): + return False + if not any(not isinstance(s, MambaSpec) for s in specs): + return False + if any(isinstance(s, (HiddenStateCacheSpec, HiSparseHotSpec)) for s in specs): + return False + parallel_config = vllm_config.parallel_config + return ( + vllm_config.cache_config.mamba_cache_mode != "all" + and vllm_config.attention_config.hisparse_config is None + and vllm_config.kv_transfer_config is None + and parallel_config.decode_context_parallel_size == 1 + and parallel_config.prefill_context_parallel_size == 1 + ) + + +def _split_mamba_pool_groups( + vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec] +) -> tuple[list[KVCacheGroupSpec], list[KVCacheGroupSpec]] | None: + """``(main-pool groups, mamba-pool groups)``, or None for one shared pool.""" + if any(group.host_resident for group in kv_cache_groups): + return None + if _glm5_next_tensor_layout(kv_cache_groups) is not None: + return None + if not _mamba_pool_eligible( + vllm_config, (group.kv_cache_spec for group in kv_cache_groups) + ): + return None + mamba_groups = [ + group for group in kv_cache_groups if isinstance(group.kv_cache_spec, MambaSpec) + ] + main_groups = [ + group + for group in kv_cache_groups + if not isinstance(group.kv_cache_spec, MambaSpec) + ] + return main_groups, mamba_groups + + +def _mamba_pool_num_blocks( + vllm_config: VllmConfig, mamba_groups: list[KVCacheGroupSpec] +) -> int: + """Blocks for every running request's state in every mamba group, plus + the null block. A request holds at most ``max_memory_usage_bytes`` of a + group's state whatever its length, and at most ``max_num_seqs`` run.""" + per_request = sum( + cdiv( + group.kv_cache_spec.max_memory_usage_bytes(vllm_config), + group.kv_cache_spec.page_size_bytes, + ) + for group in mamba_groups + ) + return vllm_config.scheduler_config.max_num_seqs * per_request + 1 + + +def _sampler_reserve_bytes(vllm_config: VllmConfig) -> int: + """Headroom for the sampler warmup that runs after the KV cache is + allocated. Its O(max_num_seqs * vocab) tensors were freed before the + profiler handed over the budget, and a compressed pool leaves no + rounding slack to absorb them. ``VLLM_SAMPLER_RESERVE_MIB`` overrides + the estimate (vocab * max_num_seqs * 8 bytes * 5 tensors).""" + if envs.VLLM_SAMPLER_RESERVE_MIB: + return envs.VLLM_SAMPLER_RESERVE_MIB * 1024 * 1024 + try: + vocab = vllm_config.model_config.get_vocab_size() + except (AttributeError, ValueError): + return 512 * 1024 * 1024 + return vocab * vllm_config.scheduler_config.max_num_seqs * 8 * 5 + + +def _mamba_pool_fixed_bytes( + vllm_config: VllmConfig, mamba_groups: list[KVCacheGroupSpec] +) -> int: + """Bytes the split layout takes off the budget before attention: the + mamba pool and the sampler headroom.""" + return _mamba_pool_num_blocks( + vllm_config, mamba_groups + ) * _get_kv_cache_bytes_per_block(mamba_groups) + _sampler_reserve_bytes( + vllm_config + ) + + +def _get_split_pool_kv_cache_config( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], + main_groups: list[KVCacheGroupSpec], + mamba_groups: list[KVCacheGroupSpec], + available_memory: int, +) -> KVCacheConfig: + """Lay the main pool and the mamba pool out as two regions. + + Each region is placed exactly as one shared pool would place its groups + (``_layout_kv_cache_tensors``): the main pool's groups overlay each other + over ``num_blocks`` blocks, the mamba groups over the mamba pool's. The + mamba region follows the main one in the same backing allocation, so the + two never alias and a block id means a block of its own pool. + """ + layout = vllm_config.cache_config.get_resolved_kv_cache_layout() + validate_kv_cache_layout(layout, main_groups) + validate_kv_cache_layout(layout, mamba_groups) + main_bytes_per_block = _get_kv_cache_bytes_per_block(main_groups) + mamba_bytes_per_block = _get_kv_cache_bytes_per_block(mamba_groups) + mamba_num_blocks = _mamba_pool_num_blocks(vllm_config, mamba_groups) + + main_memory = available_memory - _mamba_pool_fixed_bytes(vllm_config, mamba_groups) + num_blocks = max(main_memory, 0) // main_bytes_per_block + num_blocks = may_override_num_blocks(vllm_config, num_blocks) + + main_bytes = main_bytes_per_block * num_blocks + size = main_bytes + mamba_bytes_per_block * mamba_num_blocks + kv_cache_tensors = _layout_kv_cache_tensors( + main_groups, num_blocks, main_bytes_per_block, layout, size + ) + _layout_kv_cache_tensors( + mamba_groups, + mamba_num_blocks, + mamba_bytes_per_block, + layout, + size, + base_offset=main_bytes, + pool_num_blocks=mamba_num_blocks, + ) + logger.info_once( + "Separate state pool for %d Mamba group(s): %d blocks (%s GiB); " + "%d blocks for the other %d group(s).", + len(mamba_groups), + mamba_num_blocks, + format_gib(mamba_bytes_per_block * mamba_num_blocks), + num_blocks, + len(main_groups), + ) + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=kv_cache_tensors, + kv_cache_groups=kv_cache_groups, + prefix_cache_retention_interval=( + vllm_config.cache_config.prefix_cache_retention_interval + ), + mamba_pool_num_blocks=mamba_num_blocks, + mamba_pool_offset=main_bytes, + ) + + def get_kv_cache_config_from_groups( vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec], @@ -1783,79 +2027,21 @@ def add_tensor(layer_name: str, spec: KVCacheSpec, offset: int) -> None: ), ) + if (split := _split_mamba_pool_groups(vllm_config, kv_cache_groups)) is not None: + return _get_split_pool_kv_cache_config( + vllm_config, kv_cache_groups, *split, available_memory + ) + layout = vllm_config.cache_config.get_resolved_kv_cache_layout() validate_kv_cache_layout(layout, kv_cache_groups) bytes_per_block = _get_kv_cache_bytes_per_block(kv_cache_groups) - interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None num_blocks = available_memory // bytes_per_block num_blocks = may_override_num_blocks(vllm_config, num_blocks) size = bytes_per_block * num_blocks - - # Groups alias from byte 0. Spec regions are laid out differently: - # - # block-outer (the same packing repeats for every block): - # group 0: | blk 0 [ A | B | pad ] | blk 1 [ A | B | pad ] | ... - # group 1: | blk 0 [ C | D ] | blk 1 [ C | D ] | ... - # |<--- bytes_per_block -->| - # - # layer-outer (only supported for uniform page sizes or single-group models): - # group 0: | A [ blk 0 | blk 1 | ... ] | B [ blk 0 | blk 1 | ... ] | - # group 1: | C [ blk 0 | blk 1 | ... ] | D [ blk 0 | blk 1 | ... ] | - - kv_cache_tensors = [] - for group in kv_cache_groups: - group_spec = group.kv_cache_spec - layers_by_spec: defaultdict[KVCacheSpec, list[str]] = defaultdict(list) - if isinstance(group_spec, UniformTypeKVCacheSpecs): - for layer_name, spec in group_spec.kv_cache_specs.items(): - layers_by_spec[spec].append(layer_name) - elif group.layer_names: - layers_by_spec[group_spec].extend(group.layer_names) - - byte_offset = 0 - for spec, layer_names in layers_by_spec.items(): - agg = spec.aggregated_layer_count - # Layers fused into one shared page take one page slot per fused - # set (see get_tensor_slots); every layer of a set views it. - num_slots = len(layer_names) // agg if agg > 1 else len(layer_names) - layer_stride, block_stride, _, _, _ = compute_layout_strides( - spec, - num_blocks, - num_slots, - layout, - fixed_strides=(None, interleaved_block_stride, None, None, None), - ) - offset = ( - byte_offset - * max(layer_stride, spec.page_size_bytes) - // spec.page_size_bytes - ) - if agg > 1: - for slot, slot_layers in enumerate( - get_tensor_slots(KVCacheGroupSpec(layer_names, spec)) - ): - kv_cache_tensors.extend( - KVCacheTensor( - size=size, - layers=[layer_name], - layer_stride=layer_stride, - block_stride=block_stride, - offset=offset + slot * layer_stride, - ) - for layer_name in slot_layers - ) - else: - kv_cache_tensors.append( - KVCacheTensor( - size=size, - layers=layer_names, - layer_stride=layer_stride, - block_stride=block_stride, - offset=offset, - ) - ) - byte_offset += num_slots * spec.page_size_bytes + kv_cache_tensors = _layout_kv_cache_tensors( + kv_cache_groups, num_blocks, bytes_per_block, layout, size + ) return KVCacheConfig( num_blocks=num_blocks, @@ -2393,16 +2579,28 @@ def get_kv_cache_groups( ] return packed_groups - # Prefer preserving each layer's cache semantics. If physical pages cannot - # be unified, try a supported allocation-only fallback before failing. - try: - filtered_spec = unify_kv_cache_spec_page_size(filtered_spec) - except NotImplementedError: - fallback_groups = _try_get_full_allocation_fallback_groups(kv_cache_spec) - if fallback_groups is None: - raise - return fallback_groups - groups = _get_kv_cache_groups_uniform_page_size(filtered_spec) + if not hidden_specs and _mamba_pool_eligible(vllm_config, filtered_spec.values()): + # The state layers get their own pool, so their pages are not unified + # with the attention pages (that unification is what sized every + # attention block like a recurrent state); only the attention layers + # are unified among themselves. + attention_spec = unify_kv_cache_spec_page_size( + {k: v for k, v in filtered_spec.items() if not isinstance(v, MambaSpec)} + ) + filtered_spec = {k: attention_spec.get(k, v) for k, v in filtered_spec.items()} + groups = _get_kv_cache_groups_uniform_page_size(filtered_spec) + else: + # Prefer preserving each layer's cache semantics. If physical pages + # cannot be unified, try a supported allocation-only fallback before + # failing. + try: + filtered_spec = unify_kv_cache_spec_page_size(filtered_spec) + except NotImplementedError: + fallback_groups = _try_get_full_allocation_fallback_groups(kv_cache_spec) + if fallback_groups is None: + raise + return fallback_groups + groups = _get_kv_cache_groups_uniform_page_size(filtered_spec) # Add hidden-state layers back with page aligned to the common page. if hidden_specs: @@ -2793,6 +2991,24 @@ def get_kv_cache_configs( _project_kv_cache_groups_to_worker(global_kv_cache_groups, worker_spec) for worker_spec in kv_cache_specs ] + # With a separate mamba pool, its fixed size comes off the top and + # everything below plans the main pool's groups against what is left. + allocated_groups_per_worker = projected_groups_per_worker + split_fixed_bytes = [0] * len(projected_groups_per_worker) + planning_groups_per_worker = [] + for i, groups in enumerate(projected_groups_per_worker): + split = _split_mamba_pool_groups(vllm_config, groups) if groups else None + if split is None: + planning_groups_per_worker.append(groups) + continue + main_groups, mamba_groups = split + planning_groups_per_worker.append(main_groups) + split_fixed_bytes[i] = _mamba_pool_fixed_bytes(vllm_config, mamba_groups) + available_memory = [ + avail_mem - fixed + for avail_mem, fixed in zip(available_memory, split_fixed_bytes) + ] + projected_groups_per_worker = planning_groups_per_worker # If `num_gpu_blocks_override` is set, the cache size that will actually # be allocated is decoupled from the profiled `available_memory`: @@ -2839,15 +3055,23 @@ def get_kv_cache_configs( ) kv_cache_configs: list[KVCacheConfig] = [] - for projected_groups, kv_cache_spec_one_worker, available_memory_one_worker in zip( - projected_groups_per_worker, kv_cache_specs, available_memory + for ( + projected_groups, + kv_cache_spec_one_worker, + available_memory_one_worker, + fixed, + ) in zip( + allocated_groups_per_worker, + kv_cache_specs, + available_memory, + split_fixed_bytes, ): assert sum(len(group.layer_names) for group in projected_groups) == len( kv_cache_spec_one_worker ), "Some layers are not assigned to any group." kv_cache_configs.append( get_kv_cache_config_from_groups( - vllm_config, projected_groups, available_memory_one_worker + vllm_config, projected_groups, available_memory_one_worker + fixed ) ) @@ -2864,7 +3088,10 @@ def get_kv_cache_configs( # strides and offsets stay consistent with the shrunken allocation. groups = kv_cache_config.kv_cache_groups kv_cache_configs[i] = get_kv_cache_config_from_groups( - vllm_config, groups, min_num_blocks * _pool_bytes_per_block(groups) + vllm_config, + groups, + min_num_blocks * _pool_bytes_per_block(planning_groups_per_worker[i]) + + split_fixed_bytes[i], ) for kv_cache_config in kv_cache_configs: diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 2ea1f5c39bf1..2ebb7c45bdbf 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -1454,6 +1454,9 @@ class KVCacheTensor: block_stride: int offset: int = 0 # byte offset of layers[0]'s block 0 host_resident: bool = False + # Blocks of the pool backing these layers, when it is not the main pool + # (KVCacheConfig.mamba_pool_num_blocks). None: the main pool. + num_blocks: int | None = None class KVCacheGroupRole(str, Enum): @@ -1515,6 +1518,34 @@ class KVCacheConfig: kv_tp_replicas: int = 1 """Consecutive TP ranks holding identical KV for every layer (1: none).""" + mamba_pool_num_blocks: int | None = None + """Blocks of the separate pool the O(1) state groups (Mamba/GDN) draw + from. None: every group shares the main pool of ``num_blocks`` blocks. + + A hybrid model's recurrent layers need a fixed number of blocks per + running request, never one per token. Sharing the attention pool sizes + their tensors for ``num_blocks`` blocks and unifies their page with the + attention page, which spends most of the KV budget on state no request + can use. A separate pool sized for ``max_num_seqs`` leaves the rest to + attention.""" + + mamba_pool_offset: int | None = None + """Byte offset of the separate pool's region in the backing allocation; + every view at or above it belongs to that pool.""" + + def uses_mamba_pool(self, group_id: int) -> bool: + """Whether cache group ``group_id`` draws from the separate pool.""" + return self.mamba_pool_num_blocks is not None and isinstance( + self.kv_cache_groups[group_id].kv_cache_spec, MambaSpec + ) + + def group_num_blocks(self, group_id: int) -> int: + """Blocks of the pool cache group ``group_id`` draws from.""" + if self.uses_mamba_pool(group_id): + assert self.mamba_pool_num_blocks is not None + return self.mamba_pool_num_blocks + return self.num_blocks + @cached_property def transfer_group_ids(self) -> tuple[int, ...]: """IDs of cache groups that participate in external KV transfer.""" @@ -1578,6 +1609,8 @@ def host_group_ids(self) -> tuple[int, ...]: def num_blocks_of(self, tensor: KVCacheTensor) -> int: """Number of blocks addressable by the pool backing ``tensor``.""" + if tensor.num_blocks is not None: + return tensor.num_blocks if not tensor.host_resident: return self.num_blocks assert self.hisparse_host_num_blocks is not None diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index e2f68ad59865..ede4bd407ec8 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -224,7 +224,7 @@ def set_dummy_context( input_batch: InputBatch, block_tables: "BlockTables", context_len: int, - num_kv_blocks: int, + num_kv_blocks: int | Sequence[int], max_model_len: int, input_block_tables: Sequence[torch.Tensor] | None = None, ) -> None: @@ -252,17 +252,22 @@ def set_dummy_context( input_batch.positions.copy_(torch.from_numpy(local_pos + context_len)) seq_len = context_len + query_len - for block_table, block_size, bpk in zip( + # One count for all groups, or each group's own pool (a separate mamba + # pool is far smaller than the main one). + if isinstance(num_kv_blocks, int): + num_kv_blocks = [num_kv_blocks] * len(input_block_tables) + for block_table, block_size, bpk, group_num_blocks in zip( input_block_tables, block_tables.kernel_block_sizes, block_tables.blocks_per_kv_block, + num_kv_blocks, ): num_blocks = min(cdiv(seq_len, block_size), block_table.shape[1]) # Spans are disjoint until the pool runs out, then they wrap and share # blocks: profiling only needs the reads to be realistic, not distinct. block_ids = torch.arange( num_reqs * num_blocks, dtype=block_table.dtype, device=block_table.device - ) % (num_kv_blocks * bpk) + ) % (group_num_blocks * bpk) block_table[:num_reqs, :num_blocks] = block_ids.view(num_reqs, num_blocks) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index c2f29972a26c..71e5a0d1e6e6 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -175,6 +175,7 @@ clear_layer_kv_caches, copy_kv_cache_blocks_inplace, get_uniform_decode_token_count, + mamba_pool_layout, ) from vllm.v1.worker.workspace import lock_workspace, use_workspace_lane @@ -1231,6 +1232,7 @@ def update_requests(self, scheduler_output: SchedulerOutput) -> None: self.kv_caches, self.kv_cache_config.num_blocks, scheduler_output.kv_cache_block_copies, + mamba_pool=mamba_pool_layout(self.kv_cache_config), ) def gather_batch_req_state( @@ -1821,7 +1823,12 @@ def execute_model( input_batch, self.block_tables, context_len, - self.kv_cache_config.num_blocks, + [ + self.kv_cache_config.group_num_blocks(group_id) + for group_id in range( + len(self.kv_cache_config.kv_cache_groups) + ) + ], self.max_model_len, input_block_tables=block_tables, ) diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 56f7d438fc3f..d2410e1d5cf7 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -121,22 +121,30 @@ def run_mixed_prefill_decode_warmup( for decode, prefill in zip(decode_block_counts, decode_prefill_block_counts) ] prefill_block_counts = [block_count(prefill_len, s) for s in kv_cache_specs] - required_blocks = sum(decode_block_counts) + sum(prefill_block_counts) - if model_runner.kv_cache_config.num_blocks <= required_blocks: - logger.warning( - "Skipping V2 mixed prefill+decode warmup because only %d KV blocks " - "are available for %d required warmup blocks.", - model_runner.kv_cache_config.num_blocks, - required_blocks, + kv_cache_config = model_runner.kv_cache_config + pool_of = [kv_cache_config.uses_mamba_pool(g) for g in range(num_kv_cache_groups)] + for pool in set(pool_of): + group_ids = [g for g in range(num_kv_cache_groups) if pool_of[g] == pool] + required_blocks = sum(decode_block_counts[g] for g in group_ids) + sum( + prefill_block_counts[g] for g in group_ids ) - return False + available = kv_cache_config.group_num_blocks(group_ids[0]) + if available <= required_blocks: + logger.warning( + "Skipping V2 mixed prefill+decode warmup because only %d KV blocks " + "are available for %d required warmup blocks.", + available, + required_blocks, + ) + return False - next_block_id = 1 + # Each pool numbers its blocks from 1 (0 is its null block). + next_block_id = dict.fromkeys(pool_of, 1) - def _alloc_blocks(num_blocks: int) -> list[int]: - nonlocal next_block_id - block_ids = list(range(next_block_id, next_block_id + num_blocks)) - next_block_id += num_blocks + def _alloc_blocks(num_blocks: int, group_id: int) -> list[int]: + pool = pool_of[group_id] + block_ids = list(range(next_block_id[pool], next_block_id[pool] + num_blocks)) + next_block_id[pool] += num_blocks return block_ids sampling_params = SamplingParams(max_tokens=2, temperature=0.0) @@ -149,7 +157,9 @@ def _alloc_blocks(num_blocks: int) -> list[int]: mm_features=[], sampling_params=sampling_params, pooling_params=None, - block_ids=tuple(_alloc_blocks(n) for n in decode_prefill_block_counts), + block_ids=tuple( + _alloc_blocks(n, g) for g, n in enumerate(decode_prefill_block_counts) + ), num_computed_tokens=0, lora_request=None, prefill_token_ids=decode_token_ids, @@ -161,7 +171,9 @@ def _alloc_blocks(num_blocks: int) -> list[int]: decode_prefill_output.total_num_scheduled_tokens = decode_prompt_len decode_prefill_output.num_common_prefix_blocks = [0] * num_kv_cache_groups - decode_new_blocks = tuple(_alloc_blocks(n) for n in decode_block_deltas) + decode_new_blocks = tuple( + _alloc_blocks(n, g) for g, n in enumerate(decode_block_deltas) + ) cached_decode_req = CachedRequestData.make_empty() cached_decode_req.req_ids = [decode_req_id] cached_decode_req.num_computed_tokens = [decode_prompt_len] @@ -179,7 +191,9 @@ def _alloc_blocks(num_blocks: int) -> list[int]: mm_features=[], sampling_params=sampling_params, pooling_params=None, - block_ids=tuple(_alloc_blocks(n) for n in prefill_block_counts), + block_ids=tuple( + _alloc_blocks(n, g) for g, n in enumerate(prefill_block_counts) + ), num_computed_tokens=0, lora_request=None, prefill_token_ids=prefill_token_ids, @@ -294,13 +308,23 @@ def _warmup_kernels( model_runner.scheduler_config.max_num_batched_tokens // max(prompt_len, decode_query_len), ) + kv_cache_config = model_runner.kv_cache_config + pool_of = [kv_cache_config.uses_mamba_pool(g) for g in range(len(kv_cache_specs))] if max_blocks_per_req > 0: # Reserve block 0 (null block) and ensure we have enough blocks. # Encoder-only models allocate no KV blocks, so this cap doesn't apply. - num_reqs = min( - num_reqs, - max(1, (model_runner.kv_cache_config.num_blocks - 1) // max_blocks_per_req), - ) + # Each pool caps the request count on its own groups' blocks. + for pool in set(pool_of): + group_ids = [g for g in range(len(pool_of)) if pool_of[g] == pool] + per_req = sum(decode_block_counts[g] for g in group_ids) + if per_req > 0: + num_reqs = min( + num_reqs, + max( + 1, + (kv_cache_config.group_num_blocks(group_ids[0]) - 1) // per_req, + ), + ) req_ids = [f"_warmup_{i}_" for i in range(num_reqs)] @@ -316,12 +340,15 @@ def _warmup_kernels( sampling_params = SamplingParams.for_sampler_warmup() pooling_params = None - # Assign distinct block IDs per request per group. 0 null block, start from 1. - next_block_id = 1 + # Assign distinct block IDs per request per group. 0 null block, start from 1; + # each pool numbers its own blocks. + next_block_id = dict.fromkeys(pool_of, 1) - def _alloc_blocks(num_blocks: int) -> list[int]: - nonlocal next_block_id - return list(range(next_block_id, next_block_id := next_block_id + num_blocks)) + def _alloc_blocks(num_blocks: int, group_id: int) -> list[int]: + pool = pool_of[group_id] + start = next_block_id[pool] + next_block_id[pool] = start + num_blocks + return list(range(start, start + num_blocks)) # The KV-block zeroing kernel is driven by the scheduler's # new_block_ids_to_zero, so none of the steps below reach it. @@ -338,7 +365,9 @@ def _alloc_blocks(num_blocks: int) -> list[int]: pooling_params, mm_features=warmup_mm_features, ), - block_ids=tuple(_alloc_blocks(n) for n in prefill_block_counts), + block_ids=tuple( + _alloc_blocks(n, g) for g, n in enumerate(prefill_block_counts) + ), prefill_token_ids=prompt_token_ids, ) for i in range(num_reqs) @@ -394,7 +423,9 @@ def _run_decode_step(indices: list[int], spec_flags: list[bool]) -> None: for spec, held in zip(kv_cache_specs, req_blocks[i]) ] cached_req_data.new_block_ids.append( - tuple(_alloc_blocks(n) for n in deltas) if any(deltas) else None + tuple(_alloc_blocks(n, g) for g, n in enumerate(deltas)) + if any(deltas) + else None ) req_blocks[i] = [ held + delta for held, delta in zip(req_blocks[i], deltas) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 109046665016..76bfe6f10a81 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -234,6 +234,7 @@ allocate_kv_cache, bind_kv_cache, copy_kv_cache_blocks_inplace, + mamba_pool_layout, prepare_kernel_block_sizes, sanity_check_mm_encoder_outputs, ) @@ -1225,6 +1226,7 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None self.kv_caches, self.kv_cache_config.num_blocks, scheduler_output.kv_cache_block_copies, + mamba_pool=mamba_pool_layout(self.kv_cache_config), ) # Free the cached encoder outputs. diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index a7e7deff7be9..96acfd42a2c4 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -935,9 +935,7 @@ def compile_or_warm_up_model(self) -> CompilationTimes: # for backends that don't define the hook. _backends = _backends_in_use(self) for _backend_cls in _backends: - _call_backend_hook( - _backend_cls, "on_kv_cache_initialized", _backends, self - ) + _call_backend_hook(_backend_cls, "on_kv_cache_initialized", _backends, self) warmup_sizes: list[int] = [] diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index a6ef84bbd3d7..444d323cc39d 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -688,13 +688,45 @@ def clear_layer_kv_caches(layers: Iterable[Any]) -> None: layer.impl._v_scale_cache = None +def mamba_pool_layout(kv_cache_config: KVCacheConfig) -> tuple[int, int] | None: + """``(byte offset, blocks)`` of the separate mamba pool's region, or None.""" + if kv_cache_config.mamba_pool_num_blocks is None: + return None + assert kv_cache_config.mamba_pool_offset is not None + return kv_cache_config.mamba_pool_offset, kv_cache_config.mamba_pool_num_blocks + + def copy_kv_cache_blocks_inplace( kv_caches: Iterable[torch.Tensor], num_blocks: int, kv_cache_block_copies: Sequence[KVCacheBlockCopy], + mamba_pool: tuple[int, int] | None = None, ) -> None: + """Apply the scheduler's block copies to every cache. + + With a separate mamba pool (``mamba_pool_layout``), copy ids at or above + ``num_blocks`` name blocks of that pool (the scheduler shifts them), and + apply only to the caches whose views lie in its region of the backing + allocation; the rest apply to the main pool's caches. + """ if not kv_cache_block_copies: return + if mamba_pool is not None: + region_start, mamba_num_blocks = mamba_pool + main_caches: list[torch.Tensor] = [] + mamba_caches: list[torch.Tensor] = [] + for cache in kv_caches: + in_mamba = cache.storage_offset() * cache.element_size() >= region_start + (mamba_caches if in_mamba else main_caches).append(cache) + main_copies = [c for c in kv_cache_block_copies if c[0] < num_blocks] + mamba_copies = [ + KVCacheBlockCopy(c[0] - num_blocks, c[1] - num_blocks) + for c in kv_cache_block_copies + if c[0] >= num_blocks + ] + copy_kv_cache_blocks_inplace(main_caches, num_blocks, main_copies) + copy_kv_cache_blocks_inplace(mamba_caches, mamba_num_blocks, mamba_copies) + return indices_np = np.array(kv_cache_block_copies, dtype=np.int64) indices: torch.Tensor | None = None