PHBA: Prefix-State Hybrid Block Attention

Ruijie Li; Jiaxi Hu; Shiyu Wang; Yuxuan Liang

cs.LGOct 6, 2026

PHBA: Prefix-State Hybrid Block Attention

Authors: Ruijie Li, Jiaxi Hu, Shiyu Wang, Yuxuan Liang

Organizations: The Hong Kong University of Science and Technology (Guangzhou) · Independent researcher

Abstract

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.

Figures & tables

Figure 1: Overview of PHBA. The RoPE query $\bm{q}_{t}^{\mathrm{rope}}$ first performs Top- $K$ Block Routing to select relevant historical blocks. PHBA attends to the retrieved exact tokens with $\bm{q}_{t}^{\mathrm{rope}}$ and their Block-Aligned Prefix States with $\bm{q}_{t}^{\mathrm{raw}}$ , and combines both memory types through Online-Softmax Token–State Attention to produce $\bm{o}_{t}$ .

Figure 2: Long-context results and ablations. (a)–(b) RULER/NIAH averages at 16K/32K/64K; (c) LongBench average. (d)–(f) Ablations of block granularity, prefix-state capacity, and gated state construction; (d) also reports throughput (1.3B, BS=1, 8K). Bars denote 16K/32K/64K results; shadows mark default settings.

Panel of Figure 3 — Figure 3: Training efficiency on a single NVIDIA H20-96GB GPU. For MoBA and PHBA, both configurations maintain a fixed $1/8$ sparse token budget: $K=4$ dynamically adjusts the block size $C$ , whereas $C=256$ dynamically adjusts the number of routed blocks $K$ . SWA ( $1/8$ ) and NHA ( $1/8$ ) similarly scale the window size with sequence length, while SWA and NHA use a fixed window size of $64$ . All state-based models use $M=128$ memory slots. Zero throughput for NHA ( $1/8$ ) denotes out-of-memory (OOM).

Explore similar work

HydraHead: From Head-Level Functional Heterogeneity to Specialized Attention Hybridization
Jun 18, 2026Zhentao Tan, Wei Chen, Jingyi Shen +4Efficient Long-Context Inference Layer-Wise
Alibaba Group
HLA: Expressive Hybrid Linear Attention via Chunk-Wise Dynamic Mixing
Oct 5, 2026Zhuokun Chen, Xi Lin, Xiyu Wu +3Kimi Delta Attention Efficient Long-Context Inference
Monash University · Zhejiang University
Tail-Replay: Escaping the Curse of Linear Attention in Prefix Caching for Hybrid LLMs
Aug 31, 2026Yirui Liu, Ruoling Qi, Xuaner Wu +2Kimi Delta Attention Linear Attention
Institute of Artificial Intelligence, China Telecom (TeleAI) · Shanghai Jiao Tong University

Jun 18, 2026, cs.CLJ/K move · Enter open · S save

HydraHead: From Head-Level Functional Heterogeneity to Specialized Attention Hybridization

Zhentao Tan, Wei Chen, Jingyi Shen, Yao Liu+3

Alibaba Group

Efficient Long-Context Inference Layer-Wise Heterogeneity

Open paper AskPDF arXiv

The quadratic complexity of attention poses a critical bottleneck for long-context processing, spurring interest in hybrid attention designs. Most open-source hybrid models adopt a layer-wise strategy. Yet, prior work has noted the inherent difficulty of integrating Linear Attention (LA) with Full Attention (FA), suggesting that the design space of attention hybridization remains underexplored. To probe this space, we conduct interpretability analysis and observe that layers exhibit block-wise functional similarity, while individual heads within the same layer display distinct functional specialization despite sharing input features. This head-level heterogeneity suggests that the head dimension provides a natural and principled granularity for fusing heterogeneous attention signals. Building on this insight, we introduce HydraHead, a novel architecture that hybridizes FA and LA along the head axis. HydraHead features two key innovations: (1) an interpretability-driven selection strategy that identifies retrieval-critical heads and preserves FA only for them, and (2) a scale-normalized fusion module that reconciles the distributional gap between FA and LA head outputs. By leveraging a three-stage transfer pipeline with parameter reuse and distillation, we achieve high-performance hybrid models with minimal training overhead. Under a unified training setup, HydraHead outperforms other hybrid designs in long-context tasks while maintaining strong general reasoning. With interpretability-driven head selection, it matches a 3:1 layer-wise hybrid's long-context performance at a 7:1 LA-to-FA ratio. Crucially, trained on only 15B tokens, HydraHead achieves over 69% improvement over the baseline at 512K context length, approaching Qwen3.5, a leading model of comparable size with a native context length of 256K. This highlights the significant scaling potential of head-level hybridization.

Abstract

Explore similar work

HydraHead: From Head-Level Functional Heterogeneity to Specialized Attention Hybridization

HLA: Expressive Hybrid Linear Attention via Chunk-Wise Dynamic Mixing

Tail-Replay: Escaping the Curse of Linear Attention in Prefix Caching for Hybrid LLMs