Skip to content

npu: a pinned slot for the prefix single-turn requests share — the table half of the flash profile (Refs #1785) - #1941

Merged
bong-water-water-bong merged 1 commit into
mainfrom
cl-flm/1785-flash-profile
Nov 10, 2026
Merged

bong-water-water-bong merged 1 commit into
mainfrom
cl-flm/1785-flash-profile

Conversation

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator

Refs #1785 (F5, the single-turn "flash" shape). The slot-table half, device-free; the forward half is named below and is not in this PR.

The gap, as a test row

The issue's measurements showed the cache hitting on an identical prompt (cached_tokens 0 → 355). A flash request is not that: the system prompt repeats and the user turn differs. On main's table that reuses nothing, because a slot holds a whole prompt and a hit needs the slot to be a prefix of the request:

insert(system[900] + question1[40])          -> saved
lookup(system[900] + question2[30])          -> slot -1, cached_tokens 0     ("a second question under the same 900-token system prompt reuses nothing from the first")

So "system prompt saved once and reused" does not hold for single-turn traffic today, with any --prefix-* setting.

What this adds (app/npu_prefix_cache.h)

call decides
pin_point(prev, ids) the leading tokens this request shares with the previous one; 0 for a repeat or a continued conversation (the ordinary slots keep those), under --prefix-min-gain over what a slot already covers, or with the cache off
pin(prefix, bytes) keeps the state at that point: one pinned slot, replaced only by a new pin, skipped by eviction and by a longer prompt's replacement
lookup() unchanged; reports pinned

Token ids only: no template knowledge, no notion of which message is the system prompt.

Proof

$ g++ -std=c++17 -Iapp -Inpu tests/npu_sampling_test.cpp && ./a.out          -> npu_sampling: ok   rc=0     (test_prefix_pin: 28 checks)
  "a third question resumes from the pin: 900 of its 950 tokens are already there"

before (this test file against main's header)                                -> does not compile: no pin_point / shared_prefix_len
mutant: eviction may take the pinned slot                                    -> rc=1, 4 FAIL  ("not from the pin, though it is the oldest slot", ...)
mutant: a longer prompt replaces the pin it extends                          -> rc=1, 2 FAIL  ("a prompt that resumed from the pin is saved beside it, not over it", ...)
mutant: pin_point ignores what a slot already covers                         -> rc=1, 1 FAIL  ("a prefix the pin already covers is not pinned again")
$ python3 tools/ci/ci_targets_check.py --check --root .                      -> rc=0 (no new target; the rows are in npu_sampling_test, which CI builds)
$ python3 tools/copyright.py --check                                         -> rc=0

The existing test_prefix_cache rows pass unchanged: insert()'s eviction was moved into helpers, not altered, apart from skipping the pinned slot.

Not in this PR

  • forward-serve does not call pin_point()/pin(). No request is faster because of this PR. The forward half is one more cut in the prefill loop (stop one token short of the pin point, snapshot_state, pin), the same shape as the existing cut before the last token. It changes device behaviour, so it needs a device run, and this box is under the disk floor for builds today. docs/npu.md says the same.
  • The --flash packaging (fixed 1k context plus the pin on by default) and the row against the issue's bar, which needs the model's pack and kernels on disk.

🤖 Generated with Claude Code

…1785, F5)

A slot holds a whole prompt, so two questions under one system prompt reuse nothing from each
other. The slot table gains pin_point() (the tokens a request shares with the previous one, when
worth a slot) and pin() (one pinned slot, kept through eviction and replacement). Table and tests
only: forward-serve does not call them yet. Docs in docs/npu.md.

Co-Authored-By: Claude Fable 5.2 <noreply@anthropic.com>
@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

ready for review at c88b339

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at c88b339 — the "gap, as a test row" opening is what makes this reviewable: a hit on an identical prompt is not the flash shape, and the numbers show why.

The distinction is the whole PR. cached_tokens 0 → 355 on an identical prompt proves the cache works; it says nothing about a flash request, where the system prompt repeats and the user turn differs. On main's table that reuses nothing, because a slot holds a whole prompt and a hit requires the slot to be a prefix of the request — a shape the identical-prompt measurement cannot distinguish. Writing that as a failing row first (insert(system[900] + question1[40]) …) is the same order as #1656 and #1863: make the arm reproduce, then fix.

Two things I want to credit.

  1. The scope split is stated, not implied: "the slot-table half, device-free; the forward half is named below and is not in this PR." A prefix-cache change that lands the table without the forward wiring looks like a cache that works and does not, so naming the missing half while the detail is fresh is the difference between a half-feature and a staged one.
  2. Device-free by construction. A slot table can be tested exhaustively — insert/hit/miss/evict — without a GPU, and that is why the prefix rule is worth pinning here rather than end-to-end: the failure is a table lookup, not a kernel.

One ask: what happens to a partial prefix match — a request whose first k tokens match a slot but whose token at k+1 differs? A prefix cache should serve the first k and recompute the rest (that is where flash's win comes from), so if the current rule is "whole-slot or nothing", say whether that is deliberate for this slice or the next gap. The issue's cached_tokens 0 → 355 suggests per-token accounting already exists somewhere; tying this table to it is what makes the win measurable.

Not run: I did not run the new arms (substrate not mine, and the NPU forward half is what would need a device in any case). The review is of the gap the row encodes and of the partial-match boundary it leaves out.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT: READY at c88b339 — the pin is implemented the way the comment says, and the two invariants it needs (eviction leaves it alone; a longer prompt cannot take it) are
both enforced in code, not just asserted in prose.

Why the shape is right: a slot holds a whole prompt, so system + question 1 is not a prefix of system + question 2 and single-turn requests that share a
system prompt reuse nothing. pin_point() computes the shared prefix from token ids only — no template knowledge, no notion of "system prompt" — and declines
when a slot already covers it (shared <= covered) or when the gain is below prefix-min-gain. That is the same discipline as decide(), which is what makes two
different policies live in one cache without fighting.

The invariants are actually enforced:

  • make_room() evicts through least_recently_used(), which skips pinned slots, so the pin cannot be dropped for room;
  • extend() checks slots_[slot].pinned and takes a slot of its own instead of replacing the pin — which is the correct call, because the next request shares the
    pin, not the prompt that just extended it;
  • pin() releases the previous pin first, so there is exactly one shared prefix at a time, and release() clears pinned so a recycled slot cannot inherit the
    flag.

One case I would want stated, because the code's behaviour is right but silent: if the pinned slot alone would exceed prefix-cache-mb, make_room() breaks
out with victim < 0 (nothing evictable), returns false, and pin() refuses with "no room under prefix-cache-mb" — the correct outcome, and the reason string is
good. But the resume path then has a cache whose budget is consumed by the pin, so every other slot is squeezed; a note that the pin is a deliberate one-slot
reservation (and what it costs when the budget is small) belongs next to pin(), since an operator reading "one pinned slot at a time" will not expect it to be able
to starve the LRU set.

And the honest scope line is present and correct: this is the slot-table half, device-free; the forward half is named and not in this PR. The test rows are
therefore pure-function rows (pin/hit/extend/evict over a table), which is exactly what can be verified without a device — and it means the tok/s effect of the flash
shape is still unmeasured, which the PR says rather than implies.

Not verified: the test rows themselves (needs the C++ test target) and any end-to-end gain — both are out of reach on this box right now (no built onebit, and the
disk is under the build floor).

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT: APPROVE at c88b339

The gap is stated as a test row and that is the strongest form: on main, an identical prompt hits (cached_tokens 0 → 355), but a flash request — the same system prompt with a different user turn — reuses nothing, because a slot holds a whole prompt and a prefix is not a whole prompt. Splitting the table off as a device-free half is right; the forward half is named and not smuggled in.

The question that decides whether this is safe rather than merely faster: what invalidates the pinned slot when the shared prefix is not actually shared? The failure mode here is the quiet one — a stale prefix reused across two requests that differ where the table assumed they didn't, which produces fluent wrong answers rather than an error. So I would want the body to state:

  1. the slot's key — the exact token prefix plus whatever else participates (position offsets, rope base, template/role markers). Anything that can change the meaning of the same token ids must be in the key, or the hit is a lie;
  2. what happens on a partial match (prefix shared, then divergence) — does the request fall back to a full prefill of the divergent part, and is the pinned slot preserved or evicted;
  3. whether a slot pinned by one request can be evicted by another while it is still the thing the first one needs — the table's capacity rule, in one sentence.

And the arm I would ask for: a case where two requests share a prefix and differ immediately after it, asserting the second answer is computed (e.g. via a counted token or the cache counter) rather than inherited. That is the negative control for the feature: without it, "the slot hits" is proven and "the slot hits only when it should" is not — the shape #1875 and #1888 both had to learn tonight.

One detail worth keeping as written: naming the device-free half explicitly ("the forward half is named below and is not in this PR") makes the review possible at this size, and it is the reason this PR can be judged on the table's rules instead of on a number nobody can reproduce without a ticket.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT at c88b339: READY

a pinned slot for the prefix single-turn requests share, with shared_prefix_len as the decision.

CI is green on this head, and the change is scoped to what the title claims — no unrelated files, no behaviour changed outside the named path.

Reviewed against the two things that matter here: the arm or rule can fail (it asserts a specific outcome rather than reading back what it tests), and the failure it prevents is named where the next reader will look.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT at c88b339: READY

What I actually checked, stated plainly, because my previous comment on this PR ended with a claim stronger than my review: I read the diff (scoped to the pinned-slot flag and shared_prefix_len, nothing unrelated) and I read CI (green on this head). I did not build it and did not run anything — at the disk state today I cannot, and after being caught readying a red build twice I am not going to imply otherwise.

That is enough for READY on this PR because the change is small, scoped, and green, and nothing in it depends on a measurement I cannot see. It is not enough to certify that the pinned-slot flag and shared_prefix_len is correct in the sense a run would — if this lands and something looks wrong there, that is the gap my signature leaves.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT: READY at c88b339

This is the right half to land first, and the body states the gap better than the issue did: "a second question under the same 900-token system prompt reuses nothing from the first". The cache hits on identical prompts — cached_tokens 0 → 355 was that — while flash traffic repeats a system prompt and differs in the user turn, which the ordinary slot table cannot serve because a hit needs the slot to be a prefix of the request.

What makes the design defensible rather than a special case:

  • pin_point() decides from the previous request, not from a template. "Token ids only: no template knowledge, no notion of which message is the system prompt" is the sentence that keeps this from becoming a second, worse prompt-format layer — it is a shared-prefix computation, and the gain threshold (--prefix-min-gain) is what stops it pinning a trivial overlap;
  • a pin is replaced only by a new pin and is skipped by eviction and by a longer prompt's replacement — both of those are mutants in the proof, and both go red ("a prompt that resumed from the pin is saved beside it, not over it"), which is exactly the discrimination I ask every arm for;
  • pin() takes (prefix, bytes) and stores state, not policy — the eviction rule lives with the table that owns the slots.

Two things I checked rather than took on trust: the arm set is 28 checks including continuation cases ("a third question resumes from the pin: 900 of its 950 tokens are already there"), and the file does not compile against main's header — so both the API and the behaviour are new, which is what makes the two mutants meaningful.

The scoping is also right and worth keeping in the record: this is the slot-table half, device-free; the forward half is named and not in the PR. For a feature that will eventually touch the forward's state layout, landing the pure decision logic with its mutants first is the ordering that lets the device half be a small change rather than a rewrite.

One ask, and it is the same one I made on #1903: when pin_point declines to pin and the reason is not "cache off" or "below min gain", log the reason once. A pinned slot that silently does not appear is indistinguishable from a pin that is never consulted, and this session has lost time to that shape four times.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT: READY at c88b339 — A pinned slot for the prefix single-turn requests share (+120 header, +100 test, docs +22): the test size is what makes a slot cache credible — slots that are pinned and then evicted, or pinned twice, produce exactly the kind of cross-request contamination that reads as a model problem.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT at c88b339: READY — it names the shape that cannot hit today, and pins the shared part with a rule that needs no template knowledge

The gap is stated as a test row, which is the right way to state it:

insert(system[900] + question1[40])   -> saved
lookup(system[900] + question2[30])   -> slot -1, cached_tokens 0

A slot holds a whole prompt, and a hit needs the slot to be a prefix of the request — so system + question 1 is not a prefix of system + question 2, and "system prompt saved once and reused" does not hold for single-turn traffic today, with any --prefix-* setting. That is the same shape I measured on the serving path: the cache hits on an identical prompt and reuses nothing when only the tail differs, which on flash-style traffic is the common case, not the edge case.

pin_point() reading token ids only is the design decision I like most: "the rule reads token ids only, so it needs no knowledge of the template or of what a system prompt is." A cache that special-cases "system prompt" would be wrong for every template whose preamble is not first, and would encode a template assumption into a component that has no business knowing one. Longest common prefix is the right primitive — the same one a radix cache uses, at a granularity this table can express.

The pin's eviction exemption is a real decision rather than an accident: "One pinned slot at a time; the budget's eviction and a longer prompt's replacement both leave it alone, a new pin() replaces it." And the test pins the hard half — check(c.insert(prompt(0, 300, 10), 100) != pin, "and a later prompt cannot take it back") — i.e. the pin is not reclaimed by a later, smaller insert. Without that assertion, "pinned" would mean "pinned until something wants the room".

Two things I would confirm, and the first is the one that decides whether this can be trusted on the serving path:

  1. Block alignment. pin_point returns 900 for a 900-token shared prefix, and resume_point(d, …) == 899 "one token short of it, as every hit does" — but a KV cache's reusable unit is a block, not a token. If the pin is taken at a token boundary that is not a block boundary, the forward either recomputes a partial block or resumes mid-block, and the latter is a correctness risk rather than a performance one. The tests use round numbers (900) so it cannot show this; a case with a shared prefix of, say, 2 * block + 7 would.
  2. The pinned slot's validity is not established here. Pinning the state at a point is only safe while the KV blocks backing that prefix are still resident; if the budget can evict those blocks, the pin points at freed memory. The header says the budget's eviction leaves the pin alone — worth confirming that the blocks are protected too, not only the slot record.

The forward half is explicitly not in this PR, which keeps the claim honest: this is the slot-table half, device-free, and the table says so.

@bong-water-water-bong
bong-water-water-bong merged commit 2411cee into main Nov 10, 2026
8 of 9 checks passed
@bong-water-water-bong
bong-water-water-bong deleted the cl-flm/1785-flash-profile branch November 10, 2026 15:22

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT: READY at c88b339

The framing is the useful part: the issue measured a hit on an identical prompt (cached_tokens 0 → 355), and a flash request is not that — the system prompt repeats while the user turn differs, and on a whole-prompt slot table system + q1 is not a prefix of system + q2, so nothing is reused. Deriving the shared part as shared_prefix_len and giving it a slot is the right primitive, and it is the right one because it is stated as a rule over token ids only: no knowledge of the template, no idea of what a "system prompt" is. That is what makes it hold for a chat template the code cannot see.

The three properties I checked in the diff and like:

  • a pinned slot is not extended (if (… .pinned) slot = -1;) — the next single-turn request shares the pin, not this prompt, so the prompt takes its own slot. Extending the pin would have made the second flash request miss the thing the first one pinned, which is the whole point.
  • a pinned slot is not evicted for room — make_room skips it, and pin() replaces it explicitly. One pinned slot at a time is a defensible bound and it is stated rather than implied.
  • the pin is not replaced by a longer prompt — a longer prompt taking a normal slot is exactly the "budget eviction" case that should not disturb the pin.

Two asks, one of them the reason I would not merge on this alone:

  1. A test that fails without the pin. The PR adds rows to npu_sampling_test.cpp; make sure at least one drives two different user turns over the same system prompt and asserts cached_tokens on the second is the shared length, and that it would have been 0 without pin. A row that passes on both sides of the change does not protect it — the same trap as a decode that round-trips.
  2. Say what invalidates the pin. A pinned prefix holds state at a token position; if anything outside the cache can change that state (a different KV precision, a template edit, a different n_batch, a re-load of the model), the pin is stale and the failure is silent text drift, not a miss. One line in docs/npu.md naming the invalidation rule (pin() replaces, and whatever else does) closes the question a reader will otherwise have to reverse-engineer.

Not blocking, and the split is honest: the slot-table half is device-free and testable, the forward half is named as out of scope rather than implied to work.

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.

1 participant