Fix #148: walk annotation chain when final buffer is 2+ chunks ahead

[mmjerge]

Summary

Fixes #148 (chunked backward fails when a generated region spans 3+ chunks).

_unpack_from_annotation assumed the final buffer for a (layer, kind) node is at most one chunk ahead of the annotation being unpacked (final_idx in [chunk_idx, chunk_idx + 1]). With long generated regions spanning 3+ backward chunks, autograd's backward for the intermediate chunks can be served without unpacking their scatter-*/cat-* annotations, so the gap grows beyond one chunk and backward dies with e.g. ValueError: Annotation scatter-value (20,1): final chunk_idx = 3, must be in [1, 2].

Changes

1. autograd_hooks.py: walk the annotation chain across multi-chunk gaps. All scatter/cat annotations are guaranteed to be present in _packed_arg_for_id after _flush_remaining_pack_arguments (they are inserted under fresh IDs even when unmatched, precisely to keep the reconstruction chain complete). _unpack_from_annotation now applies the intermediate annotations early, in descending chunk order, and parks each reconstructed intermediate state (as a copy, since the buffer is mutated in place) in _id_to_unpacked, so autograd can still fetch those IDs later. _find_prior_annotation is generalized to _find_chain_annotation(annotation, chunk_idx); the legacy one-step ext-before-scatter behavior is preserved unchanged.

2. rl/grpo/loop.py: pass may_match_twice=may_match_twice_flex_attention_sdpa. With the new training replay cache, ext-* annotations legitimately match two pack arguments. The finetune path (may_match_twice_factory) and test_gradient.py both pass this predicate; the GRPO loop did not, so the second save of an ext buffer was left as an unmatched raw pack argument. Autograd's backward for that chunk can then be served from the raw copy without advancing the annotation chain — the likely trigger for the adjacency violation. This is the root-cause fix; change (1) makes the unpack robust even if the chain stalls for another reason.

Why it only showed with long completions

With at most 2 buffer states per cell the gap can never exceed one chunk. QA-style runs (32-token completions) never hit it; long-generation runs (e.g. 2600-token completions, chunk_size=1024) hit it stochastically, dependent on per-row eos raggedness — consistent with the reports in #148.

Testing

• New regression test test_unpack_walks_multi_chunk_annotation_chain (parametrized over chunk gaps 1, 2, 4): builds a ground-truth chain of buffer states linked by scatter-value annotations, sets the final buffer several chunks ahead, unpacks the earliest annotation, and verifies (a) correct reconstruction, (b) parked intermediate states are served when their IDs are unpacked later. The gap ≥ 2 cases fail with the exact Chunked backward fails when a generated region spans 3+ chunks (autograd_hooks annotation adjacency) #148 error before the fix and pass after.
• test/kvcache/test_autograd_hooks.py: 5 passed
• test/kvcache/test_gradient.py: 10 passed
• test/kvcache/test_gradient_main.py + test/rl: pass (CUDA-only cases skipped; verified on CPU, macOS)

Not verified: an end-to-end GPU reproduction of the original failing GRPO config (long-generation LongProc/longmath runs). We will requeue the previously crashed 32k jobs with this fix and report back on the issue.

[mmjerge]

@mseeger Fair questions — here is the full picture, including a concrete example. Short version: you are right that chunking is unaffected by generation, and right that there can be no holes in the reconstruction. Neither is violated. What breaks is the order in which autograd's backward requests the saved buffer states, and (in the GRPO path) which saved tensors get matched to annotations at all.

What is different about our runs (it is not the data)

Two things, both specific to the RL path:

1. The GRPO loop constructed the hooks without may_match_twice (config bug, fixed in this PR). finetune/longcontext_full.py passes may_match_twice=may_match_twice_factory(...), and with the new replay cache the ext-* annotations legitimately match two pack arguments. keys_values/rl/grpo/loop.py passed nothing, so the predicate defaulted to "never". Consequence: the second save of an ext buffer stays unmatched and is flushed as a raw tensor. Autograd can then run that chunk's SDPA backward entirely from the raw copy — the chunk's annotation is never unpacked at that point, so the final buffer does not advance down the chain. When an earlier chunk's annotation is later unpacked, the final buffer is now 2+ chunks ahead of it, and the old adjacency check final_idx in [chunk_idx, chunk_idx+1] raises. This is why you never see it: your path configures the predicate correctly.

2. The backward traversal order deviates more in the GRPO graph. The class docstring already documents that "the backward traversal ordering is not exactly the reverse of the forward traversal ordering", and the existing code special-cases exactly one deviation (an ext-* unpacked before its same-chunk scatter-*, handled via prior_annotation). With the GRPO head (per-token log-prob gather over the completion region + eos mask + ratio/clip, instead of the LM cross-entropy head), we observe larger deviations: after fixing (1), a 32k run still requested scatter-value (27,15) after the buffer had already been walked down to chunk 14 by an earlier request. Same phenomenon your prior_annotation mechanism handles, just at distance > 1 and in both directions. The PR generalizes that mechanism rather than adding a new one.

How generation "plays a role": it does not change chunking — chunk boundaries stay a function of cache length and chunk size only. But a gap of size g between the final buffer and a requested annotation needs g+1 saved buffer states between checkpoint columns whose unpack order deviates. With 32-token completions the loss-bearing region touches at most 2 chunks, so the reordering the GRPO graph induces has nothing to reorder across; with 2600-token completions it spans 3+ chunks. That is why hundreds of short-completion QA runs never hit it and long-generation runs hit it stochastically (the exact backward order also depends on which of the doubly-saved tensors matched, which is value-dependent).

No holes — all deltas are applied, in order

The fix does not skip any deltas. Every scatter/cat annotation is already guaranteed to be present in _packed_arg_for_id when backward starts: matched ones under their pack-argument IDs, unmatched ones inserted by _flush_remaining_pack_arguments under fresh IDs — that is pre-existing behavior, added precisely so the chain is complete ("Even if a cat or scatter annotation is not matched to a pack argument, it may be required...").

When annotation (l, c) is requested with final buffer at f > c+1, the walk applies the deltas for chunks f-1, f-2, ..., c+1, then c — strict descending order, every link applied exactly once. If a link were genuinely missing, the new code raises naming the missing chunk (Cannot reconstruct, missing 'scatter'/'cat' annotation for chunk k); we have never seen that fire.

The one genuinely new piece of state: an intermediate annotation applied early has its reconstructed buffer parked (on CPU) under that annotation's ID. If autograd later asks for that ID — the "mirror" case above — it is served the exact state instead of hitting the walked-past buffer. The values are bit-exact: a scatter undo writes back the stored delta (the overwritten old values), so the parked state equals what a strictly-ordered unpack would have produced.

Small example

This is literally what the two new unit tests construct (test/kvcache/test_autograd_hooks.py, runnable on CPU in seconds):

Take one layer, values only, cache length 32, three states linked by two scatters:

V1 --scatter(idx2, new2)--> V2 --scatter(idx3, new3)--> V3      final = V3 (chunk 3)
annotation (l,1): shape, idx2, delta2 = V1.gather(idx2)         # undoes chunk 2's scatter
annotation (l,2): shape, idx3, delta3 = V2.gather(idx3)         # undoes chunk 3's scatter

Backward asks for chunk 1's state first (its request for chunk 2's ID comes later, or is satisfied by a raw unmatched copy and never comes):

• Old code: final_idx = 3, must be in [1, 2] → ValueError. This is the exact failure from the issue logs (scatter-value (20,1): final chunk_idx = 3).
• New code: apply (l,2): V3.scatter_(idx3, delta3) → V2, park a copy of V2 under (l,2)'s ID; apply (l,1): → V1; return V1. Later request for (l,2)'s ID → served the parked V2 (second test: this request arrives after the buffer is already at V1, which the old code — and the first version of this PR — could not serve).

test_unpack_walks_multi_chunk_annotation_chain checks reconstruction for gaps of 1, 2, and 4 (gap 1 = unchanged legacy behavior); test_unpack_out_of_order_request_after_walking_past checks the parked-state serving. Both fail with the exact issue-#148 errors on the pre-PR code.

Validation status (honest)

• The failures and both mechanisms are taken from real job logs (Qwen2.5-7B, L40S, h2o-torch-quantized8, 16k/8k cache, 2600-token and 32k-tier generations).
• An 8k long-generation GRPO run has since completed all 200 steps on the fixed code (previously crashed stochastically).
• 32k runs are in flight now; the first post-fix attempt OOMed because the parking initially cloned on-device — parked states now go to CPU.
• A KV_DEBUG_ANNOTATIONS=1 traced 32k job is queued; I will attach the create/match/unpack sequence for the failing cell here when it lands, which will show the exact ordering deviation rather than my reconstruction of it.

[mmjerge]

Traced run landed (KV_DEBUG_ANNOTATIONS=1, 32k-tier GRPO, Qwen2.5-7B, cache 16384, chunk 1024, layers_per_cell 1). Full log: s3://keys-values-rl-results/runs/lm32k_trace/job.log. It replaces my earlier reconstruction with observed behavior, and the observed behavior is simpler than what I described before.

What the trace shows

Within each cell (4 chunks here), the backward's unpack requests arrive in strictly descending chunk order — no exotic reordering. Representative excerpt (layer 27, cell = chunks 21..24):

_unpack_from_annotation: ext-key   (27,24)  final_idx=24
_unpack_from_annotation: ext-key   (27,23)  final_idx=24   --> Doing scatter-key (27,23) first
_unpack_from_annotation: ext-value (27,23)  final_idx=24   --> Doing scatter-value (27,23) first
_unpack_from_annotation: ext-key   (27,22)  final_idx=23   --> Doing scatter-key (27,22) first
...

Two facts matter:

1. The prior_annotation path fires for every chunk. In this graph, the ext-* request for chunk c always arrives before the scatter-* ID for chunk c — so every scatter is applied early via your existing distance-1 mechanism. That happens ~40 times in this fragment alone. This is your design working as intended.

2. The scatter IDs are then requested late — or the "on time" case never occurs at all. There is not a single --> Skip (already done) line in 4k lines of trace. The skip branch assumes that when a scatter's own ID is unpacked, the final buffer still sits exactly at that chunk. In this graph the request arrives after later chunks' exts have already walked the buffer further down. That is precisely the failure from the real run: scatter-value (27,15): final chunk_idx = 14, must be >= 15. The buffer had legitimately moved to 14; the annotation asking for state 15 arrived afterwards.

So the invariant that breaks is narrower than "backward order deviates wildly": it is only the lifetime assumption behind the skip branch. The fix parks the state at the moment the scatter is applied early (one CPU copy per early-applied matched annotation, freed when served or at cell clear), so the late request is served bit-exactly.

The original gap-2 errors (scatter-value (20,1): final chunk_idx = 3) are the second variant: without may_match_twice (missing in the GRPO loop, fixed in this PR), the second save of an ext buffer is unmatched and served raw, so the ext→prior chain skips a link at that chunk, and the next chunk's ext arrives with a gap of 2. The chain walk covers that case.

Why the RL path and not finetune ("why would a loss head do this")

I no longer believe the loss head is the mechanism. Inside a cell the graph structure is head-independent (the head enters only through the injected head_gradients_top values, and values don't affect engine ordering). The concrete differences in the RL path are:

• Hooks config: may_match_twice was missing (gap-2 variant above). Your paths always set it.
• SDPA implementation: the trace has 624 padded-query annotation events — these runs go through the zero-padded-query SDPA, not flex_attention (FlashInfer disabled on these boxes, and the qpadded fallback is what the RL example ends up on). Different SDPA = different consumers of the cache buffers = different backward request order for the scatter IDs. Your long-sequence runs are on the flex path.

That second point is a hypothesis I can test directly: run a finetune job under the same KV_DEBUG_ANNOTATIONS=1 flag on the flex path and diff the request orderings. Happy to do that if it would settle it for you.

On the efficiency invariants you raised

The change does not touch the matching side: annotations are still removed on first match (or marked once for may_match_twice), the flush is unchanged except for recording which IDs are orphans (a set of ints), and no pack decision changes — so no new unmatched pack args and no extra large nodes written into the graph. The walk adds work only when a request arrives out of order, and each delta is still applied exactly once overall (earlier, not additionally). Parked states go to CPU, not device.

Remaining issue (separate from this PR)

The traced job — and the 32k reruns — now die in an eager-SDPA temp allocation (sdpa_attention_weights softmax, 2.31 GiB) with 40.9 GiB already allocated, i.e. a plain capacity problem in the qpadded path at cache 16384. The GRPO examples don't set attention_forward_temp_size_gb / backward temp limits the way finetune_long_lora does; we're fixing our recipe. Not an annotations problem.

[mmjerge]

Ran the two discriminating experiments discussed above. Traces in s3://keys-values-rl-results/runs/x148_*.

Experiment B: flex_attention path (your own test_gradient_row_of_cells, CUDA, h2o parametrizations, debug_print_annotations=True)

32 tests passed; trace has 1072 unpack requests. On the flex path your design assumptions hold exactly:

• 352 prior_annotation applications — all distance 1 (the ext-before-scatter case you built the mechanism for; it is common, not rare).
• 336 direct scatter unpacks, all on time (final_idx == chunk_idx + 1 in every single one).
• 576 second matches (Matched again) — may_match_twice is doing real work on this path.
• Zero adjacency/lifetime violations. This is consistent with these tests having always been green: the invariant genuinely holds on the flex path.

The qpadded RL traces, for contrast

• 32k longmath GRPO (lm32k_trace): every chunk's scatter is applied via the ext-prior path; there is not one "on time" scatter unpack (Skip (already done): 0 hits in 4k lines). The scatter IDs' own requests arrive after the buffer has moved on — that is the real-run failure scatter-value (27,15): final chunk_idx = 14 and the reason the parking half of this PR exists.
• lp5 longproc GRPO (2600-token completions, cache 8192, x148_mm_off_trace): different again — scatters are requested directly and on time (final = chunk+1 throughout), and the ext IDs are never requested at all. So even within the qpadded path, the request ordering is configuration-dependent. This is, to me, the strongest argument for the generalized walk: the invariant is a property of a specific graph, not of the machinery, and three configurations already exhibit three different orderings.

Experiment A: may_match_twice removed (attribution test) — inconclusive so far, being honest

A 2-step run on the lp5 config with the kwarg deliberately removed completed cleanly (trace shows 0 second matches, 0 unmatched, no gaps). So I have not yet demonstrated that the missing kwarg produces the original gap-2 errors — the original failures were stochastic (step ~10-60), and 2 steps is evidently not enough to hit the value-dependent conditions. A 60-step traced run without the kwarg is in flight; I will post the result either way. Two things remain true regardless of how it lands: the kwarg was missing from the GRPO loop while your paths and test_gradient set it (576 second matches in Experiment B show it matters on flex), and the chain walk + parking fix is required for the qpadded late-request behavior independent of the kwarg question.

One process suggestion from Experiment B: since may_match_twice=None silently degrades, would you take a follow-up that either makes the predicate required or logs a loud warning when a scatter/cat annotation goes unmatched while a same-shape pack argument went raw? That would have surfaced the GRPO config bug immediately.

[mmjerge]

Update: we ran the discriminating experiment for the SDPA hypothesis, and the hypothesis is wrong. Correcting the record here.

We added flex support to the RL example (two omissions had kept it off: flexatt_args was never passed, and it must go into cache_kwargs because the gradient cells use kv_cache.mha, not gpt_model.mha). We then ran the identical 32k config, same seed, traced, on both paths:

                              qpadded          flex
unpack requests                  864             96 (partial; both OOMed later)
scatters applied early           432             72
  via ext-prior mechanism
direct scatter unpacks,            0              0
  on time
on-time skips                      0              0
second matches                   576             96
padded-query events              624              0  <- proves flex engaged

Same ordering on both. On our RL workload, flex also applies every scatter early through the prior_annotation path, with zero on-time scatter requests. So the SDPA variant is not what changes the backward ordering, and to answer your question directly: removing the qpadded variant would not remove the problem (and qpadded is also the only CPU path, so it cannot be removed anyway).

The corrected picture across all five traced setups:

• your test_gradient_row_of_cells (flex, tiny model): mixed mechanisms, every direct scatter request on time. Invariant holds.
• our longproc GRPO (qpadded, 7B): all direct, all on time. Invariant holds. 60 steps clean, including with may_match_twice removed — so we also cannot yet attribute the original gap-2 errors to the missing kwarg; that attribution stays unproven.
• our 32k longmath GRPO (qpadded AND flex, 7B): all scatters applied early, none on time, and the real run requested one after the buffer moved past it. Invariant broken.

Same machinery everywhere; the ordering is a property of the specific workload's graph, and we now have direct evidence it is not selected by the SDPA implementation. Which concrete property of the 32k graph causes autograd to consume the ext saves before the scatter saves, we have not isolated — but the measured claim the PR rests on is unchanged and now better supported: multiple orderings exist in practice, the current code assumes one, and the walk + parking make the reconstruction ordering-independent at zero cost to matching/packing (accounting in the earlier comment).

Separately: both 32k runs OOM later in the backward (identical 2.31 GiB allocation on both paths, with the forward attention temp limit set). That is a capacity/tuning issue in our 32k recipe, independent of the annotation machinery — the 8k and longproc configs train to completion on the fixed code.