[KV] Pack recurrent depth planes for partially recurrent models

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.
 

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.

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

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"]

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: