Submission/hybrid RWKV token shift

[dillon-blake]

Non-record submission exploring hybrid transformer architectures that replace most attention layers with a lightweight RWKV-inspired token-shift mixing mechanism. The core idea is that most layers in a transformer only need local context, so full quadratic attention is wasteful for them. Instead, 8 of 11 layers use a simple token-shift operation that blends adjacent tokens via learned per-dimension interpolation weights, while only 3 layers retain quadratic attention with short (128-token) windows (except the final attention layer which keeps full context).
The architecture achieves a 3-seed mean val_bpb of 1.2252 with 17.0M parameters, int6 quantized and zlib compressed to ~15.86 MB. While this does not beat the current SOTA, I believe the token-shift approach is promising for its efficiency — particularly for inference, where the reduced attention overhead could significantly speed up decoding.
Beyond the hybrid architecture, the submission stacks several techniques from the leaderboard: SmearGate, bigram hash embeddings, value embeddings, XSA (cross-head suppression), partial RoPE (16/64 dims), LeakyReLU squared activation, Muon optimizer, EMA with late QAT, and logit softcapping. Full details and ablation notes are in the README.

[MatoTeziTanka]

Community Review — Submission/hybrid RWKV token shift

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 'kernels'

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 'kernels'. 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.