Non-record: Triton KV-cache backend for autoregressive eval

[LucasErcolano]

Summary

Stacked on top of #1149.

This PR adds Triton-backed KV-cache evaluation paths for the existing non-record submission:

• int8_triton: fused grouped-int8 score/apply kernels
• qjl_triton: fused QJL sign-score kernel plus grouped value apply kernel
• backend selftests and a dedicated GPU benchmark script
• prewarm logic so short evals do not get dominated by first-use Triton JIT compile cost
• benchmark-side peak CUDA memory reporting to sanity-check allocator behavior

This is a systems PR only. Training/ML behavior stays intentionally close to #1149.

Local Validation

RTX 3090 synthetic benchmark (seq_len=1024, head_dim=64):

• int8_backend -> int8_triton: 2.68x faster for score+apply
• qjl -> qjl_triton: 2.60x faster for score+apply

Benchmark-side peak allocator stats from benchmark_kv_backends.py:

• int8_backend: peak_alloc=13.71 MiB, peak_over_baseline=3.00 MiB
• int8_triton: peak_alloc=10.77 MiB, peak_over_baseline=0.03 MiB
• qjl: peak_alloc=13.76 MiB, peak_over_baseline=3.00 MiB
• qjl_triton: peak_alloc=10.83 MiB, peak_over_baseline=0.03 MiB

End-to-end autoregressive smoke eval over 512 validation tokens:

• qjl: 63.53 tok/s
• qjl_triton: 70.83 tok/s

run_kv_backend_selftests(...) passes with the Triton path enabled.

Review Notes

This branch is stacked on #1149. The intended review focus is the commit sequence on top of that PR:

• b58d15c Add Triton KV-cache backends and benchmarks
• 2da8d23 Report CUDA peak memory in Triton benchmark

[MatoTeziTanka]

[RETRACTED 2026-04-11] — This IMPORT_FAIL was a false positive. Root cause: sibling module exists in same records/ folder; runner sys.path bug. Your code is not broken. See correction below: #1153 (comment)

---

Community Review — Non-record: Triton KV-cache backend for autoregressive eval

Compliance: NEEDS AUTHOR ACTION — train_gpt.py fails to import on CT2038 (Python 3.10 / torch 2.10.0+cpu)

What I found: The CPU smoke test on CT2038 (proteus-engine, 128 GB RAM, Triton 3.6.0, flash_attn stub, cutlass_evt_fusion stub) failed at the import step with:

ModuleNotFoundError: No module named 'triton_kv_ops'

A few of the common patterns I've seen for this class of error in the 2026-04-11 sweep:

• PEP 701 f-string nesting — e.g. log(f" {cat}: {", ".join(...)}") is valid Python 3.12+ but invalid Python 3.10 because the inner ", " re-enters the outer double-quote context. One-character fix: ', ' instead of ", ". See PR Record: SP8192 + Improved Parallel Residuals + Muon 0.97 + LR 0.03 + Legal TTT — val_bpb 1.07785 (3-seed mean) #1541 / Record: SP8192 + Triple Recurrence + Banking + Fused MLP + Muon 0.97 — val_bpb 1.0778 (3-seed mean) #1523 for reference.
• Missing flash_attn variants — e.g. from flash_attn_interface import flash_attn_varlen_func when the wrapper script only stubs flash_attn_func. Not a PR defect on H100s, but the eval image / CPU preflight path needs a guarded import.
• Local compiled extension — e.g. import cutlass_evt_fusion from a records/*/cutlass_evt_fusion/ subfolder that isn't on the import path at smoke time. Usually an import-order issue inside the script.
• Actual syntax error — typo, missing bracket, etc.

Recommendation: Could you run python3 -c "import py_compile; py_compile.compile('train_gpt.py')" on your records-folder train_gpt.py under Python 3.10 specifically? The eval image is Python 3.10 per Issue #17 / the README, so any parse error on 3.10 blocks the submission at import time before any of the scored-eval logic runs.

Once the parse/import issue is fixed, I'll re-run the compliance audit through the normal pipeline. No other flags identified yet because the audit halts at the import step.

---

Reviewed by @MatoTeziTanka — The Agora. CPU smoke test (CT2038 proteus-engine, 2026-04-11): IMPORT_FAIL — ModuleNotFoundError: No module named 'triton_kv_ops'. Classification via classify_prs.py AST-based classifier; full compliance audit deferred until the import issue is resolved. Auto-drafted from a template and spot-checked before posting.

[MatoTeziTanka]

Retraction — this IMPORT_FAIL was a bug in my smoke runner

Sorry @LucasErcolano, this one's on me. I re-audited the IMPORT_FAIL I posted above and it was a false positive — the fault is in how my CPU smoke runner set up sys.path, not in your code.

What happened:

The runner imported your records/track_non_record_16mb/2026-03-30_KVCache_QJL_Polar_Turbo_1x3090/train_gpt.py with only the script's folder implicitly on sys.path, so when your file did from triton_kv_ops import ... it couldn't resolve the sibling triton_kv_ops.py that lives in the same 2026-03-30_KVCache_QJL_Polar_Turbo_1x3090/ directory. The error I reported — ModuleNotFoundError: No module named 'triton_kv_ops' — looked like a missing file, but I re-checked the head SHA 2da8d23 and records/track_non_record_16mb/2026-03-30_KVCache_QJL_Polar_Turbo_1x3090/triton_kv_ops.py is right there, committed to the PR, next to train_gpt.py.

Verified at head 2da8d23:

records/track_non_record_16mb/2026-03-30_KVCache_QJL_Polar_Turbo_1x3090/triton_kv_ops.py   ← sibling module, exists
records/track_non_record_16mb/2026-03-30_KVCache_QJL_Polar_Turbo_1x3090/train_gpt.py   ← imports it

On the real eval image (Python 3.10, records/*/ as the working dir), this import resolves correctly because the records folder ends up on sys.path via the standard cwd-driven import or via the eval harness's per-record entry point.

Your PR is not broken by this error. I'm retracting the IMPORT_FAIL classification. I'll re-queue the full compliance audit (BPB check, n-gram / TTT / SLOT flags, etc.) on the current head and post findings separately.

Again — sorry for the noise. These community reviews only work if I actually read what I'm reviewing, and I didn't in this case.

[MatoTeziTanka]

Community Review — Non-record: Triton KV-cache backend for autoregressive eval

BPB: (not parsed — see PR title) | Compliance: LOOKS CLEAN — pure-neural submission, no TTT/SLOT/n-gram-cache

What I found in the code (head SHA 2da8d23bb66f, file records/track_non_record_16mb/2026-03-30_KVCache_QJL_Polar_Turbo_1x3090/train_gpt.py):

Static code review found no TTT adaptation function, no SLOT optimization loop, no n-gram-cache class, and no pre-quant val-token fine-tune. The eval path uses the standard sliding-window stride-64 pattern. The submission is a pure-neural architecture iteration on the standard SP1024/SP4096/SP8192 baseline.

CPU smoke test (CT2038 proteus-engine, 2026-04-11): import OK in 0.22s, dim=512, layers=9, vocab=1024, code=107777 B, SMOKE_TEST_PASS

Verdict: LOOKS CLEAN.

Recommendation to @cocohearts @valerio-oai @0hq @yuzhougu-oai @notapplica: MERGE pending the usual record-track checks (3-seed validation, under-16MB artifact cap, ≤600s train + ≤600s eval on 8×H100 SXM). No compliance flags from the classification pass — this looks like a clean pure-neural iteration on the standard baseline.

Auto-classification caveat: this review was drafted by the AST-based classifier. If there's a non-standard eval mechanism (logit postprocessing, hedge mixing, etc.) that I missed because it's factored into a helper file or a non-standard function name, please flag it and I'll re-run the audit manually.

---

Reviewed by @MatoTeziTanka — The Agora. CPU smoke test (CT2038 proteus-engine, 2026-04-11): import OK in 0.22s, dim=512, layers=9, vocab=1024, code=107777 B, SMOKE_TEST_PASS. Classification via deterministic AST-based classify_prs.py (pattern bank derived from ~65 manually-reviewed PRs earlier in the 2026-04-11 sweep). This review was auto-drafted from a template and spot-checked before posting — if the template misread your code, please call it out so I can iterate the classifier.