Cool idea from Nous Research. What if you could speed up long-context pretraining with a subquadratic wrapper that you remove before deployment? That is the idea behind Lighthouse Attention. The method wraps ordinary SDPA with a hierarchical, gradient-free selection layer that compresses and decompresses queries, keys, and values symmetrically, preserving left-to-right causality. Crucially, it can be removed near the end of training in a short recovery phase, so the deployed model still runs vanilla attention with no architectural cost at inference. Preliminary LLM experiments report faster total training time and lower final loss than full-attention baselines. Why does it matter? Most efficient-attention work either changes the deployment-time architecture or pays a quality tax to do so. A training-only wrapper that survives a clean recovery phase sidesteps both. If it scales, this becomes an important training-time speedup for long-context pretraining. Paper: https://arxiv.org/abs/2605.06554 Learn to build effective AI agents in our academy: https://academy.dair.ai

Long Context Pre-Training with Lighthouse Attention

Source: https://arxiv.org/html/2605.06554 Bowen Peng* Nous Research [email protected] &Subho Ghosh* Nous Research [email protected] &Jeffrey Quesnelle Nous Research [email protected]

(February 2026)

Abstract

Training causal transformers at extreme sequence lengths is bottlenecked by the quadratic time and memory of scaled dot-product attention (SDPA). In this work, we propose Lighthouse Attention, a training-only symmetrical selection-based hierarchical attention algorithm that wraps around ordinary SDPA and can be easily removed towards the end of the training. Our hierarchical selection is also gradient-free, which exempts us from dealing with a complicated and potentially inefficient backward pass kernel. Our contribution is three-fold: (i) A subquadratic hierarchical pre- and post-processing step that does adaptive compression and decompression of the sequence. (ii) A symmetrical compression strategy that pools queries, keys and values at the same time, while preserving left-to-right causality, which greatly improves parallelism. (iii) A two stage training approach which we pre-train for the majority of the time with Lighthouse Attention and recover a full attention model at the end with a short training. We run preliminary small scale LLM pre-training experiments that show the effectiveness of our method compared to full attention training with all other settings matched, where we achieve a faster total training time and lower final loss after the recovery phase.

Full code is available at: https://github.com/ighoshsubho/lighthouse-attention.

11footnotetext:Equal contribution.## 1Introduction

The frontier of language modeling has moved toward contexts of 128K, 1M, and longer, pushed by agentic multi-step reasoning, long-document understanding, and interleaved multimodal inputs[25,1,11,22,27,8,23]. Training at these scales is the dominant hardware bottleneck: scaled dot-product attention hasΘ​(N2)\Theta(N^{2})compute and memory, a wall that FlashAttention[29]pushes back but does not remove.

A growing body of work replaces dense attention with selection: each query attends only to a small subset of keys. Block-level methods such as MoBA[20]and Native Sparse Attention[36]select contiguous blocks, while token-level methods such as DeepSeek Sparse Attention (DSA;9) score every past token via a learned indexer and forward the top-kkinto a sparse attention operator; HISA[40]adds a hierarchical indexer to keep scoring from becoming the new bottleneck. These methods produce meaningful inference speedups but inherit two design decisions that fit long-context pretraining poorly.(i) Asymmetry: queries stay at full resolution while keys and values are pooled, so the hierarchy serves only as a compressed addressable memory rather than a multi-scale representation.(ii) Architectural entanglement: selection lives inside the attention kernel, so the carefully optimized dense-attention kernels that modern tensor-core GPUs accelerate cannot be reused; every sparse method ships its own kernel.

There is also a concern specific to training. An inference-time sparse method[40,28,31,38,32]is by construction as good as its dense backbone, since the sparse substitution is evaluated only against the dense forward. A training-time sparse method must survive a harder test: once training is done, will the resulting model still be a competent dense-attention model?

We take this last question as our central correctness criterion. We introduce Lighthouse Attention: a selection-based hierarchical attention that poolsQ,K,VQ,K,Vsymmetricallyacross a multi-level pyramid, scores every pyramid entry bidirectionally with a parameter-free scorer, and selects the top-KKentries with a fused chunked-bitonic kernel. The selected entries form a dense, causally consistent sub-sequence attended to with stock FlashAttention; outputs are scattered back through a deterministic kernel. The top-KKstep is non-differentiable, with no straight-through estimator: gradients flow through scatter, FlashAttention, and gather intoWQ,WK,WVW_{Q},W_{K},W_{V}, which learn to produce values that are useful when selected. No auxiliary parameters or losses are added. Two consequences follow: the symmetric pyramid is a full multi-scale representation rather than a compressed context, and because selection sits outside the attention path, the expensive step is stock FlashAttention on a sub-sequence of size𝒪​(L​p​K+N/pL−1)\mathcal{O}(LpK+N/p^{L-1}), which reduces to𝒪​(N​log⁡N)\mathcal{O}(N\log N)atL=logp⁡(N/K)L=\log_{p}(N/K).

Our central empirical finding addresses the training-correctness concern directly: after a brief dense-SDPA resumption, Lighthouse-trained models match or beat a fully dense-SDPA baseline trained from scratch on the same token budget. The hierarchical training signal does not hollow out the model’s ability to use full attention at inference, a property inference-only sparse methods cannot claim because they never touch the training loop. We summarize our contributions:

• •A selection-based hierarchical attention designed for long-context pretraining with symmetricQ/K/VQ/K/Vpooling, bidirectional top-KKselection, and stock FlashAttention on the gathered sub-sequence, keeping sparse logic entirely outside the attention kernel.
• •Fused GPU kernels (chunked-bitonic top-KKand a custom scatter-back) that make this design fast at very large contexts.
• •The strongest empirical criterion for a training-time hierarchical method to our knowledge: dense-SDPA resumption after Lighthouse pretraining matches a dense-from-scratch baseline on training loss.

2Related works

Compression and pruning.

A first response to quadratic attention abandons softmax for a bounded-size state: linear attention[katharopoulos2020transformers,4], state-space and gated variants[12,6,34,30], and log-linear attention[13]: which gives strong asymptotics but compresses the entire past and limits long-range recall[2]. A second keeps softmax and prunes at block granularity, either training-free (MInference, FlexPrefill, XAttention, SpargeAttention[15,16,33,37]) or end-to-end (MoBA, NSA[20,36]); these map cleanly onto tiled matmul but force a single retain/discard decision per block and pool only the key–value side. A third prunes at token granularity, mostly at inference for KV-cache eviction (H2O, TOVA, SnapKV, LazyLLM, Quest, SparQ[39,26,17,10,31,28]), or via a learned indexer trained end-to-end (DSA[9]). The defining property of this family is that once selection is identified it is welded into the attention operator as a custom sparse matmul or per-query gather, foreclosing reuse of stock dense kernels.

Hierarchies and training-time correctness.

Multi-resolution attention[35]has returned to sparse LLM attention in two flavors. NSA[36], InfLLM-V2[41], Twilight[18], and DoubleP[24]build hierarchies that the attention itself reads from compression branches, centroid summaries, or quantized proxies. HISA[40]is a training-free, plug-in replacement for DSA’s indexer that runs a block-to-token two-stage score and forwards the selected tokens unchanged to the same Sparse MLA operator DSA already uses. In every case the hierarchy applies only to keys and values, and the selection that emerges still feeds a custom sparse attention kernel. Lighthouse differs on three axes: it pools queries symmetrically with keys and values into coherent multi-resolution(Q(ℓ),K(ℓ),V(ℓ))(Q^{(\ell)},K^{(\ell)},V^{(\ell)})triples; the pyramid is used purely to rank and select, so the attention that follows is stock FlashAttention on a dense sub-sequence with no sparse indexing inside the kernel; and it is trained end-to-end through a non-differentiable top-kkwrapped by a differentiable scatter, with no auxiliary loss or straight-through estimator. Inference-only sparse methods (including HISA) inherit a correctness floor from their underlying dense model, but training-time sparse methods (MoBA, NSA) must answer whether the weights they produce remain competent dense models. We take a brief dense-SDPA resumption recovering the quality of a dense-from-scratch baseline as our central correctness criterion.

Hierarchical SelectorHtH_{t}ProjectionsWQ,WK,WVW_{Q},W_{K},W_{V}PyramidPoolDenseGatherSDPAScatter-backOtO_{t}ScoreTop-KKQ~,K~,V~\widetilde{Q},\widetilde{K},\widetilde{V}O~\widetilde{O}ℐ\mathcal{I}

Figure 1:Lighthouse Attention architecture.Forward(black):HtH_{t}is projected toQ,K,VQ,K,V, passed through the symmetricPyramid Poolon the trunk, and guided by indicesℐ\mathcal{I}from theHierarchical Selector, is fed to a dense gather which topographically sorts the gathered hierarchies into a single contiguous and causal sequence, then stock SDPA, and scatter-back to produceOtO_{t}.Selection(green): the selector taps the pooled summaries off the trunk; the parameter-free scorer ranks them byℓ2\ell_{2}norm and a top-KKkernel keeps the largest entries, emitting integer indicesℐ\mathcal{I}that merge back into the trunk at the dense gather.Gradient(red, dashed):∇L\nabla Ltravels along the trunk (Ot→O_{t}\!\to\!scatter→\!\to\!FA→\!\to\!gather→\!\to\!pyramid pool→WQ,K,V→Ht\!\to\!W_{Q,K,V}\!\to\!H_{t}); the selector branch is non-differentiable and is bypassed.

3Method

We present Lighthouse Attention, a selection-based hierarchical attention mechanism for long-context pretraining. Lighthouse replaces a standard Transformer attention layer with a four-stage pipeline that surrounds, but does not modify, the attention kernel: a pre-attention selection stage drives a contiguous gather, stock FlashAttention[7]runs on the gathered sub-sequence, and a post-attention scatter writes the result back to the original positions. Selection is driven by a parameter-free scoring functional over a multi-resolution pyramid of the layer’s own queries, keys, and values, so Lighthouse introducesno new learnable parametersbeyond those of the underlying attention block.

3.1Preliminaries

LetX∈ℝN×dmodelX\in\mathbb{R}^{N\times d_{\text{model}}}be the input,WQ,WK,WV∈ℝdmodel×dW_{Q},W_{K},W_{V}\in\mathbb{R}^{d_{\text{model}}\times d}projection matrices for one head, andM∈ℝN×NM\in\mathbb{R}^{N\times N}a causal mask. Standard scaled dot-product attention[5]is

Q=X​WQ,K=X​WK,V=X​WV,Attn​(Q,K,V)=softmax​(Q​K⊤d+M)​V,Q=XW_{Q},\quad K=XW_{K},\quad V=XW_{V},\qquad\mathrm{Attn}(Q,K,V)=\mathrm{softmax}\!\left(\tfrac{QK^{\top}}{\sqrt{d}}+M\right)V,(1)with both time and memory costΘ​(N2​d)\Theta(N^{2}d). FlashAttention reduces constants but not asymptotics; atN≥105N\geq 10^{5}this term dominates. Lighthouse replaces Eq. (1) with:(i)symmetric average-pooling ofQ,K,VQ,K,Vinto anLL-level pyramid (factorpp);(ii)parameter-free scoring and a fused chunked-bitonic top-kkselection over all levels jointly;(iii)stock FlashAttention on a contiguous sub-sequence ofS≪NS\ll Nselected entries;(iv)a scatter-back that distributes each output to thepℓp^{\ell}base positions it represents. Stages (ii) and (iv) are custom kernels (Sec.5); stage (iii) is the same FlashAttention call as the dense baseline. The top-kkis treated as discrete and non-differentiable: indices carry no gradient and the scoring functional is not trained. Gradients reachWQ,WK,WVW_{Q},W_{K},W_{V}only through stages (iv), (iii), and the gather: the projections learn to produce values that areuseful when selectedrather than scores that are good at selecting, sidestepping the optimization fragility of learnable selectors.

Hierarchical Selector1. Pyramid Poolℓ=2\ell{=}2ℓ=1\ell{=}1ℓ=0\ell{=}0mean-pool bypℓp^{\ell}2. Norm Scoreℓ=0\ell{=}0ℓ=1\ell{=}1ℓ=2\ell{=}2ℓ2\ell_{2}norm + max-pool3. Chunked Bitonic Top-KKℓ=2\ell{=}2ℓ=1\ell{=}1ℓ=0\ell{=}0parentkeptbitonic argsort + tree-prune𝐐,𝐊,𝐕\mathbf{Q,K,V}[B,S,H,D][B,S,H,D]ℐ\mathcal{I}[B,H,K][B,H,K]pyramidtokensscoresqkscoreskq

Figure 2:Pyramid Pool and the Hierarchical Selector.The Pyramid Pool is a fixed pre-selection stage that livesoutsidethe selector.(1) Pyramid Poolmean-poolsQ,K,VQ,K,Vbypℓp^{\ell}; lines show which tokens feed each summary. The pooled tokens enter the selector, where(2) Norm Scorecomputes parameter-freeℓ2\ell_{2}norms‖Q(ℓ)‖2,‖K(ℓ)‖2\|Q^{(\ell)}\|_{2},\|K^{(\ell)}\|_{2}(coarser levels reuse finer norms via max-pool) and(3) Chunked Bitonic Top-KKkeeps top-KKparents(dark) that descend, while rejected-kept entries (light) emit toℐ\mathcal{I}without descending. Pruning is implicit: children of non-parents are never expanded. Per-tile bitonic argsort runs in-register; uint64 packed keys let a singletorch.sortproduceℐ\mathcal{I}.

3.2Overview

A Lighthouse attention layer replaces standard scaled dot-product attention (Eq. (1)) with a four-stage pipeline that surrounds, but does not modify, the attention kernel. LetQ,K,V∈ℝN×dQ,K,V\in\mathbb{R}^{N\times d}be the per-head projections from the layer’s ownWQ,WK,WVW_{Q},W_{K},W_{V}(Sec. 3.1).

1. (i)Pyramid.Average-poolQ,K,VQ,K,Vsymmetrically into anLL-level pyramid with pooling factorpp, producing coherent triples(Q(ℓ),K(ℓ),V(ℓ))(Q^{(\ell)},K^{(\ell)},V^{(\ell)})forℓ=0,…,L−1\ell=0,\dots,L-1.
2. (ii)Score and top-kk.Assign each pyramid entry parameter-free query and key scores and select thekkentries with the highest combined relevance across all levels via a fused chunked-bitonic top-kkkernel.
3. (iii)Dense sub-sequence attention.Gather the selected triples into a contiguous sub-sequence of lengthSSand compute softmax attention over it with stock FlashAttention.
4. (iv)Scatter-back.Distribute each entry’s output to thepℓp^{\ell}base positions it represents via a deterministic integer-atomic scatter kernel.

Stages (ii) and (iv) are custom kernel (Sec.5); stage (iii) is the same FlashAttention call used by the dense baseline. Lighthouse adds no learnable parameters or losses: the pyramid is a fixed pooling, the scorer is parameter-free, and gather/scatter are data-flow primitives. Gradients flow from the loss through stages (iv) and (iii) into the gatheredQ,K,VQ,K,Vand on intoWQ,WK,WVW_{Q},W_{K},W_{V}; the top-kkstep is discrete and non-differentiable, so its indices carry no gradient and we use no straight-through estimator. The projections therefore learn to be usefulwhen selected, not to score well at selecting.

3.3Pyramid Construction

GivenQ,K,V∈ℝN×dQ,K,V\in\mathbb{R}^{N\times d}, Lighthouse Attention constructs anLL-level pyramid whoseℓ\ell-th level is a non-overlapping window pooling of the previous level. Forℓ=0,…,L−1\ell=0,\dots,L-1, define theii-th window at levelℓ\ellas

𝒲i(ℓ)=[i​pℓ,(i+1)​pℓ−1],i=0,…,Npℓ−1,\mathcal{W}^{(\ell)}_{i}\;=\;\bigl[\,i\,p^{\ell},\;(i+1)\,p^{\ell}-1\,\bigr],\qquad i=0,\dots,\tfrac{N}{p^{\ell}}-1,(2)whereppis the pooling factor. The pyramid entries are then

Qi(ℓ)=Poolμ​{Qj∣j∈𝒲i(ℓ)},Ki(ℓ)=Poolμ​{Kj∣j∈𝒲i(ℓ)},Vi(ℓ)=Poolμ​{Vj∣j∈𝒲i(ℓ)},Q^{(\ell)}_{i}=\mathrm{Pool}_{\mu}\!\bigl\{Q_{j}\mid j\in\mathcal{W}^{(\ell)}_{i}\bigr\},\quad K^{(\ell)}_{i}=\mathrm{Pool}_{\mu}\!\bigl\{K_{j}\mid j\in\mathcal{W}^{(\ell)}_{i}\bigr\},\quad V^{(\ell)}_{i}=\mathrm{Pool}_{\mu}\!\bigl\{V_{j}\mid j\in\mathcal{W}^{(\ell)}_{i}\bigr\},(3)withPoolμ\mathrm{Pool}_{\mu}denoting mean pooling over the window. Level0is the original full-resolution sequence (𝒲i(0)={i}\mathcal{W}^{(0)}_{i}=\{i\}), and each subsequent level summarizesppconsecutive entries of the level below. We requirepL−1∣Np^{L-1}\mid N. Unlike prior hierarchical sparse designs (NSA, HISA, InfLLM-V2), which pool only the context side, Lighthouse appliesPoolμ\mathrm{Pool}_{\mu}symmetricallyto all three projections. Symmetry buys two properties used in subsequent stages: a pooled queryQi(ℓ)Q^{(\ell)}_{i}and a pooled keyKj(ℓ)K^{(\ell)}_{j}live in the same representation space, and each pyramid entry is a coherent(Q,K,V)(Q,K,V)triple summarizing the samepℓp^{\ell}-token span. The total number of pyramid entries is∑ℓ=0L−1N/pℓ≤N⋅p/(p−1)\sum_{\ell=0}^{L-1}N/p^{\ell}\leq N\cdot p/(p-1), so pyramid construction costsΘ​(N)\Theta(N)time and memory.

3.4Scoring and Selection

Each pyramid entry receives two scalar scores — one as a query, one as a key. At level0we use per-headℓ2\ell_{2}norms,

s0,iQK=‖Qi‖2,s0,iKQ=‖Ki‖2,i=0,…,N−1,s^{\mathrm{QK}}_{0,i}=\|Q_{i}\|_{2},\qquad s^{\mathrm{KQ}}_{0,i}=\|K_{i}\|_{2},\qquad i=0,\dots,N-1,(4)and at coarser levels wemax-poolfrom level0rather than recomputing from pooled projections,

sℓ,iQK=max0≤j<pℓ⁡s0,i​pℓ+jQK,sℓ,iKQ=max0≤j<pℓ⁡s0,i​pℓ+jKQ.s^{\mathrm{QK}}_{\ell,i}=\max_{0\leq j<p^{\ell}}s^{\mathrm{QK}}_{0,\,ip^{\ell}+j},\qquad s^{\mathrm{KQ}}_{\ell,i}=\max_{0\leq j<p^{\ell}}s^{\mathrm{KQ}}_{0,\,ip^{\ell}+j}.(5)Max-pooling lets a coarse span inherit the importance of its strongest token. Selection runs jointly over the concatenatedQK\mathrm{QK}andKQ\mathrm{KQ}streams across all levels via the chunked-bitonic kernel of Sec.D.2:

ℐ=TopK​({sℓ,iQK,sℓ,iKQ:(ℓ,i)∈𝒫},k),\mathcal{I}\;=\;\mathrm{TopK}\!\left(\bigl\{s^{\mathrm{QK}}_{\ell,i},\;s^{\mathrm{KQ}}_{\ell,i}:(\ell,i)\in\mathcal{P}\bigr\},\;k\right),(6) where𝒫\mathcal{P}is the full set of pyramid indices. An entry chosen via itsKQ\mathrm{KQ}score still enters the gather as its own(Qi(ℓ),Ki(ℓ),Vi(ℓ))(Q^{(\ell)}_{i},K^{(\ell)}_{i},V^{(\ell)}_{i})triple. The coarsest level is always retained in full — it is cheap and guarantees at least one contributor at every base position; the remaining budget is spent on finer levels.

3.5Gathered-Sequence Attention

Givenℐ\mathcal{I}, Lighthouse assembles a contiguous sub-sequence

Q~m=Qim(ℓm),K~m=Kim(ℓm),V~m=Vim(ℓm),(ℓm,im)∈ℐ,m=1,…,S,\widetilde{Q}_{m}=Q^{(\ell_{m})}_{i_{m}},\quad\widetilde{K}_{m}=K^{(\ell_{m})}_{i_{m}},\quad\widetilde{V}_{m}=V^{(\ell_{m})}_{i_{m}},\qquad(\ell_{m},i_{m})\in\mathcal{I},\;\;m=1,\dots,S,(7)of length

S=NpL−1+(L−1)​p​k,S\;=\;\frac{N}{p^{L-1}}\;+\;(L-1)\,p\,k,(8)because the coarsest level contributes allN/pL−1N/p^{L-1}entries while each of the remainingL−1L-1levels contributes at mostp​kpk(the factor ofppis causal-boundary bookkeeping; Sec.D.2). AtN=106,L=4,p=4,k=4096N=10^{6},L=4,p=4,k=4096,S≈6.5×104≪NS\approx 6.5\times 10^{4}\ll N. The sub-sequence is then attended to via stock SDPA or FlashAttention,

O~=Attn​(Q~,K~,V~;M~),\widetilde{O}\;=\;\mathrm{Attn}(\widetilde{Q},\,\widetilde{K},\,\widetilde{V};\,\widetilde{M}),(9)whereAttn​(⋅)\mathrm{Attn}(\cdot)is standard masked softmax attention. The causal maskM~\widetilde{M}derives from the pyramid coordinates(ℓm,im)(\ell_{m},i_{m})so each entry attends only to entries whose base positions are no greater than its own; the gather is topologically sorted, soM~\widetilde{M}reduces to a standardS×SS{\times}Scausal mask and Eq. (9) contains no sparse indexing.

Due to the hierarchical decomposition, this gathering process guaranties that there are no ”holes” or empty spaces in the sequence, which is especially important as we also compress queries Q; a hole could cause training instabilities as those missing tokens would be cut out during the forward pass and have no gradients during the backward pass. This is unlike asymmetrical methods that do not compress queries.

3.6Scatter-Back Reconstruction

The attention outputO~∈ℝS×d\widetilde{O}\in\mathbb{R}^{S\times d}is redistributed to the fullNN-token outputOO. A selected entry at levelℓ\ell, positioniisummarized window𝒲i(ℓ)\mathcal{W}^{(\ell)}_{i}during pooling but its output is written to ashiftedrange

ℛ​(ℓ,i)=[i​pℓ+pℓ−1,i​pℓ+2​pℓ−2],\mathcal{R}(\ell,i)\;=\;\bigl[\,ip^{\ell}+p^{\ell}-1,\;\;ip^{\ell}+2p^{\ell}-2\,\bigr],(10)that starts at the last summarized token. The shift ofpℓ−1p^{\ell}-1preserves causality: a base positionjjnever receives a summary that contains its own future. Within a level, consecutive windows write to disjoint adjacent ranges; across levels, contributions are summed,

Oj=∑m:j∈ℛ​(ℓm,im)O~m,O_{j}\;=\;\sum_{m\,:\,j\in\mathcal{R}(\ell_{m},i_{m})}\widetilde{O}_{m},(11)so the per-position fan-in is bounded byLLregardless ofkk.

Similarly to the gathering pass, the scattering process also has no empty spaces. This final scattered sequence is fully dense, albeit a compressive approximation of full attention.

4Design Choices

The Lighthouse pipeline of Sec.3makes four design choices that distinguish it from prior selection-based sparse attention. First,QQis pooledin lockstepwithK,VK,Vinstead of leaving queries dense as in NSA[36], HISA[40], and InfLLM-v2[41]; this is the choice that turns the dense kernel call from𝒪​(N​S​d)\mathcal{O}(NSd)to𝒪​(S2​d)\mathcal{O}(S^{2}d)at training time, and keeps pooled queries and pooled keys in the same representation space at every level. Second, the scorer isparameter-freeper-headℓ2\ell_{2}norms of the layer’s ownQ(ℓ),K(ℓ)Q^{(\ell)},K^{(\ell)}— rather than a learned scoring head as in NSA[36]or DSA[9]; this is the cheaper option and is strictly weaker than any attention- or QK-interaction-based scorer, so any positive result is a lower bound on what richer scorers can extract. The natural QK-interaction alternative we ablate against is adilated softmax-attentionscorer that runs softmax attention over the pyramid with dilation factorδ\deltaat𝒪​(N2​d/δ)\mathcal{O}(N^{2}d/\delta)per layer sub-quadratic but still super-linear inNN, and an order of magnitude more expensive than the projection-norm scorer at long context (Sec.6.4).

Third, selection isdecoupledfrom attention: top-KKproduces a contiguous, dense sub-sequence and attention is a stock SDPA or FlashAttention[5,7]call on it, with no custom sparse-attention kernel coupling the two steps as in NSA[36], DSA[9], or HISA[40]. The same kernel runs at training and inference, and disabling selection cleanly recovers the dense baseline exactly the SDPA-resume test in Sec.6.2. Fourth, we donotmake the top-KKdifferentiable: no straight-through estimator, no Gumbel softmax, no auxiliary scorer loss. Gradients flow only through the gatheredQ~,K~,V~\widetilde{Q},\widetilde{K},\widetilde{V}intoWQ,WK,WVW_{Q},W_{K},W_{V}, so the projections learn to be usefulwhen selectedrather than to game a learnable scorer. We motivate each choice and discuss alternatives in AppendixC.

5Complexity Analysis and Kernel Design

Algorithm1summarizes one Lighthouse attention layer as a sequence of GPU primitives. Most stages are standard operations executed viatorch.compile’d PyTorch code; only the top-kkselection (stage 2c) and the scatter-back (stage 5) are custom kernels.

Input:

X∈ℝN×dmodelX\in\mathbb{R}^{N\times d_{\text{model}}}, projections

WQ,WK,WVW_{Q},W_{K},W_{V}, pyramid params

(L,p)(L,p), budget

kk Output:

O∈ℝN×dO\in\mathbb{R}^{N\times d} Q,K,V←X​WQ,X​WK,X​WVQ,K,V\leftarrow XW_{Q},\;XW_{K},\;XW_{V}

//projections (fused GEMM)

{(Q(ℓ),K(ℓ),V(ℓ))}ℓ=0L−1←Poolμ​(Q,K,V)\bigl\{(Q^{(\ell)},K^{(\ell)},V^{(\ell)})\bigr\}_{\ell=0}^{L-1}\leftarrow\mathrm{Pool}_{\mu}(Q,K,V)

//pyramid (view+mean)

sℓ,iQK←Poolmax​{‖Qj‖2},sℓ,iKQ←Poolmax​{‖Kj‖2}s^{\mathrm{QK}}_{\ell,i}\leftarrow\mathrm{Pool}_{\max}\!\bigl\{\|Q_{j}\|_{2}\bigr\},\;s^{\mathrm{KQ}}_{\ell,i}\leftarrow\mathrm{Pool}_{\max}\!\bigl\{\|K_{j}\|_{2}\bigr\}

//scoring (norm+max)

ℐ←ChunkedBitonicTopK​(sQK,sKQ,k)\mathcal{I}\leftarrow\textsc{ChunkedBitonicTopK}(s^{\mathrm{QK}},s^{\mathrm{KQ}},k)

//custom kernel (§D.2)

Q~,K~,V~←Gather​(Q(⋅),K(⋅),V(⋅);ℐ)\widetilde{Q},\widetilde{K},\widetilde{V}\leftarrow\mathrm{Gather}\bigl(Q^{(\cdot)},K^{(\cdot)},V^{(\cdot)};\,\mathcal{I}\bigr)

//torch.gather

O~←FlashAttention​(Q~,K~,V~)\widetilde{O}\leftarrow\mathrm{FlashAttention}(\widetilde{Q},\widetilde{K},\widetilde{V})

//stock FA-3/FA-4

O←ScatterBack​(O~,ℐ)O\leftarrow\textsc{ScatterBack}(\widetilde{O},\mathcal{I})

//custom kernel)

1return

OO

Algorithm 1Lighthouse attention (single layer, single head).### 5.1Asymptotic Complexity

Table3decomposes per-layer cost by stage. The only super-linear term inNNis the dense sub-sequence attention,Θ​(S2​d)\Theta(S^{2}d), withS=N/pL−1+(L−1)​p​kS=N/p^{L-1}+(L-1)\,p\,kfrom Sec.3.5. ChoosingL=logp⁡(N/k)L=\log_{p}(N/k)balances the two terms inSS, givingS=Θ​(k​logp⁡(N/k))S=\Theta(k\log_{p}(N/k))and an attention cost ofΘ​(k2​log2⁡N⋅d)\Theta(k^{2}\log^{2}N\cdot d)— polylogarithmic inNNat fixedkk. Combined with the linear scoring andΘ​(N​log⁡k)\Theta(N\log k)selection passes, total per-layer compute is linear inNNup to alog⁡k\log kfactor for boundedkk. App.Bderives this and compares against dense softmax, log-linear attention, and linear/SSM families.

5.2Kernel Design and Parallelism

Of the seven stages in Algorithm1, only top-KKand scatter-back are custom kernels in CUDA and triton; the rest reduce to PyTorch primitives thattorch.compilefuses into single device passes. Ourchunked-bitonictop-KKpartitions the score stream, maintains an in-register top-mmbuffer per chunk via bitonic merge, and dispatches chunks as independent CTAs avoiding the shared-memory blow-up of textbook bitonic atk=4096k{=}4096while producing a stratified selection that resists span collapse. Crucially, gather is decoupled from attention: where NSA[36], DSA[9], HISA[40], and MoBA[20]embed selection inside a custom sparse kernel, Lighthouse hands a contiguous dense sub-sequence to stock FlashAttention[29]— making forward/backward bit-for-bit identical to a dense Transformer’s, letting context parallelism rotate the gather through standard ring attention[19]without any sparsity-aware collective, and enabling 1M-token training across 32 Blackwell GPUs (full details in App.D).

6Experiments

We evaluate Lighthouse along three axes: (1)recoverability: whether lighthouse pretraining damages the model’s ability to use full attention at inference Sec.6.2); (2)design ablations and throughputover the four knobs (scorer,pp,LL,kk) and the resulting wall-clock cost (Sec.6.4); and (3)scaling vs. dense attentionas a function of context length, including the long-context regime that requires context parallelism (Sec.6.3). All runs share the architecture and recipe of Sec.6.1.

6.1Experimental Setup

Architecture, data, optimizer.

A530530M-parameter Llama-3-style decoder (dmodel=1024d_{\text{model}}{=}1024,3030layers,H=8H{=}8, head dim128128, FFN15361536, byte-level tokenizer). Layers{0,1,28,29}\{0,1,28,29\}retain dense SDPA — PyTorch2.11.02.11.0+cu128’storch.nn.attention.sdpa_kernelrouted to cuDNN9.19.09.19.0on CUDA12.812.8; the other 26 use Lighthouse with the same cuDNN-SDPA kernel as the inner attention call on the gathered sub-sequence. Training on C4 at sequence length98,30498{,}304, global batch3232, AdamW2×10−32\!\times\!10^{-3},β1=0.9,β2=0.95\beta_{1}{=}0.9,\beta_{2}{=}0.95, weight decay0.10.1, linear warmup over 2k steps, gradient-norm clip 1, bfloat16, FSDP only.

Two-stage recipe.

Stage 1 trains with Lighthouse; stage 2 resumes the stage-1 checkpoint underdenseSDPA (same cuDNN backend), with the same optimizer state and dataloader continuation. The total budget is held at16,00016{,}000steps (≈50\approx 50B tokens); we vary the stage-1 length to test sensitivity to the switch point.

Hardware.

A single NVIDIA BGX 8×\timesB200 node is used for 98K-context runs; multi-node configurations are used with intra-node CP for 256K (Table1). We report training and validation loss, tokens/s per GPU in steady state, and total B200 hours.

6.2SDPA Recoverability

We test whether a hierarchical-trained Lighthouse model can be restored to dense attention by a brief continuation under stock SDPA. Holding the budget at16,00016{,}000steps (≈50\approx 50B tokens), we vary the stage-1 length (1010k /1111k /1212k) and resume the remainder under dense SDPA, against an full SDPA reference at matched architecture, data, and tokens (Table1, top block). At each resume the training loss transiently spikes (1.121.12–1.571.57) as the model is first run through attention it was not trained against, then recovers within≈1\approx 1–1.51.5k SDPA steps and crosses below the dense baseline; by step16,00016{,}000all three resume schedules match or beat dense-from-scratch (0.69800.6980–0.71020.7102vs.0.72370.7237), with longer dense-resume tails giving lower final loss. Recovery is robust across resume points (the recipe doesn’t pivot on a precise schedule), supporting our load-bearing claim thathierarchical training does not compromise the model’s ability to use full attention at inference, at no additional token cost over dense-from-scratch.

ConfigurationScorerParamsLHTotalTotalB200-HrsTok/s (k)Final LossStepsStepsTokens(↓\downarrow)(↑\uparrow)(↓\downarrow)SDPA Baseline (ctx=98​k=98k)—530M—16k50.3B303.245.60.7237SDPA recoverability (L=3,p=2,k=6144L{=}3,\ p{=}2,\ k{=}6144, ctx=98​k=98k)LH→\toSDPA (12k+4k)Dilated530M12k16k50.3B214.774.70.7102LH→\toSDPA (11k+5k)Dilated530M11k16k50.3B219.675.40.7001LH→\toSDPA (10k+6k)Dilated530M10k16k50.3B228.075.00.6980Hyperparameter ablations (ctx=98​k=98k)L=3,p=2,k=1536L{=}3,\ p{=}2,\ k{=}1536Dilated530M10k16k50.3B203.993.90.6825L=3,p=4,k=1536L{=}3,\ p{=}4,\ k{=}1536Dilated530M10k16k50.3B197.299.50.6881L=3,p=8,k=1536L{=}3,\ p{=}8,\ k{=}1536Dilated530M10k16k50.3B206.292.10.6828L=4,p=2,k=1536L{=}4,\ p{=}2,\ k{=}1536Dilated530M10k16k50.3B200.296.40.6978L=5,p=2,k=1536L{=}5,\ p{=}2,\ k{=}1536Dilated530M10k16k50.3B201.596.30.6991L=3,p=2,k=2048L{=}3,\ p{=}2,\ k{=}2048Dilated530M10k16k50.3B208.190.90.6880L=3,p=2,k=4096L{=}3,\ p{=}2,\ k{=}4096Dilated530M10k16k50.3B215.783.50.6951CP training(L=3,p=4L{=}3,\ p{=}4)k=1536k{=}1536, ctx=98​k=98k, CP=2=2, DP=4=4Norm530M10k16k100.7B208.391.80.6903k=2048k{=}2048, ctx=98​k=98k, CP=2=2, DP=4=4Norm530M10k16k100.7B210.989.20.6928k=4096k{=}4096, ctx=256​k=256k, CP=8=8, DP=1=1Norm530M10k16k1.07T1300.348.90.6721

Table 1:Lighthouse ablation summary.530M Llama-3,1616k optimizer steps;LH Stepsrun Lighthouse,SDPA Stepsrun dense SDPA-resume on a single8×8\!\times\!B200 node (the final block adds context parallelism).B200-Hrs: combined wall-clock×\times\,8 GPUs.Tok/s (k): Lighthouse-stagethroughput(tps)from torchtitan, aggregated over all ranks (the baseline SDPA shows the throughput when training without LH). The full set of ablations is provided in Table2, AppendixA.

6.3Scaling Laws vs. Dense Attention

We benchmark single-layer attention latency on a single B200 for contexts from88K to512512K (bf16,B=1B{=}1,H=8H{=}8,d=128d{=}128,L=3L{=}3,p=4p{=}4, sparsity1:641{:}64, medians of1010steady-state iterations), comparing Lighthouse against cuDNN-backed SDPA. SDPA scales asΘ​(N2​d)\Theta(N^{2}d)while Lighthouse scales asΘ​(S2​d)\Theta(S^{2}d)withSSdefined in Eq.8, so the gap widens withNN(Fig.3). AtN=512N{=}512K, Lighthouse is𝟐𝟏×\boldsymbol{21\times}faster on the forward pass and17.3×\boldsymbol{17.3\times}faster on forward++backward; equivalently, SDPA needs∼\sim113113K (fwd) /∼\sim122122K (fwd++bwd) of context to reach the runtime Lighthouse takes at512512K. Full-model training tells a similar story but requires care: with our530530M-parameter architecture a single B200 OOMs beyond∼\sim100100K on activations, gradients, and optimizer state regardless of attention method, so we implement context parallelism (Sec.D.4) where pyramid pooling, scoring, and top-KKrun shard-locally and the gathered sub-sequence rotates through stock ring attention[19]with no sparse-aware collectives. CP introduces a small ring-rotation overhead, costing∼\sim10%10\%in per-rank throughput vs. the single-device extrapolation, but the Lighthouse-vs-SDPA speedup is preserved under matched CP geometry (Lighthouse-CP retains the same multiplicative gain over SDPA-CP that we see in the non-CP comparison), carrying the advantage cleanly to the11M-token /3232-GPU regime.

6.4Design Ablations and Throughput

We sweep four design axes (scorer variant, pooling factorpp, number of levelsLL, top-KKbudgetkk), each varied independently while the others stay at the defaults of Sec.6.1; comparisons use the post-resume training loss at step16,00016{,}000. The full grid in Table1establishes three things. First, every Lighthouse configuration matches or beats the dense-SDPA-from-scratch baseline of0.72370.7237, so recoverability is not specific to any one hyperparameter setting.

Second, the projection-norm scorer is within∼\sim0.010.01of dilated softmax in either direction (no uniform winner) but is parameter-free and roughly9%9\%cheaper in B200-hours (179.6179.6–180.9180.9vs.197.2197.2–199.7199.7atL=3,p=4L{=}3,p{=}4).

Third, smallerpp, shallowerLL, and smallerkkall help slightly: the lowest-loss cell across the grid isL=3,p=2,k=1536L{=}3,\,p{=}2,\,k{=}1536(dilated, loss0.68250.6825), Pareto-best on every metric within its Top-KKblock. The smaller-kkdirection is the most counter-intuitive: loss decreases monotonically askkshrinks over{1,536,2,048,3,072,4,096}\{1{,}536,2{,}048,3{,}072,4{,}096\}(0.6825→0.6880→0.6890→0.69510.6825\to 0.6880\to 0.6890\to 0.6951) before dipping again atk=6,144k{=}6{,}144(0.68310.6831), plausibly because hierarchical selection regularises at our token budget; whether this reverses at substantially larger budgets is left to future work.

(a)Forward Latency (b)Backward Latency

Figure 3:Attention latency vs. context lengthfor SDPA (cuDNN) and Lighthouse (w/ cuDNN) on a single B200,L=3,p=4L{=}3,\ p{=}4, sparsity≈1:64\approx 1{:}64. SDPA scales asΘ​(N2​d)\Theta(N^{2}d); Lighthouse scales asΘ​(S2​d)\Theta(S^{2}d)withS≪NS\ll N. AtN=512N{=}512K, Lighthouse is21×21{\times}faster on the forward pass and17.3×17.3{\times}faster on the backward pass, equivalently Lighthouse at 512k takes same runtime as if training SDPA at∼\sim113K /∼\sim122K context.The throughput story is consistent. Lighthouse stage-1 sustains8484–126126k tok/s/GPU across the ablation grid against∼\sim4646k for dense SDPA, a roughly2×2\timesper-step advantage that holds across selection budgets; the projection-norm scorer atL=3,p=4,k=1536L{=}3,\,p{=}4,\,k{=}1536tops the range at126126k by skipping the dilated-attention pass entirely.

End-to-end on the1010k++66k recipe, total runtime ranges from22.522.5h (179.6179.6B200-h, normk=1536,p=4k{=}1536,p{=}4) to27.027.0h (215.7215.7B200-h,k=4096,p=2,L=3k{=}4096,p{=}2,L{=}3) against37.937.9h (303.2303.2B200-h) for dense-SDPA-from-scratch on the same1616k-step /50.350.3B-token budget: a1.40×1.40\timesto1.69×1.69\timeswall-clock speedup at matched or lower final loss. The saving comes entirely from stage-1; the SDPA-resume tail uses the same kernel as the baseline and matches its throughput. App.Egives the per-axis breakdowns, asymptotic-cost predictions, and per-stage timing decompositions.

7Conclusion

We introduce Lighthouse Attention, a selection-based hierarchical attention for long-context pretraining that poolsQ,K,VQ,K,Vsymmetricallyacross a multi-resolution pyramid and places selectionoutsidethe attention kernel, reducing the attention step to stock FlashAttention on a dense sub-sequence. The design is parameter-free, trains end-to-end with no auxiliary losses or straight-through estimators, and inherits upstream FlashAttention improvements unchanged. A brief dense-SDPA resumption after Lighthouse pretraining matches or beats dense-from-scratch at matched tokens on training loss and long-context retrieval, with1.41.4–1.7×1.7\timesend-to-end speedups against cuDNN SDPA at≥100\geq 100K context on B200 and clean scaling to11M tokens on multi-node Blackwell.

Limitations.

SymmetricQ/K/VQ/K/Vpooling presumes all queries co-occur in one forward pass, which autoregressive decoding violates; we rely on dense-SDPA resumption for an inference-ready model, and every downstream evaluation is run after that resume rather than on the hierarchical forward directly. The inner attention isΘ​(S2​d)\Theta(S^{2}d)on the gathered sub-sequence: sub-quadratic inNNat fixedkkbut not strictly linear, so regimes wherekkmust scale withNNremain uncharacterised.

Future directions.

Swapping the dense-SDPA resume for an asymmetric sparse target (DSA, NSA, HISA, MoBA) would yield a natively serveable checkpoint; per-layer or per-head adaptivekkmay outperform a fixed budget; the multi-scale pyramid extends naturally to vision, audio, and video; and serving integration (continuous batching, speculative decoding, KV-cache management) is needed to translate the training speedups into deployment.

References

• [1](2024)The Claude 3 model family.External Links:LinkCited by:§1.
• [2]S. Arora, S. Eyuboglu, A. Timalsina, I. Johnson, M. Poli, J. Zou, A. Rudra, and C. Ré(2024)Zoology: measuring and improving recall in efficient language models.InICLR,Note:arXiv:2312.04927Cited by:§2.
• [3]Y. Bengio, N. Léonard, and A. Courville(2013)Estimating or propagating gradients through stochastic neurons for conditional computation.arXiv preprint arXiv:1308.3432.Cited by:§C.4.
• [4]K. Choromanski, V. Likhosherstov, D. Dohan, X. Song, A. Gane, T. Sarlos, P. Hawkins, J. Davis, A. Mohiuddin, L. Kaiser, D. Belanger, L. Colwell, and A. Weller(2021)Rethinking attention with Performers.InICLR,Cited by:§2.
• [5]T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Ré(2022)FlashAttention: fast and memory-efficient exact attention with IO-awareness.InNeurIPS,Cited by:§C.1,§C.3,§C.4,§3.1,§4.
• [6]T. Dao and A. Gu(2024)Transformers are SSMs: generalized models and efficient algorithms through structured state space duality.InICML,Cited by:§2.
• [7]T. Dao(2024)FlashAttention-2: faster attention with better parallelism and work partitioning.InICLR,Note:arXiv:2307.08691Cited by:§C.3,§3,§4.
• [8]DeepSeek-AI(2024)DeepSeek-V3 technical report.arXiv preprint arXiv:2412.19437.Cited by:§1.
• [9]DeepSeek-AI(2025)DeepSeek-V3.2-Exp: boosting long-context efficiency with DeepSeek sparse attention.arXiv preprint.Cited by:§C.2,§C.3,§C.4,§D.3,§1,§2,§4,§4,§5.2.
• [10]Q. Fu, M. Cho, T. Merth, S. Mehta, M. Rastegari, and M. Najibi(2024)LazyLLM: dynamic token pruning for efficient long context LLM inference.arXiv preprint arXiv:2407.14057.Cited by:§2.
• [11]Gemini Team, Google DeepMind(2024)Gemini 1.5: unlocking multimodal understanding across millions of tokens of context.arXiv preprint arXiv:2403.05530.Cited by:§1.
• [12]A. Gu and T. Dao(2023)Mamba: linear-time sequence modeling with selective state spaces.arXiv preprint arXiv:2312.00752.Cited by:§2.
• [13]H. Guoet al.(2025)Log-linear attention.arXiv preprint.Cited by:Table 4,§2.
• [14]E. Jang, S. Gu, and B. Poole(2017)Categorical reparameterization with Gumbel-softmax.InInternational Conference on Learning Representations (ICLR),Cited by:§C.4.
• [15]H. Jiang, Y. Li, C. Zhang, Q. Wu, X. Luo, S. Ahn, Z. Han, A. H. Abdi, D. Li, C. Lin, Y. Yang, and L. Qiu(2024)MInference 1.0: accelerating pre-filling for long-context LLMs via dynamic sparse attention.arXiv preprint arXiv:2407.02490.Cited by:§2.
• [16]X. Laiet al.(2025)FlexPrefill: a context-aware sparse attention mechanism for efficient long-sequence inference.arXiv preprint.Cited by:§2.
• [17]Y. Li, Y. Huang, B. Yang, B. Venkitesh, A. Locatelli, H. Ye, T. Cai, P. Lewis, and D. Chen(2024)SnapKV: LLM knows what you are looking for before generation.arXiv preprint arXiv:2404.14469.Cited by:§2.
• [18]C. Linet al.(2025)Twilight: adaptive attention sparsity with hierarchical top-pppruning.arXiv preprint.Cited by:§2.
• [19]H. Liu, M. Zaharia, and P. Abbeel(2023)Ring attention with blockwise transformers for near-infinite context.arXiv preprint arXiv:2310.01889.Cited by:§D.4,§5.2,§6.3.
• [20]E. Lu, Z. Jiang, J. Liu, Y. Du, T. Jiang, C. Hong, S. Liu, W. He, E. Yuan, Y. Wang,et al.(2025)MoBA: mixture of block attention for long-context LLMs.arXiv preprint arXiv:2502.13189.Cited by:§D.3,§1,§2,§5.2.
• [21]C. J. Maddison, A. Mnih, and Y. W. Teh(2017)The concrete distribution: a continuous relaxation of discrete random variables.InInternational Conference on Learning Representations (ICLR),Cited by:§C.4.
• [22]Meta AI(2024)The Llama 3 herd of models.arXiv preprint arXiv:2407.21783.Cited by:§1.
• [23]Moonshot AI(2025)Kimi K1.5: scaling reinforcement learning with llms.arXiv preprint arXiv:2501.12599.Cited by:§1.
• [24]Y. Niet al.(2026)DoubleP: hierarchical cluster-and-refine attention with centroid approximation.arXiv preprint.Cited by:§2.
• [25]OpenAI(2024)GPT-4 technical report.arXiv preprint arXiv:2303.08774.Cited by:§1.
• [26]M. Oren, M. Hassid, Y. Adi, and R. Schwartz(2024)Transformers are multi-state RNNs.arXiv preprint arXiv:2401.06104.Cited by:§2.
• [27]Qwen Team(2025)Qwen2.5 technical report.arXiv preprint arXiv:2412.15115.Cited by:§1.
• [28]L. Ribar, I. Chelombiev, L. Hudlass-Galley, C. Blake, C. Luschi, and D. Orr(2024)SparQ attention: bandwidth-efficient LLM inference.InInternational Conference on Machine Learning (ICML),Cited by:§1,§2.
• [29]J. Shah, G. Bikshandi, Y. Zhang, V. Thakkar, P. Ramani, and T. Dao(2024)FlashAttention-3: fast and accurate attention with asynchrony and low-precision.InAdvances in Neural Information Processing Systems (NeurIPS),Cited by:§C.3,§1,§5.2.
• [30]Y. Sun, L. Dong, S. Huang, S. Ma, Y. Xia, J. Xue, J. Wang, and F. Wei(2023)Retentive network: a successor to transformer for large language models.arXiv preprint arXiv:2307.08621.Cited by:§2.
• [31]J. Tang, Y. Zhao, K. Zhu, G. Xiao, B. Kasikci, and S. Han(2024)Quest: query-aware sparsity for efficient long-context LLM inference.InInternational Conference on Machine Learning (ICML),Cited by:§1,§2.
• [32]G. Xiao, Y. Tian, B. Chen, S. Han, and M. Lewis(2024)Efficient streaming language models with attention sinks.InInternational Conference on Learning Representations (ICLR),Cited by:§1.
• [33]R. Xuet al.(2025)XAttention: block sparse attention with antidiagonal scoring.arXiv preprint.Cited by:§2.
• [34]S. Yang, B. Wang, Y. Shen, R. Panda, and Y. Kim(2024)Gated linear attention transformers with hardware-efficient training.arXiv preprint arXiv:2312.06635.Cited by:§2.
• [35]Z. Yang, D. Yang, C. Dyer, X. He, A. Smola, and E. Hovy(2016)Hierarchical attention networks for document classification.InNAACL,Cited by:§2.
• [36]J. Yuan, H. Gao, D. Dai, J. Luo, L. Zhao, Z. Zhang, Z. Xie, Y. Wei, L. Wang, Z. Xiao,et al.(2025)Native sparse attention: hardware-aligned and natively trainable sparse attention.arXiv preprint arXiv:2502.11089.Cited by:§C.1,§C.2,§C.3,§C.4,§D.3,§1,§2,§2,§4,§4,§5.2.
• [37]J. Zhanget al.(2025)SpargeAttention: accurate and training-free sparse attention accelerating any model inference.arXiv preprint.Cited by:§2.
• [38]Z. Zhang, Y. Sheng, T. Zhou, T. Chen, L. Zheng, R. Cai, Z. Song, Y. Tian, C. Ré, C. Barrett, Z. Wang, and B. Chen(2023)H2O: heavy-hitter oracle for efficient generative inference of large language models.InAdvances in Neural Information Processing Systems (NeurIPS),Cited by:§1.
• [39]Z. Zhang, Y. Sheng, T. Zhou, T. Chen, L. Zheng, R. Cai, Z. Song, Y. Tian, C. Ré, C. Barrett, Z. Wang, and B. Chen(2024)H2O: heavy-hitter oracle for efficient generative inference of large language models.InNeurIPS,Note:arXiv:2306.14048Cited by:§2.
• [40]L. Zhaoet al.(2026)HISA: efficient hierarchical indexing for fine-grained sparse attention.arXiv preprint arXiv:2603.28458.Cited by:§C.1,§C.3,§D.3,§1,§1,§2,§4,§4,§5.2.
• [41]L. Zhaoet al.(2026)InfLLM-V2: dense–sparse switchable attention for seamless short-to-long adaptation.arXiv preprint.Cited by:§C.1,§2,§4.

Appendix AAblations

ConfigurationScorerParamsLHTotalTotalB200-HrsTok/s (k)Final LossStepsStepsTokens(↓\downarrow)(↑\uparrow)(↓\downarrow)SDPA Baseline (ctx=98​k=98k)—530M—16k50.3B303.245.60.7237SDPA recoverability (L=3,p=2,k=6144L{=}3,\ p{=}2,\ k{=}6144, ctx=98​k=98k)LH→\toSDPA (12k+4k)Dilated530M12k16k50.3B214.774.70.7102LH→\toSDPA (11k+5k)Dilated530M11k16k50.3B219.675.40.7001LH→\toSDPA (10k+6k)Dilated530M10k16k50.3B228.075.00.6980Scorer ablation (L=3,p=4L{=}3,\ p{=}4, ctx=98​k=98k)k=1536k{=}1536Dilated530M10k16k50.3B197.299.50.6881k=1536k{=}1536Norm530M10k16k50.3B179.6126.00.6946k=2048k{=}2048Dilated530M10k16k50.3B199.797.10.6969k=2048k{=}2048Norm530M10k16k50.3B180.9122.40.6921Pooling-factor ablation (L=3L{=}3, ctx=98​k=98k)p=2,k=1536p{=}2,\ k{=}1536Dilated530M10k16k50.3B203.993.90.6825p=4,k=1536p{=}4,\ k{=}1536Dilated530M10k16k50.3B197.299.50.6881p=8,k=1536p{=}8,\ k{=}1536Dilated530M10k16k50.3B206.292.10.6828p=2,k=2048p{=}2,\ k{=}2048Dilated530M10k16k50.3B208.190.90.6880p=4,k=2048p{=}4,\ k{=}2048Dilated530M10k16k50.3B199.797.10.6969Number-of-levels ablation (p=2p{=}2, ctx=98​k=98k)L=3,k=1536L{=}3,\ k{=}1536Dilated530M10k16k50.3B203.993.90.6825L=4,k=1536L{=}4,\ k{=}1536Dilated530M10k16k50.3B200.296.40.6978L=5,k=1536L{=}5,\ k{=}1536Dilated530M10k16k50.3B201.596.30.6991L=3,k=2048L{=}3,\ k{=}2048Dilated530M10k16k50.3B208.190.90.6880L=4,k=2048L{=}4,\ k{=}2048Dilated530M10k16k50.3B202.294.50.6983L=5,k=2048L{=}5,\ k{=}2048Dilated530M10k16k50.3B206.592.30.7043Top-KKbudget ablation (L=3,p=2L{=}3,\ p{=}2, ctx=98​k=98k)k=1536k{=}1536Dilated530M10k16k50.3B203.993.90.6825k=2048k{=}2048Dilated530M10k16k50.3B208.190.90.6880k=3072k{=}3072Dilated530M10k16k50.3B214.986.10.6890k=4096k{=}4096Dilated530M10k16k50.3B215.783.50.6951k=6144k{=}6144Dilated530M10k16k50.3B208.188.30.6831CP training(L=3,p=4L{=}3,\ p{=}4)k=1536k{=}1536, ctx=98​k=98k, CP=2=2, DP=4=4Norm530M10k16k100.7B208.391.80.6903k=2048k{=}2048, ctx=98​k=98k, CP=2=2, DP=4=4Norm530M10k16k100.7B210.989.20.6928k=4096k{=}4096, ctx=256​k=256k, CP=8=8, DP=1=1Norm530M10k16k1.07T1300.348.90.6721

Table 2:Full Lighthouse ablation summary.530M Llama-3,1616k optimizer steps;LH Stepsrun Lighthouse,SDPA Stepsrun dense SDPA-resume on a single8×8\!\times\!B200 node (the final block adds context parallelism).B200-Hrs: combined wall-clock×\times\,8 GPUs.Tok/s (k): Lighthouse-stagethroughput(tps)from torchtitan, aggregated over 8 ranks (the Dense-SDPA-from-scratch row has no LH stage and reports its dense values).Final Loss: training loss at step16,00016{,}000. Bold marks the per-block best on each metric (skipped for throughput in the CP block since context length varies inside it).

Appendix BComplexity Derivation

Table 3:Per-layer complexity of a Lighthouse attention layer, withS=N/pL−1+(L−1)​p​kS=N/p^{L-1}+(L-1)\,p\,k.We derive the per-layer compute complexity of Lighthouse and show that, at bounded selection budgetkk, total compute is linear in the sequence lengthTT.

Setup.

Lighthouse with sequence lengthTT, pooling factorpp,LLpyramid levels, top-kkbudgetkk, and head dimensiondd.

Per-stage cost.

Table3decomposes one layer into its stages.

Sub-sequence size.

By construction the gathered sub-sequence has size

S=TpL−1+(L−1)​p​k,S\;=\;\frac{T}{p^{L-1}}\;+\;(L-1)\,p\,k,(12)where the first term is the coarsest level (kept whole) and the second is the contribution of the finerL−1L-1levels (each of which contributes at mostp​kp\,kentries).

Choice ofLL.

SettingL=logp⁡(T/k)L=\log_{p}(T/k)givespL=T/kp^{L}=T/kand thereforepL−1=T/(p​k)p^{L-1}=T/(p\,k). Substituting into Eq. (12),

TpL−1=p​k,\frac{T}{p^{L-1}}\;=\;p\,k,and the two terms combine into

S=p​k+(L−1)​p​k=p​k⋅L=p​k⋅logp⁡(T/k)=Θ​(k​log⁡T)S\;=\;p\,k+(L-1)\,p\,k\;=\;p\,k\cdot L\;=\;p\,k\cdot\log_{p}(T/k)\;=\;\Theta(k\log T)(13)(treatingppas constant).

Attention cost in terms ofTT.

The dense FlashAttention call onSStokens costsΘ​(S2⋅d)\Theta(S^{2}\cdot d). Substituting Eq. (13),

S2⋅d=Θ​(k2​log2⁡T⋅d),S^{2}\cdot d\;=\;\Theta\!\left(k^{2}\log^{2}T\cdot d\right),which is polylogarithmic inTTfor boundedkk.

Total per-layer compute.

Summing the contributions in Table3after substitutingS=Θ​(k​log⁡T)S=\Theta(k\log T),

Tlayer=Θ​(T⋅d)+Θ​(T​log⁡k)+Θ​(k2​log2⁡T⋅d).T_{\text{layer}}\;=\;\Theta(T\cdot d)\;+\;\Theta(T\log k)\;+\;\Theta\!\left(k^{2}\log^{2}T\cdot d\right). For boundedkk, the polylog term is sub-linear inTTand theΘ​(T⋅d)\Theta(T\cdot d)term from the projection, pooling, scoring, and scatter stages dominates. Therefore

Tlayer=Θ(T⋅d)at boundedk.\boxed{\,T_{\text{layer}}\;=\;\Theta(T\cdot d)\;\;\text{at bounded }k.\,}(14)

Two distinct quantities.

We emphasize that two “log” factors appear at different points and should not be conflated:

• •Sub-sequence size:S=Θ​(k​log⁡T)S=\Theta(k\log T)the size of the data tensor passed to FlashAttention.
• •Per-layer compute:Θ​(T⋅d)\Theta(T\cdot d)the total flops to run one layer.

The logarithmic factor lives inSS; the total compute is linear inTTbecause onlySStokens (notTT) attend to one another, while the linear-cost stages remain linear inTT.

Comparison to other attention families.

Table4places Lighthouse alongside dense softmax, log-linear attention, and linear-attention / SSM families.

Table 4:Per-layer compute across attention families. Lighthouse and linear/SSM families share the same asymptotic class; log-linear attention is strictly worse asymptotically; dense softmax is quadratic.

Appendix CDesign Choices (extended)

The Lighthouse pipeline of Sec.3admits several non-trivial design decisions whose rationale is not obvious from the equations alone.

C.1Symmetric Q/K/V Pooling

Prior hierarchical sparse designs[36,40,41]compress only the key–value side and leave queries at full resolution. This is natural for inference, where autoregressive decoding presents one query at a time. Training, however, exposes every query in parallel: we can compress the query side too, then recover dense behavior at inference by briefly resuming with stock SDPA[5]. Lighthouse poolsQQin lockstep withK,VK,Vat every level, reducing the dense FlashAttention call from𝒪​(N​S​d)\mathcal{O}(N\,S\,d)to𝒪​(S2​d)\mathcal{O}(S^{2}\,d)and yielding coherent(Q(ℓ),K(ℓ),V(ℓ))(Q^{(\ell)},K^{(\ell)},V^{(\ell)})triples that share a representation space across levels — pooled queries route to pooled keys, producing summary–summary interactions an asymmetric pyramid cannot express. Sec.6.2verifies the symmetric design is invariant under the recovery test.

C.2Parameter-Free Scoring

The scoring functional has the widest design space; two natural candidates are (a) a dilated softmax attention over the pyramid (most faithful to “what would softmax do,” as used by the learned selectors in NSA[36]and DSA[9], but𝒪​(N2/r)\mathcal{O}(N^{2}/r)) and (b) the per-headℓ2\ell_{2}norms of the layer’s own projections,‖Qi(ℓ)‖2\|Q^{(\ell)}_{i}\|_{2}and‖Ki(ℓ)‖2\|K^{(\ell)}_{i}\|_{2}(no parameters, no Q–K interaction,Θ​(N)\Theta(N)). We adopt the projection-norm scorer. It is the cheaper of the two and the moreconservativebenchmark: a dilated-attention scorer strictly adds Q–K interaction information and can only help. Any positive result with projection-norms is therefore a lower bound on what Lighthouse can deliver, and our ablations in Appendix.Aconfirm the dilated scorer matches projection-norms within noise — evidence that theselection structure, not the scoring function, drives the long-context behavior.

C.3Selection–Attention Decoupling

Every competing selection method we are aware of fuses its selection machinery into the attention kernel: NSA[36], DSA[9], and HISA[40]each ship a custom sparse-attention kernel that reads indices from a selection step. Lighthouse places that interfaceoutsidethe kernel: selection produces a dense, contiguous sub-sequence and attention is a stock FlashAttention[5,7,29]call on it. Two consequences follow: the same kernel runs at training and inference, so there is no train-vs-serve kernel divergence; and correctness of the attention step can be checked against the dense baseline by running attention with selection disabled, which is exactly what the SDPA-resume evaluation in Sec.6.2exercises.

C.4Gradient Flow

The top-KKstep is discrete and we do not approximate it with a straight-through estimator[3]or Gumbel softmax[14,21]. Gradients flow back from the loss through the scatter, FlashAttention[5], and gather intoWQ,WK,WVW_{Q},W_{K},W_{V}via the differentiable valuesQ~,K~,V~\widetilde{Q},\widetilde{K},\widetilde{V}; selection indices and the scoring functional carry no gradient. The projections therefore learn to produce values that areuseful when selectedrather than scores that are good at selecting — avoiding the optimization pathologies (scorer collapse, scorer–attention misalignment, auxiliary-loss tuning) that learnable selectors[36,9]are prone to.

Appendix DKernel Design and Parallelism (extended)

D.1Pyramid and Scoring

The pyramid and scoring stages are not custom kernels: pooling is aview + mean, scoring is a per-tokennormfollowed byview + maxper coarser level. All ops are pointwise or reshape-plus-reduce and are fused bytorch.compileinto a single device pass.

D.2Chunked-Bitonic Top-kkKernel

Selectingkkpyramid entries out ofΘ​(N)\Theta(N)candidates is the first stage that warrants a custom kernel. A textbook bitonic top-kksorts the entire score stream in shared memory or registers, which fails at our budgets:k=4096k=4096over a pyramid of≈2​N\approx 2Nentries cannot fit in a single thread block, and a global sort serializes across the sequence. Lighthouse instead uses achunked-bitonicdesign: the score stream is partitioned into fixed-size chunks ofNchunk=2048N_{\text{chunk}}=2048scores, each chunk maintains a running top-mmbuffer (m=128m=128) updated through an in-register bitonic merge, and theN/NchunkN/N_{\text{chunk}}chunks dispatch as independent CTAs. No thread block ever holds more thanmmscores; the work is fully parallel.

The concatenated indices arrive in score-sorted order within each chunk and chunk-major order across chunks — not the causal order the attention step requires. We apply a singletorch.sortpass to re-order them by pyramid position, after which the gatheredQ~,K~,V~\widetilde{Q},\widetilde{K},\widetilde{V}form a contiguous, causally consistent sub-sequence indistinguishable in shape from a standard dense FlashAttention input. This is what makes the rest of the pipeline kernel-agnostic.

This design does not produce the same index set as a theoretical global top-kk: if thekkglobally highest-scoring entries cluster in one chunk, some are replaced by lower-scoring entries from other chunks. Read positively, this is astratifiedtop-kkthat guarantees every region of the sequence contributes some tokens, which empirically yields more balanced attention coverage than strict global top-kkand avoids selection collapse onto a narrow span.

D.3Gather and FlashAttention Dispatch

Onceℐ\mathcal{I}is produced, gather is a stocktorch.gatherfollowed by a single FlashAttention call on the gathered tensors. This is where placing selectionoutsidethe attention kernel pays off. Prior selection-based methods[36,9,40,20]embed the selection–attention interface inside the kernel: each tile reads a sparse index list and performs an indirect KV load, which forces (i) a custom forward kernel per GPU architecture, (ii) a matching custom backward that inverts the hierarchical pattern, and (iii) ongoing tile/schedule re-tuning as hardware evolves. Lighthouse’s dense sub-sequence design sidesteps all three: forward and backward are bit-for-bit identical to a standard dense Transformer’s.

D.4Context-Parallel Execution

At sequence lengths beyond 128K we train with context parallelism acrossWWdevices, each rank holding a contiguous slice and using standard ring attention. Lighthouse extends cleanly without custom collectives because its pre-attention primitives are local:*(i)the coarsest pool windowpL−1p^{L-1}(e.g. 64) is orders of magnitude smaller than the shard size (N/W=128​KN/W=128\text{K}atN=1​M,W=8N{=}1\text{M},W{=}8), so pooling and scoring need no inter-rank communication;(ii)each rank runs the chunked-bitonic top-kkon its own pyramid, producingℐr⊆𝒫r\mathcal{I}_{r}\subseteq\mathcal{P}_{r}from tokens it already owns;(iii)*the gathered sub-sequence is dense, so FlashAttention runs under standard ring attention[19]— KV shards rotate through the ring as in a fully dense long-context run, and each rank’s queries see the cross-shard context selected by every other rank’s Lighthouse pipeline. This last property is only possible because Lighthouse’s selection output is a contiguous tensor; sparse-selection kernels cannot express ring rotation without engineering specific to the sparse layout. The combined design supports 1M-token pretraining across 32 Blackwell GPUs (4 nodes × 8 GPUs, CP degree 8) with no changes to the attention kernel itself.

Appendix EDesign Ablations and Throughput (extended)

This appendix gives the per-axis ablation results in detail, the per-stage throughput decomposition, and the asymptotic-cost predictions backing each design choice. All loss numbers are the step-16,00016{,}000post-resume training loss of the two-stage Lighthouse-then-SDPA recipe (Sec.6.2); throughput numbers are aggregated over88ranks of one B200 node at98,30498{,}304-token context (matching Table1).

E.1Scorer Variants

We compare the parameter-free projection-norm scorer against the dilated softmax-attention scorer at fixedL=3,p=4L{=}3,\,p{=}4. Atk=1536k{=}1536, dilated reaches0.68810.6881while projection-norm reaches0.69460.6946; atk=2048k{=}2048the order reverses, with0.69690.6969for dilated and0.69210.6921for projection-norm. The two are within∼\sim0.010.01of each other in both directions and neither is uniformly better. Combined with projection-norm’s lower compute cost (179.6179.6–180.9180.9B200-h vs.197.2197.2–199.7199.7for dilated, an∼\sim9%9\%saving) and absence of additional learnable parameters (Sec.C.2), projection-norm is the throughput-sensitive default; dilated remains a defensible alternative when the lowest absolute loss matters.

E.2Pooling Factor

HoldingL=3L{=}3and varyingp∈{2,4,8}p\in\{2,4,8\}atk=1536k{=}1536gives final losses of0.68250.6825,0.68810.6881,0.68280.6828respectively; atk=2048k{=}2048,p=2p{=}2andp=4p{=}4give0.68800.6880and0.69690.6969. Pooling factor has a small effect:p=2p{=}2andp=8p{=}8are essentially indistinguishable, withp=4p{=}4marginally worse by0.0050.005–0.0150.015. We adoptp=2p{=}2as the default;p=4p{=}4paired with the projection-norm scorer is the wall-clock-favoured alternative.

E.3Number of Levels

Holdingp=2p{=}2and varyingL∈{3,4,5}L\in\{3,4,5\}produces a monotonically increasing final loss: atk=1536k{=}1536,L=3L{=}3gives0.68250.6825,L=4L{=}4gives0.69780.6978, andL=5L{=}5gives0.69910.6991; the ordering carries tok=2048k{=}2048(0.6880→0.6983→0.70430.6880\to 0.6983\to 0.7043). Deeper pyramids spread the same selection budget over more coarse levels, leaving fewer entries at the finest level where the model is most sensitive.L=3L{=}3is best across both selection budgets and we adopt it as the default.

E.4Top-KKBudget

HoldingL=3,p=2L{=}3,\,p{=}2and varyingk∈{1,536,2,048,3,072,4,096,6,144}k\in\{1{,}536,2{,}048,3{,}072,4{,}096,6{,}144\}gives final losses of0.6825,0.6880,0.6890,0.6951,0.68310.6825,\,0.6880,\,0.6890,\,0.6951,\,0.6831. The trend is monotonic overk≤4096k\leq 4096and dips back atk=6144k{=}6144. A larger selection budget doesnottranslate to lower post-resume loss within the range we tested. Sparser configurations may regularise against our relatively small training-token budget; investigating whether this trend reverses at much larger budgets is left to future work. The practical implication is thatk=1536k{=}1536is preferred from both a loss and a wall-clock standpoint.

Figure 4:Needle-in-a-Haystack at9898K training, step16,00016{,}000.Four Lighthouse→\toSDPA configurations (varyingk∈{1536,2048}k\in\{1536,2048\}and scorer∈{\in\{dilated, norm}\}atL=3,p=4L{=}3,\,p{=}4) and the dense SDPA-from-scratch baseline (bottom). Each cell is the mean retrieval rate over1010single-digit passkeys at the given (context, depth); the per-panel mean is shown in each title. Random chance is10%10\%.

E.5Throughput Decomposition

Stage-1 throughput.

The recoverability runs (k=6144,L=3,p=2k{=}6144,\,L{=}3,\,p{=}2, dilated) sustain74.774.7–75.475.4k tok/s/GPU. Across the ablation grid (varyingk,p,Lk,p,Lat the9898K context) the stage-1 throughput range is83.583.5–126.0126.0k tok/s/GPU, vs.∼\sim4646k for dense SDPA. The intra-grid trends follow theS=N/pL−1+(L−1)​p​kS=N/p^{L-1}+(L-1)pkprediction: raisingppfrom22to44atk=2048,L=3k{=}2048,L{=}3lifts throughput from90.990.9to97.197.1k; raisingLLfrom33to44atk=2048,p=2k{=}2048,p{=}2lifts it from90.990.9to94.594.5k (deeper pyramid spreads the budget across more coarse levels). Swapping the dilated scorer for projection-norm atL=3,p=4,k=1536L{=}3,\,p{=}4,\,k{=}1536accelerates stage-1 from99.599.5to126.0126.0k tok/s/GPU because the scorer no longer runs an attention pass over the dilated pyramid.

End-to-end speedup.

Aggregating both stages of the1010k++66k recipe across the ablation grid, end-to-end runtime ranges from22.522.5h (179.6179.6B200-h; normk=1536,p=4,L=3k{=}1536,\,p{=}4,\,L{=}3) to27.027.0h (215.7215.7B200-h;k=4096,p=2,L=3k{=}4096,\,p{=}2,\,L{=}3) against37.937.9h (303.2303.2B200-h) for dense-SDPA-from-scratch on the same16,00016{,}000-step /50.350.3B-token budget: a1.40×1.40\timesto1.69×1.69\timeswall-clock speedup at matched or lower final loss. The saving comes entirely from stage-1; the SDPA-resume tail runs the same kernel as the baseline and matches its throughput.

Appendix FLong-Context Retrieval (NIAH)

At 530M parameters and only16,00016{,}000training steps (∼\sim50.3B tokens), full prose Needle-in-a-Haystack scores near-zero across the board, so we adopt a simplified single-digit variant that isolates the retrieval signal. To complement the loss-based recoverability evaluation in Sec.6.2, we run this test over five step-16,00016{,}000checkpoints: four Lighthouse→\toSDPA two-stage runs (varyingk∈{1536,2048}k\in\{1536,2048\}and scorer∈{\in\{dilated, norm}\}atL=3,p=4L{=}3,\,p{=}4) and the dense-SDPA-from-scratch baseline at matched compute and tokens. Inference uses dense SDPA in every case (Sec.6.1).

Protocol.

A single passkey digit (one of{0,1,…,9}\{0,1,\dots,9\}) is hidden in random alphanumeric filler at depths{0,15,30,50,70,85,100}%\{0,15,30,50,70,85,100\}\%across context lengths{4,8,16,32,64,96}\{4,8,16,32,64,96\}K. For each cell we run one forward pass over the full prompt and take an argmax restricted to the 10 digit tokens at the last position. We average the 0/1 score over the full digit sweep, so each cell reports the mean retrieval rate over1010trials; random chance is10%10\%.

Findings.

Three of the four Lighthouse runs are at or above the SDPA-from-scratch baseline (mean retrieval0.720.72):k=2048k{=}2048dilated wins overall at0.760.76,k=1536k{=}1536dilated reaches0.730.73, andk=2048k{=}2048norm matches the baseline at0.720.72; onlyk=1536k{=}1536norm dips, to0.650.65(Fig.4). Two patterns emerge. First, largerkkis the dominant axis:k=2048≥k=1536k{=}2048\geq k{=}1536for both scorers, with a gap of0.030.03(dilated) or0.070.07(norm). Second, the norm scorer hurts retrieval more than it hurts training loss: at fixedkk, switching dilated to norm costs0.040.04atk=2048k{=}2048and0.080.08atk=1536k{=}1536, the largest single-axis gaps in the grid. Combined with the loss-side finding that smallerkkregularises better, the right default depends on whether the downstream task is loss- or retrieval-driven.