Repository navigation
Conversation
`RingKVCache.update_and_fetch` wrote the incoming K/V into the ring and
then attended over the result. With `capacity == window_size == W`, a
prefill chunk of `Q` tokens at `offset` writes slots
`[offset%W, offset%W+Q)`, which hold positions `offset-W .. offset-W+Q-1`.
The chunk's first query, at position `offset`, needs positions
`offset-W+1 .. offset`. So the chunk destroys `Q-1` positions of its own
in-window history before its earliest queries ever read them.
Decode is exempt -- at `Q == 1` the only evicted slot is the one that
falls outside the window anyway -- which is why the existing parity
tests never caught this. They exercise decode, or prefill 48 tokens into
a 64-slot ring where the buffer never fills.
Sliding attention now concatenates the new K/V onto the cached window,
attends over `W + Q` keys, and writes to the ring afterwards. Split
`update_and_fetch` into `fetch_and_concat` (read-only) and `update` (the
write); `update_and_fetch` stays as their composition, documented as
decode-only, since nothing needs to change for `Q == 1`.
The mask needed a companion, because `ring_window_causal_mask` describes
the ring *after* the write. `concat_window_causal_mask` describes the
layout before it: `capacity` slots holding positions `< offset`, then
the `Q` new keys. It reconstructs slot positions by walking back from
`offset-1` with the same no-`aten.remainder` trick, and takes the
modulus of `offset + capacity - 1` so the operand stays non-negative at
`offset == 0`, where sympy would otherwise see `Mod(-1, W)`.
Writes still must not wrap: `(offset % W) + Q <= W`. That cannot be
branched on here because `offset` is a symint under `torch.export`, so
it stays a caller obligation -- pick a chunk size dividing the window.
The drafter carries the same fix. Its layers are all sliding, so it was
affected on every layer rather than three in four.
Verified: the new mask is bit-exact against a brute-force replay of the
ring across three wraps for every offset, capacity in {8,16,32}, window
in {capacity, capacity/2}. End-to-end, chunked prefill now matches
token-by-token decode (~3e-06, fp32 noise) where the old path diverged
by 4.6 on the target and 0.22-1.24 on the drafter, growing with chunk
size as expected. Both models still export under `torch.export` with
dynamic shapes, and the exported program matches eager at 0.0 on logits
*and* on the mutated K/V buffers -- which is what confirms
functionalization keeps the read ahead of the write.
The VLM `MuseGlimmerForCausalLMEmbeddings` path is untouched: it forces
`is_sliding = False` on every layer, and its exported graph contains no
ring ops at all.
(cherry picked from commit 3509a1a)
The concat parity tests trust token-by-token decode as their reference. This adds an independent brute-force sliding-window oracle that recomputes each query's attention over the true last-W key/value positions from a full non-ring K/V log, and compares the fetch_and_concat ring path against it. Covers a multi-token chunk (Q>1) written after the ring is full (offset >= window), including a chunk whose read window straddles the physical wrap, plus a Q=1 decode control. On the old write-first path these diverge from the oracle by 0.49-1.41; the read-before-write path matches within fp32 noise. (cherry picked from commit d8d13ff)
The concat prefill path writes into the ring at offset % capacity, where offset = seq_len - query_len is derived from tensor shapes. That write column must never come from tensor data: the slice_update kernel crashes on the GPU backend when its write index is a runtime data tensor instead of a shape symint. Add an export-based test that pins the property: it exports the fetch_and_concat -> update path with a dynamic sequence length, finds the ring write node, and asserts its begin index traces back to a shape query (sym_size) and never to a value pulled out of a tensor. A future refactor that reintroduced a data-dependent write index would flip this test red before it reached a device. (cherry picked from commit 1c712bc)
… decode Gate the sliding-window ring access on prefill_mode: the prefill graph reads the ring before writing (concat), correct for any chunk after wrap; the main/decode graph keeps the fast write-first path, correct for Q==1. Add a fail-loud Swift guard so a Q>1 call on the write-first main path (spec-decode verify, PPL) throws instead of silently corrupting, keyed on the presence of a sliding-ring state.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #314.
RingKVCachewrote each prefill chunk into the ring before attention ran, so once the ring was full a multi-token chunk (Q>1) evicted the oldest in-window positions its own earliest queries still needed — silent wrong output on sliding-window models (Muse-Glimmer). Decode (Q=1) and short prefill were unaffected.Fix: gate the ring access on
prefill_mode.fetch_and_concat+concat_window_causal_mask) — correct for any query length, including after the ring wraps.maingraph, query_len == 1): keep the cheaper write-first path.mainpath with a sliding ring throws (spec-decode verify, all-logits/PPL), so it can never silently corrupt.Verified on M3-Ultra (Muse-Glimmer-30B 4bit): guarded decode 35.3 vs unguarded concat 34.4 tok/s (+2.5%), prefill unchanged. Oracle tests pin write-first correct at Q=1 and divergent at Q>1 (the regime the guard refuses); concat correct throughout.
Note: the gate assumes the eager compression path (Muse's 4bit default), which traces
main(write-first) andprefill(concat) separately. A graph-mode preset would deriveprefillby trimming the write-first graph → corruption; a graph-mode Muse export should be rejected until that's handled.