Lighthouse Attention Technique

Long-context pretraining is bottlenecked by attention's O(N^2) cost. Lighthouse Attention is a selective hierarchical attention that applies symmetric pooling to Q, K, V at the same ratio to build an L-level pyramid, scores each entry with L2 norm, and keeps only the top-K. Without a learned scorer head, Gumbel-softmax, STE, or auxiliary loss, it groups only the selected tokens into a contiguous dense sub-sequence and runs FlashAttention on it as-is, so training and inference use the same kernel and automatically benefit from upstream FlashAttention improvements. The implementation is layered on top of upstream torchtitan with two files and about 600 lines of modification.

There are three key points.

• Selection logic is handled outside the kernel. It attaches ordinary FlashAttention after torch.gather and torch.sort, aligning it with the dense attention path.
• Symmetric pooling lets the pooled query and pooled key share the same representation space.
• The sorted sub-sequence becomes a hole-free causal sequence, so the standard lower-triangular mask can be used as-is.

Training happens in two stages. First, most steps are trained with Lighthouse, then a short SDPA-resume run is done with the same optimizer state and dataloader continuation to verify that dense attention can be recovered. Across all three splits — 10k+6k, 11k+5k, and 12k+4k — loss spiked by 1.12–1.57 nats immediately after resuming but recovered within 1k–1.5k SDPA steps, and the final loss of 0.6980–0.7102 was lower than the dense-from-scratch baseline of 0.7237. At 50B tokens / 16k steps, this saved 75–106 B200-hours, confirming that sparse training does not damage dense attention capability.

The performance gap is also clear. On a single B200, forward+backward at 512K context was about 17× faster than standard attention, and 98K context pretraining was 1.4–1.7× faster end-to-end. In experiments with 530M Llama-3, 16k optimiser steps, on 8×B200 single-node, stage 1 achieved 84–126k tokens/s/GPU, while dense SDPA achieved about 46k. Beyond 100K, the 530M architecture OOMs on a single B200, so it was scaled with context parallelism (CP), and 1M-token training was performed on 32 B200s. There is also a deterministic int-atomic kernel for scatter-back for reproducibility, but it is 1.2–2× slower, so fp-atomic is the default.