PaperlayerView exact revision on arXiv

Source-linked AI summary

PHBA: Prefix-State Hybrid Block Attention

Ruijie Li, Jiaxi Hu, Shiyu Wang, Yuxuan Liang

arXiv:2610.08527v1cs.LG

TL;DR

Hybrid models pair efficient sequence processing with precise retrieval, but fixed-window exact attention limits token-level access to local context. PHBA replaces that window with block retrieval paired with prefix states, and experiments report consistent retrieval and long-context gains while preserving competitive language-modeling quality and training efficiency.

  • Problem

    Hybrid sequence models balance efficient long-context processing with precise retrieval, but sliding-window exact attention restricts token-level access to a fixed local window.

    Source: §1 · p. 1

  • Method

    PHBA pairs content-dependent retrieval of historical token blocks with gated, block-aligned prefix states summarizing the context preceding each block.

    Source: §1 · p. 2

  • Results

    PHBA consistently improves retrieval and long-context understanding while preserving competitive language-modeling quality and training efficiency.

    Source: §5 · p. 9

  • Takeaways & Limitations

    Block-aligned prefix memories provide broader historical context around retrieved content without maintaining a separate recurrent state for every token.

    Source: §1 · p. 2

Abstract

from arXiv · show

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.

1 INTRODUCTION

PHBA combines content-dependent retrieval of distant token blocks with block-aligned prefix states that summarize the context preceding each block. Its implementation streams retrieved tokens and states efficiently, and experiments report stronger retrieval and long-context modeling than recurrent, sparse-attention, and hybrid baselines.

1 INTRODUCTION

  • Dense attention offers precise access to history but scales quadratically, while fixed-size recurrent states are efficient but can limit precise retrieval; NHA’s exact attention is restricted to a fixed local window (Du et al., 2026).NHA’s sliding-window attention leaves key–value pairs outside the window accessible only through compressed recurrent state.Source: §1 · p. 1
  • PHBA replaces fixed sliding-window attention with content-dependent block retrieval, pairing each selected exact-token block with a gated-recurrence prefix state summarizing its preceding causal history.The prefix state is aligned to the same block structure used for sparse routing.Source: §1 · p. 2
  • Aligning recurrent memory to routed blocks exposes broader historical context around retrieved content without maintaining a separate recurrent state for every token.PHBA selects block-aligned prefix memories together with distant exact-token blocks.Source: §1 · p. 2
  • A hardware-aware Triton implementation streams routed token blocks and prefix states from original layouts, using compact routing metadata, chunkwise prefix construction, and online-softmax accumulation.The implementation avoids materializing packed sparse payloads and supports long-context training and inference.Source: §1 · p. 2
  • Experiments report improved retrieval and long-context modeling over strong recurrent, sparse-attention, and hybrid baselines while maintaining favorable efficiency at long sequence lengths.Source: §1 · p. 2

2 PRELIMINARY

This section reviews block-sparse attention and gated slot memory, the two components PHBA builds on, and introduces notation for the method.

2.1 BLOCK-SPARSE ATTENTION

  • Content-dependent routing selects relevant historical blocks, and each query attends to tokens in its selected blocks rather than the full sequence.PHBA adopts this block-routed sparse-attention setting; its routing procedure is described in Section 3.Source: §2.1 · p. 2

2.2 GATED SLOT MEMORY

  • PHBA stores gated key/value slots only at block boundaries as prefix states, then jointly attends to them and routed exact-token blocks.The shared slot-wise gate keeps key and value memories temporally aligned; this recurrence is a gated linear attention specialization with a chunkwise-parallel formulation.Source: §2.2 · p. 3

3 PREFIX-STATE HYBRID BLOCK ATTENTION

PHBA combines top-k retrieval of exact historical token blocks with gated recurrent prefix states aligned to those blocks. It jointly normalizes token and state evidence through an online-softmax merge without materializing a combined attention matrix.

3 PREFIX-STATE HYBRID BLOCK ATTENTION

  • A single PHBA layer combines full-resolution historical token blocks with compact recurrent prefix states aligned to those blocks.The two memory types share a common softmax normalization, evaluated without explicitly materializing a combined attention matrix.Source: §3 · p. 3

3.1 TOP-K BLOCK ROUTING

  • For each query, PHBA routes to historical blocks by similarity between the query and each block’s mean RoPE-transformed key, then attends to the selected blocks’ original tokens.The current block supplies local causal attention, while query-dependent routing provides access to distant historical blocks.Source: §3.1 · pp. 3–4

3.2 PREFIX-STATE CONSTRUCTION

  • Gated recurrences build separate prefix-key and prefix-value states, which are stored immediately before each block and summarize only its preceding causal context.The key and value recurrences use the same slot-wise retention and write gates, giving corresponding slots identical temporal weighting.Source: §3.2 · p. 4

3.3 ONLINE-SOFTMAX TOKEN–STATE ATTENTION

  • The online-softmax merge is equivalent to one softmax over exact tokens and prefix-state slots, allocating probability mass directly between them without concatenation or external fusion weights.Token attention uses RoPE-transformed queries and keys; prefix-state attention uses the raw query and raw prefix-state keys.Source: §3.3 · pp. 4–5

3.4 HARDWARE-EFFICIENT CHUNKWISE AND SPARSE EXECUTION

  • Chunkwise updates parallelize writes within each chunk while carrying only boundary states recurrently; sparse kernels load routed queries, token blocks, and prefix states directly from HBM.This avoids token-wise recurrent execution within each chunk and avoids gathering sparse payloads into packed tensors before attention.Source: §3.4 · p. 5

4 EXPERIMENTS

PHBA delivers strong short-context quality and retrieval, leading long-context retrieval and understanding, while maintaining efficient training. Ablations show how block granularity, prefix-state capacity, and gated state construction shape its quality–efficiency trade-offs.

4.1 EXPERIMENTAL SETUP

  • Models are trained from scratch on FineWeb-Edu at 8K context in 760M/50B-token and 1.3B/100B-token settings, then evaluated on language modeling, reasoning, retrieval, and understanding.Models trained at 8K are evaluated directly at 16K–64K; LongBench inputs beyond 8K are truncated to the first and last 4K tokens. Baselines use matched data, context length, and model scale.Source: §4.1 · p. 6

4.2 SHORT-CONTEXT LANGUAGE MODELING AND RETRIEVAL

  • At 1.3B, PHBA achieves the best WikiText and LAMBADA results, a 60.07 commonsense average second only to KDA, and beats other efficient baselines on 4K/8K retrieval.On NIAH, it trails only the dense Transformer.Source: Table 1 · p. 7, §4.2 · p. 7

4.3 LONG-CONTEXT RETRIEVAL AND UNDERSTANDING

  • PHBA has the strongest overall RULER and NIAH performance from 16K to 64K and the strongest LongBench average, with gains across document QA, summarization, and few-shot tasks.Its retrieval advantage becomes more pronounced beyond the 8K training length, while it remains competitive on code.Source: Figure 2 · p. 8, §4.3 · p. 8

4.4 ABLATION AND DESIGN ANALYSIS

  • Ablations favor C = 128 and K = 8 as a quality–throughput trade-off, M = 128 for stable performance, and learned gated recurrence over ungated or mean aggregation.Finer blocks generally improve retrieval but reduce throughput, and performance does not improve monotonically with state size. Unlike MoBA or NHA, PHBA pairs routed historical blocks with aligned prefix states, producing stronger retrieval and long-context performance.Source: Table 1 · p. 7, Figure 2 · p. 8, §4.4 · p. 8

4.5 TRAINING EFFICIENCY.

  • Under a matched 1/8 budget, PHBA is 18.4–34.6% faster than NHA at 2K; NHA runs out of memory from 4K onward in the evaluated implementation.With C = 256, adding prefix states costs 6.7–11.9% throughput versus MoBA across 2K–16K; with K = 4, PHBA reaches 8.79 Kt/s at 16K, 12.8% above full attention.Source: Figure 3 · p. 9, §4.5 · p. 9

4.6 RELATED WORK.

  • The related-work discussion situates PHBA among efficient sequence models, sparse attention, and hybrid architectures by pairing content-routed historical blocks with block-aligned prefix states.Source: §4.6 · p. 9

5 CONCLUSION

PHBA combines block-sparse retrieval with recurrent prefix states and hardware-aware execution for efficient long-context modeling, with experiments showing retrieval and long-context improvements while preserving competitive language modeling quality and training efficiency.

5 CONCLUSION

  • PHBA combines content-dependent block-sparse retrieval with block-aligned recurrent prefix states, and its hardware-aware implementation supports efficient chunkwise construction and sparse execution.The model attends jointly to routed historical tokens and their preceding compressed context within a unified softmax.Source: §5 · p. 9

AI USE STATEMENT

  • Generative AI tools were used only to aid and polish writing, not for research design, methodology, experiments, data analysis, or generating results.Authors reviewed all AI-assisted content and take full responsibility for the final work.Source: AI USE STATEMENT · p. 10

ETHICS STATEMENT

  • The work reports no human subjects or sensitive personal data collection and identifies no specific ethical concerns beyond those generally associated with large language models.Experiments use publicly available training and evaluation datasets.Source: ETHICS STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

  • The appendix provides derivations, pseudocode, experimental settings, throughput protocol, and full results to support reproduction; source code and implementation will be public upon acceptance.The main paper specifies routing, prefix-state construction, token–state attention, and efficient execution.Source: REPRODUCIBILITY STATEMENT · p. 10

A APPENDIX · A.1 PHBA OPERATOR PSEUDOCODE

The appendix defines PHBA’s block and tile terminology and describes kernels that route directly to original HBM layouts, with training metadata reformatted for sparse backward computation.

A.1 PHBA OPERATOR PSEUDOCODE

  • C denotes the token block size in the appendix’s PHBA kernel notation.Source: §A.1 · p. 12
  • M denotes the number of prefix-state slots.Source: §A.1 · p. 12
  • H denotes the number of heads.Source: §A.1 · p. 12
  • A logical block can be either a token block or a prefix-state block.Source: §A.1 · p. 12
  • A physical tile is a smaller query, memory, or feature tile that one Triton program loads from HBM into SRAM or registers.Source: §A.1 · p. 12
  • PHBA avoids packed sparse payload tensors, consuming query-centric historical routes to load routed token and prefix-state blocks from their original HBM layouts.The forward kernel uses this direct-loading approach.Source: §A.1 · p. 12
  • During training, routing decisions are reformatted as block-major integer metadata for efficient sparse backward computation.The generic symbols are ρ, Offset, and Cblk.Source: §A.1 · p. 12

A.1.1 CHUNKWISE PARALLEL PREFIX-STATE MEMORY CONSTRUCTION

PHBA constructs boundary prefix key/value memories with a gated recurrence, then uses sparse readout to select among them. Its chunkwise form enables parallel work within chunks while carrying boundary states across chunks.

A.1.1 CHUNKWISE PARALLEL PREFIX-STATE MEMORY CONSTRUCTION

  • PHBA separates memory formation from sparse readout: a gated recurrence builds boundary prefix key/value memories that the attention layer later selects.Source: §A.1.1 · p. 13
  • For packed inputs, the recurrence resets at each sequence boundary, giving the first block of every sequence a zero incoming state.Source: §A.1.1 · p. 13
  • For parallel training, only boundary states are sequential across chunks; within-chunk writes use matrix products, yielding one prefix state per boundary and log-A prefix sums for stability.Source: §A.1.1 · p. 13

A.1.2 ROUTING METADATA KERNELS

PHBA’s routing kernels build per-block centroids, select historical top-K routes, and invert them into block-major metadata without materializing dense query–centroid scores in HBM. They also mark whether each block has a nonempty causal prefix.

A.1.2 ROUTING METADATA KERNELS

  • The three-stage routing pipeline computes block centroids, selects top-Kh historical block ids with tiled scores, and reformats routes as block-major metadata; it also records PrefixValid for nonempty causal prefixes.The dense query–centroid score matrix is never materialized in HBM.Source: §A.1.2 · p. 13
  • Centroids are computed per block and attention head within each packed sequence, averaging only valid rows after streaming RoPE keys and accumulating their sum in fp32 registers.The block boundaries are induced by cu and the block size.Source: §A.1.2 · pp. 13–14
  • Historical selection scans centroid tiles and keeps top-Kh routes per query and head, restricted to strictly preceding blocks in the same packed sequence; unavailable entries remain −1.Route indices are non-differentiable metadata.Source: §A.1.2 · p. 14
  • Block-major metadata is formed by counting assignments per logical block–head pair, prefix-summing counts into contiguous ranges, and scattering query–head row ids; payloads stay in their original HBM layouts.Only integer metadata is materialized.Source: §A.1.2 · p. 14

A.1.3 QUERY-TILE FUSED SPARSE FORWARD KERNEL

The fused forward kernel processes query-centric routes by extending block-local causal FlashAttention accumulators with prefix states and routed historical blocks. It uses a fixed-capacity route schedule to share loads within each query tile while keeping payload tensors in their original HBM layouts.

A.1.3 QUERY-TILE FUSED SPARSE FORWARD KERNEL

  • The forward pass extends block-local causal FlashAttention output and log-sum-exp accumulators with the current-block prefix state, routed historical token blocks, and aligned prefix states.Algorithm 6 operates directly from the query-centric routes Rc.Source: §A.1.3 · p. 14
  • A leader mask marks the first valid occurrence of each historical block id, allowing one loaded block to serve all rows in the query tile that route to it.Historical route entries use a fixed-capacity schedule rather than dynamic compaction.Source: §A.1.3 · p. 14
  • The token branch uses (Qrope, Krope, V), while the state branch uses (Qraw, P K, P V), and payload tensors remain in their original HBM layouts.Source: §A.1.3 · p. 14

A.1.4 SPARSE BACKWARD KERNEL

The backward pass uses a block-major schedule to enumerate query–head rows for token and prefix-state blocks. It recomputes sparse scores and probabilities on chip from the final joint log-sum-exp, avoiding materialized sparse probability tensors.

A.1.4 SPARSE BACKWARD KERNEL

  • The block-major backward schedule uses ρ and Offset to enumerate query–head rows that access each token or prefix-state block.The two logical memories contain |B_j| token rows and M prefix-state slots.Source: §A.1.4 · p. 15
  • Sparse scores and probabilities are recomputed on chip from the final joint log-sum-exp L, avoiding materialized sparse probability tensors.Source: §A.1.4 · p. 15
  • The two logical memories contain |B_j| token rows and M prefix-state slots, respectively.Source: §A.1.4 · p. 15

A.2 ADDITIONAL EXPERIMENTAL DETAILS … A.2.4 EFFICIENCY BENCHMARK PROTOCOL.

The appendix specifies PHBA’s training and baseline setup, evaluation protocols, sparse-attention forward and backward implementations, and controlled efficiency benchmarking conditions.

A.2.1 MODEL ARCHITECTURE AND TRAINING.

  • Models are trained from scratch on FineWeb-Edu at 8K context length, with 760M and 1.3B scales and PHBA defaults M = 128 and C = 128.The 760M and 1.3B models train on 50B and 100B tokens; the 1.3B main comparison uses a 1/8 exact-token budget, or 1,024 tokens per query (K = 8). Training uses AdamW, cosine decay, bfloat16, and eight NVIDIA B200-192GB GPUs.Source: §A.2.1 · p. 15

A.2.2 BASELINE CONFIGURATION.

  • Baselines are Transformer (Vaswani et al., 2017), MoBA (Lu et al., 2026), GSA (Zhang et al., 2024), KDA (Team et al., 2025), and NHA (Du et al., 2026).All use standalone architectures, matched data and context length at comparable model scale, recommended configurations, and matched exact-token budgets for sparse attention where applicable.Source: §A.2.2 · p. 15
  • Algorithm 6 processes query tiles while streaming routed token blocks and current or aligned prefix-state tiles, updating online-softmax accumulators on chip.It writes the finalized outputs and log-sum-exp values to HBM.Source: §A.2.2 · p. 16

A.2.3 EVALUATION PROTOCOL.

  • Short-context evaluation covers WikiText and LAMBADA language modeling plus zero-shot ARC-Easy/Challenge, HellaSwag, PIQA, and WinoGrande reasoning.Evaluations use lm-evaluation-harness.Source: §A.2.3 · p. 16
  • Models trained at 8K are evaluated without further training at 16K, 32K, and 64K using RULER and LongBench.RULER includes common-word extraction, HotpotQA- and SQuAD-based question answering, and controlled NIAH variants; LongBench inputs over 8K tokens retain the first and last 4K tokens.Source: §A.2.3 · pp. 16–17
  • Algorithm 7 computes token and state gradients with block-local attention backward, then propagates prefix-state adjoints backward through the gated recurrence.It recomputes chunk-local retention factors and combines token-path and state-construction gradients, including gradients through retention gates.Source: §A.2.3 · p. 17

A.2.4 EFFICIENCY BENCHMARK PROTOCOL.

  • Efficiency tests use one H20-96GB GPU, matched hardware and BF16, and measure end-to-end training throughput at 2K–16K sequence lengths.Sparse methods use a matched 1/8 exact-token budget; measurements include forward, backward, and parameter updates but exclude input preparation and data loading, and are taken after warm-up. NHA’s 1/8 configuration is OOM from 4K onward in the evaluated implementation.Source: §A.2.4 · pp. 17–18
  • MoBA and PHBA are benchmarked with fixed-K or fixed-block parameterizations, while SWA and NHA use fixed-window or matched-budget configurations.MoBA/PHBA use K = 4 with sequence-length-adjusted block size or C = 256 with adjusted K; SWA/NHA use a fixed window of 64 or a window scaled to the 1/8 budget. All state-based methods use M = 128, while Transformer uses full causal attention.Source: §A.2.4 · p. 18

Related papers

Loading 2610.08527v1…