[jason-fxz]

Refs #828

Part of the speculative decoding plan in #828; it can merge independently of the seams PR. Nothing turns it on in this PR: radix_lookahead defaults to False, and the MTP PR turns it on for MTP drafts. The full stack is in the draft PoC #832.

Summary

Adds radix_lookahead: bool to CacheManager. With radix_lookahead=True, the prefix cache keys position i by the token pair (t_i, t_{i+1}) instead of t_i (a lookahead key).

Why

An MTP draft's KV at position i is computed from the target's hidden state at i and the token at i+1. With token keys, a prefix hit can reuse a draft KV slot that was written for a different next token: when two cached requests share a prefix and then fork, the slot just before the fork may come from the other branch. With lookahead keys, a hit also has to match the token after each position, so every draft KV slot it reuses was written for the same next token. The cost: where a query forks from a cached path, its hit ends one position before the fork, then rounds down to a page (a whole page when the fork is page-aligned); a query that continues a cached path, such as the next turn of a conversation, hits as far as with token keys.

Design

• CacheManager._keys(ids, n) returns the radix keys of the first n positions: the ids themselves when radix_lookahead is off, otherwise each id packed with the next one into an int64 (ids stay below 2^31, multimodal pseudo ids included). A position whose next token is not known yet has no lookahead key, so the keys can be one shorter than n.
• Every match and insert goes through _keys, except in the SWA radix cache, which keeps token keys; radix_lookahead=True with it fails an assert.
• The GDN snapshot commits that need a key of exactly L positions skip when the lookahead key is shorter, as they already skip unaligned lengths.

Tests

tests/scheduler/test_radix_lookahead.py (CPU, real CacheManager and radix trees):

• two cached requests fork at position 8; a query that follows the second one takes the second one's slot for position 7 with lookahead keys, and its hit is one position shorter than with token keys (page sizes 1 and 4)
• the key of the last prompt position pairs it with the sampled token
• multimodal pseudo ids at the top of their range keep distinct keys
• a hybrid (GDN) snapshot hits only when the token after it agrees

Test environment

• GPU: 1x NVIDIA H100 80GB HBM3, driver 580.95.05
• CPU: Intel Xeon Platinum 8480+ (26 vCPUs), 221 GB RAM, Ubuntu 22.04.5
• CUDA 13.0, torch 2.14.1+cu130, Python 3.12.12
• This branch: f1df340, on main 408e6ae; the unit and smoke runs below are from 3b0a786, and f1df340 changes only the new test file

Results

The new test file on its own (CPU):

PYTHONPATH=python pytest tests/scheduler/test_radix_lookahead.py

5 passed.

Unit tests, with radix_lookahead left off everywhere outside the new test:

FREETOKEN_QWEN36_MODEL=/path/to/Qwen3.6-27B PYTHONPATH=python pytest tests/scheduler tests/engine tests/kvcache tests/attention tests/server tests/tokenizer tests/models

1586 passed, 202 skipped, 1 failed. The failure is tests/models/test_quant_config.py::test_scheme_for_agrees_with_the_stored_tensors[apodex/Apodex-1.1-mini-NVFP4[index-only]] (a local checkpoint with quant_method: modelopt_mixed); it fails the same way on main 408e6ae.

Smoke test on this branch, passed:

FREETOKEN_TEST_MODEL=/path/to/Qwen3.6-35B-A3B FREETOKEN_AIME25_JSONL=/path/to/aime25.jsonl PYTHONPATH=python pytest tests/e2e/test_aime.py