prefix caching changes temperature-0 answers on Flash-Next; one tool round-trip goes wrong (warm replay from a Marconi anchor)

Summary

With --enable-prefix-caching, Qwen3.8-Flash-Next gives different temperature-0 output than with caching off, and in one case the difference is a wrong answer: after a tool result, the model calls the same tool again instead of answering. Saved state is restored bit-exactly. The difference comes from how the warm turn is computed. This matters for #149, which makes prefix caching the default for this model.

Setup

GB10 (DGX Spark), Qwen3.8-Flash-Next EXL3 3.05 bpw with the n-gram table, one lane, MTP with 3 drafts, 262,144 context, thinking off, temperature 0.
Build: Closed andrew-vogulkin/atlas#6 at c1796929 plus our fixes (scheduler/ssm_pool: never compact a sequence onto an SSM slot another sequence owns (fixes #147) #150 among them). Not tested on main: the EXL3 path is not on main yet. The code involved (prefill_b.rs tail split, prefix_lookup.rs) is the same on main 4405bce3. We saw the same wrong answer on our earlier build (Qwen3.8-Flash-Next EXL3: MTP + native int8 decode for GDN/attention/lm_head (M6, follow-up to #91) #100 head plus exl3 GB10: fix concurrent MTP (K=4) and lm_head row overlap; +20% decode kernels andrew-vogulkin/atlas#2–chore(deps): Bump base64 from 0.22.1 to 0.23.1 #5 and scheduler/ssm_pool: never compact a sequence onto an SSM slot another sequence owns (fixes #147) #150).

Test

A 23-case tool-call battery (26 responses: single calls, parallel calls, and three round trips where the tool result is appended and the model must answer). Each arm ran on a fresh server; every arm is compared with caching off.

Arm	Setting	Responses that differ from caching off (of 26)	Battery
A0	caching off, run twice	0 (run 1 = run 2)	23 pass
A1	caching on, `ATLAS_NO_TAIL_SPLIT=1`, `ATLAS_MARCONI_MIN_TOKENS=1000000000` (one pass, no restore)	0	23 pass
A2	caching on, `ATLAS_MARCONI_MIN_TOKENS=1000000000` (tail split, no restore)	5 (wording, preface, argument order)	23 pass
A3	caching on, default (`ATLAS_SSM_SAVE_DUMP=1` for checksums)	6	1 fails
A4	A3 plus `ATLAS_PREFIX_SUBBLOCK=0`	the same 6	same failure

The failing case: "What's the weather in Tokyo right now? Use the tool." → correct get_weather call → tool result {"city":"Tokyo","temp_c":-3,"conditions":"light snow"} → the model emits get_weather({"city":"Tokyo"}) again with empty text. With caching off it answers with the weather. On the earlier build the warm reply was "I made a typo in the city name — let me re-check with the correct spelling." plus the same repeated call.

Log line for the failing turn: Marconi intermediate hit: restored from checkpoint at token 1056 (skipping 1056 tokens, replaying 90 SSM tokens to reach 1146; …), usage cached_tokens: 1072.

What the arms show

Restore is exact. With ATLAS_SSM_SAVE_DUMP=1, all 19 restore@1056 checksum lines equal the ckpt_save@1056 line (h_sum, h_ssq, c_sum, c_ssq), and the same holds at 1104 and 1184.
The tail split alone changes output (A2 vs A1). With caching on, the last prefill chunk is always split at floor((N-1)/16)*16 - 16 (model/trait_impl/prefill_b.rs); with caching off it is one pass. The comment there says the split was made unconditional so cold and warm match each other; caching-off runs never split, so on and off cannot match.
The wrong answer needs the restore path (A3 vs A2): the replay starts at the 1056 anchor, and the KV for tokens 1056–1071 is reused from the earlier request and not rewritten (forward_layers.rs), so those 16 tokens come from a different pass layout than the rest of the replay.
Sub-block matching is not the cause (A4 = A3).
We could not measure how close the first-token choice is: logprobs.content starts at the second generated token, and every flip is at the first.

Smallest repro (2 requests, caching on, fresh server)

Any request with the same tool preamble (it plants the anchor at 1056): the Ottawa weather question with the get_weather tool.
The Tokyo round-trip history (question, assistant tool call, tool result), 1146 prompt tokens.
Request 2 repeats the tool call. The same request with caching off, or with the A1 env, answers correctly.

Why the existing gate does not catch it

ssm-state-poisoning-gate compares rounds of one caching server with each other, passes wording changes as jitter, has no tool-result turn, and has no entry for qwen3.8-flash-next. It never compares caching on against caching off.

Suggested direction

Make the prefill pass layout depend only on absolute token position, the same with caching on and off (fixed grid; checkpoints only on grid points), so a warm replay from a grid point runs the passes a cold run would.
Do not reuse KV above the snapshot depth: recompute [snapshot, matched).
Add a caching-on vs caching-off equality check on a tool-call battery for this model before strix-windows: Flash-Next launchers + prefix caching on by default (stacked on #80) #149 makes caching the default.

We are trying 1 and 2 on our build now and will report the result here.

Thanks. This is a very clean set of arms. The A1/A2/A3 split pins the two mechanisms separately, and the checksum proof that restore is exact saves a lot of guessing. Agreed: #149 should not make caching the default until this is fixed and a caching-on vs caching-off equality check exists for this model. I'm marking #149 blocked on this.

What this means for Windows (gfx1151), from today's work:

Until now Windows only had your A2 effect. On atlas_scale builds the mid-chunk tail capture produced aux-less snapshots and displaced the aux-carrying tail-split checkpoint, so no warm restore ever happened on Windows (strix-windows: Flash-Next launchers + prefix caching on by default (stacked on #80) #149's replay: every turn logged "Prefix cache hit … but no SSM snapshot"). The tail split still ran, so Windows was in your A2 state: output drift, no restore.
prefill: refuse mid-chunk tail capture on aux-carrying models (Flash-Next warm path never engaged on Windows) #157 fixes the restore, which on Windows means it would switch on your A3 path, the one that produced the wrong answer. So prefill: refuse mid-chunk tail capture on aux-carrying models (Flash-Next warm path never engaged on Windows) #157 should land together with a fix for (2) below, or behind the same equality gate. I'll note this on prefill: refuse mid-chunk tail capture on aux-carrying models (Flash-Next warm path never engaged on Windows) #157.

On your two directions:

A grid-aligned pass layout (chunk boundaries and checkpoints only on absolute-position grid points, identical with caching on and off) is the right invariant. It also removes the cold/warm asymmetry the tail-split comment was trying to fix. The cost: one more small pass when the prompt end isn't on a grid point, which the code's own measurements put at ~150 ms on this MoE.
Recompute [snapshot, matched) instead of reusing that KV is cheap: by construction the gap is ≤ 2 blocks, so ≤ 32 tokens of extra replay, and it removes the mixed-pass-layout KV. I'd do this one first; it's small and local to the restore path.

A GB10 NVFP4 repro on main: our GB10 is back today, so I'll run your battery's 2-request repro (Ottawa preamble, then the Tokyo round trip) on main with the NVFP4 checkpoint, caching on vs off. That confirms the mechanism is format-independent and gives a main-tree baseline for whichever fix lands. I'll post the result here. Please do share your results for (1)/(2) on your build, so we can avoid fixing it twice.

Reproduced on main with the NVFP4 checkpoint (GB10 job 661, main 4405bce35, Qwen3.8-Flash-Next-NVFP4, util 0.88, thinking off, temp 0, no MTP). So it isn't EXL3-specific.

Battery: 12 tool-call cases sharing one ~2k-token system preamble. Each case is question → tool call → tool result → answer, 24 responses. Each arm ran on a fresh server; everything is compared with OFF1.

Arm	Responses that differ from caching off (of 24)	Repeats the tool call instead of answering
OFF2 (caching off, control)	0	0
ON-A1 (caching on, `ATLAS_NO_TAIL_SPLIT=1`, `ATLAS_MARCONI_MIN_TOKENS=1e9`)	0	0
ON (caching on, default)	16 (6 first requests, 10 round trips)	1

ON restored on every warm request (12/12 round trips, e.g. restored from checkpoint at token 2000 … replaying 75 SSM tokens), at ~2,000 cached tokens each.
The wrong answer is your Tokyo shape. "Has shipment C-310 been delivered?" → track_shipment → result {"status":"delivered"} → caching off answers "Yes — shipment C-310 has been delivered…", and caching on emits track_shipment({"id":"C-310"}) again with empty text.
Other diffs are wording and dropped prefaces ("I'll check the current weather…" becomes an empty-content call).
First requests differ too. They warm-hit the shared preamble from earlier cases.
A1 = OFF exactly here, so one-pass + no-restore is a fully safe configuration on main today.

Scope: this is worse than your 6/26 because the shared preamble makes every request warm. A tool-calling agent with a long system prompt lives exactly in that regime, so it's the case that matters for #149. Data is in reiner:~/jobqueue/done/661-gb10-158-cache-onoff/ (bat-*.json, serve logs).

Next: this battery is now the gate for any fix. Pass = ON equals OFF on 24/24. Happy to run it against your fix branch, or against a fix here (recompute [snapshot, matched) first, then the grid-aligned layout).

Thanks — the main + NVFP4 repro is the useful part for us: it rules out anything EXL3-specific,
and your 16/24 with a shared preamble is the regime that actually matters. Agreed that your battery
should be the gate.

We have both of your directions implemented and passing on our build. Mapping onto your numbering:

(1) grid-aligned pass layout — done. Chunk boundaries and checkpoints only on absolute-position
grid points, identical with caching on and off; radix inserts floored to the grid; no leaf or
decode checkpoints under the grid; sub-block matches disabled for grid models (a sub-block match
hands out a block computed by a different pass layout). Cold cost on our build is −2.0% at 30K
(1152 vs 1175 tok/s), close to the ~150 ms your reading predicted.
(2) recompute [snapshot, matched) — done, and it is as cheap as you expected.

One correction to the sequencing, worth having before you spend time on it: (2) first will not
get you there, and on a grid build it almost never fires. Two findings:

With (1) in place, grid-floored inserts mean the radix match always equalled the anchor, so (2)
had nothing to cap (kvcap=0 in every arm). It only matters after a snapshot eviction, which we
had to construct deliberately (--ssm-cache-slots 2) to exercise at all.
With (1) + (2) alone, warm still differed from cold at 30K, with and without MTP, and the text
diverged in the first sentence. So the layout fixes were necessary but not sufficient.

Two further causes we only found by dumping per-layer hashes cold vs warm:

GDN kernel mismatch. gdn_exact_replay was set on every warm-hit pass, forcing the GDN
prefill onto the token-sequential WY4 recurrence while the cold run computed the same pass with
the FLA chunked kernel. The two give different FP32 h — the dump showed h differing while layer-0
conv matched exactly. The forcing was correct under the old layout, where a replay started at an
arbitrary offset and FLA's 64-token grid did not line up; once the replay pass is the cold pass,
FLA is the faithful kernel. Fix: gdn_exact_replay = marconi_skip && !grid_replay_is_cold().
Drafter anchor (MTP only, so invisible in your no-MTP run). A cold turn builds drafter context
from the whole-prompt capture of target hidden rows; a warm turn's capture was short, so it fell
back to the previous turn's carried drafter KV — a different drafter state. Fix: store hidden rows
[0, a) at each grid checkpoint, restore them on a hit, decline the hit if rows are absent.
Hidden-row store is capped and host-side; no GPU memory added.

After fix 6, every dumped line is identical cold vs warm: the per-layer hash after each of 48 layers,
the logits, the SSM checksums, every aux blob at prefill end, and each decode step.

Our acceptance (all byte-identical cold vs warm and caching-on vs caching-off): tools battery
26/26 twice; 30K and 100K with MTP on and off; an eviction case through the KV cap re-reusing at a
lower anchor; overlap suite identical across 5 runs; held-out 45/50 unchanged; vision 5/5 (the grid
is skipped for those). Follow-up TTFT 1.0 s at 30K and 1.6 s at 100K, against cold 25.8 s and 97 s.
Unit tests pass except a stale baked attestation that fails against any source change.

On #157 and Windows: your read matches what we found, and fix 6 is the part that bears on it.
If #157 switches Windows from the no-restore state onto the warm-restore path, it will hit the GDN
kernel mismatch regardless of the pass layout, because that flag was set on every warm-hit pass. So
#157 needs fix 6 (or the equality gate) and not only the layout work.

Yes please to running your battery against our branch. One caveat so the result is interpretable:
our commits sit on a Flash-Next-focused branch off our own tree, not on main, so they will need
adapting — the grid change is the one with real design implications for other models and is probably
worth agreeing on before anyone merges. Happy to attach the seven patches here (individually or
combined), open a PR, or split them per fix. Each one has an env switch that restores the previous
behaviour, which should make bisecting on your side easier:
ATLAS_GRID_EXACT_REPLAY=1, ATLAS_MTP_ANCHOR=0, ATLAS_PREFIX_KV_CAP=0.

Not verified on our side: concurrency above one lane (mixed and batched prefill paths, shared
capture buffer); multi-rank / EP; a mid-prefix partial match against its own cold run; long-run host
memory of the hidden-row store under real traffic. Also, separately from reuse: on our build cold
MTP-on output does not equal cold MTP-off output, because every target token comes out of the K+1-row
verify forward. Your battery ran no-MTP so it will not see that, but it means "MTP on vs off" is not
a valid equality check on this engine.

Thanks. Dumping per-layer hashes cold vs warm was the right instrument. The GDN-kernel mismatch (gdn_exact_replay set on every warm pass) and the MTP drafter anchor are two causes we hadn't found. We agree that "MTP on vs off" isn't a valid equality check on this engine; our #130 measurements say the same.

Yes, please open a PR against Atlas-Inf/atlas main, one commit per fix. That lets us bisect, and lets the grid change get its own review. That one has cross-model reach: every SSM-hybrid model that runs with caching on goes through the tail split, including the Strix Qwen3.8-27B recipes. Keep the env switches for review. Repo rule is that whatever is proven becomes the default and its switch comes out before merge.

What we'll run once it's up:

GB10: job 661's 12-case battery (reiner:~/jobqueue/done/661-gb10-158-cache-onoff/) on NVFP4. ON must equal OFF on 24/24, with the A1 control alongside.
Windows (gfx1151): the same battery with prefill: refuse mid-chunk tail capture on aux-carrying models (Flash-Next warm path never engaged on Windows) #157 in the tree, because Windows has never taken the warm-restore path. Your fix 6 is the part prefill: refuse mid-chunk tail capture on aux-carrying models (Flash-Next warm path never engaged on Windows) #157 needs.
Concurrency > 1: you listed this as unverified. Our [bug] Flash-Next (mHC highway): a decode step between another request's prefill chunks corrupts output, and the process stays bad for concurrent requests until restart #147 C=2 probe (GB10 jobs 667–669: a short request that retires mid-way through a 47k chunked prefill) can rerun on your branch.

One question that affects the Windows recipe, where each SSM snapshot slot costs ~113 MB against a small KV pool: is the grid spacing fixed (for example the FLA 64-token chunk), or is it tied to --ssm-checkpoint-interval?

PR is up: #163 — seven commits, one per fix, env switches kept. Cherry-picked cleanly onto main 4405bce; workspace cargo check and spark-runtime tests pass there. Not yet run on NVFP4/main; that's your job 661 battery. Note: at the default grid G=4096 your ~2k preamble won't be reused, so set ATLAS_PREFILL_GRID_FINE=1024 to exercise the warm path.

Great writeup — we reproduced the mechanism on main and have a fix branch up.

What's going on, confirmed from main's source

The single-stream Marconi path and the batched path disagree about what the
replay's KV means. In prefill_b/forward_layers.rs the replay sets the
attention write floor to cached_prefix_tokens so K/V for [snap_tok, matched) is not rewritten — values computed under the original pass layout
stay in the shared blocks while the recurrent state comes from the replay
layout. The batched path (batch_kernel.rs) already rewrites the replayed
range from 0. So the same warm hit produces consistent caches under batching
and a mixed-layout pair on the single-stream path — your "KV for 1056–1071
reused from a different pass layout" is exactly this floor.

The second contributor is what you measured in A2: the tail-checkpoint split
in prefill_b.rs was gated on prefix_cache.is_active() (and an enabled
snapshot pool). With caching off the prompt takes a single-pass layout; with
it on it takes [0..cut)+[cut..N). Same prompt, two pass layouts, temp-0
flips — on/off could never match while that gate existed.

Fix — fix/158-marconi-kv (small, both your suggested directions):

Marconi replay now rewrites KV for the replayed range on the
single-stream path, matching the batched behaviour — the KV and the SSM
state then describe the same trajectory.
The tail split fires on absolute token position only (any SSM model —
cache on or off), so caching-on/off and cold/warm share one pass layout.
The checkpoint save still gates on an enabled snapshot pool; the cost is
the one extra small pass on caching-off prefills.

Verification on our side is running now: build on GB10 + serve
Qwen3.8-Flash-Next-NVFP4 + your two-request repro (Ottawa anchor → Tokyo
round-trip), same arm against the unfixed binary.

What this doesn't cover yet — your suggestion 1 in full generality: a
replay anchored deeper than the tail checkpoint still starts passes at the
snapshot position rather than on the cold chunk grid (only the tail-split
anchor coincides with a cold boundary today). Whether the residual is
measurable on your battery — after (1)+(2) — is the question our repro
answers; happy to take it further into a fixed-grid checkpoint policy if it
still flips.

Also agree on (3): an on/off equality leg belongs in the gate before #149
defaults this on.

@AzeezIsh

Thanks @AzeezIsh — that matches what we saw, and both parts of fix/158-marconi-kv line up with the two causes.

We'll run your branch on our GB10 (Qwen3.8-Flash-Next EXL3 3.05 bpw, one lane) against our tools battery, cache on vs cache off: 26 prompts, two runs each, plus the 30K and 100K cold vs warm pairs. We'll test at the default grid and at G=1024, and post the numbers here.

On #163: if your branch passes, should we close #163, or cut it down to the parts your fix does not cover (the --exact-prefix-cache opt-in and the deeper-anchor case)? Happy to go either way.

We ran fix/158-marconi-kv (01cb02e + b6a5c17, applied cleanly) on our GB10 EXL3 build (Flash-Next EXL3 3.05 bpw, base f28e1589 without our #163 commits), grid 4096, one lane, cache ON vs OFF:

Tools battery: 21/26 byte-identical in both runs (run 1 = run 2 in each arm). The same 5 differ: 3 tool round-trip final answers (the prefix caching changes temperature-0 answers on Flash-Next; one tool round-trip goes wrong (warm replay from a Marconi anchor) #158 case) and 2 tool calls with the same arguments in a different key order.
30K and 100K: the first request equals cache-OFF; the warm follow-up (29,568 / 99,616 tokens cached) differs at both lengths.

G=1024 was not tested: the grid setting comes from our #163 commits and doesn't exist on this base.

So on this build the two changes are not enough by themselves. The remaining gap may be the deeper-anchor case you mentioned, or the grid alignment our #163 adds. A caveat: this is EXL3 on our fork base, not NVFP4 on main, so your repro may still pass there. Happy to share outputs, or to test a follow-up commit.

Follow-up: we found the missing piece. On top of fix/158-marconi-kv plus our two grid commits (0f21eb1a, 3b281547), adding 2673957f from #163 alone makes cache ON byte-identical to cache OFF at G=1024 on our GB10 EXL3 build:

Tools: 26/26 in both runs (the 21/22 round-trip answers now match).
30K and 100K: warm = cold = OFF.

2673957f makes the warm replay run the same FLA GDN kernel as the cold pass, instead of the forced WY4 exact-replay path. Without it, the replayed SSM state comes from a different kernel than the cold run, so warm outputs drift even with your KV-floor fix. Adding our other #163 commits one at a time (KV cap, normalize, chunk align) did not close the gap; adding 2673957f did, and it is also enough by itself.

Not tested: 2673957f on your branch without our grid commits, and G=4096. Suggestion: take 2673957f into your branch (it applied cleanly here). We can then close #163 or cut it down to the grid opt-in.

Two corrections to my last comment, after checking the code and testing further:

2673957f is not enough by itself: it depends on our grid commit 0f21eb1a (grid_replay_is_cold() needs prefill_grid(), and it is a no-op when no grid is set). The minimal set on top of your branch is 0f21eb1a + 2673957f. With both, G=4096 also passes: tools 26/26 in both runs, and 30K/100K warm = cold = OFF.
On GB10 the warm replay before the fix ran the register-resident token-sequential GDN kernel, not WY4 (WY4 only runs if that is disabled with ATLAS_NO_GDN_REGRESIDENT=1). Both differ from the cold pass; only FLA chunked gives the same bits, because under the grid the replay starts at a cold chunk boundary. Setting ATLAS_GRID_EXACT_REPLAY=1 on the fixed build brings back the exact failure (same 2 tool cases, same 30K/100K hashes). Forcing WY4 also differs from cold.

The code comment at layer.rs:356-364 and our commit message also say WY4; we'll fix that wording in #163.

strix-hip: mask the stale uc rows of the last GDN chunk in chunk_fwd_o (

Navigation Menu

Repository navigation

prefix caching changes temperature-0 answers on Flash-Next; one tool round-trip goes wrong (warm replay from a Marconi anchor) #158

Summary

Setup

Test

What the arms show

Smallest repro (2 requests, caching on, fresh server)

Why the existing gate does not catch it

Suggested direction

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

prefix caching changes temperature-0 answers on Flash-Next; one tool round-trip goes wrong (warm replay from a Marconi anchor) #158

Description

Summary

Setup

Test

What the arms show

Smallest repro (2 requests, caching on, fresh server)

Why the existing gate does not catch it

Suggested direction

Activity

AzeezIsh commented on Oct 2, 2026

AzeezIsh commented on Oct 2, 2026

maplepaladin73 commented on Oct 2, 2026

AzeezIsh commented on Oct 4, 2026

maplepaladin73 commented on Oct 4, 2026

AzeezIsh commented on Oct 12, 2026

maplepaladin73 commented on Oct 13, 2026

maplepaladin73 commented on Oct 13, 2026

maplepaladin73 commented on Oct 14, 2026

maplepaladin73 commented on Oct 14, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions