feat(mtp): make the drafter context (prefill + cross-turn carry) the default - #352
Open
tbraun96 wants to merge 24 commits into
Open
feat(mtp): make the drafter context (prefill + cross-turn carry) the default#352tbraun96 wants to merge 24 commits into
tbraun96 wants to merge 24 commits into
Conversation
`commit_accepted_prefix_dispatch` indexes the per-token SSM intermediates with `num_accepted - 1`. That is `usize` arithmetic and release builds have overflow-checks off, so `num_accepted == 0` wraps to `usize::MAX`, `ssm_pool.h_intermediate(.., usize::MAX)` returns an out-of-range device address, and it is handed straight to `copy_d2d_async` — a wild 3 MB D2D write per SSM layer. The case is unreachable today only by caller convention: position 0 of a verify batch is accepted by construction, and DFlash passes `num_accepted + 1` (verify_dflash_step.rs). Nothing inside this function enforces that, so one caller refactor turns a rollback into memory corruption. Guard it explicitly; a genuine full-reject rewind belongs in `rollback_ssm_states`, which restores the pre-verify checkpoint. No behaviour change on any live path — every current caller passes >= 1. Not unit-testable without a GPU-backed TransformerModel (the method needs a live CUDA context, an SsmStatePool and per-sequence layer states); the guard is a precondition assertion on a private dispatch method.
…ck target `rollback_ssm_states_dispatch` had three branches: `num_accepted == 0` (restore the pre-verify checkpoint), `num_accepted <= intermediates.len()` (index-select the intermediate), and `intermediates.is_empty()` (bail). The gap is `0 < intermediates.len() < num_accepted`: no branch fires, the function returns `Ok(())`, and `h_state`/`conv_state` are left ADVANCED past the last accepted token. Every subsequent decode then runs on SSM state that includes rejected draft tokens — silent corruption, no error, no log line, surfacing much later as gibberish. The author wrote a fail-fast for the empty case and left the short case falling off the end. It is reachable as soon as the verify width K exceeds the intermediate pool depth, which is `num_drafts + 1` (impl_a1.rs:148) — i.e. exactly the direction K is being pushed. Make the whole `num_accepted > len()` range bail, and report the actual pool depth so the misconfiguration is diagnosable from one log line. Also corrects the trailing comment: `num_accepted == num_tokens` does not reach it — callers guard that case, and it would otherwise be swallowed by the branch above and copy `h_state_intermediates[K-1]`, which the fused WY kernels never write. Not unit-testable without a GPU-backed TransformerModel; this is a precondition assertion on a private dispatch method.
…e_async Neither has a caller. Repo-wide grep finds only the trait declaration (traits/model.rs), the forwarder (trait_impl/mod.rs) and the impl (trait_impl/async_chkpt.rs); the remaining hits are comments. Every live verify path (K=2, K=3, K=4, DFlash) runs the kernel directly on the canonical h_state and commits through `commit_accepted_prefix`. `pre_verify_copy_async` is not merely unused, it is a landmine: `commit_accepted_prefix` deliberately does NOT refresh `*_checkpoint` after a partial-accept rewind (the next `start_checkpoint_async` does it at prefill time). So a checkpoint -> live copy run after a commit would resurrect the pre-verify state and re-apply the rejected draft tokens. Anyone re-enabling it as an optimisation would get silent corruption. `commit_verify_state_async` carried the only implementation of the "live-state invariant" full-reject restore, so that code has in fact never executed; keeping it implies a safety net that is not there. Deleting rather than commenting: a dead landmine is worse than no code. Removes 187 lines net. No behaviour change — nothing called them.
… verify
The fused K=2/3/4 verify paths snapshot the conv1d sliding window after
every verify token, including t == K-1. Nothing reads index K-1:
* `commit_accepted_prefix(n, k)` early-returns on full accept, so it only
indexes `n - 1`, and the schedulers call it with max `n = k - 1` for a
partial accept (verify_k2/k3/k4_step.rs, verify_dflash_step.rs).
* `start_rollback_and_checkpoint_async(n)` is only ever called with
n in 1..=K-1 (impl_a2.rs:450-509, spec_step.rs:340), so index <= K-2.
* `rollback_ssm_states` is the self-spec / ngram path, which does not use
these fused kernels — the sequential fallback still writes all K
intermediates and is untouched.
So it was a `conv_bytes` D2D per SSM layer per verify step for nothing.
MEASURED COST OF THE WRITE (nsys, --cuda-graph-trace=node, 25 s window,
dgx2, W4A4 27B, K=2, 294 verify steps x 48 SSM layers):
conv_state-sized D2D (163840 B): 33,090 copies, 76.0 ms total
of which dead t=K-1 writes: 14,112 copies, ~32.4 ms
total decode GPU kernel time in the same window: 23.9 s
=> the dead write is ~0.14% of decode GPU time.
Being honest about the size: 0.14% is well below the 0.5% run-to-run wall
spread measured on this box, so an end-to-end A/B CANNOT resolve it and
none is claimed. This lands as correctness hygiene with a measured upper
bound on the win, not as a throughput lever. The campaign's verify tax
lives elsewhere: the same trace puts w4a16_gemv_batch2 +
w4a16_gemv_dual_batch2 at 84.5% of decode GPU time and the entire SSM
stack (kernels + all state D2D) at ~3.4%.
The Marconi hit lines reported `matched - snap_tok` as the recompute cost. That is the anchor->match gap only. The suffix prefill resumes at `marconi_skip_to == snap_tok` and runs the SSM recurrence forward to `total`, so the true replay is `total - snap_tok`. The omitted `[matched, total)` span is precisely the new user message on a warm agentic turn — i.e. the log understated replay cost by the term that grows with the conversation, which is the term a TTFT campaign is trying to see. Both numbers are now printed: the real replay to `total`, and the anchor->match gap that is attributable to snapshot granularity. The exact full-hit line also gains its replay length (previously it printed none, so an "exact hit" looked free even when it replayed the whole generated suffix). Also records, as a comment only, that `--ssm-checkpoint-interval` is a FILTER over chunk boundaries rather than a generator of them: this function runs only at chunk ends, so effective spacing is the CHUNK size. With `--ssm-checkpoint-interval 32` (512 tok) and `--max-prefill-tokens 8192` every chunk end is already a multiple of 32, so interval checkpoints fire every 8192 tokens, not every 512. Not changed here — making the interval a real generator would split chunks and costs prefill throughput, so it needs its own measured A/B. Logging only; no behaviour change.
…of arguing it
Adversarial review of this branch (qwen3.8-max) flagged the previous commit
as its highest correctness risk: skipping the `conv_state_intermediates[K-1]`
write turns a slot into stale data on the strength of a whole-repo "no reader
exists" proof, which a later refactor can silently invalidate with no
compiler error.
The proof itself holds — I re-verified both attacks the review constructed
against the actual code:
* DFlash bonus-token off-by-one: `verify_dflash_step.rs:200-201` passes
`k_verify = drafts.len() + 1` against `total_accepted = num_accepted + 1`,
so the two agree. DFlash also cannot reach the K=2/3/4 branches at all —
it dispatches only at `drafts.len() >= 4` (mtp_step.rs:308), i.e. verify
width >= 5, which lands on K=17 or the sequential fallback, both of which
still write every intermediate.
* `rollback_ssm_states` full accept: its only caller (spec_step.rs:165) is
guarded by `if a.seq.seq_len > expected_seq_len` (:158), which is true
only when a draft was REJECTED — so `num_accepted + 1 <= K-1` and the
index is at most K-2.
But the review is right that an argument is not an invariant. So enforce it
where it can be checked locally: `commit_accepted_prefix` now bails on
`num_accepted > k` as well as on `num_accepted == 0`, and still early-returns
on `num_accepted == k`. Its reachable intermediate index is therefore exactly
[0, k-2] by construction, which is precisely what the skipped write requires.
`num_accepted > k` is also the exact shape a bonus-token off-by-one would
take (passing `gamma` where `gamma + 1` is meant), so this catches the
regression the review was worried about at the moment it is introduced rather
than as corrupted output much later.
The comment in the fused branches is rewritten to cite the enforcement and
the caller guards by file:line, so the next reader does not have to
reconstruct the proof.
The full-accept `commit_accepted_prefix` call sites logged the error and `return`ed without marking the sequence finished, so a failed SSM state commit left the request running on state the engine had just declared untrustworthy. The reject sites already set `a.finished = true`; this makes the accept sites match. Raised by adversarial review of this branch: the new fail-fast guards are only an improvement if the caller actually stops. Otherwise they convert a wild-pointer bug into logical corruption with a clean log line — the request keeps emitting coherent-looking tokens from poisoned recurrent state, which is harder to catch than a crash. Unreachable today (the full-accept path early-returns `Ok` before touching any state), so this is purely the safety net for the guards added earlier in this branch.
…g any copy Adversarial review: bailing from inside the per-layer copy loop leaves the first N SSM layers rewound and the remaining ones still advanced past the accepted boundary. That MIXED state is strictly worse than the uniform corruption the guard was added to prevent, and much harder to diagnose. Split into two passes: an immutable pre-validation sweep that checks every LinearAttention layer has enough per-token intermediates for the requested boundary and bails before a single `copy_d2d_async` is enqueued, then the existing copy loop. The in-loop arm becomes `unreachable!` rather than the original silent `Ok(())` fallthrough, so the invariant is stated where it is relied on. The error message now says "No rollback copies were enqueued" so an operator can tell an aborted rollback from a partial one. Also switches the two new Marconi replay-length terms to `saturating_sub` (review: `snap_tok` can exceed `matched` on a tier fault-in where the effective anchor depth is recomputed, and a debug build would panic on the underflow). cargo check --workspace --all-targets clean; 1069 tests pass, 0 failed.
The accept chain short-circuits: when draft 1 is rejected, draft 2 is
discarded without being scored. So the only position-2 number the engine can
report is the CONDITIONAL p(2|1). We measure p1 ~= 0.70 and p(2|1) ~= 0.53,
and the chain alone cannot distinguish
(a) position 2 is genuinely worse — the drafter conditions on its OWN
draft-1 output rather than a verified token, so error compounds; from
(b) survivorship — position 2 is only ever scored on contexts where
position 1 already succeeded, which is a biased (easy-context) sample.
Those need completely different fixes, so the ambiguity is expensive.
The verify step already computes the target argmax at every position
(v0/v1/v2) in one batched pass, so `drafts[1] == v1` is observable on
EVERY step regardless of the chain. Record it, alongside the conditional rate and
p1, and emit all three every 100 steps:
K3 positional: steps=N p1=.. p2_uncond=.. p2_cond=.. (d1=.. d2u=.. d2c=..)
Reading it: p2_uncond ~= p2_cond means position 2 is genuinely worse;
p2_uncond ~= p1 means p2_cond was survivorship.
Caveat recorded in the code: v1 is the target's argmax GIVEN drafts[0] as the
preceding token, so on steps where draft 1 was wrong this scores the drafter
against a counterfactual context. That is the intended comparison — same
position, unbiased sample of contexts — but it is not the same as running the
drafter from a verified prefix.
Three relaxed atomics and one compare on the verify path; no allocation, no
sync, no behaviour change.
K=4 twin of the K=3 counters. Emits p1, p2_uncond, p3_uncond alongside the conditional p2_cond/p3_cond every 100 verify steps, so the per-position acceptance curve can be read without the survivorship bias that a conditional chain introduces at position 3 (which only scores contexts where positions 1 AND 2 both succeeded). Measured with these on dgx2, W4A4 27B, gate disarmed, seq_len ~10k, 600 steps: p1 = 0.660 p2_uncond = 0.485 p3_uncond = 0.407 p2_cond = 0.533 p3_cond = 0.479 The conditionals overstate positions 2 and 3 by +0.05 and +0.07 — real, but small enough that the unconditional curve is still monotonic and the position-1-to-2 cliff (x0.735) dominates the position-2-to-3 step (x0.838).
…(default OFF)
ATLAS_MTP_REFEED_ACCEPTED=1 (requires ATLAS_MTP_CATCHUP=1). Default OFF, so the
production path is byte-for-byte unchanged.
THE DEFECT. The MTP head is one module run autoregressively. Draft 1 consumes
the TARGET's verified hidden (`mtp_hidden_save`); every later draft consumes
the drafter's OWN single-block residual (`current_hidden =
ctx.buffers.hidden_states()`). So the drafter KV row written for draft d >= 2
pairs the right token with the WRONG hidden. On ACCEPT that row is kept
forever: `after_verify` trims only REJECTED rows. Every accepted draft
therefore contaminates the drafter's own context permanently, and the error
feeds every later propose.
THE EVIDENCE (dgx2, W4A4 27B, gate disarmed, seq_len ~10k, n=700/config,
measured with the unconditional per-position counters added earlier in this
branch):
unconditional acceptance by draft position: 0.660 -> 0.485 -> 0.407
first autoregressive step costs x0.735; the second only x0.838
The loss is concentrated exactly at the hidden-state handoff, not spread
evenly, which is the signature of an input-distribution change at d=2 rather
than of generic drift. Survivorship was excluded first: p2_uncond 0.499 vs
p2_cond 0.501 at nd=2 (agreement to 0.002, ~8 sd below p1). Both existing
levers were excluded too: ATLAS_MTP_CATCHUP=1 is bit-identical (its ring is
written only on SERIAL decode steps, and with the throughput gate disarmed
there are none), and dropping ATLAS_MTP_DRAFTER_PREFILL costs only
0.017/0.030 (~1 sd).
THE FIX. After a verify, the target's true hidden for every accepted position
is already sitting in the verify hidden buffer. Ring those hiddens under the
same label convention the serial-decode hook uses, and have `after_verify`
additionally drop the `num_accepted - 1` accepted rows that were written with
a drafter hidden. The next propose's catch-up feed then rebuilds exactly those
pair keys from the ring, with the target's hidden, through the already
exercised `catchup_drafter` batch path. No new kernel and no new state
machine: this reuses the gap-fill machinery for a gap that was never being
detected.
Label derivation, spelled out because it is the part most likely to be wrong:
the serial hook calls `save_hidden_for_catchup(0, seq.seq_len - 1)` AFTER
appending, so label n holds the hidden that PRODUCED the token at position n
(= hidden_{n-1} in pair-key space). In verify_k3_step, `a.seq.seq_len` has
already advanced by k=3, so the pre-verify base is L = seq_len - 3; verify row
t is the forward pass at position L+t and its hidden produced the token at
L+t+1. Hence row t is ringed under label L+t+1, for t in 0..num_accepted.
Wired for K=3 (nd=2) only — the configuration this branch measured and the one
dgx2 recommends. K=2 has no d>=2 row to fix. K=4 is deliberately left for
after this is validated.
FALSIFIABLE PREDICTION, registered before measuring. Same rig (8-turn session,
context to ~11k, nd=2, gate disarmed, n>=700):
baseline p1 = 0.653, p2_uncond = 0.499, 18.66-18.71 tok/s
IF the accepted-row contamination is the mechanism, p2_uncond rises by
>= 0.03 AND p1 also rises (p1 is contaminated too, because the bad rows
persist into later proposes) — that second part is the discriminator, and
it is what would make this worth more than the ~23 tok/s ceiling implied by
fixing position 2 alone.
IF only p2_uncond moves and p1 is flat, the contamination is confined to
within-propose drift and the ceiling arithmetic stands.
IF neither moves, the mechanism is wrong -- most likely the ring label is
off by one, which would show up as acceptance DROPPING, not staying flat.
A wrong feed cannot corrupt output: verification rejects bad drafts. The
stake is acceptance only.
Unit tests cover the trim arithmetic (`mtp_rows_to_trim`), including that the
flag OFF is exactly the legacy behaviour, that full-reject is identical with
and without the flag, and that it never trims more rows than were drafted.
…ented Measured on dgx2 (W4A4 27B, nd=2, gate disarmed, 8-turn session to ~11k, 700 verify steps per arm). The control -- this same binary with the flag OFF -- is BIT-IDENTICAL to the pre-change baseline, so the arms are clean. arm delivered feeds p1 p2_uncond OFF (baseline) -- 0.653 0.499 ON, 0..num_accepted 458 fed / 231 missed 0.669 (+0.62 sd) 0.520 (+0.80 sd) ON, 0..=num_accepted 713 fed / 5 missed 0.633 (-0.78 sd) 0.476 (-0.86 sd) The third row was a dose-response test. The exclusive bound left exactly one label unwritten per step -- a verify step advances the sequence by 1 + num_accepted, so the next step's first label sits two past the last one written -- and that single hole collapses the catch-up ring's contiguous (start, count) window, losing a third of the feeds. Closing the off-by-one took delivery from 67% to 99%, and the acceptance delta REVERSED SIGN. Feeding more of these hiddens makes acceptance worse. That is positive evidence the hidden is misaligned with the pair key it is fed under, and it retro-actively explains the +0.021 at 67% delivery as noise (0.80 sd) rather than signal. My pre-registered threshold was +0.03 on p2_uncond; neither arm met it, and the amplification test went the wrong way. The flag stays default OFF and should not be enabled. Also worth recording for whoever picks this up: throughput was NOT comparable across these arms and no tok/s delta is quoted. The arms emit different text (per-turn token counts 152/139/146/154/158/137/151/160 vs 150/140/145/156/158/138/156/159), because which token arrives through which verify batch position changes with acceptance and the verify and serial paths differ in the low bits. Acceptance rates are the trajectory-robust measure here; tok/s is not. What survives: the defect is real and measured (unconditional per-position acceptance 0.660 -> 0.485 -> 0.407, with the loss concentrated at the d=1 -> d=2 hidden-state handoff rather than spread evenly), and the catch-up machinery demonstrably reaches the drafter -- a wrong feed moved acceptance by ~0.9 sd, so a correct one should move it too. What is wrong is the label/row mapping in verify_k3_step. Next session should re-derive which verify hidden row belongs to which pair key from first principles and confirm it with a dumped hidden checksum before trusting any acceptance number. The one-label coverage fix itself is NOT included here: it is mechanically correct but it amplifies a wrong mapping, so it should land together with the corrected mapping, not before it.
…apping
The pair-key -> hidden mapping used by ATLAS_MTP_REFEED_ACCEPTED was recorded as
"wrong somewhere" after a dose-response test reversed sign at n=700. Re-derived it
from the source and then verified it directly with dumped hidden fingerprints
(FNV-1a per BF16 row) on dgx2, W4A4 27B, nd=2, gate disarmed:
ring D2D landed (fp_src == fp_dst) 658 / 658
fed hidden == live ring content at that label 422 / 422
label == key + 1 and RoPE == key + 1 on every feed always
fp(ring[position]) == fp(mtp_hidden_save) at each propose 302 / 304
feed(key k) == mtp_hidden_save at the propose whose
position was k+1 93 / 93
The last check is the non-tautological one: it compares the hidden this feature
feeds for pair key k against the hidden the drafter own forward_one consumed as
target_hidden when it wrote pair key k -- two different code paths, bit-identical
on every checkable case. The mapping is correct. The earlier inference is withdrawn
in the flag doc comment (the two arms it compared differed by ~1.7 sd, on unpaired
samples over different generated text).
What the earlier session did find is real and is fixed here: the exclusive
0..num_accepted bound left exactly one ring label unwritten per verify step (a step
advances the sequence by 1 + num_accepted), which resets the ring meta contiguous
window every step and lost a third of the feeds (458 fed / 231 missed). The bound
is now 0..=num_accepted.
K=4 gets the same ring write. mtp_rows_to_trim extra trim is K-agnostic, so with
the flag on at nd=3 the accepted drafter rows were being dropped with nothing
rebuilding them.
Two diagnostics, both default-off and both gated behind the catch-up path so they
cost nothing when it is disabled:
ATLAS_MTP_REFEED_DEBUG=1 fingerprints every hidden entering and leaving the ring
ATLAS_MTP_REFEED_SHIFT=N deliberately perturbs the ring label, so the mapping can
be falsified behaviourally (no self-consistent checksum
can -- see the doc comment for why)
Default behaviour is unchanged and proven so: the new binary with the flags OFF
reproduces the pre-change baseline BIT-IDENTICALLY over 700 verify steps (all seven
K3 positional windows: 77/71/62, 63/58/31, 64/43/30, 64/44/27, 70/37/25, 61/46/26,
58/50/28). cargo check --workspace --all-targets clean.
The catch-up reader (impl_b3.rs) feeds drafter pair key k from ring label k+1,
because pair key k is (embed(t_{k+1}), hidden_k). So ring label n must hold
hidden_{n-1} -- the hidden that PREDICTED token n.
step_decode_only forwards last_token at the OLD seq_len and only then pushes that
INPUT token and increments (decode_a2.rs:485-489, decode_b.rs:477-480: tokens.push
+ seq_len += 1). So after the step, hidden row 0 is hidden_{seq_len-1}, whose label
is seq_len -- but the hook wrote seq_len - 1. Every serially-fed drafter row was
therefore paired with the NEXT position's hidden.
The same quantity is labelled base + t + 1 for verify row t at position base + t by
the K=3 re-feed, and THAT convention is verified end-to-end by dumped hidden
fingerprints: the hidden fed for pair key k is bit-identical to the hidden the
drafter's own forward_one consumed as target_hidden when it wrote pair key k, on
93/93 checkable cases. The serial hook was the side that disagreed.
Found by adversarial review (qwen3.8-max) of the mapping derivation, then confirmed
against decode_a2.rs / decode_b.rs.
NOT measured: ATLAS_MTP_CATCHUP only fires during serial-decode stretches, which
exist only when the throughput gate is ARMED, and every leg tonight ran with
ATLAS_MTP_GATE_FORCE=1. It is derived and default-off. Whoever evaluates
ATLAS_MTP_CATCHUP next must A/B it with the gate ARMED -- and note the previously
recorded "CATCHUP: +1.2% n.s." was measured with this off-by-one in place.
Powered A/B on dgx2 (W4A4 27B, nd=2, gate disarmed, 16 documents x 8 turns, ~10k verify steps per arm -- n is raised by NEW CONTENT because GATE_FORCE=1 makes repeated legs bit-identical): arm n p1 p2_uncond tokens/verify step OFF 10,400 0.6100 0.4182 1.882 ON 10,100 0.6262 (+2.4 sd) 0.4452 (+3.9 sd) 1.926 (+2.3%) Pre-registered criterion, written before the run: p2_uncond up by >= 0.015 at >= 3 sd. Measured +0.027 at 3.9 sd. Met. At n=700 -- the sample behind the earlier "refuted" verdict -- this same effect is ~1.0 sd, i.e. invisible. That verdict was a power problem, not a mapping problem, and the mapping itself is now verified by dumped hidden fingerprints (93/93 cross-step, see the preceding commit). Flag stays DEFAULT OFF pending the standard gates (C2 smoke, A 35B webserver_ok 10/10, B/D ST-995). Caveat recorded in the doc comment: the arms emit different text, so the binomial sd understates the true variance; content is matched but this is one measurement, not a replication. Context, also in the doc comment: this is a +2.3% accepted-tokens effect, and a much larger one was measured the same night on the same box -- the MTP drafter is prompt-blind on every WARM turn (142 KV rows at sequence position 10,098), worth +0.086 p1 / +0.101 p2_uncond / +10.2% accepted tokens per verify step. Build that first.
The flag's doc comment now carries its powered result AND what it is worth next
to the effects that dominate it, so nobody spends another session here first.
this flag (dgx2, n ~ 10k/arm) p2_uncond +0.027 (3.9 sd), +2.3% tokens/step
drafter INPUT hidden at d >= 2 p2_cond 0.5265 -> 0.7196, McNemar z = +18.4
(dgx1 ATLAS_MTP_ORACLE_P2) = +0.193, recovers 1.40x the p1-p2 gap, ~7x this flag
drafter context blindness on p1 +0.086 / p2_uncond +0.101, +10.2% tokens/step
WARM turns (dgx2) (142 drafter KV rows at sequence position 10,098)
The oracle probe refutes "exposure bias": the drafter is not mis-calibrated at
draft position 2, it is fed the wrong vector. And a drafter that can see the
prompt is a different drafter, so the blindness fix has to land before this
flag's +0.027 can be taken at face value.
Doc-only; no behaviour change.
The +0.086 p1 / +0.101 p2_uncond figure for prompt-prefilling the drafter on every turn was measured by removing --enable-prefix-caching, which also removes KV/SSM warm restore. A de-confounding pair settles the split: two more 10k-step arms with the drafter prefill DISABLED in both (coverage pinned at zero, prefix caching the only variable) give a warm-restore effect of only +0.0070 p1 (1.0 sd) / +0.0120 p2_uncond (1.8 sd), not significant. drafter coverage = +0.0794 p1 (92% of the total) / +0.0885 p2_uncond (88%) A third contrast agrees: with the prefill firing on turn 0 only -- 1 turn in 8 -- that coefficient predicts +0.0099 p1 and the measurement is +0.0157 (2.3 sd). Three contrasts, one coefficient. Doc-only; no behaviour change.
ATLAS_MTP_PREFILL_PROFILE=1 times the drafter prefill's phases separately (each phase synced, so it is diagnostic-only and must never be on in a timed leg). It exists because the whole-prompt drafter prefill turned out to cost 1136 ms over 11,947 rows on GB10 — the same order as the entire warm TTFT it would have to be paid out of — and the cost had to be attributed rather than guessed at. Measured 2026-07-21, 11,947 rows, h=5120, PREFILL_CHUNK=512: embed_loop=40.2 ms concat_loop=41.0 ms fc_gemm=874.0 ms kv_gemm=176.3 ms tail_sync=2.1 ms The obvious read was wrong. The two PER-ROW loops (one 10 KB D2D per row via `self_copy_embed_row`, one `bf16_concat` launch per row — ~23,900 tiny launches for 12k rows) look like the cost and are NOT: together they are 7.2% of it. 92% is the already-batched GPU work, and 77% is the single `fc` projection [rows, 2h] x [2h, h], running at ~1.4 TFLOP/s. So batching those loops buys ~7%, not the ~10x it appears to promise. Also measured and rejected: routing the three projections through `dense_gemm_tc` (the m16n8k16 tensor-core kernel the attention prefill uses) made `fc` 21% SLOWER at this shape — 874 -> 1060 ms — while helping k/v by 22%. Net +13%, so it is not the drop-in its doc comment suggests.
…OFF)
The MTP drafter is prompt-prefilled only on a COLD turn. On a WARM turn the
target reuses a cached prefix, so no prefill chunk starts at 0,
`mtp_prefill_capture_len` stays 0, the propose-site guard `captured >=
prompt_len` fails, and `prefill_drafter` is SKIPPED entirely. Proposer state
is per-request, so the drafter then starts EMPTY and gains one row per
decoded token — measured 142 drafter KV rows at sequence position 10,098,
while the target sees all 10,098. 987 of 1007 scored MLPerf-edge samples are
warm, so the drafter is nearly blind on ~98% of the benchmark.
Measured cost of that blindness (dgx2, ~10k verify steps/arm, de-confounded
from SSM warm restore, which is only +0.0070 p1 on its own): +0.079 p1 /
+0.089 p2_uncond, about +10% accepted tokens per verify step. Reproduced on
dgx1: cold turns decode at 45.1-49.3 ms/token, warm turns at 54.7-58.4.
The obvious fix -- rebuild the whole-prompt drafter prefill on every warm
turn -- CANNOT PAY FOR ITSELF, and this is why it is not what this commit
does. `prefill_drafter` costs 1136 ms over 11,947 rows (see the phase
attribution in the preceding commit) against a warm TTFT of 1134 ms. From
the scored e2e (TTFT median 1661.70 ms, TPOT median 52.12 ms, ~45 output
tokens/turn): +10% accepted tokens saves ~211 ms of decode per turn and the
rebuild costs 1140 ms, i.e. a net loss of ~927 ms/turn, and TTFT would go
1662 -> 2802 ms against vLLM's 2985 -- a 1.80x lead collapsed to 1.07x. The
verdict holds at every plausible turn length (-770 ms at 71 tokens, -680 ms
at 97). Trading TTFT for decode at this exchange rate is a loss.
Instead: a warm turn's prompt is a strict extension of the previous turn's
sequence -- that is exactly why the prefix cache hits -- so the drafter rows
the previous turn already built ARE the rows this turn needs. Keep them.
* `mtp_carry` holds ONE finished turn's drafter KV. One slot is enough
and makes ownership unambiguous (blocks are owned by the slot XOR by a
live sequence): every spec path is gated `active.len() == 1`.
* A turn adopts the carried rows iff its prompt begins with exactly the
tokens that produced them. `hidden_i` is a pure function of
`tokens[0..=i]`, so prefix equality IS the validity condition.
* It then appends only the missing pair keys, `[last_pair_key+1,
prompt_len-2]`, whose hiddens are exactly the span this warm turn's own
prefill computed. Cost scales with NEW tokens, not context length.
* A COLD turn releases any carried entry before its own prefill: the
drafter KV pool holds `max_seq_len / block_size + 1` blocks -- one
sequence's worth -- so a stale entry would starve `alloc_block`.
`try_mtp_prefill_capture` now also writes warm chunks at their ABSOLUTE
position. On its own that is the known trap (a warm turn computes only the
new user message, so coverage would go 0% -> ~0.3%); it is useful only
because the carried rows supply the rest.
The position interval this adds is reset per `alloc_sequence`, so a warm
append can only read hiddens THIS sequence computed. That is strictly
tighter than the legacy `captured >= prompt_len` guard, which has a latent
cross-sequence hole: `mtp_prefill_hidden` is model-level and nothing checks
that its rows came from the same token prefix. This path does check.
Conventions, since this is where the code kills people: pair key k is
(embed(t_{k+1}), hidden_k) at RoPE k+1; `mtp_prefill_hidden` row i is
hidden_i; the catch-up ring uses the OTHER convention (label n holds
hidden_{n-1}). Rows are compacted while RoPE stays in sequence space, so a
skipped key leaves no hole -- only a missing row, which is this row space's
steady state anyway. And a wrong or missing drafter row cannot corrupt
output: the target verifies every draft.
Default OFF (PCND). 10 unit tests cover the validity predicate, the append
plan (including both refusal cases and the clamp), and the interval merge's
refusal to claim a gap. Gate C2 NVFP4 tool-call + coherence smoke: 8/8.
The mechanism gate caught this before it cost an A/B. Requiring the WHOLE
carried entry to be a prefix of the next turn's prompt refused EVERY warm
turn on the 27B rig:
MTP_CARRY adopt: prompt_len=12139 ... -> prefix mismatch
MTP_CARRY store: rows=74 last_pair_key=Some(11686) seq_tokens=11689
i.e. the carry was rejected, the drafter stayed blind, and it limped to 74
rows again. The cause is benign: the chat template re-tokenizes the
assistant/user boundary, so the tail of a finished turn's sequence does not
reappear verbatim at the head of the next turn's prompt. Roughly 170 tokens
out of 12,000 differ — and the all-or-nothing rule threw away the 11,900
that were identical.
`usable_by` replaces `adoptable_by`: it takes the common prefix, keeps every
pair key the prompt actually agrees with, and truncates the rest. Rows are
append-only in increasing key order, so dropping d rows from the tail drops
the d highest keys; clamping `last_pair_key` to `common - 2` can only
OVERSTATE the surviving row's key when the tail had gaps, which starts the
append later and costs coverage, never correctness. Truncated rows are then
overwritten by the append or never read (the drafter reads `seq_len` rows).
Measured after the fix, same rig, same turn:
MTP_CARRY store: rows=12036 last_pair_key=Some(12114) seq_tokens=12117
MTP_CARRY adopt: prompt_len=12139 -> adopted rows=11864 appended=195
first_key=11943
11,864 carried + 195 appended = 12,059 drafter rows for a 12,139-token warm
prompt: 99.3% coverage, up from 0%, for an append of 195 rows (~19 ms at the
measured 0.0954 ms/row) instead of a 1136 ms rebuild.
Also logs the common-prefix length on refusal, so the next person who sees a
mismatch can tell a tokenizer boundary from a genuinely unrelated session.
`run_mtp_propose_inner` had grown a 54-line block deciding how the drafter
gets its prompt context, and impl_b3.rs was already 518 lines — over the
500-line budget before this branch touched it, and 547 after. Extracting the
block into `ensure_drafter_context` next to the code it calls leaves
impl_b3.rs at **497 lines: under budget, and smaller than it was at the
branch point.**
Pure extraction, no behaviour change. The helper destructures
`SequenceState` for disjoint field borrows rather than cloning the token
vector — a clone would have allocated ~48 KB on every propose, i.e. on every
decode step.
Also reverts formatting-only churn in four files this branch does not
otherwise touch (async_chkpt.rs, scheduler/{mod,verify_k3_step,
verify_k4_step}.rs), which `cargo fmt` had picked up from pre-existing drift
on the base commit. The diff is now exactly the change.
`cargo fmt` reformats pre-existing drift in async_chkpt.rs and three scheduler files. That drift is present on the base commit and is not this branch's to fix; carrying it made the diff look like it touched the verify paths, which it does not.
…default `ATLAS_MTP_DRAFTER_PREFILL` and `ATLAS_MTP_CARRY_DRAFTER` were opt-in and were never set, so the shipping recipe drafted CONTEXT-BLIND: 0 drafter-prefill events across a whole 1007-sample benchmark, with the drafter seeing 0-1.4% of a context the target sees all of on the ~98% of scored samples that are warm. Measured, dgx3 2026-07-21, MLPerf-edge compliance OVERALL PASS, 1007/1007, `dpcarry` vs the matched same-binary control `ctrl04f`: wall 5839.26 s -> 4984.53 s (-14.6%) TPOT p50 52.04 ms -> 39.60 ms (-23.9%) TTFT p50 1665.5 ms -> 1557.1 ms (BETTER, not a regression) IoU 0.6281 -> 0.6285 (noise floor 0.0028) The IoU cost that justified keeping drafter prefill off does not exist, and TTFT improved because the carry pays the context as a ~21.5 ms append instead of an 887 ms rebuild. The two halves are now ONE feature with ONE kill switch, because they are not independent: carry is INERT without prefill (its call site is nested inside `!mtp_prefill_hidden.is_null()`, and that buffer only exists when prefill is on), while prefill WITHOUT carry is a measured -927 ms/turn net loss. Resolving them separately makes an actively harmful configuration reachable by accident, so `drafter_context::resolve` enforces `carry => prefill` and a test asserts it over every value either switch can hold. (unset) -> prefill ON, carry ON ATLAS_NO_MTP_DRAFTER_CONTEXT=1 -> prefill OFF, carry OFF ATLAS_MTP_DRAFTER_CONTEXT_PREFILL_ONLY_UNSAFE=1 -> prefill ON, carry OFF Only `ATLAS_NO_*` disables, only at exactly "1". `ATLAS_*=0` disables nothing here on purpose: several unrelated flags in this tree are presence-checked and are ENABLED by `=0`, and that trap has cost whole sessions. The old opt-in names are obsolete and ignored, and their presence is WARNED at startup so a stale launch script cannot leave an operator believing a value took effect. Also closes an observability gap this campaign hit repeatedly: neither half was visible in the serve log, only in the launcher's `-e` flags. Startup now prints `MTP drafter context: prefill=.. carry=..`, and the prefill-only arm prints a warning naming its measured cost. The 336 MB prompt-hidden capture (max_seq_len x hidden, BF16 — 2.7 GB at 256k) is additionally gated on `MtpQuantization::supports_drafter_prefill()`, so an NVFP4/FP8 MTP head no longer allocates a buffer the batched prefill can never write now that the feature is on by default.
…g config `ATLAS_SSM_TAIL_PROTECT=1` and `ATLAS_SSM_TAIL_LEASE_TTL=128` appear in the frozen MLPerf-edge launch script and read as if they are doing work. They are not. The lease only ever shields an entry with `is_tail == true`; the sole production writer of that flag is `insert_tail_snapshot`, reached only from `finalize_midchunk_capture`, which returns early when `ssm_tail_midchunk_enabled()` is false — and the frozen config sets `ATLAS_SSM_TAIL_MIDCHUNK=0`. So the TTL governs an empty set. Behaviour is deliberately unchanged; this is a warning to the next reader, who would otherwise infer from the launch script that tail protection is active.
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
1 participant
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.
Makes the MTP drafter context on by default. Measured end-to-end on GB10 (DGX Spark), MLPerf-edge agentic, 1007 samples, temp 0 / seed 42.
Why
The shipping config was drafting blind:
ATLAS_MTP_DRAFTER_PREFILLwas never set, so the drafter never received the prompt — 0 drafter-prefill events across a full 1007-sample benchmark, and 142 drafter KV rows at sequence position 10,098. Every prior e2e number was a blind-drafter baseline.Result (five e2e runs, all MLPerf compliance PASS)
TPOT −24%, wall −14.6% vs control. TTFT improved rather than regressed. IoU held.
The
IoU costthat justified keeping drafter-prefill off does not exist: 0.6285 vs 0.6281 against a 0.0028 noise floor.Design
New SSOT
crates/spark-model/src/model/drafter_context.rs. One switch, not two, because they are not independent: carry is inert without prefill (its call site is nested inside!mtp_prefill_hidden.is_null()), and prefill without carry does the work and discards it.resolve()enforcescarry => prefill; disabling carry alone is unreachable.ATLAS_NO_MTP_DRAFTER_CONTEXT=1ATLAS_MTP_DRAFTER_CONTEXT_PREFILL_ONLY_UNSAFE=1Everything is
== Some("1"), soATLAS_*=0disables nothing (pinned by a test). Obsolete opt-in names are ignored and their presence is warned at startup.Startup now logs
MTP drafter context: prefill=ON carry=ON (default)— previously neither flag was verifiable from the serve log, only from the container env.The prompt-hidden capture is 336 MB at 32k / 2.7 GB at 256k, so the allocation is additionally gated on the MTP head supporting batched prefill.
Verification
cargo check --workspace --all-targetsclean; release build 158 kernels.spark-model213/213,spark-runtime145 pass / 2 ignored.centml/Qwen3.6-27B-NVFP4-W4A4-mlpinf: 3/3 tool calls, 3/3 coherent, 0 degenerate — run before any perf number.adopted rows=1785 appended=108, same ~10.5 ms append); the kill switch yields 0 prefill events and 0 adopts.Not covered
Gates A / B / D were not run. Clippy reports 4 errors that are pre-existing on the parent commit at lines this branch does not modify.