# Log-Depth Chunked Linear-Attention Scan

- ID: 1070
- Canonical URL: https://synthcore.org/idea/1070/log-depth-chunked-linear-attention-scan
- API JSON: https://synthcore.org/api/idea/1070.json
- API Markdown: https://synthcore.org/api/idea/1070.md
- Verification status: mech_ok_no_baseline
- Source: [arXiv:2607.17419](https://arxiv.org/abs/2607.17419)
- Category: architecture
- Solves: speedup, scalability, memory
- ML areas: attention, transformer, inference-speedup, batching
- Math tags: linear-algebra, computation, optimization
- Ratings: usefulness 8/10; difficulty 6/10; novelty 6/10

## 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.

## 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)\).

## Key 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.$$$

## 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

- Status: mech_ok_no_baseline
- Mechanism evidence: yes
- Mechanism confirmed: yes
- Practical verdict: no_effect
- Verdict: 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.

### Mechanism check

- Verdict: 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.
- Confidence: 8/10
- 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.

### Practical benchmark

- Paired seeds: 8
- Baseline mean: 0
- Idea mean: 0
- p-value: 0.78825
- Paired wins: 3/8
- Benchmark verdict: no significant win

## Artifacts

- [bench_report.json](https://synthcore.org/code/191/bench_report.json)
- [report.md](https://synthcore.org/code/191/report.md)
- [report_bench_2026-08-30T194641.md](https://synthcore.org/code/191/report_bench_2026-08-30T194641.md)
- [results.json](https://synthcore.org/code/191/results.json)
- [run_experiment.py](https://synthcore.org/code/191/run_experiment.py)
- [stage2_bench.py](https://synthcore.org/code/191/stage2_bench.py)
- [Download all files as ZIP](https://synthcore.org/download/191)

## Disclaimer

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