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..2ed22680 100644 --- a/tests/test_huginn.py +++ b/tests/test_huginn.py @@ -381,3 +381,294 @@ 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 + + +def test_huginn_prefix_cache_experiment_reuses_prompt_pages(monkeypatch): + """With the experiment gate, Huginn publishes completed prompt pages and reuses them. + + Publication requires written coverage only for the layers each plane holds live + KV for (prelude/coda write at depth zero only). A cached-prefix run is + deterministic and token-identical to another cached-prefix run, but not to a + cold run: the recomputed suffix draws its recurrent-state noise from the start + of the request's RNG stream, where the cold run drew noise for the first prompt + tokens. Position-keyed state noise would close that gap; until then the + engine guard stays and this path is an experiment. + """ + import pytest as _pytest + + torch.manual_seed(42) + shared = HuginnForCausalLM(tiny_huginn_config(mean_recurrence=4)).eval() + params = SamplingParams(max_tokens=4, ignore_eos=True) + prompt = [4, 7, 3, 9, 2, 6, 1, 8] + with _pytest.raises(ValueError, match="prefix caching"): + LLMEngine( + shared, + cache_config=CacheConfig(num_blocks=128, block_size=2, enable_prefix_caching=True), + ) + monkeypatch.setenv("VLLM_RLT_HUGINN_PREFIX_CACHE", "1") + + def engine(): + return LLMEngine( + shared, + cache_config=CacheConfig(num_blocks=128, block_size=2, enable_prefix_caching=True), + scheduler_config=SchedulerConfig( + max_num_seqs=2, max_num_batched_tokens=8, prefill_chunk_size=4 + ), + ) + + def run(eng, name, tokens): + eng.add_request(name, tokens, params) + out = {} + while eng.has_unfinished_requests(): + for o in eng.step(): + if o.finished: + out[o.request_id] = (list(o.token_ids), list(o.exit_depths)) + return out[name] + + warm, other = engine(), engine() + first = run(warm, "first", prompt) + assert warm.cache_manager.lookup_prefix(prompt), "the finished prompt should be published" + assert first == run(other, "first", prompt) # cold runs agree with each other + hits = warm.cache_manager.prefix_hits + second = run(warm, "second", prompt) + assert warm.cache_manager.prefix_hits > hits # prefill resumed after the cached pages + run(other, "second", prompt) + assert second == run(other, "second-again", prompt) # cached runs are deterministic + assert second[1] == first[1] # fixed depth either way + # A longer prompt sharing the prefix reuses only complete prompt blocks. + hits = warm.cache_manager.prefix_hits + run(warm, "third", prompt + [5, 3, 7, 7]) + assert warm.cache_manager.prefix_hits - hits >= len(prompt) - 2 + assert warm.cache_manager.num_used_blocks == 0 + + +def test_position_keyed_state_noise_is_a_truncated_normal_keyed_by_seed_and_position(): + from vllm_rlt.models.huginn import position_keyed_state_noise + + noise = position_keyed_state_noise( + torch.arange(64), torch.arange(64) * 37, 5280, 0.02, torch.float32 + ) + z = noise / 0.02 + assert z.abs().max() <= 3.0 + assert ( + abs(z.mean().item()) < 0.01 and abs(z.std().item() - 0.987) < 0.01 + ) # trunc N(0,1) on [-3,3] + assert 0.0005 < (z.abs() > 2.9).float().mean().item() < 0.002 + # Same (seed, position) gives the same channel vector regardless of batch composition. + alone = position_keyed_state_noise( + torch.tensor([21]), torch.tensor([21 * 37]), 5280, 0.02, torch.float32 + ) + assert torch.equal(alone[0], noise[21]) + # Different seed or position gives different noise. + other = position_keyed_state_noise( + torch.tensor([21, 22]), torch.tensor([21 * 37 + 1, 21 * 37]), 5280, 0.02, torch.float32 + ) + assert not torch.equal(other[0], noise[21]) and not torch.equal(other[1], noise[21]) + # Adjacent seeds must not be channel-shifted copies of each other (a linear + # seed/channel combination before the mix would alias them exactly). + a = position_keyed_state_noise(torch.tensor([5]), torch.tensor([100]), 512, 0.02, torch.float32) + b = position_keyed_state_noise(torch.tensor([6]), torch.tensor([100]), 512, 0.02, torch.float32) + assert not torch.equal(a[0][:-1], b[0][1:]) and not torch.equal(a[0][1:], b[0][:-1]) + + +def test_position_keyed_noise_makes_huginn_independent_of_chunking_and_prefix_cache(monkeypatch): + """With position-keyed state noise a cached-prefix run reproduces the cold run exactly.""" + import vllm_rlt.models.huginn as huginn_module + + monkeypatch.setattr(huginn_module, "POSITION_NOISE", True) + monkeypatch.setenv("VLLM_RLT_HUGINN_PREFIX_CACHE", "1") + torch.manual_seed(42) + shared = HuginnForCausalLM(tiny_huginn_config(mean_recurrence=4)).eval() + params = SamplingParams(max_tokens=4, ignore_eos=True, seed=9) + prompt = [4, 7, 3, 9, 2, 6, 1, 8] + + def engine(chunk, cache): + return LLMEngine( + shared, + cache_config=CacheConfig(num_blocks=128, block_size=2, enable_prefix_caching=cache), + scheduler_config=SchedulerConfig( + max_num_seqs=2, max_num_batched_tokens=8, prefill_chunk_size=chunk + ), + ) + + def run(eng, name, tokens): + eng.add_request(name, tokens, params) + out = {} + while eng.has_unfinished_requests(): + for o in eng.step(): + if o.finished: + out[o.request_id] = (list(o.token_ids), list(o.exit_depths)) + return out[name] + + cold_chunk2 = run(engine(2, False), "a", prompt) + cold_chunk8 = run(engine(8, False), "b", prompt) + assert cold_chunk2 == cold_chunk8 # chunking no longer changes the noise + warm = engine(4, True) + assert run(warm, "first", prompt) == cold_chunk2 + hits = warm.cache_manager.prefix_hits + assert run(warm, "second", prompt) == cold_chunk2 # cached prefix, identical tokens + assert warm.cache_manager.prefix_hits > hits + longer = prompt + [5, 3, 7, 7] + assert run(warm, "third", longer) == run(engine(4, False), "c", longer) + assert warm.cache_manager.num_used_blocks == 0 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..54e1eeeb 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: @@ -254,15 +373,24 @@ def publish_prefix(self, request_id, tokens, length, event=None): if key in self._prefixes: continue end = (index + 1) * self.block_size - # Every depth and layer must have a contiguous written prefix. + # Every depth must have a contiguous written prefix in each layer it + # holds live KV for: all layers at depth zero, only the recurrent + # layers deeper (Huginn's prelude/coda write at depth zero only). # GPU completion alone does not establish this coverage. - if any(w.prefix < end for plane in allocation.written for w in plane): + if not self._prefix_complete(allocation, end): break blocks = tuple(t[index] for t in allocation.block_tables) for b in blocks: self._refs[b] += 1 self._prefixes[key] = blocks + def _prefix_complete(self, allocation: _Allocation, end: int) -> bool: + for plane, written in enumerate(allocation.written): + layers = range(self.num_layers) if plane == 0 else self.recurrent_layers + if any(written[layer].prefix < end for layer in layers): + return False + return True + def lookup_prefix(self, tokens): if not self.enable_prefix_caching: return () @@ -296,14 +424,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 +442,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 +455,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 +527,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 +801,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 +809,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 +846,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 +873,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 +920,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..608d94a3 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -29,6 +29,19 @@ class SignalHandle: row: int +def _huginn_prefix_cache_experiment() -> bool: + """Opt-in: let Huginn share immutable prompt prefixes (VLLM_RLT_HUGINN_PREFIX_CACHE=1). + + Huginn prefill runs every recurrence depth for every prompt token and writes the + prelude/coda KV at depth 0, so a completed prompt prefix holds valid KV in every + plane a later request can read. Combine with VLLM_RLT_HUGINN_POSITION_NOISE=1 so a + cached run draws the same recurrent-state noise as a cold one. + """ + import os + + return os.environ.get("VLLM_RLT_HUGINN_PREFIX_CACHE", "0") == "1" + + class StageEngine: """Single-device loop-level engine with synchronous or pipelined scheduling.""" @@ -55,7 +68,7 @@ def __init__( if getattr(model, "requires_boundary_kv", False): if cache_config.layout != "last_exited": raise ValueError("Huginn requires last_exited KV") - if cache_config.enable_prefix_caching: + if cache_config.enable_prefix_caching and not _huginn_prefix_cache_experiment(): raise ValueError("Huginn prefix caching is not yet supported") self.exit_config = exit_config or ExitConfig() self.execution_config = execution_config or ExecutionConfig() @@ -161,6 +174,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/models/huginn.py b/vllm_rlt/models/huginn.py index 9312df67..f43ef7b4 100644 --- a/vllm_rlt/models/huginn.py +++ b/vllm_rlt/models/huginn.py @@ -8,6 +8,7 @@ import json import math +import os from dataclasses import asdict, dataclass, fields from pathlib import Path from typing import TYPE_CHECKING @@ -198,6 +199,50 @@ def forward(self, hidden, freqs, batch, cache): return self.norm_4(self.mlp(self.norm_3(hidden)) + hidden) +# Opt-in: draw each token's recurrent-state noise from a counter-based hash of +# (request seed, position, channel) instead of the request's sequential RNG stream. +# The sample is then independent of batching, prefill chunking and prefix-cache +# hits, so a cached-prefix generation reproduces the cold run token for token. +POSITION_NOISE = os.environ.get("VLLM_RLT_HUGINN_POSITION_NOISE", "0") == "1" + + +_GOLDEN = 0x9E3779B97F4A7C15 - (1 << 64) # SplitMix64 increment as a signed int64 + + +def _lsr(x, bits): + """Logical right shift of an int64 tensor.""" + return (x >> bits) & ((1 << (64 - bits)) - 1) + + +def position_keyed_state_noise(seeds, positions, width, std, dtype): + """Truncated normal (mean 0, ``std``, within 3 std) keyed by (seed, position, channel). + + Three chained SplitMix64-style mixes (seed, then position, then channel) produce + 24 uniform bits per element; the inverse normal CDF maps them onto the truncated + range. int64 arithmetic wraps on every backend, so CPU and CUDA agree bitwise. + """ + rows = seeds.shape[0] + channel = torch.arange(width, dtype=torch.int64, device=seeds.device) + + def mix(v): + v = (v ^ _lsr(v, 30)) * -4658895280553007687 + v = (v ^ _lsr(v, 27)) * -7723592293110705685 + return v ^ _lsr(v, 31) + + # Stage-wise: mix the seed, then fold in the position, then the channel. Each + # stage is a bijection, so (seed, position, channel) triples never alias the + # way a single linear combination of the three could. + x = mix(seeds[:, None] * -7046029254386353131 + _GOLDEN) + x = mix(x + positions[:, None] * 6364136223846793005 + _GOLDEN) + x = mix(x + channel[None, :] * 1442695040888963407 + _GOLDEN) + uniform = _lsr(x, 40).to(torch.float32) * (1.0 / 16777216.0) # 24 bits -> [0, 1) + uniform = uniform * (1.0 - 2.0 / 16777216.0) + 1.0 / 16777216.0 # (0, 1), away from the edges + lower = 0.5 * (1.0 + math.erf(-3.0 / math.sqrt(2.0))) + probability = lower + uniform * (1.0 - 2.0 * lower) + z = torch.special.ndtri(probability).clamp_(-3.0, 3.0) + return (z * std).to(dtype).reshape(rows, width) + + class HuginnForCausalLM(nn.Module): """Huginn-0125 with eager KV boundaries and a graphable recurrent core.""" @@ -257,21 +302,32 @@ def prelude_prepared( freqs = self._freqs(batch) for block in self.transformer.prelude: hidden = block(hidden, freqs, batch, cache) - state = torch.empty_like(hidden) - start = 0 - while start < len(state): - end = start + 1 - generator = generators[start] if generators is not None else None - while end < len(state) and (generators is None or generators[end] is generator): - end += 1 - torch.nn.init.trunc_normal_( - state[start:end], - std=self.config.initializer_range, - a=-3 * self.config.initializer_range, - b=3 * self.config.initializer_range, - generator=generator, + if POSITION_NOISE and generators is not None: + state = position_keyed_state_noise( + torch.tensor( + [g.initial_seed() for g in generators], dtype=torch.int64, device=hidden.device + ), + batch.position_ids.to(torch.int64), + self.config.n_embd, + self.config.initializer_range, + hidden.dtype, ) - start = end + else: + state = torch.empty_like(hidden) + start = 0 + while start < len(state): + end = start + 1 + generator = generators[start] if generators is not None else None + while end < len(state) and (generators is None or generators[end] is generator): + end += 1 + torch.nn.init.trunc_normal_( + state[start:end], + std=self.config.initializer_range, + a=-3 * self.config.initializer_range, + b=3 * self.config.initializer_range, + generator=generator, + ) + start = end state = state * math.sqrt(self.config.n_embd) return torch.cat((state, hidden), dim=-1) 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()