[jaredpalmer]

The small Kevs (0.8B, 4B, 9B) cannot score long documents on a GPU. The 2026-09-30 family longdoc reads (/bench/fam-*-longdoctest on kev-runs) all stopped with OutOfMemoryError. Kev-4B and both 9B versions failed at the first 32k record, and Kev-0.8B at the first 64k record even on an H200. Kev-27B scores the same records without trouble, and agents-ood / guardrails-ood development states (up to 51k / 39k tokens) are in the same range, so this blocks every long read of the small family.

Why fp32 fell to the L x L kernel

#149 already handles long rows. A record whose longest row is over ROW_PASS_TOKENS (16,384) runs its state once through kev.shared_prefix under SDPA's flash and memory-efficient kernels, and is labelled kernels: efficient. For Kev-27B (bf16) that means flash. The small family is evaluated in fp32, and in fp32 the state pass never reached an efficient kernel:

• kev.shared_prefix passes no mask for an unpadded state, so SDPA can run it causal.
• transformers' sdpa_attention_forward sees no mask and asks SDPA for grouped-query attention (enable_gqa=True, use_gqa_in_sdpa). Grouped-query attention means 16 query heads share 4 key/value heads.
• Of the allowed kernels, only flash and math accept enable_gqa. Flash has no fp32, so SDPA quietly uses math, which materialises a 16 x L x L fp32 score matrix: 69 GB at 32k and 275 GB at 64k, per attention layer.

The profile of one 4k pass shows this (runs/long-state-4b-h100/report.json["kernels_4k"]). Under #149's context the state pass ran cutlass_80_simt_sgemm + cunn_SoftMaxForward, the math kernel. Only the masked branch pass ran fmha_cutlassF_f32 (memory-efficient). One consequence: every row #149 labelled efficient on an fp32 backbone that did fit had its state attention computed by math.

The change

kev.predictors.long_row_kernels() is #149's context with one addition. While the context is active, an fp32 SDPA call repeats its keys and values per query head instead of asking for enable_gqa. transformers already does exactly this whenever a mask is given. With the repeat, the memory-efficient kernel takes the call: fp32 inputs and output, and memory linear in L. Mechanically, it swaps transformers.integrations.sdpa_attention.use_gqa_in_sdpa for the duration of the context and restores it afterwards. A bf16 / fp16 call keeps enable_gqa and flash, so Kev-27B's long rows take exactly the path they took before. Nothing else changes:

• Rows under the threshold never enter the context, so their logits stay bit for bit.
• The labels stay as Train and serve 64k-token states: caps raised to 65,536, unpadded shared-prefix states run SDPA's causal kernel, the benchmark scores long rows through the shared prefix #149 defined them: kernels: efficient on long rows and long_rows in report.json.
• CPU / MPS stay on eager attention.

The repeated K/V cost 1 GiB each at 64k for Kev-9B.

Measurements

scripts/long_state_memory.py (new) loads a checkpoint the way LocalPredictor does (fp32, TF32 off) and runs the longest longdoc-v1 development record per bucket under three settings: math (the global exact policy alone), legacy (#149's context as it was), and long (this PR). Peak is GPU memory over the loaded weights.

	state tokens	math	legacy (#149)	long (this PR)
Kev-4B, H100	15,231	33.6 GiB, 3.7 s	33.6 GiB, 3.7 s	4.3 GiB, 3.0 s
	30,468	OOM	OOM	8.6 GiB, 6.8 s
	60,926	OOM	OOM	17.2 GiB, 17.2 s
Kev-9B, H100	15,231	33.7 GiB, 5.8 s	33.7 GiB, 5.8 s	4.6 GiB, 5.1 s
	30,468	OOM	OOM	9.1 GiB, 10.9 s
	60,926	OOM	OOM	18.2 GiB, 25.5 s
Kev-4B, H200	15k / 30k / 61k	33.6 GiB / OOM / OOM	same	4.3 / 8.6 / 17.2 GiB, 2.9 / 6.7 / 16.6 s
Kev-9B, H200	15k / 30k / 61k	33.7 GiB / OOM / OOM	same	4.6 / 9.1 / 18.2 GiB, 5.1 / 10.9 / 25.1 s

At 16k, math's 33.6 GiB is two 16 x 15,231² fp32 matrices (scores and softmax). That matches the cause above.

Parity against math where math fits: the 8k and 16k buckets, cuad and synthetic records, every question. "rows" is the row form that records under the threshold actually get.

	questions	long vs math: max |dp| raw / at the shipped T, flips	for scale: legacy vs math	rows vs math
Kev-4B 8k, H100	40	2.2e-4 / 2.2e-4, 0	1.5e-4	7.4e-4
Kev-4B 16k, H100	40	3.0e-4 / 1.9e-4, 0	3.4e-4	–
Kev-9B 8k, H100	79	3.9e-4 / 3.3e-4, 0	2.8e-4	1.5e-3
Kev-9B 16k, H100	79	6.4e-4 / 5.4e-4, 0	3.3e-4	–

So the fp32 memory-efficient kernel moves answers by at most 6.4e-4 with no flips in 238 questions. That is the same order as what #149's branch pass already moved (legacy column), and smaller than the gap between the shared prefix and the row form, which comes from fla's TF32 DeltaNet kernels. The bar I set before measuring was "no flips, max |dp| within the TF32 DeltaNet noise (~3e-3, AGENTS.md)". The H200 runs repeat the H100 values exactly where they overlap (same sm80 kernel).

Reports: runs/long-state-{4b,9b}-{h100,h200}/report.json (numbers and kernel names only, no document text).

End to end

Kev-4B on longdoc-v1 development (1,200 records, 4,654 questions, states 4k-64k), H100, kev.benchmark through modal_app.py::benchmarks, runs/longmem-kev-4b-longdoc-dev-r2/report.json:

• Main (d1faa41, same image, same GPU type) stops at longdoc-v1/cuad/development/32k/0000 with OutOfMemoryError: Tried to allocate 55.00 GiB. That is one 16 x 30k² fp32 score matrix, the cause above, reproduced on the development partition.
• This branch reads all 1,200 records with 0 rejected. 1,866 rows in 480 records are labelled kernels: efficient (long_rows.threshold 16,384). Those are exactly the 32k and 64k buckets; under the Qwen3.5 tokenizer no 16k-bucket row crosses the threshold.

bucket	records	median / max s per record	accuracy
4k	240	2.6 / 8.6	0.922
8k	240	5.7 / 6.2	0.915
16k	240	13.4 / 14.6	0.896
32k (long)	240	6.3 / 17.1	0.868
64k (long)	240	16.8 / 27.0	0.834

Clean overall: acc 0.887, Brier 0.180, ECE 0.066 (n = 4,654). These are report-only probe numbers, not a release claim. The 16k bucket is slower than 32k because it is on the exact row form (the state once per question); the long buckets run the state once.

Short rows are bit for bit. The main read's 720 records before the OOM (every 4k / 8k / 16k record) have logits identical to this branch's in all 720 records (predictions.jsonl compared field for field).

What this changes for existing reads

A long row short enough to have fit in math before (just over 16,384 tokens; none in longdoc-v1) now takes the memory-efficient kernel for its state too, so it moves by up to the amounts above. Pair those rows only with reads from the same side of this PR; AGENTS.md says so next to the long-row rule. Kev-27B's reads and all short rows are unchanged.

Not in this PR

• kev.serve with KEV_DTYPE=fp32 on CUDA has the same grouped-query fallback on long states. It serves bf16 by default, and the fp32 server is not on any read path, so I left it alone.
• The kernels: efficient label still does not say which efficient kernel ran (flash for bf16, memory-efficient for fp32). report.json["environment"]["dtype"] already tells them apart.
• No bf16 fallback. fp32 fits with a wide margin (18 GiB over the weights at 61k for Kev-9B on an H100), so Train and serve 64k-token states: caps raised to 65,536, unpadded shared-prefix states run SDPA's causal kernel, the benchmark scores long rows through the shared prefix #149's bf16 path is not needed here.

Test plan

• tests/test_unit.py::test_long_row_kernels_repeat_fp32_keys_instead_of_grouped_attention: inside the context an unmasked fp32 call repeats K/V and a bf16 call keeps enable_gqa; outside it nothing changes; the patch is restored; the outputs match.
• tests/test_unit.py::test_local_predictor_scores_long_rows_through_the_shared_prefix, extended: the context is entered only for a long row on CUDA (not on CPU, not under the threshold), and the labels are unchanged.
• tests/test_model.py::test_long_rows_run_fp32_attention_in_linear_memory on Modal H100 (modal_app.py::gpu_tests): Kev-0.8B, a 22.5k-token state. The long pass stays under 1/4 of one fp32 L x L score matrix per head, the math pass exceeds it, the answers match within 2e-3, and the argmax is the same.
• CI suites locally: test_unit, test_conventions and the rest of the CI list pass (one flaky torchrun test_resume_is_bit_identical case passed on rerun; it also passes on main).
• Short rows: 720 of 720 longdoc-v1 development records under the threshold have bit-identical logits on main and on this branch (Kev-4B, H100).
• End to end: Kev-4B longdoc-v1 development, 1,200 of 1,200 records on H100, where main runs out of memory at the first 32k record.