△ Mechanism confirmed, baseline not beaten 2026

Log-Depth Chunked Linear-Attention Scan

Implementation & benchmark of arXiv:2607.17419 — Kernelized Linear Attention: Breaking the Capacity Wall with Symmetric Cones

Usefulness8/10
Difficulty6/10
Novelty6/10

Source paper: Kernelized Linear Attention: Breaking the Capacity Wall with Symmetric Cones arXiv:2607.17419 ⓘ · analyzed Aug 30, 2026

AI-generated research hypothesis, automatically tested. Not peer-reviewed.

Idea description

Implement causal linear attention in chunks and combine chunk summaries with an associative scan instead of carrying the recurrent state through all chunks sequentially. This preserves the exact causal computation while reducing inter-chunk dependency depth from the number of chunks to its logarithm, enabling substantially more GPU parallelism for long-context training and prefill.

Formulas

$$\mathbf S_c^{\mathrm{loc}}=\sum_{i\in I_c}\psi(\mathbf k_i)\mathbf v_i^\top,\qquad \mathbf Z_c^{\mathrm{loc}}=\sum_{i\in I_c}\psi(\mathbf k_i).$$
$$(\mathbf S_1,\mathbf Z_1)\oplus(\mathbf S_2,\mathbf Z_2)=(\mathbf S_1+\mathbf S_2,\mathbf Z_1+\mathbf Z_2),\qquad (\mathbf S_{<c},\mathbf Z_{<c})=\bigoplus_{j<c}(\mathbf S_j^{\mathrm{loc}},\mathbf Z_j^{\mathrm{loc}}).$$
$$t_{\mathrm{total}}=\mathcal O(Cn_\psi d_v)+\mathcal O(\log_2N_C\cdot d_v)+\mathcal O\!\left(Cn_\psi d_v+C^2(n_\psi+d_v)\right),\qquad N_C=\left\lceil\frac{T}{C}\right\rceil.$$$

Mathematical statement

For a feature map \(\psi\), linear attention has numerator state \(\mathbf S_t=\sum_{i\le t}\psi(\mathbf k_i)\mathbf v_i^\top\) and denominator state \(\mathbf Z_t=\sum_{i\le t}\psi(\mathbf k_i)\). For chunk \(c\), with token index set \(I_c\), define \(\mathbf S_c^{\mathrm{loc}}=\sum_{i\in I_c}\psi(\mathbf k_i)\mathbf v_i^\top\) and \(\mathbf Z_c^{\mathrm{loc}}=\sum_{i\in I_c}\psi(\mathbf k_i)\). Chunk summaries form an associative monoid under \((\mathbf S_1,\mathbf Z_1)\oplus(\mathbf S_2,\mathbf Z_2)=(\mathbf S_1+\mathbf S_2,\mathbf Z_1+\mathbf Z_2)\). Therefore the exclusive prefix state before chunk \(c\), \((\mathbf S_{<c},\mathbf Z_{<c})=\bigoplus_{j<c}(\mathbf S_j^{\mathrm{loc}},\mathbf Z_j^{\mathrm{loc}})\), can be computed in \(\lceil\log_2N_C\rceil\) parallel scan rounds, where \(N_C=\lceil T/C\rceil\), \(T\) is sequence length, and \(C\) is chunk size. Within each chunk, the output combines the prefix state with the causal local state. The resulting cost model is \(\mathcal O(Cn_\psi d_v)+\mathcal O(\log_2N_C\cdot d_v)+\mathcal O(Cn_\psi d_v+C^2(n_\psi+d_v))\), replacing the sequential middle term \(\mathcal O(N_Cd_v)\).

Implementation notes

Integrate this into the training-time and prefill implementation of causal linear attention, replacing the token-by-token recurrent loop. Choose a chunk size C, initially 32 or 64, and reshape each sequence into N_C = ceil(T/C) chunks. Kernel 1 computes the causal local outputs within every chunk and emits S_loc[c] and Z_loc[c]; padded positions must contribute zero summaries. Kernel 2 performs an exclusive Blelloch or Hillis-Steele scan over pairs (S_loc, Z_loc). Every scan combination is exactly S = S_left + S_right and Z = Z_left + Z_right, so the scan is mathematically identical to sequential accumulation. Kernel 3 broadcasts each exclusive prefix pair to its chunk, adds it to each token's local prefix state, and evaluates z_t = (phi_q[t] @ S_total[t]) / (phi_q[t] @ Z_total[t] + eps). Use FP32 accumulation for summaries and fused Triton kernels to reduce HBM traffic. The paper's mathematics determines the associative operation and logarithmic scan depth; empirically tune C for occupancy, register pressure, and shared memory. Validate against a slow sequential reference on random tensors before benchmarking. Test sequence lengths 2K, 8K, 32K, and 131K on a small Transformer, comparing sequential linear attention and FlashAttention-2 for forward and backward latency, peak memory, and numerical output error. Success is near-identical outputs with several-fold faster long-context prefill or training and linear memory growth.

Verification

Mechanism confirmed, baseline not beaten

Stage 1 · Toy mechanism gate: Passed ✓

Stage 2 · Mechanism transferred to benchmark: Confirmed ✓

Stage 2 · Practical benchmark result: No measurable effect

Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works

Stage 1 — Mechanism check agent confidence 8/10

Built an exact chunked causal linear-attention MVP with FP32 chunk summaries and a Hillis-Steele associative scan. Across all tested chunk sizes and sequence lengths, maximum error versus the sequential reference was 2.63e-7, and observed scan rounds exactly matched ceil(log2(number of chunks)) in every sweep. The MVP was slower at T=2048 but faster at T=8192 and T=16384, providing a promising signal but not definitive production-kernel evidence.

Agent confidence
8/10
Baseline
Sequential reference CUDA forward latency: 1.57 ms at T=2048, 5.62 ms at T=8192, and 11.20 ms at T=16384.
Idea
Chunked scan CUDA forward latency: 3.34 ms at T=2048, 4.08 ms at T=8192, and 3.47 ms at T=16384; maximum absolute reference error was 2.63e-7; scan-round predictions matched exactly.

Limitations: This is a readable PyTorch MVP rather than a fused Triton implementation. Backward latency, FlashAttention-2, training convergence, longer sequences, multi-sequence batching, peak memory, and kernel occupancy were not tested. Timing includes Python and PyTorch launch overhead.

How to run: python3 run_experiment.py

Stage 2 — Benchmark vs. tuned baseline (arXiv:2607.17419)

No significant difference vs. the tuned baseline (+1.0%, p=0.7883, wins 3 / 8 paired seeds).

Benchmark
Sequence modeling
Model
transformer_tiny
Paired seeds
8
Baseline mean
0.2409
Idea mean
0.2434
Effect (Δ)
0.0025 (+1.0%; negative = idea better)
Wins
3 / 8 paired seeds
p-value
0.7883 (permutation test, 20 000 shuffles)
Smallest detectable effect
±9.0%
Mechanism
Confirmed ✓
Practical effect
No effect
Baseline tuning
swept over 3 configs
Limitations:

Only the registered built-in sequence forecasting track was tested, with a tiny 32-token window, 800 training examples, 300 test examples, 8 epochs, and a readable PyTorch implementation. Fused Triton kernels, long-context throughput, backward latency, peak memory, FlashAttention-2, and production-scale training were not tested.

How to run: python3 stage2_bench.py

Verdict computed by deterministic test code from paired-seed statistics — not by the language model.

Stage 2 — Benchmark latest Failed ✗

Agent confidence: 9/10

The registered sequence track was structurally matched to causal attention. The chunked linear-attention system confirmed the predicted logarithmic scan depth and near-exact trained-model outputs, but its task MSE was slightly worse than the tuned softmax-attention baseline and the paired permutation test was non-significant.

Baseline
8-seed MSE mean 0.2409101166; best lr=0.003.
Idea
8-seed MSE mean 0.2434249818; best lr=0.003, chunk=8; paired delta +0.0025148652, p=0.78825.

Benchmark result

No significant difference vs. the tuned baseline (+1.0%, p=0.7883, wins 3 / 8 paired seeds).

Benchmark
Sequence modeling
Model
transformer_tiny
Paired seeds
8
Baseline mean
0.2409
Idea mean
0.2434
Effect (Δ)
0.0025 (+1.0%; negative = idea better)
Wins
3 / 8 paired seeds
p-value
0.7883 (permutation test, 20 000 shuffles)
Smallest detectable effect
±9.0%
Mechanism
Confirmed ✓
Practical effect
No effect
Baseline tuning
swept over 3 configs
Limitations:

Only the registered built-in sequence forecasting track was tested, with a tiny 32-token window, 800 training examples, 300 test examples, 8 epochs, and a readable PyTorch implementation. Fused Triton kernels, long-context throughput, backward latency, peak memory, FlashAttention-2, and production-scale training were not tested.

How to run: python3 stage2_bench.py

Artifacts

Implementation overview ⬇ Download all as ZIP 6 files · code, reports and structured results