[Huginn] Opt-in prefix caching with position-keyed recurrent-state noise (stacked on #147)

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,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