Search papers, labs, and topics across Lattice.
PHBA introduces a novel hybrid attention mechanism that replaces fixed local windows with top-k block-sparse retrieval, augmenting each retrieved block with a compact prefix state summarizing preceding context via a gated linear recurrence. This design enables precise long-range evidence retrieval and compressed historical context integration within a single layer, outperforming existing linear and hybrid attention models on long-context and retrieval tasks. A hardware-aware Triton implementation further optimizes efficiency by streaming routed blocks and states, avoiding large intermediate tensor materialization.
Precise long-range evidence retrieval is unlocked by coupling block-sparse attention with learned prefix states, overcoming the fixed local window limitations of prior hybrid architectures.
Hybrid architectures combining linear sequence models with softmax attention provide an effective balance between efficient long-context modeling and precise token retrieval. Existing designs such as Native Hybrid Attention (NHA) combine compressed long-term states with sliding-window attention, but their exact attention is restricted to a fixed local window. In this work, we introduce Prefix-State Hybrid Block Attention (PHBA), which replaces local sliding-window attention with top-k block-sparse retrieval and couples each retrieved block with a compact prefix state summarizing its preceding context. The prefix states are constructed by a gated linear recurrence at block boundaries and retrieved together with the corresponding token blocks, allowing the model to combine precise long-range evidence with compressed historical context within a unified layer. We further develop a hardware-aware Triton implementation that streams routed token blocks and prefix states without materializing large intermediate tensors. Experiments show that PHBA improves long-context and retrieval performance over strong linear and hybrid baselines while retaining efficient training and inference.
