cs.LGFeb 26, 2026

S2O: Early Stopping for Sparse Attention via Online Permutation

Authors: Yu Zhang, Songwei Liu, Chenqian Yan, Sheng Lin, Beichen Ning, Fangmin Chen, Xing Wang

Organizations: ByteDance · Xiamen University

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Attention scales quadratically with sequence length, fundamentally limiting long-context inference. Existing block-granularity sparsification can reduce latency, but coarse blocks impose an intrinsic sparsity ceiling, making further improvements difficult even with carefully engineered designs. We present S2O, which performs early stopping for sparse attention via online permutation. Inspired by virtual-to-physical address mapping in memory systems, S2O revisits and factorizes FlashAttention execution, enabling inference to load non-contiguous tokens rather than a contiguous span in the original order. Motivated by fine-grained structures in attention heatmaps, we transform explicit permutation into an online, index-guided, discrete loading policy; with extremely lightweight preprocessing and index-remapping overhead, it concentrates importance on a small set of high-priority blocks. Building on this importance-guided online permutation for loading, S2O further introduces an early-stopping rule: computation proceeds from high to low importance; once the current block score falls below a threshold, S2O terminates early and skips the remaining low-contribution blocks, thereby increasing effective sparsity and reducing computation under a controlled error budget. As a result, S2O substantially raises the practical sparsity ceiling. On Llama-3.1-8B under a 128K context, S2O reduces single-operator MSE by 3.82×\times at matched sparsity, and reduces prefill compute density by 3.31×\times at matched MSE; meanwhile, it preserves end-to-end accuracy and achieves 7.51×\times attention and 3.81×\times end-to-end speedups.

Figures & tables

Figure 1: Speedup at 128K context length (Llama-3-8B) with sparsity–error trade-off. We report end-to-end latency speedup over FlashAttention and break down the major components, including sparse preprocessing time, attention compute time, and other overheads. We also show each method’s sparsity ratio and the corresponding mean squared error (MSE; lower is better), highlighting that our method achieves lower error at higher sparsity.
Figure 2
Figure 4: S2O workflow. Step 1 (Permutate) builds two lightweight index arrays without moving tensors: (i) an intra-segment query permutation index QpermQ_{\mathrm{perm}} (token-level, segment-local), and (ii) a prefix key/value logical map KVpermKV_{\mathrm{perm}} (token-level, global indices) obtained by retrieving segment representatives. Step 2 (Sparse attention) runs attention in two passes: Pass-1 computes a dense intra-segment causal window to initialize online-softmax states; Pass-2 resumes the states and processes the historical prefix in the retrieved order with early stopping.
Figure 5: Evaluation overview on Llama-3.1-8B . (a) Operator-level accuracy under matched sparsity. (b) Attention speedup: dark bars denote attention time and light bars denote preprocessing time; the corresponding sparsity(s) and MSE are also reported. (c) End-to-end prefill latency breakdown across context lengths: dark bars denote total attention-related time (including preprocessing) and light bars denote all remaining overheads.
Hyperparameter Ablations (128K)
SSτ\tauSparsity ↑\uparrowavg. MAE ↓\downarrowavg. MSE ↓\downarrow
20480.0010.7210.005360.00058
20480.0020.7980.008730.00131
20480.0040.8610.013530.00265
20480.0050.8770.015350.00323
20480.0100.9180.021630.00560
Table 2: Hyperparameter ablations. We report mean sparsity and the average approximation error (MAE/MSE, averaged over heads) on Qwen3-8B at 128K. Top block fixes S=2048S{=}2048 and varies τ\tau ; bottom block fixes τ=0.005\tau{=}0.005 and varies SS .
MethodSparsity ↑\uparrowMAE ↓\downarrowMSE ↓\downarrow
local window only0.9830.017901.8395 ×10−3\times 10^{-3}
Ours (w/o QpermQ_{\mathrm{perm}} )0.7960.001331.86 ×10−5\times 10^{-5}
Ours (w. QpermQ_{\mathrm{perm}} )0.8520.001291.87 ×10−5\times 10^{-5}
Table 3: Component ablation. We report mean sparsity and average approximation error (MAE/MSE) on Qwen3-8B at 128K, and additionally present one representative head.
Figure 6: Comparison with BSA under matched sparsity. Speedup over FA2 dense under different TopK block ratios for non-causal attention. Left: sequence length N=32N=32 K. Right: sequence length N=128N=128 K. Under matched sparsity settings, S2O achieves runtime comparable to BSA, and the measured speedup remains largely aligned with the actual sparsity level. This suggests that sparsity beyond block granularity is practical and can be effectively translated into wall-clock acceleration.
LLRetrieval (ms)Flash latency (ms)
S=512S{=}512S=1024S{=}1024S=2048S{=}2048S=4096S{=}4096
4K0.2540.2550.2560.2460.756
8K0.3740.3110.2660.3151.720
16K0.7850.5450.3860.3285.712
32K2.4151.4100.9200.63621.09
64K8.6444.6462.6671.71381.48
Table 9: Measured retrieval latency (ms) as a function of segment length SS and sequence length LL , with Flash latency shown on the right.
Figure 7: Additional attention heatmaps on Qwen3-8B under different reordering strategies.
Figure 8: Additional attention heatmaps on LLaMA-3.1-8B under different reordering strategies.
More in this topic

LLM Inference Acceleration

LLM inference acceleration develops algorithms and systems that reduce language-model latency, memory use, and serving cost while preserving output quality.

See all 1,084 papers

Explore similar work

CardsList
Dec 7, 2025

Block Sparse Flash Attention

Modern large language models increasingly require long contexts for reasoning and multi-document tasks, but attention's quadratic complexity creates a severe computational bottleneck. We present Block Sparse Flash Attention (BSFA), a drop-in replacement that accelerates long-context inference while preserving model quality. Unlike methods that predict importance before computing scores, BSFA computes exact query-key similarities to select the top-k most important value blocks for each query. By comparing per-block maximum scores against calibrated thresholds, we skip approximately 50% of the computation and memory transfers for pruned blocks. Our training-free approach requires only a one-time threshold calibration on a small dataset to learn the per-layer and per-head attention score distributions. We provide a CUDA kernel implementation that can be used as a drop-in replacement for FlashAttention. On Llama-3.1-8B, BSFA achieves up to 1.13x end-to-end speedup on LongBench with only a 1.1% accuracy drop, and up to 1.24x on Needle-in-a-Haystack retrieval at a 1% accuracy drop. The attention kernel itself accelerates by up to 1.38x. We compare BSFA against five recent sparse attention baselines (SpargeAttention, MInference, FlexPrefill, XAttention, and BLASST), and verify the method on Qwen2.5-7B and on A6000 and H100 GPUs. The implementation is available at https://github.com/Danielohayon/Block-Sparse-Flash-Attention.

Daniel Ohayon, Itay Lamprecht, Itay Hubara +3

Technion – Israel Institute of Technology, Haifa, Israel · Intel - Habana Labs · Intel – Habana Labs, Tel Aviv, Israel

GPU AccelerationSelf-AttentionLong-Context Language Model InferenceHigh-Performance ComputingLLM Inference AccelerationSparse Attention+3
AskPDF
Jun 3, 2026

SparDA: Sparse Decoupled Attention for Efficient Long-Context LLM Inference

Sparse attention reduces compute and memory bandwidth for long-context LLM inference. However, two key challenges remain: (1) KV cache capacity still grows with sequence length, and offloading to CPU memory introduces a PCIe transfer bottleneck; (2) the sparse selection step itself retains O(T2)O(T^2) complexity and can dominate attention cost at long contexts. We propose SparDA, a decoupled sparse attention architecture that introduces a fourth per-layer projection, the Forecast, alongside Query, Key, and Value. The Forecast predicts the KV blocks needed by the next layer, enabling lookahead selection that overlaps CPU-to-GPU prefetch with current-layer execution. Because Forecast is decoupled from the attention query, our GQA implementation uses one Forecast head per GQA group, reducing selection overhead versus the original multi-head selector. SparDA adds <<0.5% parameters and trains only the Forecast projections by matching the original selector's attention distribution. On two sparse-pretrained 8B models, SparDA matches or slightly improves accuracy and delivers up to 1.25×\times prefill speedup and 1.7×\times decode speedup over the sparse-attention offload baseline. By enabling larger feasible batch sizes on a single GPU, SparDA further reaches up to 5.3×\times higher decode throughput than the non-offload sparse baseline. Our source code is available at https://github.com/NVlabs/SparDA.

Yaosheng Fu, Guangxuan Xiao, Xin Dong +2

1NVIDIA · 2Thinking Machines Lab · 3ByteDance Seed +1

KV-Cache OffloadingLong-Context Language Model InferenceLLM Inference AccelerationSparse Attention
AskPDF
Oct 22, 2026

LazyFlash: Forward-Dense, Backward-Sparse Attention for Accelerated Long-Context Fine-Tuning

As large language models (LLMs) process increasingly long sequences, attention becomes a major computational bottleneck. Sparse attention mostly targets inference, yet many methods still train densely or cannot support sparse training, and their added complexity makes sparsity hard to turn into wall-clock speedups, leaving dense attention widely used in practice. This paper introduces LazyFlash, a forward-dense, backward-sparse attention method built on the observation that attention training is dominated by the backward pass. LazyFlash keeps the forward pass exactly dense, so the trained model remains a standard dense model that existing systems such as vLLM can serve without any change. For the backward pass, LazyFlash proposes a gradient-aware block sparsification design that ranks query blocks by their attention-output gradient norms and skips those contributing negligibly to parameter updates, turning the dense backward into a block-sparse one. A gradient-mass budget adapts the amount of skipped computation to each head and training step, while importance sampling preserves the dense gradient in expectation. These designs are implemented inside the tiled FlashAttention backward kernel, where skipped blocks are never loaded or computed, translating gradient sparsity into real wall-clock acceleration. Experiments on long-context text and speech LLMs show that LazyFlash maintains downstream performance while achieving up to 5.7x attention-backward speedup over FlashAttention and 1.7x end-to-end training speedup.

Keqi Deng, Ruchao Fan, Jinyu Li

Microsoft, USA

Sparse Fine-TuningLLM Fine-TuningBlock-Sparse AttentionLong-Context Language Modeling
AskPDF