Efficient fixed-length KV-cache for recurrent Transformer sampling

[josephdviviano]

Summary

torchgfn's recurrent Transformer sampling path was both incorrect and slow. The estimator re-fed the whole committed prefix each step while the module threaded a growing carry, so the KV-cache and absolute positions diverged (cache grew 0,1,2,4,7,… instead of 0,1,2,3,4; positions likewise). It survived CI only because the sole integration test asserts isfinite(loss).

This PR makes the recurrent Transformer carry correct and adds an opt-in in-place KV-cache for fast autoregressive sampling.

Changes

• TransformerDiscreteSequenceModel.init_carry(max_len=…) preallocates [B, n_heads, max_len, head_dim] K/V buffers plus a write cursor; the block writes new K/V in place (narrow().copy_(), TorchScript-safe) instead of torch.cat. The two carry modes are selected by an Optional[int] cursor — None keeps the existing grow-by-cat carry, so this is fully backward compatible.
• RecurrentDiscretePolicyEstimator(use_kv_cache=…):
  • False (default): corrected, grad-bearing full-prefix baseline.
  • True: feed only the newest committed token against the persistent in-place cache under no_grad — O(L²) rollout cost vs the O(L³) re-forward. Ideal for inference/rollout.
  • cache_max_len (defaults to the module's max_position_embeddings) is validated eagerly at construction, so an undersized cache fails fast rather than mid-rollout.

Both paths build the canonical [BOS, w0, …, w_{t-1}] stream and are numerically identical to a single full-sequence forward.

Scope: fixed-length trajectories

The cached path assumes every active row shares a committed length (true for e.g. BitSequence). Variable-length carry slicing is out of scope — see the Gap-2 TODO in prob_calculations.py. Enforcement is layered so the hot path stays free of per-step host↔device syncs:

• always on — the module's batch-size carry check (a masked, unsliced carry mismatches the active batch);
• under debug=True — the stronger equal-length and cursor-sync checks.

Why no_grad?

In-place cache mutation invalidates the autograd graph (a buffer slice saved for an earlier step's backward is overwritten by a later step). Sampling doesn't need gradients, so the cached path runs under no_grad. The stacked follow-up (#526) adds an efficient teacher-forced loss recompute that turns this sampling win into an end-to-end training speedup.

Tests

testing/test_kv_cache.py — 7 tests (12 with parametrization):

• test_inplace_cache_matches_full_forward — dense-equivalence vs one full forward, across shape configs
• test_inplace_cache_is_preallocated_and_mutated_in_place — horizon-length buffers, same tensor object across steps, no realloc
• test_decode_position_is_load_bearing — perturbing the decode position by ±1 breaks equivalence
• test_estimator_cache_equivalence — use_kv_cache OFF == ON == teacher-forced
• test_estimator_fixed_length_guard — ragged active batches rejected (under debug)
• test_cache_max_len_validated_eagerly — undersized cache raises at construction
• test_sampling_runs_in_both_modes — end-to-end rollout in both modes

Benchmark

tutorials/examples/benchmark_kv_cache.py is an instructional example (it explains why KV-caching helps and how it is wired here) and gates every reported speedup behind a |Δ logits| equivalence check.

Measured on this branch — single CPU run (Apple Silicon), batch 32, dim 64 / 4 heads / 2 layers, word_size=1:

seq_len	sampling, no cache	sampling, cached	speedup	|Δ logits|
8	15.74 ms	10.57 ms	1.5×	3.0e-07
32	107.25 ms	43.99 ms	2.4×	4.8e-07
64	308.05 ms	99.54 ms	3.1×	4.8e-07

Absolute figures are CPU- and machine-dependent and vary run to run; the robust claim is the trend — the gap widens with sequence length, as the O(L³) → O(L²) argument predicts.

Unaffected

Stateless policies (hypergrid, etc.), DiscretePolicyEstimator, and the Sampler API are byte-for-byte unaffected. The full testing/ suite passes on this branch: 1939 passed, 735 skipped.

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 86.36364% with 12 lines in your changes missing coverage. Please review.
✅ Project coverage is 74.49%. Comparing base (654b828) to head (d9cfd34).

Files with missing lines	Patch %	Lines
src/gfn/estimators.py	87.50%	4 Missing and 2 partials ⚠️
src/gfn/utils/modules.py	85.00%	2 Missing and 4 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #525      +/-   ##
==========================================
+ Coverage   74.39%   74.49%   +0.09%     
==========================================
  Files          59       59              
  Lines       10636    10692      +56     
  Branches     1546     1558      +12     
==========================================
+ Hits         7913     7965      +52     
- Misses       2219     2222       +3     
- Partials      504      505       +1     

Files with missing lines	Coverage Δ
src/gfn/estimators.py	86.11% <87.50%> (+0.19%)	⬆️
src/gfn/utils/modules.py	73.93% <85.00%> (+0.77%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.