From cc59ce9a672206ffaee9c56770761b73bb79bfe7 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Sun, 25 Oct 2026 17:20:05 +0100 Subject: [PATCH] [KV] Pack recurrent depth planes for partially recurrent models LAST_EXITED keeps one KV plane per recurrence depth and reserves every layer's slot in every plane. For a model whose prelude and coda layers run once per token (Huginn), only plane 0 ever holds those layers: of Huginn-0125's 8 layers x 32 planes = 256 layer slots per token, 124 are never written. The pool fills with dead slots, so few requests fit and throughput stops growing with concurrency. With CacheConfig.pack_recurrent_planes (default on) each plane maps only the layer slots it can hold: plane 0 all layers, deeper planes the recurrent ones. A page is a run of blocks of num_layers slots, entry = block * num_layers + slot in a flattened view of the pool, so a Huginn page needs 17 blocks instead of 32. Writes, attention, early-exit plane copies and the async finalize kernel take an entry stride and first-layer offset; PD sends each plane's live slot range (plane_slots, checked on connect) and the coordinator prices requests in blocks per page. Values never change, only their addresses. Fully recurrent models (Ouro) keep the old layout. --no-pack-recurrent-planes restores the unpacked layout. GH200, Huginn-0125 (32 loops, CUDA graphs, sync), same KV pool of 24,165 blocks, 512-token prompts, 64 output tokens, one run per level: tokens in the pool: 12,080 unpacked -> 22,736 packed (1.88x) throughput: 0.99x / 0.99x / 1.14x / 1.16x / 1.06x at 1 / 8 / 32 / 64 / 128 median TTFT at 64 requests: 15.7 -> 10.9 s 2048-token prompts: 1.13x / 1.11x at 8 / 32 token identity: 233/233 requests PD: 2.79 MB per prompt token sent instead of 5.41 MB (only live slots). Tests: CPU unit tests for the packed layout, routing, PD segments and slot checks; GPU test comparing packed and unpacked Huginn (tiny model, CUDA graphs, early-exit copies) token for token. Signed-off-by: 0z5a Co-Authored-By: Claude Opus 5.5 --- docs/kv_cache_walkthrough.md | 6 +- docs/kv_layout_computation.md | 19 +++ tests/test_huginn.py | 161 ++++++++++++++++++++++ tests/test_pd.py | 41 ++++++ vllm_rlt/config.py | 4 +- vllm_rlt/core/kv_cache_manager.py | 198 ++++++++++++++++++++++----- vllm_rlt/core/memory.py | 11 +- vllm_rlt/core/scheduler.py | 2 +- vllm_rlt/engine/stage_engine.py | 1 + vllm_rlt/engine/worker.py | 1 + vllm_rlt/entrypoints/runtime_args.py | 7 + vllm_rlt/kernels/routing.py | 16 ++- vllm_rlt/pd/engine.py | 9 +- vllm_rlt/pd/transport.py | 23 +++- vllm_rlt/worker/async_speculative.py | 2 +- vllm_rlt/worker/async_state.py | 3 +- vllm_rlt/worker/cuda_graph.py | 19 +-- vllm_rlt/worker/metadata.py | 4 +- vllm_rlt/worker/model_runner.py | 16 ++- 19 files changed, 465 insertions(+), 78 deletions(-) diff --git a/docs/kv_cache_walkthrough.md b/docs/kv_cache_walkthrough.md index 4ea2661e..95e20bf9 100644 --- a/docs/kv_cache_walkthrough.md +++ b/docs/kv_cache_walkthrough.md @@ -27,7 +27,11 @@ Code: [manager](../vllm_rlt/core/kv_cache_manager.py), ## Capacity, validity and lifetime `required_blocks(tokens)` counts physical blocks across storage planes. -LAST_EXITED has one plane per recurrent depth; SHARED has one plane. A request's +LAST_EXITED has one plane per recurrent depth; SHARED has one plane. With plane +packing (see [KV layout computation](kv_layout_computation.md#7-plane-packing-for-partially-recurrent-models)) +several deeper planes share one physical block, `block_tables` has +`physical_planes` entries and `_entries(allocation, depth)` maps a depth to +layer-flattened page entries; written state stays per logical plane. A request's maximum input extent is prompt length plus `max_tokens - 1`: the last sampled output is not fed back into the model. diff --git a/docs/kv_layout_computation.md b/docs/kv_layout_computation.md index 15e8ceae..7df557c2 100644 --- a/docs/kv_layout_computation.md +++ b/docs/kv_layout_computation.md @@ -304,3 +304,22 @@ shared: ceil(T/B) physical blocks The cache pool is preallocated using `num_blocks`. With the same `num_blocks`, both layouts allocate the same KV tensor bytes; `shared` allows that pool to cover more token positions. Early exit under `last_exited` does not automatically release deeper storage because those positions retain the copied history. `shared` reduces block demand and eliminates finalization copies, but sequential prefill positions within a request can limit performance. The layout alone does not establish which implementation is faster. Because historical attention inputs differ, the layouts are not interchangeable as a lossless memory optimization, and this example does not establish an accuracy ranking. + +## 7. Plane packing for partially recurrent models + +Huginn runs a prelude and a coda once per token and repeats only the recurrent block. Prelude and coda attention use depth 0. A `last_exited` block in any deeper plane therefore holds live KV only in the recurrent layers' slots; the other layer slots are allocated but never written or read. For Huginn-0125 (2 prelude, 4 recurrent, 2 coda layers, 32 depths) 124 of 256 layer slots per logical page are dead. + +When the recurrent layers are contiguous and occupy at most half of the layers, `CacheConfig(pack_recurrent_planes=True)` (the default) stores `k = num_layers // num_recurrent_layers` deeper planes in one physical block. Each packed plane uses its own slot range. Plane 0 still holds every layer: + +```text +physical planes per page = 1 + ceil((D - 1) / k) # Huginn-0125: 1 + ceil(31 / 2) = 17 +plane 0 page entry = physical_block * num_layers +plane d >= 1 page entry = physical_block * num_layers + member * num_recurrent - first_recurrent +member = (d - 1) % k, physical plane = 1 + (d - 1) // k +``` + +`get_block_table(request_id, depth)` returns these page entries. A layer's K/V for a row lives at page `entry + layer` of the layer-flattened cache `key_cache.view(num_blocks * num_layers, block_size, kv_heads, head_dim)`. Writes, attention, CUDA graphs, the device-resident routing tables and finalization (the Triton `finalize_kernel` copies `finalize_layers` slots from each entry with `entry_stride`) add the layer to the entry instead of indexing the layer axis. The stored values and the attention arithmetic are unchanged; only addresses move, so packed and unpacked execution produce bitwise-identical tokens and KV. Nothing is copied or recomputed. + +`plane_block_tables` and the allocation views count physical planes, so `required_blocks` shrinks by `D / physical planes` (1.88x for Huginn-0125) and the same KV pool holds that many more tokens. `transfer_slots` reports each physical plane's live slot range. PD sends only those slots, packed or not: 132 of 256 layer slots per Huginn-0125 page, and the coordinator prices reservations in `blocks_per_page` physical blocks. + +Models whose layers all recur (Ouro, Nanbeige) are unaffected: `k` is 1 and every path keeps the unpacked addressing. A non-recurrent layer may be accessed only at depth 0 in a packed cache; the manager rejects any other access because those slots belong to another plane. diff --git a/tests/test_huginn.py b/tests/test_huginn.py index 2a606835..391820d6 100644 --- a/tests/test_huginn.py +++ b/tests/test_huginn.py @@ -381,3 +381,164 @@ def test_huginn_graph_states_logits_and_kv_match_eager(): for cache in caches: cache.free("request") assert cache.num_used_blocks == 0 + + +def test_plane_packing_geometry_and_live_transfer_slots(): + from vllm_rlt.core.kv_cache_manager import physical_planes, plane_packing + + # Huginn-0125: 2 prelude + 4 recurrent + 2 coda layers, 32 recurrences. + assert plane_packing(8, range(2, 6), 32) == 2 + assert physical_planes(8, range(2, 6), 32) == 17 + assert physical_planes(8, range(2, 6), 32, enabled=False) == 32 + assert physical_planes(8, range(2, 6), 32, layout="shared") == 1 + # Fully recurrent models (Ouro, Nanbeige) and non-contiguous cores never pack. + assert plane_packing(24, None, 4) == 1 + assert plane_packing(8, (1, 2, 4, 5), 32) == 1 + assert plane_packing(8, range(1, 6), 32) == 1 # Five of eight layers recur. + assert physical_planes(9, range(3, 6), 8) == 1 + 3 # Three planes per block. + + def cache(pack, loops=5): + return KVCacheManager(4, 1, 8, 64, 2, loops, recurrent_layers=(1, 2), **pack) + + packed, unpacked = cache(dict(pack_recurrent_planes=True)), cache({}) + assert (packed.plane_packing, packed.physical_planes) == (2, 3) + assert packed.transfer_slots == ((0, 4), (0, 4), (0, 4)) + assert unpacked.transfer_slots == ((0, 4),) + ((1, 3),) * 4 + assert cache(dict(pack_recurrent_planes=True), loops=4).transfer_slots == ( + (0, 4), + (0, 4), + (0, 2), # The last group holds one plane. + ) + assert packed.required_blocks(5) == 3 * 3 and unpacked.required_blocks(5) == 3 * 5 + # The finalize kernel copies recurrent slots of flattened pages when packed. + assert packed.finalize_layers == (1, 2) and unpacked.finalize_layers == (0, 4) + assert packed.entry_stride == packed.key_cache.stride(1) + assert unpacked.entry_stride == unpacked.key_cache.stride(0) + + +def test_packed_entries_address_disjoint_live_slots(): + cache = KVCacheManager(4, 1, 8, 64, 2, 5, recurrent_layers=(1, 2), pack_recurrent_planes=True) + assert cache.allocate("a", 5) and cache.allocate("b", 3) + assert cache.num_used_blocks == 3 * 3 + 2 * 3 + addresses = {} + for rid in ("a", "b"): + for depth in range(5): + layers = range(4) if depth == 0 else (1, 2) + for entry in cache.get_block_table(rid, depth): + for layer in layers: + page = entry + layer + assert 0 <= page < cache.num_blocks * cache.num_layers + assert page not in addresses, "two live (depth, layer) slots alias" + addresses[page] = (rid, depth, layer) + physical = {b for rid in ("a", "b") for table in cache.plane_block_tables(rid) for b in table} + assert {page // cache.num_layers for page in addresses} == physical + # Growth keeps entries in step with the physical tables. + assert cache.allocate("c", 9, initial_tokens=2) and cache.ensure_capacity("c", 9) + grown = cache._get_allocation("c") + assert len(grown.block_tables[0]) == 5 and all(len(t) == 5 for t in grown.entry_tables) + assert grown.entry_tables[0] == tuple(b * 4 for b in grown.block_tables[0]) + assert grown.entry_tables[2] == tuple(b * 4 + 2 - 1 for b in grown.block_tables[1]) + batch = cache._prepare_batch(["a"], [1], [0]) + with pytest.raises(ValueError, match="depth zero"): + cache._write_prepared(0, batch, torch.zeros(1, 1, 8), torch.zeros(1, 1, 8)) + for rid in ("a", "b", "c"): + cache.free(rid) + assert cache.num_used_blocks == 0 + + +@pytest.mark.parametrize("max_loops", [None, 2]) +@pytest.mark.parametrize("temperature", [0, 0.7]) +def test_packed_planes_match_unpacked_outputs_and_kv_bitwise(model, max_loops, temperature): + """Packing changes addresses only: tokens, exit copies and every live KV byte agree.""" + model = HuginnForCausalLM(tiny_huginn_config(mean_recurrence=5)).eval() + prompts = {"first": [4, 7, 3, 9, 2], "second": [9, 2, 5]} + runs = [] + for pack in (False, True): + engine = LLMEngine( + model, + cache_config=CacheConfig(num_blocks=128, block_size=2, pack_recurrent_planes=pack), + scheduler_config=SchedulerConfig( + max_num_seqs=2, max_num_batched_tokens=8, prefill_chunk_size=3 + ), + ) + cache = engine.cache_manager + assert cache.plane_packing == (2 if pack else 1) + params = SamplingParams( + max_tokens=4, max_loops=max_loops, ignore_eos=True, seed=11, temperature=temperature + ) + for name, prompt in prompts.items(): + engine.add_request(name, prompt, params) + outputs, kv, peak = {}, {}, 0 + while engine.has_unfinished_requests(): + for output in engine.step(): + if output.finished: + outputs[output.request_id] = (output.token_ids, output.exit_depths) + peak = max(peak, cache.num_used_blocks) + for name in prompts: + if not cache.has_allocation(name): + continue + allocation = cache._get_allocation(name) + for depth in range(5): + for layer in range(4) if depth == 0 else model.recurrent_kv_layers: + length = allocation.written[depth][layer].prefix + if length: + kv[name, depth, layer, length] = cache.read(layer, name, depth, length) + assert cache.num_used_blocks == 0 + runs.append((outputs, kv, peak)) + (unpacked, unpacked_kv, unpacked_peak), (packed, packed_kv, packed_peak) = runs + assert packed == unpacked + if max_loops is not None: + # Decode exits early, so finalization copies depth-2 KV into packed planes. + assert all(depths[1:] == [2] * 3 for _, depths in packed.values()) + assert packed_kv.keys() == unpacked_kv.keys() and packed_kv + for key, (k, v) in packed_kv.items(): + torch.testing.assert_close(k, unpacked_kv[key][0], rtol=0, atol=0) + torch.testing.assert_close(v, unpacked_kv[key][1], rtol=0, atol=0) + # Five planes need three physical blocks per page instead of five. + assert packed_peak * 5 == unpacked_peak * 3 + + +@pytest.mark.gpu +def test_packed_planes_match_unpacked_on_gpu(): + """CUDA graphs and exit copies address packed planes (Huginn runs synchronously).""" + torch.manual_seed(42) + model = ( + HuginnForCausalLM( + tiny_huginn_config(n_embd=384, n_heads=4, intermediate_size=768, mean_recurrence=5) + ) + .to(device="cuda", dtype=torch.bfloat16) + .eval() + ) + prompts = ([4, 7, 3], [9, 2], [3, 8, 7, 6], [4]) + runs = [] + for pack in (False, True): + engine = LLMEngine( + model, + attention_backend="triton", + cache_config=CacheConfig(num_blocks=256, block_size=4, pack_recurrent_planes=pack), + scheduler_config=SchedulerConfig( + max_num_seqs=4, max_num_batched_tokens=8, prefill_chunk_size=2 + ), + execution_config=ExecutionConfig(cuda_graphs=True, cuda_graph_max_graphs=16), + ) + cache = engine.cache_manager + assert cache.plane_packing == (2 if pack else 1) + torch.manual_seed(123) + for i, prompt in enumerate(prompts): + # Exit at depth 2 of 5: finalization copies into three deeper planes. + engine.add_request( + str(i), + prompt, + SamplingParams(max_tokens=4 + i, min_loops=2, max_loops=2, ignore_eos=True), + ) + finished, peak = {}, 0 + while engine.has_unfinished_requests(): + for output in engine.step(): + if output.finished: + finished[output.request_id] = (output.token_ids, output.exit_depths) + peak = max(peak, cache.num_used_blocks) + assert len(finished) == 4 and cache.num_used_blocks == 0 + runs.append((finished, peak)) + (unpacked, unpacked_peak), (packed, packed_peak) = runs + assert packed == unpacked + assert packed_peak * 5 == unpacked_peak * 3 diff --git a/tests/test_pd.py b/tests/test_pd.py index 22ced40c..8f426644 100644 --- a/tests/test_pd.py +++ b/tests/test_pd.py @@ -628,6 +628,7 @@ def test_pd_reservations_are_reported_before_worker_acknowledgments(): num_blocks=64, block_size=2, depths=4, + blocks_per_page=4, watermark_blocks=0, info={"agent": "test"}, ) @@ -642,3 +643,43 @@ def test_pd_reservations_are_reported_before_worker_acknowledgments(): assert engine.stats() == dict(requests=1, reserved_blocks=20) assert len(messages) == 1 and messages[0][0] == ("decode", "reserve") assert not hasattr(engine, "cache_manager") and not hasattr(engine, "scheduler") + + +@pytest.mark.parametrize("pack", [False, True]) +def test_segments_send_exactly_live_huginn_slots(pack): + # Prelude 0, recurrent 1-2, coda 3; five depths. + c = KVCacheManager(4, 2, 8, 64, 4, 5, recurrent_layers=(1, 2), pack_recurrent_planes=pack) + c.allocate("a", 9) + k = c.key_cache + info = dict( + block_size=4, + num_blocks=64, + num_layers=4, + device=0, + block_bytes=k.stride(0) * 4, + layer_bytes=k.stride(1) * 4, + token_bytes=k.stride(2) * 4, + key_ptr=0, + value_ptr=10**9, + plane_slots=c.transfer_slots, + ) + expected = set() + for depth in range(5): + for token in range(2, 9): + for layer in range(4) if depth == 0 else (1, 2): + page = c.get_block_table("a", depth)[token // 4] + if not pack: + page = page * 4 + layer + else: + page += layer + for base in (0, 10**9): + pos = base + page * info["layer_bytes"] + (token % 4) * info["token_bytes"] + expected.update(range(pos, pos + info["token_bytes"])) + actual = set() + segments = list(kv_segments(info, c.plane_block_tables("a"), 2, 9)) + for address, size, _ in segments: + assert not actual & set(range(address, address + size)) + actual.update(range(address, address + size)) + assert actual == expected + # Huginn-0125 shape: live bytes are 132 of 256 layer planes either way. + assert len(expected) == 2 * 7 * (4 + 4 * 2) * info["token_bytes"] diff --git a/vllm_rlt/config.py b/vllm_rlt/config.py index 8872d018..0a0d5b7c 100644 --- a/vllm_rlt/config.py +++ b/vllm_rlt/config.py @@ -20,9 +20,11 @@ class CacheConfig: incremental_allocation: bool = False # Admission headroom for incremental allocation; otherwise effective watermark is zero. watermark_ratio: float = 0.0 + # Share physical blocks among deeper planes when only some layers recur. + pack_recurrent_planes: bool = True def __post_init__(self): - for name in ("enable_prefix_caching", "incremental_allocation"): + for name in ("enable_prefix_caching", "incremental_allocation", "pack_recurrent_planes"): if type(getattr(self, name)) is not bool: raise ValueError(f"{name} must be a boolean") if not 0 <= self.watermark_ratio < 1: diff --git a/vllm_rlt/core/kv_cache_manager.py b/vllm_rlt/core/kv_cache_manager.py index 3bb6bf02..267e7177 100644 --- a/vllm_rlt/core/kv_cache_manager.py +++ b/vllm_rlt/core/kv_cache_manager.py @@ -4,9 +4,17 @@ SHARED overwrites a single plane and needs no exit copies. Optional incremental allocation separates logical capacity from physical pages. Immutable full-depth prompt blocks can be shared through the prefix cache. + +Plane packing: when only a contiguous subset of layers recurs (Huginn's core +between a prelude and coda), planes deeper than zero leave the other layer +slots of every block unused. Packing stores several such planes in one physical +block, each in its own slot range, and addresses a row's layer as page +``entry + layer`` of the layer-flattened cache. Nothing is copied or recomputed; +only block tables change. """ import hashlib +import math import struct from collections import OrderedDict from collections.abc import Callable, Sequence @@ -53,10 +61,32 @@ def __contains__(self, position: int) -> bool: @dataclass class _Allocation: max_tokens: int + # Physical blocks, one table per physical plane. block_tables: tuple[tuple[int, ...], ...] written: list[list[_WrittenPositions]] transfer_leases: set[str] = field(default_factory=set) release_requested: bool = False + # Packed layouts only: per logical plane, the layer-flattened page entry of + # each logical page (physical block * layers + plane slot base). + entry_tables: tuple[tuple[int, ...], ...] | None = None + + +def plane_packing(num_layers, recurrent_layers, max_loops, layout="last_exited", enabled=True): + """How many depth planes share one physical block (1 means no packing).""" + layers = sorted(range(num_layers) if recurrent_layers is None else recurrent_layers) + contiguous = layers == list(range(layers[0], layers[0] + len(layers))) + factor = num_layers // len(layers) + if not enabled or layout != "last_exited" or max_loops < 2 or not contiguous or factor < 2: + return 1 + return factor + + +def physical_planes(num_layers, recurrent_layers, max_loops, layout="last_exited", enabled=True): + """Physical blocks per logical page across every depth plane.""" + if layout != "last_exited": + return 1 + factor = plane_packing(num_layers, recurrent_layers, max_loops, layout, enabled) + return max_loops if factor == 1 else 1 + math.ceil((max_loops - 1) / factor) @dataclass(frozen=True) @@ -103,6 +133,7 @@ def __init__( recurrent_layers: Sequence[int] | None = None, *, attention: AttentionBackend | None = None, + pack_recurrent_planes: bool = False, ): for name, value in ( ("num_layers", num_layers), @@ -133,6 +164,13 @@ def __init__( ): raise ValueError("recurrent_layers must be valid cache layer indices") self.storage_depths = max_loops if layout == "last_exited" else 1 + self.plane_packing = plane_packing( + num_layers, self.recurrent_layers, max_loops, layout, pack_recurrent_planes + ) + self.physical_planes = physical_planes( + num_layers, self.recurrent_layers, max_loops, layout, pack_recurrent_planes + ) + self._recurrent_set = frozenset(self.recurrent_layers) self.device = torch.device(device) self.dtype = dtype self.attention = attention or create_backend( @@ -143,6 +181,9 @@ def __init__( shape = (num_blocks, num_layers, block_size, num_kv_heads, head_dim) self.key_cache = torch.empty(shape, device=self.device, dtype=dtype) self.value_cache = torch.empty_like(self.key_cache) + # Layer-flattened views: page (block * num_layers + layer) is one layer's block. + self._flat_keys = self.key_cache.view(-1, block_size, num_kv_heads, head_dim) + self._flat_values = self.value_cache.view(-1, block_size, num_kv_heads, head_dim) # Resolve an implicit CUDA index to the actual storage device once. self.device = self.key_cache.device self._free_blocks = list(reversed(range(num_blocks))) @@ -186,7 +227,85 @@ def bytes_per_block(self) -> int: def required_blocks(self, max_tokens: int) -> int: if not isinstance(max_tokens, Integral) or isinstance(max_tokens, bool) or max_tokens <= 0: raise ValueError("max_tokens must be a positive integer") - return ((max_tokens + self.block_size - 1) // self.block_size) * self.storage_depths + return ((max_tokens + self.block_size - 1) // self.block_size) * self.physical_planes + + @property + def transfer_slots(self) -> tuple[tuple[int, int], ...]: + """Per physical plane, the layer slot range holding live KV. + + Plane zero holds every layer. Deeper planes hold only recurrent layers, + in their own slots when unpacked or packed from slot zero when packed. + """ + full = (0, self.num_layers) + if self.physical_planes == 1: + return (full,) + count = len(self.recurrent_layers) + if self.plane_packing > 1: + deeper = self.storage_depths - 1 + return (full,) + tuple( + (0, min(self.plane_packing, deeper - group * self.plane_packing) * count) + for group in range(self.physical_planes - 1) + ) + layers = sorted(self.recurrent_layers) + if layers != list(range(layers[0], layers[0] + count)): + return (full,) * self.physical_planes + return (full,) + ((layers[0], layers[0] + count),) * (self.physical_planes - 1) + + @property + def entry_stride(self) -> int: + """Element stride between consecutive page entries of ``get_block_table``.""" + return self.key_cache.stride(1 if self.plane_packing > 1 else 0) + + @property + def finalize_layers(self) -> tuple[int, int]: + """Layer slot range an early-exit copy moves: ``(first, count)``.""" + if self.plane_packing > 1: + return min(self.recurrent_layers), len(self.recurrent_layers) + return 0, self.num_layers + + @staticmethod + def entry_tables(allocation: _Allocation) -> tuple[tuple[int, ...], ...]: + """Per logical plane, the page entries rows address (physical blocks unless packed).""" + if allocation.entry_tables is None: + return allocation.block_tables + return allocation.entry_tables + + def _entries(self, allocation: _Allocation, depth: int) -> tuple[int, ...]: + """Page entries a row at ``depth`` addresses: physical blocks unless packed.""" + return self.entry_tables(allocation)[self._plane(depth)] + + def _refresh_entries(self, allocation: _Allocation) -> None: + if self.plane_packing == 1: + return + tables = [tuple(b * self.num_layers for b in allocation.block_tables[0])] + count, first = len(self.recurrent_layers), min(self.recurrent_layers) + for plane in range(1, self.storage_depths): + group, member = divmod(plane - 1, self.plane_packing) + base = member * count - first + physical = allocation.block_tables[1 + group] + tables.append(tuple(b * self.num_layers + base for b in physical)) + allocation.entry_tables = tuple(tables) + + def _layer_kv(self, layer: int, tables: torch.Tensor): + """K/V block views and tables for one layer's attention call.""" + if self.plane_packing == 1: + return self.key_cache[:, layer], self.value_cache[:, layer], tables + return self._flat_keys, self._flat_values, tables + layer + + def _store_layer(self, layer, blocks, offsets, k, v) -> None: + if self.plane_packing == 1: + self.key_cache[blocks, layer, offsets] = k + self.value_cache[blocks, layer, offsets] = v + else: + self._flat_keys[blocks + layer, offsets] = k + self._flat_values[blocks + layer, offsets] = v + + def _require_packed_layer(self, layer, rows) -> None: + # A boundary layer exists only in plane zero; elsewhere its slots belong + # to another packed plane. + if self.plane_packing > 1 and layer not in self._recurrent_set: + if any(self._plane(row[1]) for row in rows): + raise ValueError("non-recurrent layers have KV only at depth zero") def _drop_refs(self, blocks): for b in blocks: @@ -296,14 +415,14 @@ def allocate(self, request_id: str, max_tokens: int, *, initial_tokens=None, pre claimed = [b for group in prefix for b in group] for b in claimed: self._refs[b] += 1 - fresh = self._claim((pages - len(prefix)) * self.storage_depths) + fresh = self._claim((pages - len(prefix)) * self.physical_planes) if fresh is None: self._drop_refs(claimed) return False tail = pages - len(prefix) tables = tuple( tuple(g[d] for g in prefix) + tuple(fresh[d * tail : (d + 1) * tail]) - for d in range(self.storage_depths) + for d in range(self.physical_planes) ) length = len(prefix) * self.block_size self._allocations[request_id] = _Allocation( @@ -314,6 +433,7 @@ def allocate(self, request_id: str, max_tokens: int, *, initial_tokens=None, pre for _ in range(self.storage_depths) ], ) + self._refresh_entries(self._allocations[request_id]) self.prefix_queries += 1 self.prefix_hits += length return True @@ -326,13 +446,14 @@ def ensure_capacity(self, request_id, tokens): extra = pages - len(allocation.block_tables[0]) if extra <= 0: return True - blocks = self._claim(extra * self.storage_depths) + blocks = self._claim(extra * self.physical_planes) if blocks is None: return False allocation.block_tables = tuple( t + tuple(blocks[d * extra : (d + 1) * extra]) for d, t in enumerate(allocation.block_tables) ) + self._refresh_entries(allocation) return True def truncate_suffix(self, request_id: str, frontier: int) -> None: @@ -397,18 +518,20 @@ def has_allocation(self, request_id: str) -> bool: def allocated_blocks(self, request_id: str) -> int: """How many physical blocks this request holds, counting every plane.""" - return len(self._get_allocation(request_id).block_tables[0]) * self.storage_depths + return len(self._get_allocation(request_id).block_tables[0]) * self.physical_planes def get_block_table(self, request_id: str, depth: int) -> tuple[int, ...]: + """Page entries for ``depth``: physical blocks, or flattened pages when packed.""" self._validate_depth(depth) - return self._get_allocation(request_id).block_tables[self._plane(depth)] + return self._entries(self._get_allocation(request_id), depth) def plane_block_tables(self, request_id: str) -> tuple[tuple[int, ...], ...]: - """One block table per plane. + """One physical block table per physical plane. - This tuple is indexed by plane, not by loop depth. LAST_EXITED has - ``max_loops`` planes and SHARED has one, so ``tables[2]`` fails under - SHARED while ``get_block_table(rid, 2)`` works under both layouts. + This tuple is indexed by physical plane, not by loop depth. LAST_EXITED + has ``max_loops`` planes (fewer when packed) and SHARED has one, so + ``tables[2]`` fails under SHARED while ``get_block_table(rid, 2)`` works + under both layouts. ``transfer_slots`` gives each plane's live layers. """ return self._get_allocation(request_id).block_tables @@ -669,8 +792,7 @@ def _write_prepared( if isinstance(batch, DeviceKVMetadata): self._validate_tensor(k, len(batch.position_ids), "k") self._validate_tensor(v, len(batch.position_ids), "v") - self.key_cache[batch.write_blocks, layer, batch.write_offsets] = k - self.value_cache[batch.write_blocks, layer, batch.write_offsets] = v + self._store_layer(layer, batch.write_blocks, batch.write_offsets, k, v) return if not batch.writable: raise ValueError("a read-only prepared KV batch cannot be written") @@ -678,8 +800,9 @@ def _write_prepared( self._validate_tensor(v, len(batch.position_ids), "v") if not batch.rows: return - self.key_cache[batch.write_blocks, layer, batch.write_offsets] = k[: len(batch.rows)] - self.value_cache[batch.write_blocks, layer, batch.write_offsets] = v[: len(batch.rows)] + self._require_packed_layer(layer, batch.rows) + rows = len(batch.rows) + self._store_layer(layer, batch.write_blocks, batch.write_offsets, k[:rows], v[:rows]) for allocation, depth, position in batch.rows: allocation.written[self._plane(depth)][layer].add(position) @@ -714,34 +837,25 @@ def _attend_prepared( self._require_live_batch(batch) self._validate_tensor(q, len(batch.position_ids), "q", query=True) if isinstance(batch, DeviceKVMetadata): - return self.attention( - q, - self.key_cache[:, layer], - self.value_cache[:, layer], - batch.block_tables, - batch.context_lengths, - ) + keys, values, tables = self._layer_kv(layer, batch.block_tables) + return self.attention(q, keys, values, tables, batch.context_lengths) if not batch.rows: return torch.empty_like(q) + self._require_packed_layer(layer, batch.rows) for allocation, depth, position in batch.rows: self._require_prefix(allocation, layer, depth, position + 1) + keys, values, tables = self._layer_kv(layer, batch.block_tables) if batch.cu_seqlens_q is not None: return self.attention.prefill( q, - self.key_cache[:, layer], - self.value_cache[:, layer], - batch.block_tables, + keys, + values, + tables, batch.context_lengths, batch.cu_seqlens_q, batch.max_seqlen_q, ) - return self.attention( - q, - self.key_cache[:, layer], - self.value_cache[:, layer], - batch.block_tables, - batch.context_lengths, - ) + return self.attention(q, keys, values, tables, batch.context_lengths) @torch.no_grad() def finalize_token(self, request_id: str, position: int, exit_depth: int) -> None: @@ -750,13 +864,23 @@ def finalize_token(self, request_id: str, position: int, exit_depth: int) -> Non if self.layout == "shared": return logical_page, offset = divmod(position, self.block_size) - source = allocation.block_tables[self._plane(exit_depth)][logical_page] + source = self._entries(allocation, exit_depth)[logical_page] destinations = [ - allocation.block_tables[self._plane(depth)][logical_page] + self._entries(allocation, depth)[logical_page] for depth in range(exit_depth + 1, self.max_loops) ] # Basic-index copies avoid a blocking host->device index tensor on the # boundary stream. That transfer would serialize final core and routing. + if self.plane_packing > 1: + # Packed recurrent layers are contiguous pages of the flattened cache. + first, last = min(self.recurrent_layers), max(self.recurrent_layers) + 1 + for destination in destinations: + for cache in (self._flat_keys, self._flat_values): + cache[destination + first : destination + last, offset].copy_( + cache[source + first : source + last, offset] + ) + self._record_finalized(allocation, position, exit_depth) + return for destination in destinations: if len(self.recurrent_layers) == self.num_layers: self.key_cache[destination, :, offset].copy_(self.key_cache[source, :, offset]) @@ -787,9 +911,11 @@ def read(self, layer: int, request_id: str, depth: int, length: int | None = Non raise ValueError("read length exceeds the request's reserved token budget") self._require_prefix(allocation, layer, depth, length) positions = torch.arange(length, device=self.device) - table = torch.tensor( - allocation.block_tables[self._plane(depth)], device=self.device, dtype=torch.long - ) + table = torch.tensor(self._entries(allocation, depth), device=self.device, dtype=torch.long) blocks = table[positions // self.block_size] offsets = positions % self.block_size + if self.plane_packing > 1: + self._require_packed_layer(layer, ((allocation, depth, 0),)) + pages = blocks + layer + return self._flat_keys[pages, offsets], self._flat_values[pages, offsets] return self.key_cache[blocks, layer, offsets], self.value_cache[blocks, layer, offsets] diff --git a/vllm_rlt/core/memory.py b/vllm_rlt/core/memory.py index 220643d9..aa3781bd 100644 --- a/vllm_rlt/core/memory.py +++ b/vllm_rlt/core/memory.py @@ -4,7 +4,7 @@ import torch -from vllm_rlt.core.kv_cache_manager import KVCacheManager +from vllm_rlt.core.kv_cache_manager import KVCacheManager, physical_planes from vllm_rlt.worker.buffers import execution_buffer_bytes @@ -67,6 +67,7 @@ def plan_cache(model, cache, scheduler, execution, backend, *, staged_prefill=No dtype=parameter.dtype, backend=backend, recurrent_layers=getattr(model, "recurrent_kv_layers", None), + pack_recurrent_planes=cache.pack_recurrent_planes, ) ids, positions = [], [] for index, length in enumerate(lengths): @@ -131,7 +132,13 @@ def plan_cache(model, cache, scheduler, execution, backend, *, staged_prefill=No # memory, particularly for tiny-model GPU tests with very small pages. useful_blocks = ( math.ceil(config.max_position_embeddings / cache.block_size) - * (config.total_ut_steps if cache.layout == "last_exited" else 1) + * physical_planes( + config.num_hidden_layers, + getattr(model, "recurrent_kv_layers", None), + config.total_ut_steps, + cache.layout, + cache.pack_recurrent_planes, + ) * scheduler.max_num_seqs ) blocks = min(budget_blocks(budget, per_block), useful_blocks) diff --git a/vllm_rlt/core/scheduler.py b/vllm_rlt/core/scheduler.py index f443d742..3e2ffa47 100644 --- a/vllm_rlt/core/scheduler.py +++ b/vllm_rlt/core/scheduler.py @@ -398,7 +398,7 @@ def _plan_admission(self, request: Request, active_count: int) -> AdmissionPlan: if cache.incremental_allocation and not reserve_outputs else capacity ) - required = cache.required_blocks(budget_tokens) - len(prefix) * cache.storage_depths + required = cache.required_blocks(budget_tokens) - len(prefix) * cache.physical_planes reserved = self._reserved_growth_blocks(reserve_outputs=reserve_outputs) # Allow an otherwise idle engine to start a large request. admission_headroom_blocks = cache.watermark_blocks if active_count else 0 diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index 40a6c9a8..a4fbca0d 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -161,6 +161,7 @@ def __init__( incremental_allocation=cache_config.incremental_allocation, watermark_ratio=cache_config.watermark_ratio, recurrent_layers=getattr(model, "recurrent_kv_layers", None), + pack_recurrent_planes=cache_config.pack_recurrent_planes, ) self.scheduler = Scheduler( scheduler_config, diff --git a/vllm_rlt/engine/worker.py b/vllm_rlt/engine/worker.py index ce15bf9d..db3cfabe 100644 --- a/vllm_rlt/engine/worker.py +++ b/vllm_rlt/engine/worker.py @@ -71,6 +71,7 @@ def __init__( num_blocks=cache.num_blocks, block_size=cache.block_size, depths=cache.storage_depths, + blocks_per_page=cache.physical_planes, max_active_requests=active_limit, watermark_blocks=cache.watermark_blocks, ) diff --git a/vllm_rlt/entrypoints/runtime_args.py b/vllm_rlt/entrypoints/runtime_args.py index a166705b..b7a0b1bb 100644 --- a/vllm_rlt/entrypoints/runtime_args.py +++ b/vllm_rlt/entrypoints/runtime_args.py @@ -35,6 +35,12 @@ def add_runtime_args(parser): help="Copy metadata to the GPU instead of reading mapped host memory", ) parser.add_argument("--kv-layout", choices=["last_exited", "shared"], default="last_exited") + parser.add_argument( + "--no-pack-recurrent-planes", + action="store_false", + dest="pack_recurrent_planes", + help="Keep one physical block per depth plane even when only some layers recur", + ) parser.add_argument("--gpu-memory-utilization", type=float, default=0.9) parser.add_argument("--kv-cache-memory-bytes", type=int) parser.add_argument("--memory-reserve-bytes", type=int, default=256 * 1024 * 1024) @@ -94,6 +100,7 @@ def runtime_configs(args): num_blocks=args.num_blocks, block_size=args.block_size, layout=args.kv_layout, + pack_recurrent_planes=args.pack_recurrent_planes, gpu_memory_utilization=args.gpu_memory_utilization, kv_cache_memory_bytes=args.kv_cache_memory_bytes, memory_reserve_bytes=args.memory_reserve_bytes, diff --git a/vllm_rlt/kernels/routing.py b/vllm_rlt/kernels/routing.py index 9a641327..3cc82443 100644 --- a/vllm_rlt/kernels/routing.py +++ b/vllm_rlt/kernels/routing.py @@ -153,13 +153,17 @@ def finalize_kernel( WIDTH: tl.constexpr, PLANES: tl.constexpr, PAGE: tl.constexpr, + FIRST_LAYER: tl.constexpr, LAYERS: tl.constexpr, CHANNELS: tl.constexpr, - BLOCK_STRIDE: tl.constexpr, + ENTRY_STRIDE: tl.constexpr, LAYER_STRIDE: tl.constexpr, TOKEN_STRIDE: tl.constexpr, TILE: tl.constexpr, ): + # Tables hold page entries: physical blocks (ENTRY_STRIDE = block stride) or, + # for packed planes, layer-flattened pages (ENTRY_STRIDE = layer stride). The + # copy covers LAYERS slots starting at FIRST_LAYER relative to the entry. row, target = tl.program_id(0), tl.program_id(1) slot = tl.load(D + row * 3) depth = tl.load(D + row * 3 + 1) @@ -168,9 +172,13 @@ def finalize_kernel( source = tl.load(Tables + (slot * PLANES + depth) * WIDTH + pos // PAGE) dest = tl.load(Tables + (slot * PLANES + target) * WIDTH + pos // PAGE) x = tl.program_id(2) * TILE + tl.arange(0, TILE) - offset = (x // CHANNELS) * LAYER_STRIDE + (pos % PAGE) * TOKEN_STRIDE + x % CHANNELS - src = source.to(tl.int64) * BLOCK_STRIDE + offset - dst = dest.to(tl.int64) * BLOCK_STRIDE + offset + offset = ( + (FIRST_LAYER + x // CHANNELS) * LAYER_STRIDE + + (pos % PAGE) * TOKEN_STRIDE + + x % CHANNELS + ) + src = source.to(tl.int64) * ENTRY_STRIDE + offset + dst = dest.to(tl.int64) * ENTRY_STRIDE + offset k = tl.load(K + src, x < LAYERS * CHANNELS, 0) v = tl.load(V + src, x < LAYERS * CHANNELS, 0) tl.store(K + dst, k, x < LAYERS * CHANNELS) diff --git a/vllm_rlt/pd/engine.py b/vllm_rlt/pd/engine.py index ce2b3b7d..7372353b 100644 --- a/vllm_rlt/pd/engine.py +++ b/vllm_rlt/pd/engine.py @@ -257,11 +257,10 @@ def _message(self, peer, m): del self.transfers[w.tid] def _blocks(self, peer, tokens): - return ( - (tokens + peer.client.info["block_size"] - 1) - // peer.client.info["block_size"] - * peer.client.info["depths"] - ) + info = peer.client.info + pages = (tokens + info["block_size"] - 1) // info["block_size"] + # Packed planes need fewer physical blocks per page than depths. + return pages * info.get("blocks_per_page", info["depths"]) def add_request(self, request_id, prompt_token_ids, sampling_params=None, *, trace_id=None): if self.closed or self.failure: diff --git a/vllm_rlt/pd/transport.py b/vllm_rlt/pd/transport.py index fca044a1..1ab55130 100644 --- a/vllm_rlt/pd/transport.py +++ b/vllm_rlt/pd/transport.py @@ -18,13 +18,18 @@ class Transfer: def kv_segments(info, tables, start, end): """Map valid token ranges to contiguous byte segments, preserving every depth. - A full physical block holds all layers. Partial blocks need one range per - layer; unused tail bytes are never sent or advertised as initialized. + ``tables`` has one physical block table per physical plane and + ``plane_slots`` each plane's live layer slots (default: every layer). A full + page sends its live slots as one range; partial pages need one range per + slot. Unused tail bytes and dead layer slots are never sent. """ page = info["block_size"] if not 0 <= start < end: raise ValueError("invalid KV range") - for table in tables: + slots = info.get("plane_slots") or ((0, info["num_layers"]),) * len(tables) + if len(slots) != len(tables): + raise ValueError("one live slot range is required per physical plane") + for table, (first, last) in zip(tables, slots): if end > len(table) * page: raise ValueError("KV range exceeds allocation") for logical in range(start // page, (end - 1) // page + 1): @@ -36,9 +41,13 @@ def kv_segments(info, tables, start, end): for base in (info["key_ptr"], info["value_ptr"]): address = base + block * info["block_bytes"] if lo == 0 and hi == page: - yield (address, info["block_bytes"], info["device"]) + yield ( + address + first * info["layer_bytes"], + (last - first) * info["layer_bytes"], + info["device"], + ) else: - for layer in range(info["num_layers"]): + for layer in range(first, last): yield ( address + layer * info["layer_bytes"] + lo * info["token_bytes"], (hi - lo) * info["token_bytes"], @@ -90,6 +99,7 @@ def __init__(self, name, cache, hidden, config): block_size=cache.block_size, num_layers=cache.num_layers, depths=cache.storage_depths, + plane_slots=tuple(tuple(slot) for slot in cache.transfer_slots), block_bytes=k.stride(0) * element, layer_bytes=k.stride(1) * element, token_bytes=k.stride(2) * element, @@ -111,6 +121,9 @@ def connect(self, peer): for key in ("block_size", "num_layers", "depths", "block_bytes", "hidden_bytes"): if peer[key] != self.info[key]: raise ValueError(f"PD KV geometry mismatch: {key}") + # Slot ranges may arrive as lists after serialization. + if [list(s) for s in peer["plane_slots"]] != [list(s) for s in self.info["plane_slots"]]: + raise ValueError("PD KV geometry mismatch: plane_slots") self.peers[name] = peer def submit(self, transfer_id, sequence, peer, local, remote, size): diff --git a/vllm_rlt/worker/async_speculative.py b/vllm_rlt/worker/async_speculative.py index dd31ebae..7022c423 100644 --- a/vllm_rlt/worker/async_speculative.py +++ b/vllm_rlt/worker/async_speculative.py @@ -195,7 +195,7 @@ def _prepare_round(self, batch): table_view = host_tables.numpy() table_view.fill(0) for i, ((_, allocation), pages) in enumerate(zip(allocations, page_counts)): - for depth, table in enumerate(allocation.block_tables): + for depth, table in enumerate(self.cache.entry_tables(allocation)): table_view[i, depth, :pages] = table[:pages] buffer.leased = True return buffer, allocations, host_tables, tables diff --git a/vllm_rlt/worker/async_state.py b/vllm_rlt/worker/async_state.py index be0d423a..3d5e0ae7 100644 --- a/vllm_rlt/worker/async_state.py +++ b/vllm_rlt/worker/async_state.py @@ -175,7 +175,8 @@ def ensure_slot(self, request, routing_buffers): was_reserved = rid in self.slots slot = self.reserve_slot(request) previous = self.table_versions.get(rid) - current = allocation.block_tables + # Page entries per logical plane: physical blocks, or flattened pages when packed. + current = self.cache.entry_tables(allocation) if previous != current: for plane, pages in enumerate(current): old = previous[plane] if previous is not None else () diff --git a/vllm_rlt/worker/cuda_graph.py b/vllm_rlt/worker/cuda_graph.py index db1a7519..a129f96c 100644 --- a/vllm_rlt/worker/cuda_graph.py +++ b/vllm_rlt/worker/cuda_graph.py @@ -39,29 +39,22 @@ def __init__(self, cache): self.cache = cache def _write_prepared(self, layer, batch, k, v): - cache = self.cache - cache.key_cache[batch.write_blocks, layer, batch.write_offsets] = k - cache.value_cache[batch.write_blocks, layer, batch.write_offsets] = v + self.cache._store_layer(layer, batch.write_blocks, batch.write_offsets, k, v) def _attend_prepared(self, layer, batch, q): cache = self.cache + keys, values, tables = cache._layer_kv(layer, batch.block_tables) if batch.cu_seqlens_q is not None: return cache.attention.prefill( q, - cache.key_cache[:, layer], - cache.value_cache[:, layer], - batch.block_tables, + keys, + values, + tables, batch.context_lengths, batch.cu_seqlens_q, batch.max_seqlen_q, ) - return cache.attention( - q, - cache.key_cache[:, layer], - cache.value_cache[:, layer], - batch.block_tables, - batch.context_lengths, - ) + return cache.attention(q, keys, values, tables, batch.context_lengths) def _graph_entry(hidden, batch, cache, max_position_embeddings): diff --git a/vllm_rlt/worker/metadata.py b/vllm_rlt/worker/metadata.py index c82a75cd..8c752b0d 100644 --- a/vllm_rlt/worker/metadata.py +++ b/vllm_rlt/worker/metadata.py @@ -144,7 +144,7 @@ def _write_addresses(self, rows, *, offsets=None, for_write=True): """Resolve physical addresses and reject aliased writes across query rows.""" cache = self.cache blocks = [ - allocation.block_tables[cache._plane(depth)][position // cache.block_size] + cache._entries(allocation, depth)[position // cache.block_size] for allocation, depth, position in rows ] if offsets is None: @@ -197,7 +197,7 @@ def build( blocks, _ = self._write_addresses(rows, offsets=offsets, for_write=for_write) attention = plan_attention_metadata( AttentionRows( - block_tables=[a.block_tables[cache._plane(d)] for a, d, _ in rows], + block_tables=[cache._entries(a, d) for a, d, _ in rows], positions=positions, sequence_keys=[(id(a), d) for a, d, _ in rows] if packed else None, ), diff --git a/vllm_rlt/worker/model_runner.py b/vllm_rlt/worker/model_runner.py index ada23d20..37fd9e4a 100644 --- a/vllm_rlt/worker/model_runner.py +++ b/vllm_rlt/worker/model_runner.py @@ -655,7 +655,9 @@ def _prepare_shared_loop_inputs(self, batch): for row, span in enumerate(active): position = span.token_start + offset positions.append(position) - blocks.append(base.rows[row][0].block_tables[0][position // cache.block_size]) + blocks.append( + cache.entry_tables(base.rows[row][0])[0][position // cache.block_size] + ) tokens.append(span.input_tokens[offset]) contexts.append(position + 1) size = self._size(len(active)) @@ -1323,9 +1325,9 @@ def finalize_many(self, requests): self.boundary_stream.wait_event(event) routing_buffers.descriptor = routing_buffers.routing.stage(routing_buffers.count) channels = cache.num_kv_heads * cache.head_dim - finalize_kernel[ - (len(copies), cache.max_loops, math.ceil(cache.num_layers * channels / 256)) - ]( + # Packed planes copy only the recurrent slots of flattened page entries. + first_layer, layers = cache.finalize_layers + finalize_kernel[(len(copies), cache.max_loops, math.ceil(layers * channels / 256))]( routing_buffers.descriptor, self.async_state.tables, cache.key_cache, @@ -1333,9 +1335,11 @@ def finalize_many(self, requests): self.async_state.width, cache.max_loops, cache.block_size, - cache.num_layers, + first_layer, + layers, channels, - *cache.key_cache.stride()[:3], + cache.entry_stride, + *cache.key_cache.stride()[1:3], 256, ) routing_buffers.record_done()