Block Sparse Flash Attention

Momentum

8 papers in the last four weeks, up 60% on the four weeks before. 0.1% of all new papers.

Jul 6Week of Sep 21

Latest papers 29

Sep 30, 2026cs.LG

LampAttention: Look-Ahead Mixed-Precision FlashAttention for Dedicated Accelerators

While most attention logits can be computed in low precision without degrading numerical stability, current attention kernels fail to exploit this phenomenon. We introduce a novel hardware-algorithm co-design in the form of mixed-precision FlashAttention. Our method accumulates key-query products and evaluates their exponentials in 8-bit formats, then adaptively identifies sensitive sub-blocks and recomputes them in 16-bit formats. We propose the specifications for a dedicated accelerator capable of executing this pipeline efficiently. Simulated experiments with Qwen3 and Gemma 3 show that rerouting a selective minority of sub-blocks to high precision is sufficient to recover the baseline model performance.

Stanislav Budzinskiy, Marian Gloser, Tolunay Yilmaz +5

Faculty of Mathematics University of Vienna, Austria · Huawei Heisenberg Research Center, Munich, Germany · Huawei Technologies Co. Ltd

Block Sparse Flash Attention

Ask PDF

Sep 29, 2026cs.CV

Parameterized Stripe Attention for Efficient Video Generation

Diffusion Transformers (DiTs) enable high-quality video generation but suffer from substantial inference latency, primarily attributable to the computationally expensive full spatio-temporal attention. While sparse attention methods offer potential solutions, existing approaches face an inherent flexibility--efficiency dilemma: predefined masks lack the flexibility to capture diverse attention patterns, while runtime-determined masks introduce overheads and sacrifice hardware efficiency. We identify the lack of a unified structural characterization of DiT attention as a key limitation of existing methods, and establish that video DiT attention exhibits \textbf{periodic diagonal stripe structures} along both temporal and spatial dimensions. To formally encode these structured patterns within a single efficient kernel, we present {\bf PSA}, a parameterized stripe attention that formalizes the observed stripe regularity, unifying diverse attention patterns for efficient mask generation. This unified representation enables a single hardware-efficient CUDA kernel to process all sparse patterns, achieving FlashAttention-3-level Model FLOPs Utilization. To determine optimal sparsity configurations, we propose a training-free offline search algorithm that automatically maximizes sparsity under a specified error tolerance for each attention head. Experiments on HunyuanVideo and Wan~2.1 demonstrate that PSA achieves 1.57

\times

and 1.37

\times

end-to-end speedups over FlashAttention-3 baselines, with acceptable visual quality degradation.

Xingyu Jia, Baole Ai, Ang Wang +2

Alibaba Group

Spatio-Temporal Attention Video Diffusion Models Block Sparse Flash Attention Video Generation Cuda Step-Level Credit

Ask PDF

Sep 28, 2026cs.CV

DORA: Dynamic Online Reinforcement Agent for Token Pruning in Vision Transformers

Vision Transformers (ViTs) incur quadratic self-attention cost in the number of tokens. Most token-reduction methods adapt token identities within a prescribed layer-wise compression schedule, or search a static mask offline, and thus limit online adaptation of when and how much to prune. We propose DORA (Dynamic Online Reinforcement Agent), which learns an input-adaptive pruning policy itself for frozen ViTs. At each eligible block, a hierarchical actor decides whether to prune, how many tokens to remove, and which tokens to remove from each image's evolving representation. Because early deletions change the states observed by later decisions, DORA formulates pruning as a finite-horizon Markov decision process. Complete-prefix shadow evaluations convert final-prediction fidelity into localized per-step credit, while closed-loop accuracy feedback adjusts the fidelity penalty toward a shared accuracy-drop target. A privileged critic and all shadow computations are training-only. Deployment retains the frozen backbone and a lightweight actor that applies hard deletion and packed variable-length FlashAttention, converting token reduction into measured speedups. On ImageNet-1K with DeiT-Base, DORA reduces FLOPs by 38.4% relative to the uncompressed backbone within one percentage point of accuracy loss. Averaged across four ViT-type backbones at matched accuracy, DORA uses 13.2% fewer FLOPs and achieves 32.4% higher throughput than the corresponding per-backbone baseline means. Under zero-shot transfer to ImageNet-A, these gains widen to 20.3% and 45.6%, respectively.

Kaixuan He, Song Chen, Yi Kang

University of Science and Technology of China, Hefei, China · Institute of Artificial Intelligence, Hefei Comprehensive National Science Center, Hefei, China

Visual Token Pruning Vision Transformer Block Sparse Flash Attention Linear Attention Offline Reinforcement Learning

Ask PDF

Sep 28, 2026cs.LG

Broken Symmetry in BF16 Attention: Why FlashAttention Gradients Blow Up Late in Training

BF16 is now standard in large-scale pretraining, including in fused attention kernels such as FlashAttention, and these kernels are widely trusted. When we used FlashAttention-3 to pretrain a 450M-parameter transformer on 50B tokens, however, we ran into a problem: training was healthy for 25B tokens, then the gradient norm grew a thousandfold and the loss ended 0.2 nats above FP32 attention, without a single NaN. Recomputing the attention backward of just two layers in FP32 removes almost all of the excess gradient. Part of the cause is known: a fused multiply-add in the forward softmax, so far treated as an extreme-input NaN case and never fixed in FlashAttention-3. Repairing it stops the blow-up, but the query gradient is still wrong by more than its own size, and training still drives attention logits to thousands of times their size under accurate gradients. The remaining error comes from a broken conservation law. The softmax score gradient sums to zero along every row, which makes the query gradient blind to where the keys sit as a group; rounding it to BF16 leaves a small nonzero sum that leaks the mean key into the gradient, and the leak grows exactly as late training makes keys large and attention sharp. We introduce GProj (gauge projection), which restores the zero sum after the cast with two rank-one corrections per row. It cuts the remaining median query/key gradient errors from 219%/13% to 0.34%/0.37%, on par with FP32 attention, for 4.7% more time per training step. In matched from-scratch runs it trains to the same loss as FP32 attention, while FlashAttention-3 and key smoothing both destabilize.

Junlin Chen, Daize Dong, Huanwei Di +9

Rutgers University · Carnegie Mellon University · Oracle +2

Block Sparse Flash Attention Matched Fp16 Intermediate Pretraining Symmetry

Ask PDF

Sep 27, 2026cs.LG

FoldAttention: Declared-Reference Softmax for Fast Decode and Deterministic Backward

Autoregressive decode repeatedly streams a growing KV cache, making attention a major cost at long context. Existing high-performance kernels use online softmax, which discovers a row's normalization reference as it scans keys. Earlier contributions therefore remain provisional and may require rescaling. We argue that the reference need not be discovered: softmax is invariant to a common shift, so the reference only has to keep the weights in range. We present FoldAttention, an additive formulation of softmax attention that fixes a finite reference

Z_i

before scanning the KV cache. Each weight

2^{s_{ij}-Z_i}

is then final when computed, so contributions add across disjoint key ranges and their quotient equals softmax attention in real arithmetic. We use this property to develop two techniques for Hopper decode: (1) final weights gate key and value reads before the bytes are fetched, and a per-call depth

T

cuts keys below

2^{-T}

while keeping their mass, and (2) additive partials compose split KV and shared-prefix cascades without rescaling. On H100 at

T=16

, FoldAttention decodes seven real-model generations 1.36-2.30

\times

faster than the fastest BF16 baseline, and up to 3.09

\times

faster across MHA and GQA shapes, at an error within 1.5% of the lowest BF16 error on six of the seven; reading every key, it is 1.14-1.30

\times

faster at matched error. We validate on Qwen3-8B that a whole decode step is up to 1.46

\times

faster while likelihood and long-context accuracy match those under BF16 kernels. The same principle makes the backward deterministic: CTAs round bounded partial gradients onto an integer grid declared before the reduction and add them in any order. FoldAttention thereby removes the determinism tax: its deterministic backward is up to 1.84

\times

faster than deterministic FlashAttention-3/4 and 1.05

\times

faster than the fastest nondeterministic kernel.

Sriman Achanta

Virginia Commonwealth University

Block Sparse Flash Attention Autoregressive Decoding Gumbel-Softmax Relaxation Key-Value Cache Activation Functions

Ask PDF

Sep 24, 2026cs.LG

FlashLoop: Fast and Memory-Efficient Looped Transformers via Lazy Updates

Looped Transformers have attracted substantial attention as a parameter-efficient approach to increasing computational depth through repeated application of shared Transformer blocks. However, their practical advantages over conventional Transformers remain under debate: each additional loop incurs another Transformer pass and requires caching another set of KV states, causing inference FLOPs and KV-cache memory to grow continuously with loop depth. This overhead becomes particularly severe at large loop counts and long context, preventing the parameter efficiency of Looped Transformers from translating into practical inference efficiency. In this paper, we find that much of the additional computation and storage introduced by looping is redundant. As recurrence proceeds, state changes become increasingly concentrated on a small subset of tokens; attention-output differences are dominated by a sparse and stable subset of key columns; and KV residuals between adjacent loops become progressively more amenable to low-bit quantization. Building on these observations, we introduce FlashLoop, a training-free inference framework that reduces cross-loop redundancy through token-sparse updates, sparse attention, and KV-residual quantization. Across several Looped Transformers models, FlashLoop delivers lossless accuracy while achieving up to 1.64

\times

end-to-end speedup and up to 6

\times

KV-cache memory reduction, substantially improving the practicality of scaling Looped Transformers to greater computational depths and longer context.

Wanqi Yang, Shiwei Liu

ELLIS Institute Tübingen Max Planck Institute for Intelligent Systems Tübingen AI Center

Block Sparse Flash Attention Transformer Architectures Flash Kv-Cache Management

Ask PDF

Sep 21, 2026cs.LG

FlashBoB: I/O-Efficient Exact Backward-over-Backward for Softmax Attention

Transformer models built on the attention mechanism have become a central building block in modern deep learning, yet softmax attention remains a major bottleneck for long-context workloads. While FlashAttention makes the forward and first backward passes I/O-efficient, it does not support backward-over-backward (BoB), which enables exact differentiation through the backward pass for applications such as second-order optimization, test-time training, gradient-based memory, and meta-learning. Existing BoB implementations either materialize large intermediate tensors or exhaust GPU memory at long sequence lengths. We present FlashBoB, an exact, I/O-efficient algorithm for BoB in softmax attention that keeps computation within on-chip tiles and avoids all

N \times N

intermediate tensors, where

N

is the sequence length. The key insight is a hierarchical affine structure in the softmax double backward: two row-wise scalars determine all outputs through affine transformations. This yields a two-pass schedule with bounded on-chip static random-access memory (SRAM) usage and minimal off-chip high-bandwidth memory (HBM) traffic. FlashBoB achieves

Θ(N^2 d^2/M)

HBM traffic (

d

is the head dimension and

M

is the memory size) and, within the standard FlashAttention-style score-recomputation model, matches the inherited large-cache lower bound for exact forward attention. Empirically, it scales exact attention BoB to

N=262\text{K}

on a single A100 80GB GPU, where prior PyTorch exact baselines fail by

N=16\text{K}

, and is up to

6.3\times

faster than FlashBack. These results make exact second-order attention practical at long-context sequence lengths where prior implementations cannot run efficiently.

Anthony Givans, Michael Crawshaw, Mingrui Liu

Block Sparse Flash Attention Gumbel-Softmax Relaxation Backward Passes Activation Functions Gradient-Based Optimization

Ask PDF

Sep 20, 2026cs.LG

LumoTree: Path-Parallel Speculative Verification for Hybrid Language Models

Tree speculative decoding for hybrid language models must preserve one coherent continuation across recurrent, convolution, and attention state. We present LumoTree, a verifier that executes recurrent paths in parallel, reuses state tiles within each path, and coordinates native recurrent replay, convolution-history gathering, and attention-cache remapping through a shared logical tree. Fused candidate selection, GPU-resident acceptance, and grouped split-K attention support the verification cycle. Component experiments show exact candidate-selection parity and recurrent agreement within paired error bounds. An exploratory Qwen3.8-27B NVFP4 deployment on a single NVIDIA DGX Spark (GB10) records 25.63 pooled tokens/s on ten SWE-bench Verified Astropy tasks. The results characterize component-level numerical agreement and coding-agent deployment; complete-model continuation and controlled application speedups remain open.

Zhiyuan Ma

Gated Deltanet Block Sparse Flash Attention

Ask PDF

Sep 17, 2026cs.CL

DeepSeek-V4.1-Flash: Pushing the Limits of KV Cache Compression

The widespread adoption of long-horizon agents has made model workloads increasingly input-heavy. Although prior work has substantially reduced the cost of long-context computation, prefill remains computationally expensive, and large KV caches continue to strain HBM and SSD capacity and data-transfer bandwidth. Together, these compute, storage, and bandwidth demands constitute the primary bottleneck to further lowering deployment costs. To address this challenge, we introduce DeepSeek-V4.1-Flash, a multimodal Mixture-of-Experts (MoE) model with 552B backbone parameters and support for contexts of up to one million tokens. With its Causal Encoder-Decoder (CED) architecture, the model activates 16B parameters per token during decode but only 8B parameters during prefill, substantially improving cost efficiency for agentic workloads. To push the limits of KV cache compression, DeepSeek-V4.1-Flash combines cross-layer KV cache reuse in Compressed Sparse Attention 2 (CSA2) with FP4 KV caching. These designs reduce its global KV cache footprint (always in HBM) to 890 bytes per token, roughly 1/4 of the corresponding footprint of DeepSeek-V4-Flash. Further, through a dedicated deployment optimization known as SWA Bounded Replay, DeepSeek-V4.1-Flash reduces its persistent KV cache footprint (always on SSD or in host memory) to roughly 1/8 of that of DeepSeek-V4-Flash. Despite its much smaller KV cache footprint, the model delivers substantially better performance than the baseline. In addition, we streamline the DeepSeek-V4 architecture and introduce several efficient architectural extensions. We pretrain DeepSeek-V4.1-Flash on a multimodal corpus comprising 45T tokens and conduct comprehensive post-training, yielding strong performance across diverse text-based and multimodal agentic scenarios. Model checkpoints are available at https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash.

DeepSeek-AI, :, Anyi Xu +590

Key-Value Cache Compression Key-Value Cache Deepseek Block Sparse Flash Attention Long-Context Reasoning Data Compression Methods+1

Ask PDF

Sep 14, 2026cs.CL

SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking

Post-training attention sparsification reduces the quadratic cumulative attention cost of pretrained Transformers by selecting a small set of context units (tokens or blocks) for each query. Existing trainable methods usually use a lightweight selector to score context units, followed by hard Top-K selection that blocks gradients from the language modeling loss. Consequently, these methods commonly distill layer-wise dense attention distributions. Although this encourages the selector to rank context units by dense attention weights in the original model, the ranking is not directly aligned with their impact on predictions under a fixed attention budget (i.e., the number of attended context units per query), potentially wasting the limited budget on less useful units. To address this misalignment, we propose Simple Attention Sparsification (SAS), a gated sparse attention mechanism that optimizes context ranking end-to-end with the language modeling loss. The key idea is to inject the selector's continuous scores into attention logits during training, allowing the loss to update the selector through standard backpropagation. We identify several choices crucial for this simple design to work well in practice: placing the gate inside the attention softmax in log form, using normalized softmax gates to calibrate historical context against the always-retained current block, and preserving continuous selector scores so the model learns relative priorities rather than only hard selections. To support long-sequence training, we implement a memory-efficient Triton kernel that integrates SAS into FlashAttention-style computation. Across reasoning, long-context understanding, and agentic tasks, SAS consistently outperforms trainable sparse attention baselines across attention budgets, with especially large gains under tight budgets, demonstrating more effective context ranking for downstream tasks.

Zhiwei Li, Lei Zhu, Hao Gu +6

Tencent HY LLM Frontier · Hong Kong University of Science and Technology (Guangzhou) · Hong Kong University of Science and Technology

Dynamic Sparse Attention Block Sparse Flash Attention Sparsity Efficient Long-Context Inference Loss Function Ranking

Ask PDF

Sep 9, 2026cs.LG

EFQ-Softmax: Exp-Free Quantization for Softmax

Low-bit attention accelerates Transformer inference by moving the

QK^\top

and

PV

matrix multiplications to FP8 or FP4 matrix engines. However, the softmax path often evaluates shifted-score exponentials in higher precision, forms a temporary probability block, and quantizes it before low-bit

PV

multiplication. This exp-then-quantize path creates a mismatch between a high-precision probability producer and a low-bit matrix consumer. We propose EFQ-Softmax (Exp-Free Quantization for Softmax), a low-bit probability-generation method that directly maps shifted attention scores to block-scaled E2M1 operands. For each microscaling block, EFQ-Softmax selects an exponent-only scale from the local maximum, maps the shifted scores to a normalized residual domain, and generates nonnegative E2M1 probability codes using a single affine rule. The resulting operand is used consistently in both the

\widetilde{P}V

numerator update and the

\widetilde{P}\mathbf{1}

denominator update. The FlashAttention-style row-maximum update, historical rescaling, high-precision accumulation, and final normalization remain unchanged. We evaluate end-to-end quality on Qwen3-8B, Qwen3-VL-8B-Instruct, and WAN2.2-TI2V-5B, and separately measure kernel-level performance on the A5 vector unit. EFQ-Softmax improves the Qwen3-8B seven-task mean from 0.6749 with MXFP4 to 0.6773 and the Qwen3-VL nine-task mean from 0.7826 to 0.8000. On WAN2.2, it maintains temporal consistency and visual quality comparable to the FP16 and MXFP4 baselines under VBench. On the A5 vector unit, EFQ-Softmax reduces the vector-stage latency of the fused probability-generation kernel by 40.33% on average across sequence lengths from 16K to 128K. These results show that direct low-bit probability generation can replace the conventional exp-then-quantize path while preserving end-to-end model quality.

Haohui Han, Yuming Wan, Hongni Wang +4

Xi'an Jiaotong University · Xi’an Jiaotong University · Huawei Technologies Co., Ltd +1

Gumbel-Softmax Relaxation Block Sparse Flash Attention Matrix Multiplication Transformer Architectures Activation Functions

Ask PDF

Sep 3, 2026cs.LG

Hardware-Aware FP4 FlashAttention-4

Blackwell's 4-bit floating-point (FP4) tensor cores do not automatically make attention faster because softmax conversion and on-chip dependencies dominate once its matrix products shrink. We address this with \emph{Direct-P} for noncausal inference and a causal path that passes the forward quantization directly into backward. Direct-P maps scores directly to FP4 probabilities and reaches up to 2.13

\times

the bfloat16 (BF16) forward throughput on an NVIDIA GB200. The causal path reconstructs probabilities from saved quantized queries and keys and uses 8-bit floating-point (FP8) gradient operands, accelerating a complete single-GPU 8-billion-parameter update by up to 1.14

\times

. Matched distributed training retains FP8 probabilities and values; every tested MXFP4 probability/value training trajectory diverges.

Robert Hu

GRAPHCORE RESEARCH

Block Sparse Flash Attention Causal Inferences Probability

Ask PDF

Aug 20, 2026cs.LG

Ask Self, Ask Others: Relation Is All You Need

Attention dominates token mixing, but it collapses relation formation and flow allocation into a single score-to-flow step. We introduce Relation, which separates them by first organizing pairwise evidence into explicit Self and Exchange relations and deriving information flow afterward. Relation first decides whether a token should rely on itself or draw from its history, and if it draws from history, where to look. This relational organization gives rise to Full Relation, FlashRelation, Linear Relation, and Hybrid Relation. Across matched decoder-only models, Full Relation achieves lower mean final-validation NLL than MHA and reaches the paired MHA final training loss with 4.5-7.3% fewer tokens. Structural diagnostics further show that Relation learns a distinct depth organization: the first layer acts as a current-token anchor and a high-rank router, while later layers shift strongly toward history. In a fixed-context reference benchmark, FlashRelation is 4.17-5.28x faster than the materialized Full Relation implementation. Across scale-matched production workloads, it reaches 89.7-92.9% of PyTorch FlashAttention throughput while executing the exact Full Relation operator. Hybrid Relation demonstrates that Full and Linear Relation layers can be composed within a single decoder. These results support a relation-first view of token mixing: ask Self, ask Others, then let Flow follow Relation.

Yuting Ge, Pengju Yang, Mingkai Nie

Department of Electrical Engineering City University of Hong Kong · Jilin University · National University of Singapore

Linear Attention Block Sparse Flash Attention Self Stance Flash Meanflow

Ask PDF

Aug 5, 2026cs.LG

Training-Free Hashing-Based Attention via Binary Principal Components

Long-context large language models (LLMs) are increasingly deployed in real-world applications, yet self-attention remains a major efficiency bottleneck -- especially during decoding -- due to the necessity of repeatedly processing ever-growing key-value (KV) caches. Existing sparse attention reduce computation by attending to fewer KV pairs, but often suffer from substantial accuracy degradation, require additional training, or rely on expensive hashing. In this work, we present BinaryPC, a training-free, data-aware hashing-based sparse attention for long-context LLMs. BinaryPC constructs compact binary hash codes and corresponding hash function by computing binary principal components of data. Unlike Locality-Sensitive Hashing (LSH) with data-independent random projections or learned non-linear hashing methods, BinaryPC constructs binary codes that explicitly preserve the structural information of data without requiring gradient-based training. Comprehensive experiments across multiple model families and long-context benchmarks show that BinaryPC preserves accuracy relative to full attention while achieving superior performance among sparse and hashing-based baselines. On modern GPUs, BinaryPC improves end-to-end decoding throughput by 3.56

\times

over the FlashAttention kernel. Our code is available at https://github.com/yudaohai666/BPC.

Daohai Yu, Zhanpeng Zeng, Keyu Chen +6

Key Laboratory of Multimedia Trusted Perception and Efficient Computing, Ministry of Education of China, Xiamen University, 361005, P.R. China. · Tencent YouTu Lab, Shenzhen, China · Sino-Russian Research Center for Digital Economy.

Dynamic Sparse Attention Block Sparse Flash Attention Efficient Long-Context Inference Locality-Sensitive Hashing Training-Free Principal Component Analysis

Ask PDF

Aug 3, 2026cs.LG

ATFlash: Per-RoPE-Wavelength Attention Windows for Compute/Memory-Efficient LLM Inference

The attention score with rotary position embeddings (RoPE) decomposes exactly into a sum over its 2D-rotation frequency pairs, and each pair's wavelength limits how far it can discriminate position. Aligned with this structure, we propose the per-RoPE-wavelength distance window: it prunes the query--key inner-product terms beyond a wavelength-proportional distance. Unlike a sliding window, every key remains reachable, at least through the low-frequency pairs. The reduction rate is input-independent, with a closed form logarithmic in the sequence length

N

, in contrast to dynamic-sparse methods like MInference. Such token-level selection is orthogonal to our frequency-level pruning. The window can therefore be applied on top of those methods. On Qwen2.5-0.5B and Llama-3.2-3B, the window prunes 37--48% of the query--key inner-product terms within each model's native context length. Relative to full attention, the top-1 match rate stays at 96--98% and the mean output-distribution KL at the

10^{-3}

-nat level on LongBench-v2 contexts. We examine absolute scores on long-context benchmarks such as RULER, OpenAI-MRCR, LongCodeQA, and

\infty

Bench: they are broadly preserved. We implement the window as a slice of the query--key contraction axis, leaving the online-softmax recurrences untouched, and port it with minimal diffs into the released FlashAttention-4 prefill and FlashInfer decode. On RTX PRO 6000 with Llama, both ports outpace stock with gains growing with context length, up to

1.29\times

at 128K. End to end on Qwen2.5-7B-1M, with 57% of the inner-product terms pruned, the speedup reaches

1.31\times

at a 1M-token context.

Shun-ichiro Hayashi, Daichi Mukunoki, Tetsuya Hoshino +1

Graduate School of Informatics, Nagoya University · Information Technology Center, Nagoya University

Block Sparse Flash Attention Rotary Position Embedding LLM Inference Optimization Flash

Ask PDF

Aug 3, 2026cs.LG

Output-Aware Rotation for INT2 KV-Cache Quantization

The key-value (KV) cache has become a major memory and bandwidth bottleneck in long-context large language model inference, making ultra-low-bit quantization increasingly important. However, existing rotation-based INT2 methods optimize cache statistics or proxy errors before the complete attention readout, even though the model is ultimately affected by the error propagated through attention and the output projection

W_O

. To address this mismatch, we propose \textit{OptR}, an output-aware rotation method that minimizes post-

W_O

attention-output error. OptR decomposes the post-

W_O

attention-output error into key- and value-induced terms and learns per-head orthogonal corrections through the full INT2 quantization and attention path. OptR further applies an attention-equivalent key reparameterization to reduce large channel-wise offsets without changing the softmax distribution. Across three models and five reasoning and coding benchmarks, OptR consistently improves both QuaRot and OSCAR and strengthens long-context retrieval, while preserving the paged KV-cache format with negligible inference overhead.

Vincent-Daniel Yun, Woosang Lim, Minsoo Cheong +4

University of Southern California · Seoul National University · Inha University

Key-Value Cache Quantization Kv-Cache Management Block Sparse Flash Attention LLM Inference Optimization Reparameterization

Ask PDF

Aug 3, 2026cs.CL

Understanding Sparse Attention Selectivity in Long-Context Foundation Models via Counterfactual Evaluation

Sparse attention is widely deployed in long-context serving stacks, yet no framework audits how discarding blocks changes the influence of specific content on model output. We first establish that the phenomenon is real and causal: Block Sparse Flash Attention (BSFA) route replay across four architectures changes output decisions in 13 of 16 cells, with zero identity-replay label flips. We then introduce a dense-calibrated counterfactual audit using matched probe cards---Gold (carrying the correct answer label), Poison (carrying a target wrong label), and Benign (filler only)---under six-layout position symmetry, isolating the sparsification-specific effect. Two patterns compete. Signal concentration: the selector preserves Gold and Poison blocks far above filler-matched Benign blocks (G

\approx

\gg

B across all model--task pairs). Integration loss: discarding blocks severs cross-block attention---confirmed by an ablation where isolating the probe block collapses its influence from 4.48 logits to zero. Compression ratio governs the balance: a full sweep from mild (

c=0.25

) to aggressive (

c=0.75

) compression across four model--task pairs reveals that three of four cells move toward stronger sparse amplification at higher compression, with two exhibiting sign reversals. Three independent arms---BSFA route replay, controlled block-top-

k

, and KV-cache eviction---converge: sparsification changes content influence in ways aggregate accuracy cannot detect. We provide an open measurement framework deployable on any model exposing block identities.

Xingyu Ren, Youran Sun, Chugang Yi +1

The Chinese University of Hong Kong, Hong Kong, China · University of Maryland, College Park, MD, USA

Dynamic Sparse Attention Block Sparse Flash Attention Sparsity Efficient Long-Context Inference Wireless Foundation Models Model Auditing+1

Ask PDF

Jul 23, 2026cs.CV

Flash EQ-Linear: Accelerating Equivariant Linear Layers via Group-wise Discrete Fourier Transform

Equivariant networks embed geometric symmetries as structural priors through weight sharing, achieving remarkable parameter efficiency across vision tasks. However, this parameter efficiency does not translate into compute efficiency: most existing implementations unroll the structured weights into dense matrices and dispatch them to generic dense kernels, so an equivariant layer costs no fewer MACs than its non-equivariant counterpart. In this paper, we observe that the equivariant linear (EQ-Linear) layer---the most fundamental and frequently used module in modern equivariant architectures---is essentially a circular convolution along the group dimension composed with a linear transform along the channel dimension. Building on this observation, we propose Flash EQ-Linear, an exact acceleration algorithm that reduces the cost to

2(T-1)/T^2

of the original dense formulation (

T

is the equivariant group size) by combining the Fourier convolution theorem along the group dimension with the conjugate symmetry of the real DFT. To translate these computational savings into wall-clock speedups, we further develop dedicated CUDA kernels for the

\mathrm{p}4

group. At the operator level, Flash EQ-Linear achieves up to

2.1\times

forward speedup over PyTorch's highly optimized F.linear; at the network level, Flash EQ-ViT achieves up to

{1.7\times}

end-to-end speedup over both equivariant and non-equivariant baselines. As an operator-level acceleration algorithm, Flash EQ-Linear provides plug-and-play acceleration for diverse pretrained equivariant models, including EQ-ViT, EQ-Swin, EQ-VMamba, and EQ-INR, without retraining or architectural changes. Code is available at https://github.com/zhongchenzhao/FlashEQLinear.

Zhongchen Zhao, Jixin Wang, Qi Xie +4

Xi’an Jiaotong University · The Hong Kong Polytechnic University · OPPO Research

Equivariant Neural Networks Rotation Equivariance Equivariance Block Sparse Flash Attention Fourier

Ask PDF

Jul 5, 2026cs.LG

HiFA4: Training-Free 4-bit FlashAttention on Ascend HIF4 NPUs for LLM Inference

We present HiFA4, a post-training operator-level design that executes both QK^T and PV in FlashAttention as 4-bit HIF4 Cube GEMMs for LLM inference on Ascend NPUs, while maintaining the online softmax state in FP16. To our knowledge, HiFA4 is the first Ascend-HIF4-targeted design of this kind evaluated on standard NLP benchmarks. HiFA4 combines two mechanisms. Smooth-QK applies a calibration-static per-channel equivalent rescaling to Q and K after RoPE, transferring quantization difficulty from K to Q without per-tile online reduction at inference. P-Reordering accumulates the softmax normalizer from the same quantized attention weights P_hat used in the PV GEMM, rather than from a higher-precision reconstruction. We show that this inconsistent formulation introduces a coherent output-scaling error, and validate the effect on a Qwen3-8B Layer-0 MMLU trace, where all 3.6M measured attention tiles exhibit net probability-mass loss with median epsilon_bar = -0.064. P-Reordering also allows the normalizer to be fused into the PV Cube GEMM. Across five LLMs, HiFA4 reduces quantization-induced decision drift. On Qwen3-8B, it recovers 37.5% of the accuracy gap introduced by direct HIF4 quantization, narrows the sample-weighted accuracy loss from 1.12 pp to 0.70 pp, reduces BF16-inconsistent MMLU predictions from 16.3% to 8.2%, and cuts MMLU accuracy regressions by 57% (1071 to 465). On Gemma2-9B, mild smoothing keeps HiFA4 within 0.7 pp of BF16 while reducing MMLU regressions by 27%. On LLaMA3.1-8B, Mistral-7B, and Phi-4B, where Smooth-QK is disabled, P-Reordering with the adopted Q-Mean auxiliary still reduces full-set MMLU regressions by 41-52%. A preliminary instruction-scheduling analysis projects a 35.4% critical-path latency reduction relative to BF16 by fusing the softmax normalizer into the PV Cube GEMM; on-hardware validation is left to future work.

Hui Dong, Yanzhao Li, Jie Gao +5

Huawei Technologies

Block Sparse Flash Attention LLM Inference Optimization Matrix Multiplication Batch Normalization Post-Training

Ask PDF

Jul 1, 2026cs.CV

RotateAttention: RoPE-Aware Rotation and Range Rectification for INT4 Quantized Attention in Video Generation

\textbf{DiT-based video generation models equipped with 3D Rotary Position Embeddings (3D RoPE)}

, the attention mechanism remains a primary computational bottleneck due to its quadratic complexity with respect to sequence length. While quantized

\textbf{FlashAttention}

offers a promising path toward hardware acceleration, existing low-bit quantization methods overlook two critical challenges in this setting:

\textbf{1)}

applying online rotation matrices -- a widely used technique for mitigating outliers in Queries (

Q

) and Keys (

K

) -- is difficult to reconcile with

\textbf{RoPE}

; and

\textbf{2)}

the non-negative attention matrix

P = \exp(QK - \max(QK))

makes symmetric quantization waste half of the 4-bit dynamic range. In this work, we observe that the outlier distributions of

Q

and

K

are strongly affected by the dimensional partitioning of

\textbf{3D RoPE}

. Based on this finding, we propose

\textbf{RotateAttention}

, an efficient

\textbf{mixed-precision INT4 FlashAttention}

framework tailored for

\textbf{DiT-based video generation models with 3D RoPE}

, using selective

\textbf{FP16 fallback}

for accuracy-sensitive attention blocks and denoising steps. RotateAttention introduces two core techniques:

\textbf{1) RoPE-aware Rotation}

, which employs either mergeable rotation matrices that can be fused into RoPE or negligible-overhead matrices to mitigate RoPE-induced outliers in

Q

and

K

; and

\textbf{2) Range-optimized

Quantization}

, which uses fixed scales and zero-points to fully exploit the

\textbf{INT4 numerical range}

with minimal computational overhead. Experiments show that

\textbf{RotateAttention}

preserves video generation quality nearly identical to full-precision baselines while achieving up to 1.68

\times

end-to-end speedup and 2.2

\times

kernel-level acceleration.

Yaofu Liu, Wanli Lan, Jinxi Li +2

The Hong Kong University of Science and Technology, Hong Kong SAR, China · Independent Researcher, China

Block Sparse Flash Attention Video Diffusion Models Video Generation Video Latents Dynamic Rope Manipulation Rectification

Ask PDF

Jun 29, 2026cs.CV

Reweighting Framewise Attention in Video Transformers for Facial Expression Understanding

Understanding facial expressions in videos requires modeling subtle and localized facial dynamics under unconstrained conditions. Although recent Vision Transformer (ViT)-based video models have shown strong performance through large-scale self-supervised pretraining, their attention mechanisms often emphasize dominant global motions and coarse temporal dynamics, limiting sensitivity to fine-grained facial variations. To address this limitation, we propose MiRA (Marginal-induced Attention Redistribution), a plug-in frame-marginal attention redistribution framework for ViT backbones that enhances spatio-temporal selectivity toward subtle facial dynamics without introducing additional trainable parameters. MiRA derives frame-level confidence and intra-frame concentration statistics from self-attention maps to estimate frame-wise marginal importance and redistribute attention toward spatiotemporally localized facial cues. We first introduce a principled exact mode based on post-softmax attention redistribution. To further improve efficiency, we propose flashLite mode, a lightweight pre-softmax approximation that integrates frame-marginal redistribution into FlashAttention kernels while preserving the effectiveness of the exact formulation. Experimental results on challenging Facial Expression Recognition (FER) benchmarks demonstrate consistent improvements over strong ViT baselines.

Seongro Yoon, Donghyeon Cho, Jinsun Park +1

Inria, Université Côte d’Azur, France · Hanyang University, South Korea · Pusan National University, South Korea

Facial Expression Recognition Self-Supervised Vision Transformers Spatio-Temporal Attention Block Sparse Flash Attention

Ask PDF

Jun 25, 2026cs.LG

EpiKV: Epiphany-Aware KV Cache Eviction Without the Attention Matrix

Reasoning models can generate chains of thought tens of thousands of tokens long, making the key--value (KV) cache that holds them a major bottleneck for inference throughput. Existing eviction policies for long reasoning traces typically rank cached tokens using attention weights, requiring access to the attention matrix and making them incompatible with fast inference kernels. In this work we study the limits of such policies under tight cache budgets. Surprisingly, we find that under the strongest of them the generations that finish are wrong about as often as without eviction; most of the accuracy loss comes from generations that enter loops and run until the length limit, and retaining more tokens according to a fixed importance score exacerbates this behavior. What stops the looping is keeping the tokens the model's recent queries point to, and the forward pass the model already runs reveals them without the attention matrix. Motivated by this observation, we introduce epiphany-aware KV cache eviction EpiKV, which combines hidden-state shifts with the model's recent query--key relevance to rank cached tokens without materializing the attention matrix. On multiple benchmarks, EpiKV matches or outperforms the strongest attention-based eviction baselines while running directly in vLLM with unmodified attention kernels.

Steven Kolawole, Virginia Smith

Language Technologies Institute Carnegie Mellon University · Machine Learning Department Carnegie Mellon University

Key-Value Cache Eviction Key-Value Cache Block Sparse Flash Attention Reasoning Traces

Ask PDF

Jun 2, 2026cs.AR

P-Cast Precision in FP8 Attention: Sink-Induced Collapse and the Optimality of S=2^8

FP8 (E4M3) acceleration for attention computation offers significant throughput gains, but the 3-bit mantissa introduces precision challenges when the softmax probability matrix~

P

is cast to FP8 before the

P \cdot V

matrix multiplication. We analyze two implementation choices that affect output precision under the \emph{Attention Sink} phenomenon: (1)~the KV block iteration order, and (2) the static scaling factor applied to

P

before casting. We show that forward KV iteration causes \emph{P-collapse} -- to leading order a fraction

Φ(Δ+ δ_k - 6.93 - \ln S)

of non-sink

P

values underflow to zero, where the small shift

δ_k \approx 1

(for

k_{\text{sink}}{=}4

) is the expected within-sink-block score maximum -- and that reverse iteration removes it, with a zero-underflow guarantee when reverse is combined with

S{=}256

. We further give a constructive characterization of

S = 256 = 2^8

as the static scale that simultaneously satisfies (i)~bit-exact IEEE 754 scaling, (ii) the lower envelope of a sawtooth function

dp(S)

over the E4M3 number line (

dp = 2^{-4}

, the minimum worst-case quantization step), and (iii)~the maximum normal-range coverage \emph{among bit-exact (

2^k

) scales} (a non-bit-exact scale such as

448

attains slightly higher coverage; sec.5}). Both optimizations are already deployed in FlashAttention-3/4 on engineering grounds; our contribution is a quantitative account of \emph{why} these choices are good and a closed-form threshold

Δ_c = 6.93 + \ln S - δ_k

for predicting kernel-level precision loss. Kernel-faithful experiments (

Q, K, V

in FP32 to isolate the P-cast effect) show

3

10\times

MSE improvement at moderate sink strengths, and paired tests confirm both fixes saturate to the same precision floor when combined -- which motivated updating the hpc-ops kernel from

S{=}1

S{=}256

Reed Lau

Tencent

Full-Precision Performance Fp8 Block Sparse Flash Attention Keys And Value Optimality Stance

Ask PDF

May 25, 2026cs.LG

Quantized Keys Steal Attention: Bias Correction for KV-Cache Compression in Video Diffusion

Chunk-wise autoregressive video diffusion models rely on a KV cache of previously generated chunks to avoid redundant computation, but this cache quickly becomes a memory bottleneck as videos grow longer. Methods that quantize the KV cache to low bitwidths reduce memory pressure but degrade video quality. We show that a key driver of this degradation is a systematic bias in attention weights: due to the convexity of the exponential in softmax attention, quantization noise inflates the contribution of cached keys, a phenomenon we call the Jensen bias. This effect causes quantized keys to steal attention mass from the unquantized current chunk. We derive a per-attention-score correction that removes this bias in expectation, computed on the fly from the quantization step sizes of the cached keys and the query norm. Using a second-order Taylor approximation, the additional computational overhead is negligible, and no additional memory is needed alongside the cache. Evaluated on MAGI-1, SkyReels-V2, and HY-WorldPlay at INT2 quantization, our correction recovers most of the quality lost to aggressive quantization, reaching near-BF16 video quality, and can outperform INT4 quantization while using 50% less memory.

Tuna Tuncer, Felix Becker, Thomas Pfeil

1Technical University of Munich · 2Tensordyne

Autoregressive Video Diffusion Models Key-Value Cache Quantization Video Quality Block Sparse Flash Attention Key-Value Cache Compression Key-Value Cache+1

Ask PDF

May 21, 2026cs.LG

ThriftAttention: Selective Mixed Precision for Long-Context FP4 Attention

Efficient attention algorithms are critical to mitigate the quadratic cost of attention in long-context workloads. Prior work utilises block-scaled quantisation techniques on Blackwell GPUs to move attention computation to 4-bit precision to accelerate inference. However, these techniques result in significant quality degradation in long-context settings. We show that the output impact of quantisation error is highly non-uniform and increases with the importance of each query-key interaction, concentrating functionally relevant error in a small number of attention blocks that contain the most important tokens. We propose ThriftAttention, a low-bit attention variant that delivers near-FP16 long-context quality at FP4 inference efficiency. This approach proceeds in two stages. First, a heuristic rapidly selects a small number of important query-key block pairs for FP16 precision. Second, the selected blocks are computed in FP16 and the remaining blocks in FP4, with both paths merged via online softmax into a single output. We demonstrate across long-context benchmarks and model families that by computing only 5% of query-key blocks in FP16, ThriftAttention recovers on average 89.1% of the FP4-to-FP16 performance gap. We show ThriftAttention's advantage grows with sequence length, mitigating the systematic FP4 quality degradation observed at longer contexts. The code is available at https://github.com/joesharratt1229/ThriftAttention.

Joe Sharratt

Block Sparse Flash Attention Efficient Long-Context Inference Quantizer LLM Inference Optimization

Ask PDF

May 14, 2026cs.LG

DualKV: Shared-Prompt Flash Attention for Efficient RL Training with Large Rollouts and Long Contexts

Modern RL post-training methods such as GRPO and DAPO train on

N

response sequences of

R

tokens sampled from a shared prompt of

P

tokens, but standard FlashAttention replicates all

P

prompt tokens

N

times across both forward and backward passes -- duplicating compute and memory on identical hidden states. In large-rollout, long-context RL training (

N{\geq}16

P{\geq}8\text{K}

), this redundancy dominates the policy update cost. We observe that in decoder-only models, causal masking makes prompt representations invariant across sequences at every layer, so all per-token operations (norms, projections, MLP) and attention can process the prompt once -- a property not yet exploited at the kernel level for training. We propose \textbf{DualKV}, the first FlashAttention kernel variant that eliminates shared-prompt replication during RL training, via (1)~fused CUDA forward and backward kernels that iterate over two disjoint KV regions -- shared context and per-sequence response -- in a single kernel launch, and (2)~a data-pipeline redesign in veRL that repacks

N(P{+}R)

tokens into

P{+}NR

tokens per micro-batch, extending the token reduction from attention to the entire model by a factor

ρ= N(P{+}R)/(P{+}NR)

. DualKV is mathematically equivalent to standard attention and introduces no approximation. On Qwen3-8B GRPO training with 8

\times

H100 GPUs (

N{=}32

, 8K-context), DualKV achieves

1.63

2.09\times

policy-update speedup, enables

2\times

larger micro-batches, and raises MFU from

36\%

76\%

. Similar gains hold for DAPO (

2.47\times

speedup,

77\%

MFU). At 30B MoE scale on 16

\times

H100, DualKV achieves

3.82\times

policy-update and

3.38\times

end-to-end step speedup over FlashAttention (which requires 4-way Ulysses sequence parallelism to avoid OOM).

Jiading Gai, Shuai Zhang, Xiang Song +2

1Amazon Web Services · 2Google · University of Minnesota

Block Sparse Flash Attention Efficient Long-Context Inference Decode Speedup Offline Reinforcement Learning Responses

Ask PDF

Apr 28, 2026cs.LG

QFlash: Bridging Quantization and Memory Efficiency in Vision Transformer Attention

FlashAttention improves efficiency through tiling, but its online softmax still relies on floating-point arithmetic for numerical stability, making full quantization difficult. We identify three main obstacles to integer-only FlashAttention: (1) scale explosion during tile-wise accumulation, (2) inefficient shift-based exponential operations on GPUs, and (3) quantization granularity constraints requiring uniform scales for integer comparison. To address these challenges, we propose \textit{QFlash}, an end-to-end integer FlashAttention design that performs softmax entirely in the integer domain and runs as a single Triton kernel. On seven attention workloads from ViT, DeiT, and Swin models, QFlash achieves up to 6.73

\times

speedup over I-ViT and up to 8.69

\times

speedup on Swin, while reducing energy consumption by 18.8% compared to FP16 FlashAttention, without sacrificing Top-1 accuracy on ViT/DeiT and remaining competitive on Swin under per-tensor quantization. Our code is publicly available at https://github.com/EfficientCompLab/qflash.

Sehyeon Oh, Yongin Kwon, Jemin Lee

University of Science and Technology, Daejeon, Republic of Korea · Electronics and Telecommunications Research Institute, Daejeon, Republic of Korea · Pusan National University, Busan, Republic of Korea +1

Block Sparse Flash Attention Quantized Vision Transformer Flash Peak Memory Activation Functions

Ask PDF

Apr 16, 2026cs.LG

Dispatch-Aware Ragged Attention for Pruned Vision Transformers

Token pruning methods for Vision Transformers (ViTs) promise quadratic reductions in attention FLOPs by dropping uninformative patches. Yet standard variable-length attention APIs -- including FlashAttention-2's varlen and PyTorch's NestedTensor SDPA -- fail to translate these savings into proportional wall-clock gains at the short post-pruning sequence lengths typical of ViTs (

\leq

197 tokens). We identify a dispatch-overhead bottleneck: at these lengths, host-side kernel dispatch consumes

{\sim}

50,

μ

s regardless of workload, exceeding the actual GPU compute time at moderate-to-high pruning rates. We present a lightweight bidirectional Triton attention kernel whose dispatch floor is

{\sim}

24,

μ

s -- roughly 2.17

\times

lower than FlashAttention-2 varlen -- allowing pruning savings to become visible in wall-clock time. Integrated into a complete pack-attend-unpack pipeline and evaluated on an NVIDIA RTX 4000 Ada Generation GPU, our system achieves 1.88

\times

end-to-end throughput over padded PyTorch SDPA at standard 224

\times

224 inputs, scaling to 2.51

\times

at 384

\times

384. Against FlashAttention-2 varlen -- the strongest baseline -- our kernel delivers 9-12% higher throughput at serving batch sizes (BS=1-4), and 2.17

\times

lower kernel latency at 80% token pruning. Numerical correctness is verified with max absolute logit difference

<

0.004 and bit-exact top-1 predictions.

Seifeldin Abdellatif, Ahmad Almasri

College of Engineering Al Ain University Abu Dhabi, United Arab Emirates

Visual Token Pruning Block Sparse Flash Attention Self-Supervised Vision Transformers Vision Transformer Dynamic Sparse Attention Pytorch

Ask PDF

Dec 7, 2025cs.LG

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

Block Sparse Flash Attention Efficient Long-Context Inference LLM Inference Optimization Bottlenecks

Ask PDF

Momentum

Related topics

Latest papers 29