Skip to content

prompt_radix: a lookahead in the tree and a planned grid, so a family like GLM's can share the store - #1659

Merged
ashhart merged 4 commits into
ashhart:mainfrom
Bizuayeu:pr/radix-lookahead-grid
Oct 30, 2026
Merged

ashhart merged 4 commits into
ashhart:mainfrom
Bizuayeu:pr/radix-lookahead-grid

Conversation

@Bizuayeu

Copy link
Copy Markdown
Contributor

No family turns these rules on here, so nothing on main changes: HIP's families and #1357's Qwen path keep their marks, pages and bytes, and a golden run pins that. This is the half of #1510's option B that prompt_radix itself needs, and #1657's shared adapter builds on it. What it enables, from GLM's port on these rules (TF_GLM_KV_SHARE=4, the branch on #1510) against main's flat store, on GB10 at TP=3 with the 1M window:

main's flat store GLM's port
64k message sent again (s) 37.07, 37.06 1.77, 1.77
128k message sent again (s) 75.63, 75.55 2.14, 2.14
decode (tok/s, bench_openai) 51.17, 50.13 51.07, 51.06

What this changes

Three commits in core/prompt_radix.zig, core/prompt_radix_tree.zig and their test file. No family code.

A golden run first. Seven passes under the rules HIP's families and #1357's Qwen path use (lookahead 0, no grid, the tail kept): two conversations over a shared system block, a growing and an edited turn, a third prompt that evicts, and one past the limits. Each pass's resume point and marks, and the tree's states, pages, bytes and counts after it, are pinned before the store changes, and the next two commits leave the run as it was.

A lookahead in the tree. GLM keeps its MTP rows on pages, and row at−1 reads token at, so a page followed by a different token holds different rows. A page is now matched and shared only while the prompt repeats the lookahead's tokens past it: each node keeps those tokens past its last page, child, common, insert and split compare and carry them, and every slice the store hands the tree grows by the lookahead. A resume leaves max(1, lookahead) tokens to prefill.

A planned grid. A planned store cuts and resumes only on its grid, a multiple of the page: marks snap to it, deepest skips states off it and keep refuses them.

Rules it cannot keep. Store.init returns error.BadPromptRadixRules for a zero page, or for a planned grid that is zero or no multiple of the page. It was only a std.debug.assert, so a release build took them, as @jschmied noted on #1510.

With lookahead 0 and no grid, every path stays as it was.

Receipt

  • Environment: TensorFold 9f805f5a0 (three commits on main f350cd3ea), host tests on Linux x86_64. The table above is GLM's port at 0cc5cd348, on cuda qwen3_5: the prompt cache keeps KV pages on the device (TF_QWEN_KV_SHARE, #1156) #1357's 65ef4ffcc and these commits (GLM-5.3-Flash XSFP4, PARALLEL=4, KV pages on, the drafts on, three GB10 in a ring on NCCL, launched unset, on, on, unset): the comment on #1510 has the method and the TP=2 pair, where main's flat store has room and the two match.
  • Exactness: with lookahead 0 and no grid the golden run pins what the store keeps, resumes and evicts. In GLM's port, every reply's tokens are the same in all four launches, resumed and concurrent ones among them.
  • Decode and prompt speed before and after: n/a on main, no family runs the new rules. The table above is the port's.
  • Kernel contract: n/a, no GPU kernel is added or changed.
  • Tests run (Linux x86_64): test-kv-pool 28/28, engine_api's tests 411/411, zig build test 441/441 steps (the first run's one failure, cuda_cost_identity's timing-dependent drafter test outside these files, passed when that step ran again), zig fmt --check, and tools/zig/lean_check.py with nothing new (426 against main's 426). Merged with cuda qwen3_5: the prompt cache keeps KV pages on the device (TF_QWEN_KV_SHARE, #1156) #1357's 65ef4ffcc (no conflicts), test-kv-pool 28/28 and test-qwen35-share-faults-host 14/14.
  • Not run, and why: Metal, HIP devices and CUDA on this tree alone, since no family's path changes here and the golden run pins the rules theirs use.

Checklist

  • Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off. (Only the store's matching changes, and no family turns the new rules on.)
  • Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. (Nothing on main changes, and the golden run pins the store's resumes. On GLM's port at TP=3 and TP=2, with the drafts on, a resumed prompt equals a fresh one and concurrent equals solo, on GLM CUDA: the KV page pool crowds out the prompt cache, so long chats prefill again every turn #1510. Drafted against plain was not run there.)
  • No precision traded for speed. If the bits change, the description says which sums change and why. (No bits change: the store decides what to keep and resume, and no family's rules change here.)
  • A new or changed GPU kernel is registered with the kernel contract (docs/kernel-contract.md); each backend's SUMMARY line is in the receipt. (No GPU kernel is added or changed.)
  • Prompt processing is no slower than the last release. (No family runs the new paths. With lookahead 0 the comparisons take no extra tokens.)
  • The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not. (Host tests on Linux; GLM's port on CUDA on GB10 at TP=3 and TP=2. Metal and HIP devices not run, as above.)
  • New tests fail before the change, pass after it, and skip cleanly without their dependency. (On main the deepest test does not build, as deepest takes no step there, and with that test left out the first lookahead test stops at the store's assert. With their property taken out here, the tree without the lookahead fails the two lookahead tests, a planned store stepping by the page fails the two planned tests and GLM's shape, and deepest off the step fails its own. The rules test fails on main's assert. The golden run passes before and after, as it pins what main does. None needs a device.)
  • Comments and docstrings are one line. No measurements or history in the source.
  • No personal data, machine names, internal hosts or local paths. No AI attribution lines.

Bizuayeu and others added 4 commits October 30, 2026 23:00
…sume and evict

Seven passes under the rules HIP's families and Qwen's CUDA path use (a page of 4 tokens here, min_gap 16, the tail
kept): two conversations over a shared system block, a growing and an edited turn, a third prompt that evicts, and one
past the limits. Each pass's resume point and marks, and the tree's states, pages, bytes and counts after it, are the
ones the store kept before it learns a lookahead and a grid.
… like GLM's can share the store

A page is shared and matched only while the prompt repeats the lookahead's tokens past it: GLM's MTP rows sit on pages,
and row at-1 reads token at, so a page whose next token differs holds other rows. Each node keeps those tokens past its
last page; child, common, insert and split compare and carry them, and every slice the store hands the tree grows by the
lookahead. A planned store cuts and resumes only on its grid, a multiple of the page: marks snap to it, deepest skips
states off it and keep refuses them. With lookahead 0 and no grid every path is as before (the golden run).
Store.init refused a zero page or a planned grid that is no multiple of the page only through std.debug.assert, so a
release build would take them. It now returns error.BadPromptRadixRules, with a test for each.
@ashhart
ashhart force-pushed the pr/radix-lookahead-grid branch from 9f805f5 to b100164 Compare October 30, 2026 23:00
@ashhart
ashhart merged commit c03b722 into ashhart:main Oct 30, 2026
8 of 9 checks passed
@ashhart

ashhart commented Oct 31, 2026

Copy link
Copy Markdown
Owner

Thank you @Bizuayeu. This is in TensorFold 1.2.0, released today: https://github.com/ashhart/TensorFold/releases/tag/v1.2.0

It landed with your commits under your own name, so you're listed as a TensorFold contributor, and it's credited to you in the release notes and the changelog.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants