Log-Depth Chunked Linear-Attention Scan
Implementation & benchmark of arXiv:2607.17419 — Kernelized Linear Attention: Breaking the Capacity Wall with Symmetric Cones
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
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
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
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
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
- bench_report.json 3.1 KB View Raw
- report.md 1.3 KB View
- report_bench_2026-08-30T194641.md 4.7 KB View
- results.json 7.2 KB View Raw
- run_experiment.py 5.6 KB View Raw
- stage2_bench.py 7.8 KB View Raw