InfLLM v2 Sparse Attention

• The paper introduces a novel two-stage blockwise sparse attention mechanism that dynamically selects top-k blocks, achieving up to 7× speedup in decoding long contexts.
• It employs blockwise semantic pooling and hardware-aware optimizations to significantly reduce memory usage while maintaining competitive model performance.
• Empirical results demonstrate that groupwise mask sharing and dynamic selection allow each query to attend to roughly 5% of tokens without retraining dense kernels.

InfLLM v2 Sparse Attention is a trainable, two-stage sparse attention framework designed to enable LLMs to efficiently process long contexts with drastically reduced computational and memory overhead. By employing dynamic top-k block selection over mean-pooled semantic kernels, and grouping query heads to share the same sparsity mask, InfLLM v2 enables both prefill (context processing) and decoding (autoregressive generation) acceleration without the need for retraining or significant modification of the dense attention kernels. This approach features hardware-aware optimizations and achieves high sparsity rates while maintaining competitive model performance across long-context benchmarks.

1. Sparse Attention Fundamentals and Theoretical Justification

Sparse attention methods aim to reduce the inherent $O(N^2)$ complexity of dense attention by restricting each query to attend only a subset of keys. The theoretical underpinning for such sparsification is that, under common assumptions (e.g., LayerNorm-induced Gaussianity of queries/keys), the softmax attention matrix is already inherently sparse: most entries are negligibly small, with only $O(n^C)$ (for some $C\in(0,1)$) significant contributions per row. This “natural attention sparsity” (Deng et al., 3 Apr 2024) provides justification for aggressive top-$k$ or blockwise selection strategies that keep approximation errors provably bounded.

InfLLM v2 employs this insight and adopts a blockwise, content-adaptive selection scheme, reinforced by theoretical results showing that keeping the largest $k=O(d)$ entries attains nearly full representational power due to the Carathéodory theorem (Sason et al., 3 Mar 2025).

2. Two-Stage Blockwise Sparse Attention Mechanism

InfLLM v2’s mechanism comprises two sequential stages:

1. Blockwise Semantic Pooling and Scoring:
   • The context is partitioned into equal-length blocks (size $m$), yielding $l/m$ blocks for an input of length $l$.
   • For each block $S$, a semantic kernel is computed by mean pooling the key vectors:

   $\text{Mean}(S) = \frac{1}{|S|} \sum_{t \in S} k_t$

• For each query $q_i$, the similarity to each semantic kernel $S$ is evaluated as

  $$r_{\text{kernel}}(q_i, S) = \text{softmax}(q_i \cdot \text{Mean}(S))$$

• For each block $B_j$, containing or overlapping kernels $\mathcal{S}$, set

  $$r_{\text{block}}(q_i, B_j) = \max_{S \in \mathcal{S}: S \cap B_j \ne \emptyset} r_{\text{kernel}}(q_i, S)$$

1. Top-$k$ Block Selection and Groupwise Mask Sharing:
   • For each query (or head/group), select the $k$ blocks with the highest $r_{\text{block}}$ scores, forming a block-level sparse set $\mathcal{K}_i$.
   • Attention is then computed only using the tokens from those selected blocks.
   • To maximize memory and compute efficiency, block selection is shared across a group of queries or heads.

This design ensures each query attends exclusively to a small, dynamically determined subset of context tokens, reducing both arithmetic operations and cache accesses.

3. Hardware-Aware Implementation and Integration

InfLLM v2 is implemented in the CUDA-optimized CPM.cu inference system within the MiniCPM4 framework (Team et al., 9 Jun 2025). The blockwise selection is fused into a dedicated kernel, statically managing memory and capitalizing on data reuse. Kernel fusion allows for groupwise sharing of memory access, while speculative sampling and quantization further minimize runtime overhead.

A table summarizing the critical steps:

Stage	Operation	Complexity / Motivation
Block pooling	Mean across block key vectors	$O(l)$
Block scoring	Query–semantic kernel dot-product, softmax	$O(l)$ per query
Top-$k$ sel.	Select $k$ blocks with highest scores	$O(b\log b)$, $b=l/m$
Attention	Compute only over selected tokens per query	$O(km)$ per query
Group share	Group queries/heads share sparsity mask	Reduces redundant block access

This approach is structurally similar to blockwise and windowed patterns in attention (cf. MoA (Fu et al., 21 Jun 2024), S2-Attention (Lin et al., 25 Jul 2024)), but distinct in its trainable, content-adaptive block selection and focus on extreme memory and compute reduction for long-context inference.

4. Comparison with Alternative Sparse Attention Approaches

InfLLM v2 distinguishes itself from static, pattern-based, or purely heuristic sparse attention methods:

• Blockwise mean pooling enables rapid, parameter-free semantic summarization of long contexts, avoiding the cost of token-level pairwise computation as in classical top-$k$.
• Dynamic block selection, reinforced at inference time, provides accuracy on par with dense attention for long sequences, supported by experimental results showing marked advantages in both throughput and resource usage over baselines such as Qwen3-8B (Team et al., 9 Jun 2025).
• Groupwise mask sharing is motivated by empirical observations of the similarity of attention patterns across nearby queries/heads, allowing substantial elimination of redundant computation—a strategy that harmonizes with findings in pattern-sharing sparse attention (Peng et al., 26 May 2025).
• Hybridization and Hardware Efficiency: Unlike schemes that require retraining (as in SEA (Lee et al., 2023)) or use hand-crafted patterns (TriangleMix (He et al., 29 Jul 2025)), InfLLM v2 maintains deployment flexibility and plug compatibility with both GPU and mobile SoC pipelines. Dynamic sparsity ratios and groupings can be tuned with minimal impact on model performance.

5. Performance Analysis and Empirical Results

Experimental results in (Team et al., 9 Jun 2025) indicate:

• Sparsity: For context lengths up to 128k tokens, each query typically attends to around 5% of the full set of tokens.
• Prefill and Decoding Speed: On end-device hardware (e.g., Jetson AGX Orin), models using InfLLM v2 achieve up to a 7× speedup in decoding relative to dense attention models such as Qwen3-8B.
• Throughput and Latency: The reduction in memory accesses and arithmetic operations leads to lower time-to-first-token (TTFT) and scalable throughput, as context lengths increase.
• Accuracy: Despite aggressive sparsification, performance on open benchmarks remains competitive, with empirical evidence of maintained or improved perplexity under long-context regimes.

These advances are achieved without retraining or model surgery, and are especially well-suited for resource-constrained or latency-sensitive deployment scenarios.

6. Relation to Broader Ecosystem and Adaptive Extensions

InfLLM v2’s architecture can be synergistically combined with other long-context and adaptive attention techniques:

• Hybrid Context Sparsity: The block selection logic can be tailored or layered with context sharding across attention heads as in S2-Attention (Lin et al., 25 Jul 2024), or with layerwise triangle patterns as in TriangleMix (He et al., 29 Jul 2025).
• Token/Block Allocation: Learned or profile-guided variants (e.g., MoA (Fu et al., 21 Jun 2024), SeerAttention (Gao et al., 17 Oct 2024)) can be integrated for per-head or per-layer adaptivity, and dynamic sparsity thresholds can be introduced based on empirical mean or separability of semantic kernels.
• Hardware Integration: The CUDA/CPM.cu implementation demonstrates the benefit of harmonizing sparsity-induced arithmetic reduction with I/O-aware kernel fusion; on mobile SoCs, dynamic sparse attention can be efficiently mapped by allocating top-k screening to NPUs and blockwise computation to CPU/GPU as in shadowAttn (Yin et al., 22 Aug 2025).

7. Future Prospects and Limitations

While InfLLM v2 achieves scalable, efficient long-context attention, certain trade-offs may arise:

• The selection of block size $m$ and number of top-$k$ blocks impacts the memory/speed/accuracy trade-off; adaptive tuning remains a subject for further paper.
• For highly nonuniform or pattern-agnostic input distributions, mean pooling may fail to capture sharp local semantics; hybrid or learnable pooling functions could address this.
• The approach is optimized for inference; further work could explore training-time block sparsity regularization (cf. condensation (Sason et al., 3 Mar 2025)) and end-to-end differentiable variants.

Overall, InfLLM v2 Sparse Attention exemplifies a scalable, blockwise, dynamically sparse attention mechanism supported by both theoretical and empirical evidence, enabling high-throughput, low-memory long-context inference without the compromises of static pattern sparsity or mandatory retraining.