Long-Context Attention Benchmark: From Kernel Efficiency to Distributed Context Parallelism
Authors: Tao Bu
Qiangang Wang
Bowen Zeng
Hanwen Sun
Yunpeng Huang
Chun Cao
Jingwei Xu
Venue: ICLR 2026 Poster (2026)
Keywords: Long Context
Dense Attention Kernel
Sparse Attention Kernel
Context Parallel Machenism
Mask Pattern
Primary area: datasets and benchmarks
TL;DR: We build a unified benchmark to evaluate attention kernels and context parallel mechanisms for ultra-long sequence training, highlighting efficiency, scalability, and trade-offs.

Abstract: Transformer-based large language models (LLMs) have achieved remarkable success, yet their standard softmax-operator-based attention mechanism incurs quadratic computation and memory costs with respect to sequence length, posing a major bottleneck for long-context training. Prior work tackles this challenge along two directions: (1) kernel-level optimizations, which accelerate dense and sparse attention operators; and (2) module-level strategies, often referred to as distributed attention or context parallel training, which scale attention across multiple devices. However, systematic evaluation still remains limited: operator-level comparisons are often incomplete, while context parallel strategies are typically framework-specific, with unclear performance analysis across contexts. To address these gaps, we propose a unified benchmark that integrates representative attention kernels and context parallel mechanisms with a modular and extensible interface for evaluation. The benchmark evaluates methods along two critical dimensions: (1) attention mask patterns, which strongly affect efficiency, scalability, and usability, and (2) sequence length and distributed scale, which determine performance under extreme long-context training. Through comprehensive experiments on the cluster of up to 96 GPUs, our benchmark enables reproducible comparisons, highlights method-specific trade-offs, and provides practical guidance for designing and deploying attention mechanisms in long-context LLM training.

Published as a conference paper at ICLR 2026
LONG-CONTEXI ATTENTION BENCHMARK: FROM KERNEL EFFICIENCY TO DISTRIBUTED CON- TEXT PARALLELISM
Tao Bu' Qiangang Wang'* Bowen Zeng? Hanwen Sun
Yunpeng Huang' Chun Cao' Jingwei Xu
State Key Laboratory for Novel Software Technology. Nanjing University, China
2 ZZhejiang University, China
Peeking University, China
(butao, qgwang, hsmals iiniiln... edu. cn, zbw.cs csez ju. ju.jjueeee edu on sunhamwen& stu. www.enesssuuuuuuuee.e.mmm edu. {caochun, jingweix) jn edu, cn
ABSTRACT
Transformer-basel large language models (LLMs) have achieved remarkable success, yet their standard softmax-operator-basee attention mechanism incurs quadratic computation and memory costs with respect to sequence length, pos- ing a major bottleneck for long-context training. Prior work tackles this challenge along two directions: (1) kernel-level optimizations, which accelerate dense and sparse attention operators; and (2) module-level strategies, often referred to as distributed attention or context parallel training, which scale attention across mul- tiple devices. However, systematic evaluation still remains limited: operator-level comparisons are often incomplete, while context parallel strategies are typically framework-specific. with unclear performance analysis across contexts. To ad- dress these gaps, we propose a unified benchmark that integrates representative attention kernels and context parallel mechanisms with a modular and extensible interface for evaluation. The benchmark evaluates methods along two critical di- mensions: (1) attention mask patterns, which strongly affect efficiency, scalability, and usability, and (2) sequence length and distributed scale, which determine per- formance under extreme long-context training. Through comprehensive experi- ments on the cluster of up to 96 GPUs, our benchmark enables reproducible com- parisons, highlights method-specine trade-offt and provides practical guidance for designing and deploying attention mechanisms in long-context LLM training.
INTRODUCTION
The Transformer architecture, powered by the attention mechanism, has become the foundation of
large language models (LLMs) (Achiam et al.] 2023 Team et al. 2222 Dubey et al., 2024). With
the guidance of the scaling law (Kaplan et al. 2020 lay al. 2022), ), current state-of-the-art LLMs, such as GPT (Agarwal et al. 2025). Gemini (Team et al. (024). and DeepSeek (Liu et al. 2024),
contain billions of parameters and are trained on trillions data using large-scale distributed GPU clusters. However, as model size and training data continue to grow, the computational and mem- ory costs of conventional attention scale quadratically with sequence length, posing a fundamental efficiency bottleneck for large-scale LLM training (Dai et al. (0224) Although the context window of LLMs has expanded dramatically from 4K tokens to I28K (Grattafiori et al.) 2024), 1M (Yang ct al. 2025a), and even 10M (Team et al. 2024) tokens, the design and performance characteris- ties of long-context attention mechanisms at these seales in distributed training remain insufficiently understood (Gao et al.; 2024b).
Recent research on efficient long-context attention at scale has progressed along two main directions. The first focuses kernel-level optimizations (Zhang et al.), such dense and sparse kernels,
Equaal contribution
1 Corresponding author

---

Published as a conference paper at ICLR 2026
reducing attention complexity single GPU. The second emphasizes module-Ivvel designs, or context parallelism (Duan et al. 2024), which partition long sequences (c.g., 32K-128K tokens) across multiple GPUs with tailored communication and scheduling for scalability. Despite these advances, comprehensive analyses of long-context attention mechanisms remain lacking. Attention operators differ significantly in their support for mask patterns, and even the exhibit substantial performance variation across masks. Currently, no unified evaluation has been established. Furthermore, existing context parallel attention mechanisms are often tightly integrated with specific training frameworks (e.g., DeepSpeed (Rasley et al., 2020) and InternEvo (Chen et al.) 2024)) which limits reusability and hinders systematic comparison. As a result, researchers lack clear understanding of the trade-off's between methods, and practitioners have reliable benchmark or reference to guide the selection of attention mechanisms long-context training.
operator can
To address these issues, we collect representative attention operators and context parallel mecha- nisms, and design unified framework to systematically benchmark their capabilities, limitations, and potential risks in ultra-long context training. In our framework, we establish unified data prepa- ration interface that supports both non-distributed kernels and context parallel attention mechanisms, enabling fair evaluation across methods. Specifically, for non-distributed kernel scenarios, we in- tegrate a variety of dense and sparse attention kernels, implementing standardized interfaces that eliminate inconsistencies in data representation and ensure comparability under the benchmark. For distributed scenarios, we reconstruct and optimize representative context parallel attention mech- anisms within the unified framework, providing efficient, scalable implementations with modular interfaces. Building on the foundation, we conduct large-scale experiments and in-depth analyses along two critical dimensions: (I) attention mask patterns (up 14 mask patterns), which are of- ten overlooked but have significant impact on efficiency, scalability, and usability, and (2) context length and distributed scale, where we systematically evaluate performance trends and capability limits as both the input length and distributed scale grow, reaching up to 768K on 96 GPUs. We hope our results offer valuable insights for research on long-context training large models, as well as for the design and development of next-generation distributed attention mechanisms. Our contributions are threefold:
1. Unified benchmarking: we provide standardized framework with consistent data prepa- ration for fair evaluation attention mechanisms across diverse long-context scenarios.
2. Modular components: unify dense and sparse kernels under high-level modular inter- face, and provide optimized distributed attention terms of context parallelism.
3. In-depth analysis: conduct extensive experiments across dense long-context scenar- ios to identify key factors affecting attention efficiency and scalability, providing valuable guidance for ultra-long context training and development.
MANGCA-BENCH
LongCA-bench a benchmark designed evaluate the efficiency of long-context attention across
both single-device kemnels and distributed context parallel mechanisms. The benchmark consists of
three core components: (1) a unified data preparation interface that standardizes preprocessing, (2)
unified input representation interface that supports 7 dense and 5 sparse attention kernels, and (3)
an optimized context parallelism framework that incorporates 5 distributed attention mechanisms. Together, these components provide a systematic and extensible platform for analyzing long- long-context attention, enabling fair comparisons across operator-level efficiency and distributed scalability.
2.1 DATA PREPARATION
We first describe the data preparation process in the benchmark. To generate inputs practical for long-context attention benchmarking, we introduce a dedicated data preparation interface. Rather than directly using the downstream datasets, our interface combines diverse mask patterns with variable length of sequence sampling, ensuring that the evaluation data accurately reflects the char- acteristics and challenges of long-context training.
2

---

Published as a conference paper at ICLR 2026
Dense
Spaars Masks
Mask Generator
Data Preparation
Input Sampling Interface
Batasets Distribution
Masks
Data Component
Long-Context Attention Benchnark Evaluation
Harbbare--WiiA
Kernel Component
Module Component
Input Representation Interface
Distributed Fav Setup Interfuce Partition Interface
Dense Kernel
Distributed Nadule
Naive Attention
Flash Attention
Flex Actention
Flash Infer
VSA
VSA Triton
Ulysses
Ring Attention (P2P, All All-Gather)
Tarch Sdpa
CuDNN Fused Attention
Flash Mask
FlashAttentian2 Spanse
Flex Actention
Loong
Figure 1: The architecture of LongCA benchmark
Dral
Gond fliing Winer 2911
Blak
Static
Figure 2: Attention mask patterns
2.1.1 INPUT MASK PATTERNS
Diffcrent tasks require specific mask types based on the training scenario. In LongCA-bench, we categorize total of 14 mask patterns into major classes (see Figure 2 12 static masks (6 regular and 6 heterogeneous), and dynamic masks. The key distinction lies in whether the mask can be predetermined before training or must be generated adaptively during the training process.
Static regular mask. The FULL and CAUSAL masks are the most widely used in training Vaswani al. 2017). Considering the document-level variants, FULL DOCUMENT and CAUSAL DOCU- MENT are employed for efficient sequence packing and in-batch/in- token processing (Krell et al. 2021 Dehghani et al. 2023). In addition, by applying the sliding window variants, FULL SLIDING WINDOW and CAUSAL SLIDING WINDOW can leverage sparsity to balance computational cost
and token coverage (Beltagy et al. 2020p.
Static heterogeneous mask. SHARED QUESTION mask used in reward models allows multiple answers to share the same question (Ouyang al. (022). GLOBAL SLIDING mask is designed to effectively capture both global context and local details İZaheer et al. 2020). CAUSAL BLOCK- WISE mask, which is widely adopted in-context learning, restricts demonstrations to local blocks while letting the test example attend globally, supporting long-context evaluation (Bertsch et al.) 2024). The PREFIX LM CAUSAL and PRÉFIX DOCUMENT masks are specifically tailored to introduce prefix for language modeling tasks (Raffel et al. 2020 BLOCK CAUSAL DOCU- MENT mask combines the block and document concepts and is widely used multimodal model training (Zewei & Yunpeng 2025).
Dynamic mask. In long-context scenarios, block sparse masks reduce computational latency and memory usage by restricting attention computation to the most salient blocks of the input. Since

---

Published as a conference paper at ICLR 2026
the selected blocks depend the contextual input, the mask pattern varies across examples. Block sparse masks have been widely adopted in both natural language processing (Lu et al. 2025 Yuan
et al. 2025 Xu et al. 2025| Guo et al. 2024; Ye et al. (2025) and visual generation Zhang et al.
2025d Zewei & Yunpeng 2025; Yang et al. 2025b). We categorize block sparse masks into two
types: uniform and variable (Figure 2). The uniform block mask applies attention blocks of fixed
size (e.g., 64x64) across the entire attention map and computes the selected blocks during attention. In contrast, the variable block mask provides greater flexibility by allowing blocks of different sizes, offering more efficient and expressive representation of sparse attention patterns.
2.1.2 INPUT DATA SAMPLING
Dense data sampling. The mask specifies the area of attention interactions within a context win- dow. As the context window expands (e.g., from 8K to 768K), input data sampling becomes crucial. To ensure that the benchmark data reflects realistic training scenarios, we analyzed several widely used public pretraining datasets, including Pile (Gao et al. 2020), ProLong64K (Gao et al. 2024a), ProLong5128 (Gao et al. (024a),
Sapaajana--Pr-SSurcce
(Yaofu 2024)) Open-
WebText (Gokaslan et al., 2019). and C4 (Raffel et al.) 2020) (see Appendix ...2 for details). Fol- lowing prior findings (Fu ct al. 2024 Gao et al. 2024b), we note that: (1) language model training typically requires datasets from diverse sources; (2) extending the context length requires maintain- ing domain diversity while upsampling long- uin-sequence samples; and (3) mixing long- long-context sources (e.g., code repositories and books) with high-quality short-text data improves long-context modeling without sacrificing overall performance. In our benchmark, the data sampling method therefore uses the Pile dataset for samples up to 8K, ProLong6-K for long-context samples up to 64K, and Pro- Long5 12K for ultra-long samples up to 768K. This combination ensures that evaluation data reflects realistic training scenarios across different context scales.
Sparse data sampling. Mask generation for block sparse attention in our benchmark follows stan- dard methodology (Xia et al. 2025 Lu et al. 2025 Yuan et al., 2025) Zhang et al.] 2025d). The attention matrix is first partitioned into wineeiinnsional grid of blocks, with either uniform or variable pre-defined block sizes. An importance score is then computed for each block in the grid, which in multi-head attention may be assigned on a per-head or per-group basis. Guided by target sparsity ratio, top-K selection is performed for each query block to identify the most salient key blocks for attention computation. In our benchmark, however, we simplify the process to specif- ically evaluate kernel performance under varying sparsity levels (c.g.. 0.2, 0.5, and 0.8). Instead of computing explicit important scores, we simulate the scoring and selection process by randomly generating block masks to achieve the desired sparsity. To create real-world workloads, we evaluate sequence lengths from 32K to 128K, sampled at 32K intervals. This range is derived from analyzing prominent benchmarks for block sparse attention's 's primary applications in video generation (c.g., VBench (Huang et al. (024)) and LLMS (e.g., RULER (Hsieh et al., 2024)). Our evaluation covers both MHA and GQA using uniform block sizes of 64x64 and 128x 128 (see Appendix A.3 A.3 for full results).
2.2 ATTENTION KERNEL.
Efficient attention kernels aim to reduce time and memory complexity without compromising ex- pressiveness. The most straightforward approach is hardware acceleration, which speeds up compu- tation without altering the original attention logic. Another common strategy leverages the inherent sparsity of attention, skipping unnecessary computations, often guided by a dynamic sparse mask.
Dense attention kernel. We integrate seven dense attention kernels and categorize their support for different mask types (see Table []. Since kernels often apply for different requirements on data structure and mask formats, we implement dedicated adapter interfaces for each. These interfaces generate kemel-speecfic input representations from unified data format, climinating inconsisten- cies in data expression across kernels. This design simplifies input preparation for diverse mask scenarios, ensures comparability within the benchmark, and provides unified solution for future kernel extensions.
As baselines in our benchmark, we include both the step-by-step naive attention and PyTorch's 's fused scaled dot product attention (SDPA) (PyTorch Contributors (024b), both construct full 2D masks and theoretically support arbitrary masking patterns. For hardware-optimizeed kernels,
4

---

Published as a conference paper at ICLR 2026
Table
Dense kernel support across mask patterns
Mask Type
Naive-Torch
SDPA
FA2
FA3
cuDNN-Fused--ae
FlexAttn
FlashMask
FULL
CAUSAL
FULL SLIDING WINDOW
CAUSAL SLIDING WINDOW
DOCUMENT
CAUSAL DOCUMENT
V
SHARE QUESTION
V
K
X
K
CAUSAL BLOCKWISE
PREFX LM CAUSAL
X
PREPI LM DOCUMENT
Table 2: Characteristics of sparse kernels
Sparse Kernels
Characteristics
VSA
Triton VSA
FA2 Sparse
FlexAttention
FlashInfer
Linform/Variahle Masks
Uniform only
Uniform only
Uniform only
Both
Both Forward only
Fonwart/Backwwrd
only
Block Size
64 only
64 only
128 only
Arbitrary
Arbitrary
Support
GPU Support
sm10
> sm88
> sm80
sm800
> sm80
Performance
High
Medium
Medium
Low
Medium
Memory Overhead
Low
Low
Low
High
Mediam
integrate the FlashAttention series, including FA (Dao al.) 2022), FA2 Dao 2023), FA3 (Shah et al., [2024), as well as cuDNN fused kernels (NVIDIA Corporation 2025). These kemnels em- ploy advanced techniques such as shared memory, block- partitioning, warp scheduling. FP8, and asynchronous processing. For flexible kernels, we integrate Flex Attention (Flex) (Dong et al.) 2.24]. a general fused operator with memory complexity close O(S²), where S denotes the sequence length, that generates specialized kernels based on per-position boolean functions and en- ables compatibility with arbitrary masks. also include FlashMask (Wang et al. 2025a), which introduces column-wise representation to optimize heterogeneous computation.
Sparse attention kernel. Sparse attention significantly reduces the computational complexity of attention for long sequences. Due to its versatile mask representation. block sparse attention is widely used in state-of-the, of-hh--att sparse attention methods (Lu et al.) 2025 Zhang et al.) 20225b] Yuan
al. 2025 Zhang al. 2025d Xu et al. 2025). Therefore, our benchmark incorporates five block sparse attention kernels to evaluate long-context sparse attention.
We categorize these kernels into two main types. The first type consists of dedicated block sparse attention kernels, which are highly optimized for sparse patterns with uniform block sizes (e.g.. 64). Representative implementations include VSA (Zhang al. 2025d), its Triton-based ver- sion (Triton VSA), and the FlashAttentiinn22buueee 2024). The second type comprises general- purppss sparse attention kernels, which offer greater flex- ibility and support arbitrary block structures. They are compatible with both uniform and variable block sparse masks. This category includes FlexAttention (Dong et al., (0024) and FlashInfer (Ye) et al., (025). These kernels exhibit different characteristics, as summarized in Table 2 (refer to Appendix A.4 for details). We evaluate performance through comparisons using two mask types: uniform block mask and variable block mask. Note that backward computation in training is sup- ported by only a limited set of block-sparse attention kernels. For comprchensiveness, we select FA2 Sparse and FlashInfer, two inference-side methods, for comparisons in benchmark.
block sparse attention (FA2 Sparse) (Guo et al.)
2.3 DISTRIBUTED ATTENTION MECHANISM
In our benchmark, we reproduce and optimize 5 representative distributed attention mechanisms under unified framework, including Ulysess, Ring P2P, Ring All-Gather, USP, and LoongTrain. Train. We establish unified infrastructure that standardizes distributed setup and sequence partitioning, ensuring a consistent invocation protocol across all methods. The integrated distributed attention mechanisms be categorized into three architectural designs:
5

---

Published as a conference paper at ICLR 2026
partitions both the sequence and head dimensions in multi-head attention, using All-to-All communication to switch parallel dimen- sions. This approach is simple. general, and numerically precise, but the scalability is constrained by the number of attention heads, particularly under GQA, MQA, or tensor parallelism.
All-to-all based design. DeepSpeed's Ulysses (Jacobs al. 2023)
Ring P2P based design. Ring P2P (Liu al. 2023) uses multi-round ring-structured point-to-poing communication, while Ring All-Gather (NVIDIA) 2025) performs a single all- all-gather of key-value tensors, relying on ring topologies. These approaches exhibit strong scalability and naturally overlap computation with communication via pipelining, However, they suffer from lower efficiency and potential numerical error accumulation.
Hybrid design. USP (Fang& & Zhao, 2024) and LoongTrain (Gu et al. 2022) extend Ulysses and ring-based designs into a two-dimensional scheme. An inner layer applies Ulysses with All-to- All for intra-node bandwidth, while an outer layer uses ring-based attention to enhance scalability and enable compute-communice communication overlap. LoongTrain further proposes DoubleRing Attention, enhancing Ring P2P with a two-level sliding window to improve communication efficiency.
In our reproduction and optimization, we draw inspiration from TransformerEnginne (NVIDIA 2025), achieving perfect load balancing through double-parallel partitioning combined with head- to-tail reordering (zhuzilin) 2024). We also incorporate optimizations such as double buffering and multi-stream overlap of computation. For each method, implement backend support for both Flash Attention v3 and cuDNN Fused Attention operators. We extend the input layout to a variable- length (varlen) format, allowing multiple sequences of different lengths to be concatenated along the sequence dimension while handling padding tokens. This ensures the flexibility and usability of varlen inputs under different distributed scales. Since varlen inputs can introduce substantial syn- chronization and waiting overhead across devices, we precompute the necessary meta-information for all distributed strategies as a one-time preprocessing step, thereby minimizing distribution- related performance degradation. Despite these extensions and optimizations, our benchmark re- mains constrained by the underlying distributed attention designs, thus currently supporting only FULL, CAUSAL, FULL/CAUSAL DOCUMENT masks.
3 EVALUATION
In this section, we present experiments evaluating the speed and memory efficiency of different attention methods under long-context scenarios. The speed is measured in TFLOPs/s metric, and peak memory usage is reported in gigabytes (GB), Kernel performance is evaluated on single GPU, while the performance of distributed context parallel attention is assessed across multi-GPU clusters of varying scales. All context parallel evaluations are performed on NVIDIA H100 GPUs (80GB HBM3). Kernel-level experiments are conducted on both NVIDIA H100 and NVIDIA A800 GPUs. We primarily present and analyze the results on H100 here, with most A800 results deferred to the appendix. The code is publicly available for the community!
3.1 DENSE ATTENTION KERNEL PERFORMANCE
We evaluate dense kernels across 12 static mask configurations to assess both expressiveness and efficiency. Sequence lengths range from 1K to 48K, with BFloat16 precision, a hidden dimension of 128, and two head settings: GQA (64:8) and MHA (64:64)) We record forward and back- ward throughput as well as peak memory usage (see Appendix A.5 for full results on H100, see Appendix A.9 for full A800 results).
Figures and report TFLOPs at 8K sequence length under GQA (64:8), where X denotes un- supported configurations. The six groups on the left correspond to static regular masks, and the remaining on the right show static heterogeneous masks. Note that the FA series and cuDNN fused kernels do not support heterogeneous masks. In particular, cuDNN fused kernel does support the FULL SLIDING WINDOW mask with GQA (64:8), though other configurations are recommended.
Although the naive implementation and Torch SDPA thcoretically support arbitrary masks, their quadratic complexity leads to severe efficiency degradation and excessive memory overhead, mak-
'The implementation is accessible at: wttp://ttw..cccooo///eeeeeeeeeee http:://gttt.ccccoou//uuuggggggg Peadd headge

---

Published as a conference paper at ICLR 2026
GQA Fwd Flops (Bk)
fs3
f=2
flxx
flashmmakk
sipa
torch
720
Speed (taraFlo
6.00
630
330
130
F3W
C3W
Figure 3: Forward TFLOPs of dense kernels with different masks (8K length)
ing them impractical in long-context settings. Under computation-intensivee dense settings (e.g.- FULL or CAUSAL), SDPA achieves performance comparable to general fused operators such as FlexAttention. These results provide baseline for fused operators without bardware-speecin mizations. FlashMask, another generic fused operator, leverages a column-wise mask representation to mitigate computational sparsity. While optimized for heterogeneous masks, its column- -wie rep- resentation cannot cover all scenarios, making less general than FlexAttention.
opti-
For regular scenarios, FA series and cuDNN fused attention are all hardware hardware-optimized kernels. On H100 GPUs, FA3, specifically optimized for the Hopper architecture, achieves the best performance. cuDNN fused attention supports multiple architectures but imposes stricter constraints on data pat- terns (e.g., GQA (64:8) with FUL.I SL.DINGG WINDOW). While some of these limitations can be circumvented by preprocessing techniques such padding, doing so introduces extra overhead. Note that although FA2 and cuDNN fused attention yield lower performance, kernel selection should be guided by the target hardware architecture. In Figures 4 and 5 we present the main experimental results on both the H100 and A800.
3.2 SPARSE ATTENTION KERNEL PERFORMANCE
To comprehensively evaluate the functionality and computational efficiency of various sparse ker- nels, we include VSA (Zhang et al. 2025d). Triton VSA, FA2 sparse (Guo al. (2024) and Flash- Infer (Ye et al. 2025) in our evaluation (see Appendix(AA6]]aa full results). We perform kernel-level evaluations across two kinds of block sizes (64 and 128), both forward and backward computation, two attention variants (MHA (64:64) and GQA (64:8)), and sequence lengths ranging from 32K to 128K. Note that Flex AtexAttenionn (Dong al. 2024) excluded due severe out-of-memory (OOM) issues originating from its mask representations.
From functionality perspective, comparisons in Figures 4 (a) and (b) reveal that FA2 sparse does not support a block size of 64, while FlashInfer lacks backward computation. As shown in Fig- ures (a), (c), and (d). VSA does not support a block size of 128. Due to its specific design for the MHA architecture in video diffuusion models, VSA does not currently support GQA. In contrast, FlashInfer is prone OOM errors at longer sequence lengths and smaller block sizes, stemming from the substantial metadata storage it requires.
These limitations highlight the need for further engineering optimizations in block-sparss kernels. Backward computation is essential for trainable sparse attention, particularly in GQA and MHA. Flexibility across block sizes required to support diverse sparse attention designs, and memory challenges in block-sparse mask representations also need to be addressed.
From a performance perspective, Figures (a) and (b) show that VSA outperforms both Triton VSA and FlashInfer, while Figures and (d) indicate that FlashInfer outperforms fa2 sparse. Across all kernels, the forward pass consistently achieves higher percentage theoretical TFLOPs than the backward pass (with theoretical TFLOPs taken from FA3 in Figure (). Additionally, Figures (c) and (d) show minimal performance differences between MHA and GQA, though GQA achieves better GPU memory efficiency. By comparing Figures 4 (a) and (c), observe that FlashInfer
7

---

Published as a conference paper at ICLR 2026
H0.00
fo2_spame spans
flashinter
Via
Speed (teraFLOPs/s)
400
300
200
100
22k
54c
36k
120k
324
54k
1,25c
52k
64k
1288
32%
648
12.5k
Sequence Length
Figure 4: Performance results (TFLOPs) of sparse kemelss with 50% sparsity ratio on H100 GPU
ABC0
fa2 sparse
Tashinfar
Via triton
Speed (teraFLoPPiN
200
1.0
50
611 MIAA, Bloccm54
SEc
6.0
0.1k
12
122-
6.02
54k
128c
120K
(6) Bwd.
Pws,
MIA
Siccb-128
Sequance Length
Figure Performance results (TFL.Ps) sparse kernels with a 50% sparsity ratio on A800 GPU
performs significantly better with block size of 128 than with 64, suggesting that larger block sizes are more effective for achieving higher performance.
Overall, these results demonstrate that performance in block sparse attention significantly im- proved through specialization, where kernels tailored to particular parameters (e.g., block size or hardware architecture) consistently outperform general implementations. Meanwhile, the backward pass remains major bottleneck, underscoring urgent need for optimization. A key future di- rection the development more flexible, comprehensive kernels that deliver high performance across wide range of block sizes. Achieving this goal requires moving beyond mangl-parrameer tuning toward deeper, hardware-level optimizations.
3.3 CONTEXT PARALLEL ATTENTION PERFORMANCE
We evaluate four mask patterns: FULL, CAUSAL, FULL DOCUMENT, and CAUSAL DOCU- MENT. The per-device sequence length is fixed at 8K with hidden size 128, validated under the GQA (64:8) setting. Experiments are conducted on NVIDIA H100 GPUs, scaling from to 96 GPUs across 12 servers, with total context windows from 64K 768K. Since Ulysses requires di- visibility constraints, GQA is converted to MHA by replicating KV heads (X indicates divisibility failure). Performance under FULL DOCUMENT shown Figures and 24 with additional details provided in Appendix A.7
Context parallel attention inevitably involves distributed communication, raising
major con- cerns: (1) whether communication effectively overlaps with computation, and if not, how efficient the communication is; and (2) whether workload balance is achieved in terms of both data vol- ume and computation across devices. Ideally, a context parallel strategy should behave close to non-distributed setting. Inter-node communication constitutes the dominant bottleneck compared to computation and intra-node communication. In experiments, we fix the large-load AllToAll groups within the mixed architecture to 8 per node, while the small-load P2P groups are placed across nodes and scale with the number of nodes. For the secondary P2P communication groups in Loong Train, we adopt a balanced configuration (e.g., 12 = 3 x 4) to maximize inter-node bandwidth utilization. All experiments performed the FA3 backend for consistency.
Ulysses' AllToAl communication is entirely exposed outside the computation. Thanks to its col- lective communication pattern (NVIDIA Corporation, 2020) with low communication overhead and
8

---

Published as a conference paper at ICLR 2026
full Hull_document fwd. RRop
620
-1:30
ulysses
ring_p2p
ring. ring_allgather
loongtrain
Speed (taraF(oPa))
4.00
320
130
8
16
24
32
40
48
56
64
80
96
(#GPUs)
Figure 6: Forward TFLOPs of Context Parallel Attention FULL DOCUMENT
Table Communication volume per device of Context Parallel Attention
Contect Parallel Attention
Stagr
Uyyees
Ring All-Gather
Ring P2P
USP
Train
Nad
Nel chuud
Alka Auld
Backward
A Bnadd
its head-sharded computation pattern leading to perfectly balanced workloads, Ulysses still delivers solid performance. However, its scalability bounded by the number of attention heads. load- balanced Ring P2P ensures that each GPU processes the same amount of computation and commu- nication per iteration. However, Ring P2P communication mask- mask-independent, always transferring in fixed ratio of D/N, where D is the total data and N is the world size, meaning performance de- pends entirely on the amount of computation workload. Ring P2P performs optimally the FULL scenario. However, in the DOCUMENT scenario, variable-length padding depends scale and sampling, leading to noticeable per-GPU computation variation and performance fluctuations.
The hybrid architecture alleviates the above issues. While the intra-node AllToAll communication group still remains exposed outside computation, its per-communication volume reduced from D (N - 1)/N in Ulysses to D x (8 1)/N per group (one-way). Meanwhile, the inter-node Ring P2P computation volume increases from D D/N in pure Ring P2P to D/K, enabling USP and LoongTrai Train to achieve optimal performance improvements, where K denotes the size of the inter- node communication group. Additionally, Loong Train introduces a secondary P2P architecture to further improve inter-node bandwidth utilization, providing modest forward speedups compared USP. However, because the secondary architecture involves extra window synchronization, the Ring backward pass cannot directly continue from the forward state, negating the overall performance gains. Overall, the experimental results demonstrate that fully leveraging MHA by first partitioning the heads yields significant performance benefits.
3.4 COMMUNICATION ANALYSIS FOR CONTEXT PARALLEL ATTENTION
We provide theoretical analysis of the communication volume of the entire pipeline per device for each context parallel attention. For each mechanism, we describe the communication operators involved, their invocation timing, call frequency, and data types, and we further report the corre- sponding theoretical communication volume per device. We further include an analysis of com- munication-compuueee computation overlap. We denote query heads, AR. as the number of key or value heads, and as the hidden dimension. Full detailed analysis provided in Appendix A.8 present only the results in Table
the total sequence length, h, the number of
RELATED WORK
Long context language modeling. Models such as BERT (Devlin al., 2019) and GPT (Brown al. 2020) can process thousands tokens, supporting document and dialogue level tasks, with full document understanding and long range retrieval emerging key challenges. Recently, model con-

---

Published as a conference paper at ICLR 2026
text windows have expanded dramatically, from 4K tokens to 128K (Dubey al. 2024). 1M (Yang et al. 2025a), and even 10M tokens (Team et al. 2024). The ability to model ultra-long contexts enables continuous reference, reasoning, and summarization over extended input sequences. This enhances advanced capabilities such as long-text reasoning (Guo et al. 2025 Muennighoft et al.)
2025). improved in-context learning et al. 2025 Team al. 2024). efficient information com-
pression Lee et al. 2024 Wang et al. 2024), and multimodal understanding (Weng et al., 2024).
Attention kernels. Attention the core component in Transformers with the time complexity of O(r²) in terms of context length. Hardware-efficien attention leverages hardware features to
reduce the time and memory costs. [ao al. (2022); ; Dao (2023); [Shah et al.) (2024) employs
matrix tiling and kernel fusion. Zhhag et al. (2024b a 2025a) uses quantization to leverage low-bit Tensor Cores. Sparse Kernels use the inherent sparsity (Child et al. 2019, Zhang et al. 2025b) of
the attention map P Softmax (Q /vd) to accelerate computation. Other directions include KV cache compression (Zhao et al.) 2023a) via weight sharing (Ainslie et al.) (203)) low-rank decomposition (Liu et al. 2024) to reduce memory overhead without extra computation.
Parallelism for distributed training. Various parallel paradigms have been developed to tackle resource challenges in large-scale distributed model training. Data parallclism (PyTorch Contrib- utors 2024a; Rajbhandari al., 2020 Zhao et al. (023b) partitions data along the batch dimen- sion. Tensor parallelism (Shoeybi bi et al. 2219; Xu & You 2023; Wang et al. (022). pipeline par- allelism (Huang al., 2019 Li & Hoefler [202]] Narayaaan
al. 2019 Qi et al. (024) [39 -2], and expert parallelism Gale et al. [2023] Hwang et al. 2023 Li et al. 2023 Liu et al. (2024) par-
tition model parameters along different dimensions. Hybrid parallel strategies (Smith et al. 2022] Ge et al. 2025 Wang et al. 2025b) are used to meet diverse needs and balance computation and memory. However, these strategies cannot fully address activation memory overhead from ultra-long sequences. Context parallelism (Korthikanti et al. [023] [i et al. (2021) partitions data by sequence, overlap, balancing, scalability, and usability; many designs remain underexplored, and near-linear scalability is still difficult.
but faces challenges in cnmpulatioo-commmiccni
5 CONCLUSION
The complexity of distributed environments is far greater than that of single-device settings. In ultra-long context training, selecting or developing appropriate kernels and context parallel strategies poses significant challenges and requires substantial effort and resources. To address this, we present fair and unified benchmark for attention mechanisms in ultra-long context training, covering the spectrum from single-device kernels to large-scale distributed context paraliel methods. Although our work has limitations, it aims improve fairness in comparing different approaches, expose their performance trade-offs and constraints, and provide objective references to guide future research and development in ultra-long context training.
ACKNOWLEDGMENTS
We thankful to the anonymous reviewers for their helpful comments. This work is supported by Jiangsu Science and Technology Major Project (#BG2025005), the Collaborative Innovation Center of Novel Software Technology and Industrialization, and Jiangsu Wukong Intelligent Computing Digital Technology. Jingwei Xu (jingweix@nju.er edd.uuuueedeee...o en) is the corresponding author.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, llge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. Gpt-4 technical report. arXiv preprint arXiv:2303. 08774, 2023.
Sandhini Agarwal, Lama Ahmad, Jason Ai, Sam Altman, Andy Applehaum, Edwin Arbus, Rahul K Arora, Yu Bai, Bowen Baker, Haiming Bao, et al. gpt-oss-120b & gpt-oss-20b model card. arXiv preprint arXiv:2508,1099 2025.
Sand. ai, Hansi Teng, Hongyu Jia, Lei Sun, Lingzhi Li, Maolin Li, Mingqiu Tang, Shuai Han, Tianning Zhang, W. Q Zhang, Weifeng Luo, Xiaoyang Kang, Yuchen Sun, Yue Cao, Yunpeng
10

---

Published as a conference paper at ICLR 2026
Huang. Yutong Lin, Yuxin Fang, Zewei Tao, Zheng Zhang, Zhongshu Wang. Zixun Liu, Dai Shi, Guoli Su, Hanwen Sun, Hong Pan, Jie Wang. Jiexin Sheng, Min Cui, Min Hu, Ming Yan, Shucheng Yin, Siran Zhang, Tingting Liu, Xianping Yin, Xiaoyu Yang, Xin Song, Xuan Hu, Yankai Zhang, and Yuqiao Li. Magi-1: Autoregressive video generation at scale, 2025. URL https::/arxiis:
13211
Joshua Ainslie, James James Lee-Thorp, Michiel De Jong, Yury Zemlyanskiy, Federico I Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head check- points. arXiv preprint arXiv:2305. 13245, 2023.
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.055 1550000.05/5,, 150, 2020.
Amanda Bertsch, Maor Ivgi, Emily Xiao, Uri Alon, Jonathan Berant, Matthew R Gormley, and Graham Neubig. In-context learning with long-context models: An in-depth exploration. arXiv preprint ar Xixx:244550000000 2024.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Qiaoling Chen, Diandian Gu, Guoteng Wang, Xun Chen, YingTong Xiong, Ting Huang, Qing- hao Hu, Xin Jin, Yonggang Wen, Tianwei Zhang, et al. Internevo: Efficient long-sequence large language model training via hybrid parallelism and redundant sharding. arXiv preprint arXiv:2401.09149 2024.
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174 2016.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arviv preprint arXiv: 10509, 2019.
Liuyao Dai, Hao Qi, Weicong Chen, and Xiaoyi Lu. High-speed data communication with advanced networks large language model training. IEEE Micro, 44(2):31 40:3,04-00000 2024.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint wtw.iiv:337700869.1 2023.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in neural information processing systems, 35:16344-16333 2022.
DeepSeck-AL. Deepseck-v3.2-ox Boosting cfff- with deepseek sparse attention, 2025.
Mostafa Dehghani, Basil Mustafa, Josip Djolonga, Jonathan Heek, Matthias Minderer, Mathilde Caron, Andreas Steiner, Joan Puigcerver, Robert Geirhos, Ibrahim M Alabdulmohsin, et al. Patch n' n'pack: Navit, a vision transformer for any aspect ratio and resolution. Advances in Neural 2023.
Information Processing Systems,
Jacob Devlin, Ming-Was Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume (long and short papers), pp. 4171-41866 2019.
Juechu Dong, Boyuan Feng, Driss Guessous, Yanbo Liang, and Horace He. Flex attention: A programming model for generating optimized attention kemnels. arXiv preprint arXhv:2412.05496, 2024.
Jiangfei Duan, Shuo Zhang, Zerui Wang, Lijuan Jiang, Wenwen Qu Qinghao Hu, Guoteng Wang, Qizhen Weng, Hang Yan, Xingcheng Zhang, et al. Efficient training of large language models on distributed infrastructures: survey. arXiv preprint arXiv:2407.2000 2024.

---

Published as a conference paper at ICLR 2026
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, et al. The llama 3 herd of models. arXiv e-prints, pp. arXiv-2407, 2024.
Jiarui Fang and Shangchun Zhao. Usp: A unified sequence parallelism approach for long context generative ai. arXiv preprint arXiv:2405.07771 2024.
Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, and Hao Peng. Data engineering for scaling language models 128k context. arXiv preprint arXiv:2402. 10171,
2024.
Zichuan Fu, Wentao Song, Yejing Wang, Xian Wu, Yefeng Zheng, Yingying Zhang, Derong Xu, Xuetao Wei, Tong Xu, and Xiangyu Zhao, Sliding window attention training for efficient large language models. arXiv preprint arXiv:2502. 18845, 2025.
Trevor Gale, Deepak Narayanan, Cliff Young, and Matei Zaharia. Megablocks: Efficient sparse training with mixture-of-experts. Proceedings of Machine Learming and Systems, 5:288-304, 2023.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang. Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00m 00027, 2020.
Tianyu Gao, Alexander Wettig. Howard Yen, and Danqi Chen. Enabling large language models to generate text with citations. 2024a.
Tianyu Gao, Alexander Wettig, Howard Yen, and Danqi Chen. How to train long-context language models (effectively). arXiv preprint arXiv:2410.02660, 2024b.
Hao Ge, Junda Feng, Qi Huang, Fangcheng Fu, Xiaonan Nie, Lei Zuo, Haibin Lin, Bin Cui, and Xin Liu. Bytescale: Efficient scaling of I1m training with a 2048k context length on more than 12,000 gpus. arXiv preprint arXiv:2502. 2/231, 2025.
Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanic Tellex. Openwebtext corpus. http: 2019.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshor, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Ko- renev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez, Austen Gregerson, Ava Spataru. Baptiste Roziere, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, Danielle Pintz, Danny Livshits, Danny Wyatt, David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia-Olano, Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Francisco Guzmán, Frank Zhang, Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Anderson, Govind That- tai, Graeme Nail, Gregoire Mialon, Guan Pang, Guillem Cucurell, Hailey Nguyen, Hannah Kore- vaar, Hu Xu, Hugo Touvron, lliyan Zarov, Imanol Arrieta Ibarra, Isabel Kloumann, Ishan Misra, Ivan Evtimov, Jack Zhang, Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Ma- hadeokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong, Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jong- soo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia. Kalyan Vasuden Alwala, Karthik Prasad, Kartikeya Upasani. Kate Plawiak, Ke Li, Kenneth Hcafield. Kevin Stone, Khalid El-Arini, Krithika Iyer, Kshitiz Malik, Kuenley Chiu, Kunal Bhalla, Kushal Lakhotia, Lauren Cantala-Yeary Laurens van der Maaten, Lawrence Chen, Liang Tan, Liz Jenkins, Louis Martin, Lovish Madaan, Lubo Malo, Lukas Blecher, Lukas Landzaat, Luke de Oliveira, Madeline Muzzi, Mahesh Pasupuleti, Mannat Singh, Manohar Paluri, Marcin Kardas, Maria Tsimpoukelli, Mathew Oldham, Mathieu Rita, Maya Pavlova, Melanie Kambadur, Mike Lewis, Min Si, Mitesh Ku- mar Singh, Mona Hassan, Naman Goyal, Narjes Torabi, Nikolay Bashlykov, Nikolay Bogoy- chev, Niladri Chatterji, Ning Zhang, Olivier Duchenne, Onur Çelebi, Patrick Alrassy, Pengchuan Zhang, Pengwei Li, Petar Vasic, Peter Weng, Prajjwal Bhargava, Pratik Dubal, Praveen Krishnan,
12

---

Published as a conference paper at ICLR 2026
Punit Singh Koura, Puxin Xu, Qing He, Qingxiao Dong, Ragavan Srinivasan, Raj Ganapathy, Ra- mon Calderer, Ricardo Silveira Cabral, Robert Stojnic, Roberta Raileanu, Rohan Maheswari, Ro- hit Girdhar, Rohit Patel, Romain Sauvestre, Ronnie Polidoro, Roshan Sumbaly, Ross Taylor, Ruan Silva, Rui Hou, Rui Wang, Saghar Hosseini, Sahana Chennabasappa, Sanjay Singh, Sean Bell, Seohyun Sonia Kim, Sergey Edunov, Shaoliang Nie, Sharan Narang. Sharath Raparthy, Sheng Shen, Shengye Wan, Shruti Bhosale, Shun Zhang, Simon Vandenhende, Soumya Batra, Spencer Whitman, Sten Sootla, Stephane Collot, Suchin Gururangan, Sydney Borodinsky, Tamar Herman, Tara Fowler, Tarek Sheasha, Thomas Georgiou, Thomas Scialom, Tobias Speckbacher, Todor Mi- haylov, Tong Xiao, Ujjwal Kam. Vedanuj Goswami, Vibhor Gupta, Vignesh Ramanathan, Viktor Kerkez, Vincent Gonguet, Virginie Do, Vish Vogeti, Vítor Albiero, Vladan Petrovic, Weiwei Chu, Wenhan Xiong, Wenyin Fu, Whitney Meers, Xavier Martinet, Xiaodong Wang. Xiaofang Wang, Xiaoqing Ellen Tan, Xide Xia, Xinfeng Xie, Xuchao Jia, Xuewei Wang, Yaelle Gold- schlag, Yashesh Gaur, Yasmine Babaci, Yi Wen, Yiwen Song, Yuchen Zhang. Yue Li, Yuning Mao, Zacharie Delpierre Coudert, Zheng Yan, Zhengxing Chen, Zoe Papakipos, Aaditya Singh, Aayushi Srivastava, Abha Jain, Adam Kelsey, Adam Shajofeld, Adithya Gangidi, Adolfo Victoria, Ahuva Goldstand, Ajay Menon, Ajay Sharma, Alex Boesenberg, Alexei Baevski, Allie Feinstein, Amanda Kallet, Amit Sangani, Amos Teo, Anam Yunus, Andrei Lupu, Andres Alvarado, An- drew Caples, Andrew Gu, Andrew Ho, Andrew Poulton, Andrew Ryan, Ankit Ramchandani, An- nie Dong, Annie Franco, Anuj Goyal, Aparajita Saraf, Arkabandhu Chowdhury, Ashley Gabriel, Ashwin Bharambe, Assaf Eisenman, Azadeh Yazdan, Beau James, Ben Maurer, Benjamin Leon- hardi, Bernie Huang. Beth Loyd, Beto De Paola. Bhargavi Paranjape, Bing Liu, Bo Wu, Boyu Ni, Braden Hancock, Bram Wasti, Brandon Spence, Brani Stojkovic, Brian Gamido, Britt Mon- talvo, Carl Parker, Carly Burton, Catalina Mejia, Ce Liu, Changhan Wang, Changkyu Kim, Chao Zhou, Chester Hu, Ching-Hsiang Chu, Chris Cai, Chris Tindal, Christoph Feichtenhofer, Cymthia Gao, Damon Civin, Dana Beaty, Daniel Kreymer, Daniel Li, David Adkins, David Xu, Davide Testuggine, Delia David, Devi Parikh, Diana Liskovich, Didem Foss, Dingkang Wang, Due Le, Dustin Holland, Edward Dowling, Eissa Jamil, Elaine Montgomery, Eleonora Presani, Emily Hahn, Emily Wood, Eric-Tuan Le. Erik Brinkman, Esteban Arcaute, Evan Dunbar, Evan Smoth- ers, Fei Sun, Felix Kreuk, Feng Tian, Filippos Kokkinos, Firat Ozgenel, Francesco Caggioni, Frank Kanayet, Frank Seide, Gabriela Medina Florez, Gabriella Schwarz, Gada Badeer, Georgia Swee, Gil Halpern, Grant Herman, Grigory Sizov, Guangyi, Zhang, Guna Lakshminarayanan, Hakan Inan, Hamid Shojanazeri, Han Zou, Hannah Wang, Hanwen Zha, Haroun Habeeb, Harri- son Rudolph, Helen Suk, Henry Aspegren, Hunter Goldman, Hongyuan Zhan, Ibrahim Damlaj, Igor Molybog. Igor Tufanov, Ilias Leontiadis, Irina-Elens Veliche. Itai Gat, Jake Weissman, James Geboski, James Kohli, Janice Lam, Japhet Asher. Jean-Baptiste Gaya, Jeff Marcus, Jeff Tang, Jen- nifer Chan, Jenny Zhen, Jeremy Reizenstein, Jeremy Teboul, Jessica Zhong, Jian Jin, Jingyi Yang, Joe Cummings, Jon Carvill, Jon Shepard, Jonathan McPhie, Jonathan Torres, Josh Ginsburg, Jun- jie Wang. Kai Wu, Kam Hou U, Karan Saxena, Kartikay Khandelwal, Katayoun Zand, Kathy Matosich, Kaushik Veeraraghavan, Kelly Michelena, Keqian Li, Kiran Jagadeesh, Kun Huang, Kunal Chawla, Kyle Huang, Lailin Chen, Lakshya Garg, Lavender A, Leandro Silva, Lee Bell, Lei Zhang, Liangpeng Guo, Licheng Yu, Liron Moshkovich, Luca Wehrstedt, Madian Khabsa, Manav Avalani, Manish Bhaut Martynas Mankus, Matan Hasson, Matthew Lennie, Matthias Reso, Maxim Groshev, Maxim Naumov, Maya Lathi, Meghan Keneally, Miao Liu, Michael L. Seltzer, Michal Valko, Michelle Restrepo, Mihir Patel, Mik Vyatskov, Mikayel Samvelyan, Mike Clark, Mike Macey, Mike Wang, Miquel Jubert Hermoso, Mo Metanat, Mohammad Rastegari, Munish Bansal, Nandhini Santhanam, Natascha Parks, Natasha White, Navyata Bawa, Nayan Singhal, Nick Egebo, Nicolas Usunier, Nikhil Mehta, Nikolay Pavlovich Laptev, Ning Dong, Norman Cheng, Oleg Chernoguz, Olivia Hart, Omkar Salpekar, Ozlem Kalinli, Parkin Kent, Parth Parekh, Paul Saab, Pavan Balaji. Pedro Rittner, Philip Bontrager, Pierre Roux, Piotr Dollar, Polina Zvyagina, Prashant Ratanchandani, Pritish Yuvraj, Qian Liang, Rachad Alao, Rachel Ro- driguez, Rafi Ayub, Raghotham Murthy, Raghu Nayani, Rahul Mitra, Rangaprabhu Parthasarathy, Raymond Li, Rebekkah Hogan, Robin Battey, Rocky Wang, Russ Howes, Ruty Rinott, Sachin Mehta, Sachin Siby, Sai Jayesh Bondu, Samyak Datta, Sara Chugh, Sara Hunt, Sargun Dhillon, Sasha Sidorov, Satadru Pan, Saurabh Mahajan, Saurabh Verma, Seiji Yamamoto, Sharadh Ra- maswamy, Shaun Lindsay, Shaun Lindsay, Sheng Feng, Shenghao Lin, Shengxin Cindy Zha, Shishir Patil, Shiva Shankar. Shuqiang Zhang, Shuqiang Zhang, Sinong Wang, Sneha Agarwal, Soji Sajuyigbe, Soumith Chintala, Stephanie Max, Stephen Chen, Steve Kehoe, Steve Satter- field, Sudarshan Govindaprasad, Sumit Gupta, Summer Deng, Sungmin Cho, Sunny Virk, Suraj Subramanian, Sy Choudhury, Sydney Goldman, Tal Remez, Tamar Glaser, Tamara Best, Thilo
13

---

Published as a conference paper at ICLR 2026
Koehler, Thomas Robinson, Tianhe Li, Tianjun Zhang, Tim Matthews, Timothy Chou, Tzook Shaked, Varum Vontimitta, Victoria Ajayi, Victoria Montanez, Vijai Mohan, Vinay Satish Ku- mar, Vishal Mangla, Vlad lonescu, Vlad Poenaru, Vlad Tiberiu Mihailescu, Vladimir Ivanov, Wei Li, Wenchen Wang, Wenwen Jiang. Wes Bouaziz, Will Constable, Xiaocheng Tang, Xiao- jian Wu, Xiaolan Wang, Xilun Wu, Xinbo Gao, Yaniv Kleinman, Yanjun Chen, Ye Hu, Ye Jia, Ye Qi. Yenda Li, Yilin Zhang, Ying Zhang, Yossi Adi, Youngjin Nam, Yu, Wang. Yu Zhao, Yuchen Hao, Yundi Qian, Yunlu Li, Yuzi He, Zach Rait, Zachary DeVito, Zef Rosnbrick, Zhao- duo Wen. Zhenyu Yang, Zhiwei Zhao, and Zhiyu Ma. The llama 3 herd of models, 2024. URL https::/aaxxa org/ als 2477.21783
Diandian Gu, Peng Sun, Qinghao Hu, Ting Huang, Xun Chen, Yingtong Xiong, Guoteng Wang, Qiaoling Chen, Shangchun Zhao, Jiarui Fang, et al. Loongtrain: Efficient training of long- sequence Ilms with head-context parallelism. arXiv preprint arXiv:2406. 18485, 2024.
Suriya Gunasekar, Yi Zhang, Jyoti Aneja, Caio César Teodoro Mendes, Allie Del Giorno, Sivakanth Gopi, Mojan Javaheripi, Piero Kauffmann, Gustavo de Rosa, Olli Saarikivi, et al. Textbooks are all you need. arXiv preprint arXiv:2306. 1644, 2023.
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint 1rxxiv:5501.1244/ 2025.
Junxian Guo, Haotian Tang, Shang Yang, Zhekai Zhang, Zhijian Liu, and Song Han. Block Sparse Attention. attps://githnn com/ mit han watnn-labBBooocco Sewpsse-ttteeteeeeeeeeeee 2024.
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models? arXiv preprint arXiv:2404.06654 2024.
Yanping Huang. Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoe V Le, Yonghui Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems, 32, 2019.
Ziqi Huang, Yinan He, Jiashuo Yu, Fan Zhang, Chenyang Si, Yuming Jiang, Yuanhan Zhang, Tianx- ing Wu, Qingyang Jin, Nattapol Chanpaisit, al. Vbench: Comprehensive benchmark suite for video generative models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 21807-21818, 2024.
Changho Hwang, Wei Cui, Yifan Xiong, Ziyue Yang. Ze Liu, Han Hu, Zilong Wang, Rafael Salas, Jithin Jose, Prabhat Ram, et al. Tutel: Adaptive mixture-of-experts at scale. Proceedings of Machine Learning and Systems, 5:269-287, 2023.
Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Shuaiwen Leon Song, Samyam Rajbhandari, and Yuxiong He. Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models. arXiv preprint arXiv: 14509, 2023.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Alec Gray, Scott ..and Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001. 08361,, 2020.
Vijay Anand Korthikanti, Jared Casper. Sangkug Lym, Lawrence McAfee, Michael Andersch, Mo- hammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. Proceedings of Machine Learning and Systems, 5:341 -3311 2023.
Mario Michael Krell, Matej Kosec, Sergio P Perez, and Andrew Fitzgibbon. Efficient sequence packing without cross cross-contamination: Accelcrating large language models without impacting per- formance. arXiv preprint arXiv:2107 02027, 1.21.
Jinhyuk Lee, Anthony Chen, Zhuyun Dai, Dheeru Dua, Devendra Singh Sachan, Michael Boratko, Yi Luan, Sébastien MR Arnold, Vincent Perot, Siddharth Dalmia, et al. Can long-contes. lan- guage models subsume retrieval, rag, sql, and more? arXiv preprint arXiv:2406. 1 13121, 2024.
14

---

Published as a conference paper at ICLR 2026
Aonian Li. Bangwei Gong, Bo Yang, Boji Shan, Chang Liu, Cheng Zhu, Chunhao Zhang, Congchao Guo, Da Chen, Dong Li, et al. Minimax-01: Scaling foundation models with lightning attention. arXiv preprint (rXXv:2501.08113. 2025.
Jiamin Li, Yimin Jiang, Yibo Zhu, Cong Wang, and Hong Xu. Accelerating distributed (MoE) training and inference with lina. In 2023 USENIX Annual Technical Conference (USENIX ATC 23), pp. 945-959, 2023.
Shenggui Li, Fuzhao Xue, Chaitanya Baranwal, Yongbin Li, and Yang You. Sequence parallelism: Long sequence training from system perspective. arXiv preprint arXiv:2105.13122 2021.
Shigang Li and Torsten Hoefler. Chimera: efficiently training large-scale neural networks with bidirectional pipelines. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-14, 2021.
Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi Dengr, Chong Ruan, Damai Dai, Daya Guo, al. Deepseek-v2. strong, economical, and efficient mixture- of-experts language model. arXiv preprint arXiv:2405.04444 2024.
Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near- infinite context. arXiv preprint arXiv:2310.0188; 2023.
Enzhe Lu, Zhejun Jiang, Jingyuan Liu, Yulun Du, Tao Jiang, Chao Hong. Shaowei Liu, Weiran He, Enming Yuan, Yuzhi Wang, al. Moba: Mixture of block attention for long-context llms. arXiv preprint arXiv:2502. 13119 13189, 2025.
Niklas Muennighoff, Zitong Yang, Weijia Shi, Xiang Lisa Li, Li Fei-Fei, Hanmaneh Hajishirzi, Luke Zetllemyer, Percy Liang, Emmanuel Candès, and Tatsunori Hashimoto. sl: Simple test-time scaling, arXiv preprint arXiv:2501. 100::2001.9993, 2025.
Deepak Narayanan. Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R Devanur, Gre- gory R Ganger, Phillip B Gibbons, and Matei Zaharia. Pipedream: Generalized pipeline par- allelism for dnn training. In Proceedings of the 27th ACM symposiun on operating systems principles, pp. 1-15, 2019.
NVIDIA. Transformer engine. 2025. Accessed: 2025-09-23.
NVIDIA Corporation. Collective communication functions.
rctps://docs wnviiiiiaa........ com/ 2020. Accessed: 2025-
09-23.
NVIDIA Corporation. fused_attn.h.
engine transformer Tnanssooreer---giinee-- re leases wew.ineerreeeeesse/eeeeeeeeeeeee
attn. html 2025. Accessed: 2025-09-23.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to fol- low instructions with human feedback. Advances in neural information processing systems, 35: 27730-27744, 2022.
PyTorch Contributors. Distributed data parallel. httpss//bocs.pytoocc.orrr/// stable/notes/adt hdp.htmlllmmmmmmmmmmmmmmmmmmmm.. 2024a.
PyTorch Contributors. forch.m.functiinallssalledd htt.aicciin...aeeeeooooooooooooe https://docs.
2024b.
Penghui Qi, Xinyi Wan, Guangxing Huang, and Min Lin. Zero bubble (almost) pipeline parallelism. In The Twelfth International Conference on Learning Representattonss, 2024.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140): 2020.
15

---

Published as a conference paper at ICLR 2026
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. SC20: International Conference for High Perfor- mance Computing, Networking. Storage and Analysis, pp. 1-16. IEEE, 2020.
Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. Deepspeed: System opti-
mizations enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD international conference on knowledge discovery & data mining, pp- 3505-3500 2020.
Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Min- jia Zhang, Dong Li, and Yuxiong He. {Zero-offload]: Democratizing (billion-scale) model training. In 2021 USENIX Annual Technical Conference (USENIX ATC 21), pp. 551-564, 2021.
Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision. Advances in Neural Information Processing Systems, 37:68658-668 1176666886665,,, 2024.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan
Catanzaro. Megatron-Im: Training multi-billion parameter language models using model par- allelism. arXiv preprint arXiv:1909.08000 2019.
Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, al. Using deep- speed and megatron to train megatron-turing nlg 530b, large-scale generative language model. arXiv preprint arXiv:2201.119901 2022.
Daria Soboleva, Faisal Al-Khateeb, Robert SlimPajama:
Myers,
Jacob R Steeves, Joel 627B token cleaned and
Hestness, and Nolan Dey.
A
deduplicated version slimpa janaajjnnn -a 627b token cleaned and wew.deeeeeceeeeeeessiiooro wwwpppiccteedeeeeeee of redpa jama 627B
of RedPajama.
ht https://orebras.ai/
2023. URL httpt:ttt//h/ttttet http::////ccciiga.e..ee///aaatat
Benjamin F Spector, Simran Arora, Aaryan Singhal, Daniel Y Fu, and Christopher Ré. Thunderkit- tens: Simple, fast, and adorable ai kernels. arXiv preprini arXiv:2410.20399, 2024.
Yi Tay, Jason Wei, Hyung Won Chung, Vinh Q Tran, David So, Siamak Shakeri, Xavier Garcia, Huaixiu Steven Zheng, Jinfeng Rao, Aakanksha Chowdhery, et al. Transcending scaling laws with 0.1 extra compute. arXiv preprint arXiv:2210. 11399, 2022.
Gemini Team, Rohan Anil, Sebastian Borgeaud, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew Dai, Anja Hauth, Katie Millican, et al. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11805, 2023.
Gemini Team, Petko Georgiev, Ving Ian Lei, Ryan Burnell, Libin Bai, Anmol Gulati, Garrett Tanzer, Damien Vincent, Zhufeng Pan, Shibo Wang, et al. Gemini 1.5: Unlocking multimodal under- standing across millions of tokens of context. arXiv preprint wrXiv:24033055344 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Boxiang Wang, Qifan Xu, Zhengda Bian, and Yang You. Tesseract: Parallelize the tensor parallelism efficiently. In Proceedings of the 51st International Conference on Parallel Processing, pp. 1-11, 2022.
Guoxia Wang. Jinle Zeng, Xiyuan Xiao, Siming Wu, Jiabin Yang, Lujing Zheng, Zeyu Chen, Jiang Bian, Dianhai Yu, and Haifeng Wang. Flashmask: Efficient and rich mask extension of flashat- tention, 2025a. htt.l:
01359
Liang Wang. Nan Yang, Xiaolong Huang, Linjun Yang, Rangan Majumder, and Furu Wei. Large search model: Redefining search stack in the era of llms. In ACM SIGIR Forum, volume 57, pp. 1-16. ACM New York, NY, USA, 2024.
16

---

Published as a conference paper at ICLR 2026
Zheng Wang, Anna Cai, Xinfeng Xie, Zaifeng Pan, Yue Guan, Weiwei Chu, Jie Wang, Shikai Li, Jianyu Huang, Chris Cai, et al. Wib-Ilm: Workload-balann 4d parallelism for large language model training. arXiv preprint arXiv:2503.17924, 2025b.
Yuetian Weng, Mingfei Han, Haoyu He, Xiaojun Chang, and Bohan Zhuang. Longvlm: Efficient long video understanding via large language models. In European Conference on Computer Vision, pp. 453-470. Springer, 2024.
Yifei Xia, Suhan Ling, Fangcheng Fu, Yujie Wang, Huixia Li, Xuefeng Xiao, and Bin Cui. Training-free free and adaptive sparse attention for efficient long video generation. arXiv preprint arXiv:2502.21079, 2025.
Qifan Xu and Yang You. An efficient 2d method for training super-large deep learning models. In 2023 IEEE International Parallel and Distributed Processing Symposium (IPDPS), pp. 222-232. IEEE, 2023.
Ruyi Xu, Guangxuan Xiao, Haofeng Huang, Junxian Guo, and Song Han. Xattention: Block sparse attention with antidiagonal scoring. arXiv preprint arXiv:2503. 16428, 2025.
An Yang, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoyan Huang, Jiandong Jiang, Jianhong Tu, Jianwei Zhang, Jingren Zhou, et al. Qwen2. 5-1m technical report. arXiv preprint arXiv:2501.1538. 2025a.
Shuo Yang. Haocheng Xi, Yilong Zhao, Muyang Li, Jintao Zhang, Han Cai, Yujun Lin, Xiuyu Li, Chenfeng Xu, Kelly Peng, et al. Sparse videogen2: Accelerate video generation with sparse attention via semantic-aware permutation. arXiv preprint arXiv 2025b.
Yaofu. Slimpajama per source length upsample, 2024. URL http:s/tttt///uggggggaaa cc/ waaasstssyaafuuu/e 09-15.
caama-pee-aae -etparrperrs-oooo-----e-e-------
Accessed: 2025-
Zihao Ye, Lequn Chen, Ruihang Lai, Wuwei Lin, Yineng Zhang, Stephanie Wang, Tianqi Chen, Baris Kasikci, Vinod Grover, Arvind Krishnamurthy, al. Flashinfer: Efficient and customizable attention engine for I1m inference serving. arXiv preprint arXiv:2501. 01005, 2025.
Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang. Zhenda Xie, Y. X. Wei, Lean Wang, Zhiping Xiao, Yuqing Wang, Chong Ruan, Ming Zhang, Wenfeng Liang, and Wangding Zeng. Native sparse attention: Hardware-aligned and natively trainable sparse attention, 2025.
www.bbbb.oogg.oo/ooooaaaa/aabaat http:////.aa...org//////222222. 11089
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences, Advances in neural information processing systems, 33:17283-17297, 2020.
Tao Zewei and Huang Yunpeng. Magiattention: A distributed attention towards linear scalability for ultra-long context, heterogeneous mask training. hitps: http://ggttub..t/// com/ SandAI- MagiAttention/ 2025.
Jintao Zhang, Rundong Su. Chunyu Liu, Jia Wei, Ziteng Wang, Pengle Zhang, Haoxu Wang, Huiqiang Jiang, Haofeng Huang, Chendong Xiang, et al. A survey of efficient attention methods: Hardware-efficient sparse, compact, and linear attention.
Jintao Zhang, Haofeng Huang, Pengle Zhang, Jia Wei, Jun Zhu, and Jianfei Chen. Sageatten- tion2: Efficient attention with thorough outlier smoothing and per-thread int4 quantization. arXiv preprint arXiv:241 1. 100958, 2024a.
Jintao Zhang, Jia Wei, Haofeng Huang, Pengle Zhang, Jun Zhu, and Jianfei Chen. Sageat- tention: Accurate 8-bit attention for -un-ann- inference acceleration. arXiv preprint arXiv:2410.02366 2024b.
Jintao Zhang, Jia Wei, Pengle Zhang, Xiaoming Xu, Haofeng Huang, Haoxu Wang, Kai Jiang, Jun Zhu, and Jianfei Chen. Sageattention3: Microscaling fp4 attention for inference and an exploration of b- training. arXiv preprint arXiv:2505. 11594, 2025a.
17

---

Published as a conference paper at ICLR 2026
Jintao Zhang, Chendong Xiang, Haofeng Huang, Jia Wei, Haocheng Xi, Jun Zhu, and Jianfei Chen. Spargeattn: Accurate sparse attention accelerating any model inference. arXiv preprint 0rXI::25002.81137.22222
Peiyuan Zhang, Yonggi Chen, Runlong Su, Hangliang Ding, lon Stoica, Zhengzhong Liu, and Hao Zhang. Fast video generation with sliding tile attention. arXiv preprini arXiv:2502.04500 2025c.
Peiyuan Zhang, Haofeng Huang, Yongqi Chen, Will Lin, Zhengzhong Liu, lon Stoica, Eric Xing, and Hao Zhang. Vsa: Faster video diffusion with trainable sparse attention. arXiv preprint 13Xi::5005.1338, 2025d.
Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yinggian Min, Beichen Zhang, Junjie Zhang. Zican Dong, et al. survey of large language models. arXiv preprint ar arXiv:2333.18223 1(2), 2023a.
Yanli Zhao, Andrew Gu Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, et al. Pytorch fsdp: experiences on scaling fully sharded data parallel. arXiv preprint arXiv:2304. / 1277, 2023b.
zhuzilin. [feature request] balancing computation with zigzag blocking. httts:////tthb
com/ wwuzziiiiin ring- flash wtt.:/attttttootssses/s////s Feb 2024.
18

---

Published as a conference paper at ICLR 2026
A APPENDIX
A.1 LLM USAGE
In this work, LLMs served supporting tools for manuscript polishing. They were mainly used to enhance textual clarity, refine grammar, and adjust minor phrasing. All conceptual development, baseline optimization, benchmark design, and experimental analysis were conducted independently by the authors, without reliance on LLMs.
A.2 DETAILS OF DENSE DATA SAMPLING
In large language model construction, the quality and diversity of the training data are crucial for enhancing model performance (Gunasekar et al. 2023). Numerous studies have explored various methods to improve data quality. Our benchmark builds upon these efforts by systematically an- alyzing several publicly available, high-quality, and widely used English datasets. Full results are presented in Figure 7
We mainly analyze the Pile and SlimPajama (Soboleva et al. 2023) datasets to study their effects on the model's short-context modeling capabilities. The Pile dataset is a large-scale, diverse English text corpus designed for training large language models, with a total size of 825 GB. It consists of 22 high-quality subsets, many drawn from academic or professional sources, including Common Crawl, Wikipedia, Open Web'ee,,,, ArXiv, and PubMed. Such diversity across multiple domains and topics substantially increases the richness and variety of the training data. SlimPajama is an open-source dataset obtained from the original RedPajama corpus through multiple preprocessing steps such as NFC normalization, cleaning, deduplication, and document interleaving, comprising a total of 627B tokens. Compared to Pile, SlimPajama contains less web data and more content from Books, ArXiv, and Wikipedia. These are sigh-quality long-form text sources that help improve the model's long-context modeling capabilities. Owing its large scale, SlimPajama is not fully sampled; instead, our benchmark samples sequences of up to 8k tokens from Pile, which is sufficient to represent realistic short-context modeling scenarios.
The above data cleaning mainly focused on a limited context window (e.g., 8k). To extend the model's context window, recent studies have begun exploring data mixing strategies for long con- texts. We follow the findings of (Fu et al. (2024) and (Gao et al., 2024b): (1) continual pretraining on long-context data can significantly improve the model's ability to accurately retrieve information in long contexts; (2) when extending the context length, oversampling long sequences while pre- serving the original domain diversity of the pretraining dataset is crucial; (3) mixing high-quality long-context sources with high-quality short-context sources is essential for enhancing long-context modeling capability while maintaining performance on short contexts. In our benchmark, we col- lected statistics on the publicly available long-text upsampled dataset caampaajaaa-p--suur---lllee upsample (referred to as Upsampled SlimPajama) (Yaofu) 2024) from (Fu et al.) 2024), as well as the datasets prolong data--4tt (ProLong64K) (Gao et al. [2024a) and prolong-data-51 512K (Pro- Long512K) (Gao et al. 2024a) from (Gao et al. (2024b), which are used to extend the model's context window to 64K and 768K tokens, respectively.
Considering the significant differences resulting from various tokenization methods, we directly split English samples by spaces in our statisties (approximately reflecting tokenized lengths). All statistics are shown in Figure 7 with short-context and long-context distributions arranged side by side. It worth noting that datasets are generally expressed in terms of the number of tokens rather than the number of samples. Although long-text text tokens in datasets such as ProLong64 and ProLong512K account for up 60%. their prominence the length distribution may still be limited.
In real training scenarios, the proportion of long- and short-text tokens seen by the model should reflect the inherent distribution of the dataset. Prior to training, data preprocessing typically shuffles the dataset and packs all samples into batches, ensuring that the model eventually traverses every sample. However, in benchmark evaluation, it is impractical to cover the entire dataset in each run. Therefore, we set a fixed number of steps per evaluation. To avoid significant deviations between the sampled data and the true dataset distribution when using only a small number of steps, we adopt stratified sampling, enforcing a 6:4 ratio of long- to short-text tokens. This achieves trade-ofi off between evaluation cost and result fidelity. The 6:4 ratio is a common choice, but for special cases, recommend defining scenario-specifie ratios for evaluation.
19

---

Published as a conference paper at ICLR 2026
Proportion
Proportion
(0.32)
0.00
100
200
000
0004
10
0,07
0.01
02
00
900
90
000
(4)
2.99%
6.40%
6.49
5738
160)
7.01%
204)
5.57%
(24 (2)
(22224)
7.374
5.377
Protong Dataset Binned Length Distribution
3.91%
HIC
DRA 4gg
3.07%
2.59%
HEEST
TRIS 480)
(R0.522)
(ML
540
1.53%
1.72%
1.30%
(M6 (54)
(RUE (60)
1.11%
1.23%
DRC 700
D74 (68)
DRE (0
(789 B22)
DE2 680
1.54%
Length Distribution
(RCAA
1.54%
(MG 1040
1.10%
1.0355
1.50%
1060)
/H000 782)
AAM 16R0)
(40)
11 1.22%
2200
14782
200%
10000
2994
PRE. 1790
1.25%
F2TAR (714)
1.54% 205
(0162204
20160
FERM maU
2908)
1.10%
19782 2080
F47M 7780)
1. 1.88 2
1.00
(OTR NOTT
1813
3 100%
(H077 51
RIZATO
1.0035
1.00%
0.27%
Proportion
Proportion
0.00
101
200
00
400
0.05
0.00
200
500
300
4.00%
1015%
4.62%
3.62%
4.65%
4.53%
A.15%
1202200
400)
1868 (20.320)
4.70%
ME
OpenWebToxt Binned Length Distribution
6.35%
Protong 512k
1.20%%
3.00%
3.03%
380
3.50%
2.85%
3.20%
2.50%
4.4%%
(42
12 70
2,3795
D20 ND
2.133
2.25%
RR
JAM
2.42%
1800
2.13%
1 1.92%
1.6354
5.71%
1872 P00
1.60%
1.76%
Binned Length Distribution
1.50%
1288 7M)
3.40%
1 1.40%
DRg
1.19%%
1.31%
1.28%
1.22%
1.15%
1.0%%
F9R 18:4
1.02%
ROO LORD
L28)
1860
1.63%
1.48%
T.30%
5.00%
1.19%
1.0%%
1.44%
1.34%
1.33%
109%
1.%
1.22%
1.21%
1.14%
Dest 4787
1.04%
1.05%6
1.04%
1.05%
..00%
1.01%
1.00%
3.09%
6.78%
20

---

Published as a conference paper at ICLR 2026
Proportion
Proportion
(0.20)
000
20'0
900
900
80
01:0
0.00
10
200
20
600
500
900
000
80'0
(22
PANN
(6)
7.41%
(6120)
(122.80))
Fres 122)
5.13%
(122220)
1.30%
222.200
(2062200
5.07%
43%
SlimPajama Dataset Binned Length Distribution
(200
2220.252)
3.00%
1.71%
Dataset Binned Length
Taee 41o)
1.21%
4416440)
7.7074
(445-800
52%
5.3%%
1..0%
1512 13125466
1.144
1.80%
1376
1.12%
L03%
fouE 6409
1.0%
1840 2041
1.07%
1.13%
1.09%
1765.
1.17%
1822 2918
5.5%
3.10%
loc. 1owe)
fowe 201
1.00%
1.11%
0,005
1.33%
1.09 1.07%
11216 (126,100) Lean
.66%%
1.01%
1.05%
1.02%
Hone 2100
1.06%
DAT241NG
12 220223222 3072)
1.02%
1.01%
3.01%
1,01%
1.00%
0.00%
0.47%
0.00%
(a) C4 Dataset Distribution
(b) Upsampled SlimPajama Dataset Distribution
Figure 7: Pretraining Dataset Length Distributions.
21

---

Published as a conference paper at ICLR 2026
A.3 DETAILS OF SPARSE DATA SAMPLING
Referencing common block sparse mask generation methods, our block sparse mask generation process is as follows:
1. Block Partitioning: For a task with an input sequence length of seglen, we can conceptualize a seglen seqlen attention map. Based on the provided p_block_lists and b.block_lista this two-dimensional attention map is partitioned into len(g halllg_llockkaa .lisss x len(k_block Jists) blocks. For uniform block mask, the block sizes within g.block lists and b.block lists are the same. In contrast, for variable block mask, these block sizes can differ.
2. Score Calculation: Common block sparse methods typically calculate a score for each block in some manner (e.g., by applying mean pooling over each block) and generate a score matrix which represents the importance of the block. Blocks with higher importance are more likely to be selected. It is important to note that scores may vary across different attention heads. Generally, a distinct mask is generated for each KV head. This means for Multi-Head Attention, the score matrix can be unique for each head. For Grouped-Quers Attention, masks are the same within a group but can differ between groups. In our experiments, we abstract away the specifics of score calculation and use randomly generated numbers, thereby focusing the final block sparse attention computation.
3. Top-k Selection: For each block in the query dimension (g), we select the top-k blocks from the key dimension (k) for computation. The overall degree of sparsity can be expressed the fraction Ten(K lists)- We define a sparsity ratio to represent this degree of sparsity, which has a direct conversion relationship with top-k and is essentially equivalent. To observe the kernel's 's performance under different sparsity levels, we select 0.2, 0.5, and 0.8 as representative sparsity ratios.
A.4 DETAILS OF SPARSE ATTENTION KERNELS
VSA (Zhang et al. 2025d) is trainable block sparse attention implementation designed specifically for video diffusion models. It employs a two-stage methodology. In the coarse-grained stage, applies token rearrangement strategy from STA (Zhang al. 2025c) to increase computational density. It then calculates inter-block scores via cube partition and meen-pooling on the QK ma- trix, which selects the top-k K blocks for each Q block. Subsequently, the fine-grained stage, utilizes ThunderKittens (Spector et al. 2024) to develop a kernel. This kernel customized for Hopper GPUs maximize hardware utilization. According to VSA's analysis, the ine-grained stage accounts for over 80% of latency at context lengths of 32K or more. This finding highlights the critical need to optimize the block sparse attention kernel, and our performance benchmarks also focus on this kernel for fair comparisons. As variant, Triton VSA implements the algorithm in the Triton language. This approach aims to enhance cross-hard hoosw-e compatibility but results in some performance degradation. However, both implementations specifically optimized for video models and they do not support common LLM attention modes like Grouped-Quee Attention (GQA). FA2 Sparse Guo et al. (2024) is an open-source implementation based on the FlashAttention-2 (Dao) 2023) codebase. It enables block sparse functionality by mod- ifying the computation logic, allowing each Q block to traverse only its designated KV blocks. The primary limitation of this implementation is its lack of support for the backward pass and without optimization for advanced GPUs like NVIDIA Hopper GPUs FlexAttention Dong et al. (2024) leverages compiler technology to introduce a more flexible mask description method. This approach enables it to support sparse attention in the form of a block mask. However, the representation for block masks is relatively complex. Its compilation technique can therefore degenerate to instanti- ating full OS22) mask, which causes significant memory overhead. FlashInfer Ye et al. (2025) is general- purpose kernel library oriented toward LLM inference. It designs a block sparse matrix general-purpose
block sparse attention
structure as a unified format for the KV cache. This design allows block sparse attention input to be converted into a paged attention format where page size equals to 1. This process enables the reuse of its efficient attention kernel and supports arbitrary block sizes. Due to its positioning as an inference library, it does not support the backward pass.
22

---

Published as a conference paper at ICLR 2026
A.5 DETAILS OF DENSE KERNEL PERFORMANCE ON H100 GPU
For single-sequence samples such as FULL/CAUSAL and SLIDING WINDOW, we conduct 2 runs of sampling, with each rum followed by 5 warm-up steps and 20 kernel computation steps. For multi- sequence data such as DOCUMENT and SHARE QUESTION, considering the sparsity differences introduced by sampling. we perform 30 runs with independent sampling in each run, also followed by 5 warm-up steps and 20 kernel computation steps. We record the median values of FLOPS and peak memory. It is worth noting that our results represent the expected outcomes under these specific settings, and occasional large deviations in individual kernel runs are considered normal.
A.5.1 PERFORMANCE METRIC: FLOPS
GQA Bwd Flops (8k)
HOre
fb3
fs2
cudrn
flex
flashumask
sdpa
torch
500
(twnFLoPPN)
600
200
300
101
30:
FSW
CSW
FD
CD
SQ
CB
65
PLC
PLD
BCD
Figure 8: Backward TFLOPs of dense kernels with different masks (8K length)
FULL and CAUSAL are the most common masks used in language model pretraining, as shown in Figure 9 SIPA and Flex approximately represent the baselines of fused operators without general, increasing sequence length improves kernel perfor- mance, which typically stabilizes around 16K. The results highlight the importance of hardware- aware optimization: on the H100 Hopper architecture, FA3 achieve significant performance gains.
hardware-speciffc optimization.
FULL/CAUSAL DOCUMENT is primarily designed for concatenating variable-length input quences to reduce unnecessary padding while preserving full or causal connectivity. It is important to note that concatenating variable-length sequences can introduce computational instability, which becomes particularly pronounced when the data contains many small fragmented chunks, as shown in Figure TO
In our experiments, we fixed the sliding window size to 1024, as shown in Figure 11 though we recommend evaluating with other window sizes as well (Fu et al. 2025).
Overall, in heterogeneous mask scenarios, Flex and FlashMask show varying performance gains depending on the context.
PREFIX LM and PREFIX LM DOCUMENT extend the standard language model regular mask by introducing a prefix, allowing the prefix to attend all tokens. In our experiments, a prefix is randomly generated for each run, and the median across multiple runs is reported to reflect expected performance in realistic scenarios with varying prefixes, as shown in Figure 12, Models trained with a prefix demonstrate advantages in handling long-text and multi-turn dialogue tasks. This approach enables the model to better leverage contextual information, improving performance in generation tasks (Raffel et al. 2020).
SHARE QUESTION and CAUSAL BLOCKWISE can be viewed as variants of DOCUMENT, as shown in Figure [3 SHARE QUESTION allows all query tokens to share the key tokens from the first document and commonly used in Reward Models (RM) as a shared-question mask, enabling multiple answers to reference the same question (Ouyang et al. 2022). This eliminates redundant computation and accelerates training. CAUSAL BLOCKWISE, on the other hand, allows all key to- kens share the query tokens from the last document and is typically applied in demonstration test tasks, where test examples can attend to all demonstrations. This facilitates studying model perfor- mance improvements in long-context tasks (Bertsch et al. 2024
23

---

Published as a conference paper at ICLR 2026
full GQA fwdd flops
000
HUSe
f63
Fa2
cuttin
flex
fash_mask mask
sdpa
torch
700
Speed (teraPLoPPaN
000
5010
1k
2k
18k
24c
32%
40k
48k
Sequence Lengtth
Fud TELOPSS
Tull GQA bwd flops
700
1006
fa3
faz
cusiin
flcx
flash_mask mask
sdpa
torch
600
(twrafLPP.N)
400
1000
1K
2k
8k
14k
24c
19c
6:k
40k
Sequence Length
GQA
full fwdd flops
000
1000
fa3
Fa2
flex
flash fish_mask
torth
700
Speed (baraPLOPPIN))
500
200
100
1k
2k
4x
16k
24
32k
40k
48k
Sequence
full bwdd
700
flex
flash
torch
600
Speed (teraf LOPaNN)
400
300
100
14k
24c
32c
40k
GLOBAL SLIDING combines global attention with sliding- attention. In each run, randomly sample a window size, treating the leftmost window size tokens in the Query and Key global tokens, which attend all Key and Query tokens, respectively. Due the increased sparsity of the mask, the performance Flex and FlashMask correspondingly decreases, as shown in Figure
24

---

Published as a conference paper at ICLR 2026
GQA,
fa33
Fa2
flex
flashmaak
torth
700
(teraPLOPais)
000
500
200
200
100
1k
2k
16k
24
32x
40k
(e) CAUSAL GQA TELOPS
flex
flah
Speed (terel LOPA)
400
2000
14k
24c
12x
4:kk
Sequence Length
(6) CAUSAL GOA Hwd TELOPS
MHA
sdpa
700
Speed (taraLoPPu))
000
500
200
100
Sequence Length
(g) CALSAL MHA Eid THLOPS
causal MHA flops
flash_mask
sdpa
600
400
200
Sequence Length
O0 CAUSAL MHA Bud TFLOPS
Figure 9: TFLOPS of FULL and CAUSAL
25

---

Published as a conference paper at ICLR 2026
full document, GQA, fwd. flos
700
f63
FA2
cutiin
flex
faash_m mash_mask
sdpa
torch
600
Spaed (teraPLoPaats
400
200
200
100
1k
2k
18k
24
32%
40k
48k
Sequence Lengtth
FULJ DOCUMENT PPS
full document GQA bwd Nops
:000
fa3
fa2
cusin
flcx
flash mask
sdpa
torch
500
Spaed (terafLOP.N)
100
1k
2k
8k
14k
24c
19c
6:k
4Bk
Sequence Length
(b)
full dol.ddocumnn MWH. flops
700
1000
fa3
Fa2
cuadinn
flex
flash_ mask
torth
Speed (baraPOPPPI))
500
000
100
1.
2k
4x
8k
16k
24
32x
40k
48k
Sequencce
(6) DOCUMENTT
full document bwd
flex
flash.
torch
Speed [teral LOPAlN)
200
100
1dkk
24x
32c
4.k
(6) DOCUMENT Rood OPS
BLOCK CAUSAL DOCUMENT can also be viewed as a variant DOCUMENT, elevating com- putation from the token level to the block level. In experiments, we fix block_sizz 1024, as shown in Figure 15 This approach commonly used training autoregressive multimodal large models (ai et al. 2025).
26

---

Published as a conference paper at ICLR 2026
document, Fwdd
fu.3
Fi2
flex
flash_mask
torth
600
Speed (taraPLOPa2)
200
200
100
1k
2k
36k
24
32X
40k
Sequence
(e) CAUSAL DOCUMENT GQA Fed
bwd
flex
flch
borch
Speed (teraP LPPP)))
200
100
14k
24
12x
45k
Sequence Length
(9 CAUSAL DOCUMENT GQA Bwd TELOPS
MHA fwd flops
600
flash_mask
sdpa
Speed (traFLoPPe)
200
600
2000
200
200
Sequencc Lengtt
(g) CAUSAL DOCUMENT MHA Pwd TELOPS
causal document bwd flop
flash_mask
sdpa
FLDPP())
200
200
100
Sequence Lengtth
do CAUSAL DOCUMENT MHA Tewd TFLOPS
Figure 10: TELOPS of FULL/CAUSAL DOCUMENT
27

---

Published as a conference paper at ICLR 2026
sliding window full GQA, fwd. flops
700
-238
f63
Fa2
flx
flash_mask mask
<dpa
tarch
Speed (twaPLoPPs))
500
603
200
300
24k
40c
ABcc
Sa Sequence Length
(o FUTL SLIDING WINDOW GQA Food TFI OPS
sliding www.wwwwwwwww.e full, GQA, bwd_flops
500
2330
Taa
faz
flxx
fassh mask
s:pa
tarch
603
Speed (teraff
003
000
300
Sequence Length
(1) FULL SLIDING WINDOW GQA Bed TELOPS
slidiing window, full MHA Ewid flop
700
Taa
fa2
flex
flssh mask
sdpaa
tarch
Speed (brraLOPPS)))
503
100
200
Se Sequence Length
FULL SLIDING WINDOW MHA Ped TELOPS
sliding, www.ow/wwwwwwwwwoo..m full, MHA bwd. flops
230
flex
flash_mask
sdpaa
tarch
603
Speed (araPLOPEN)
202
102
Sequence Length
di FULL SLIDING WINDOW MILA Bwd TELOPS
28

---

Published as a conference paper at ICLR 2026
sliding
fa3
Fi2
flex
flash_aask
torth
600
Speed (teraPLOPaIs)
200
200
100
1k
2K
36k:
24
32X
40k
(e) CAUSAL GQA TFLLOPS
sliding
GQA,
600
flex
flah
250
Speed (teraP LOPAN
200
25.
200
150
100
14k
24c
12x
Sequence
(8) CAUSAL SLIDING WINDOW GQA Bwd THLOPS
sliding, wwtp://lddwwweeee httpi//wwwwsssaaeeee MHA__al__A___eeea wew.t__d__foppppp flops
flash_mask
Speed (traFLoPPe)
500
400
000
200
300
Sequencc Lengtth
CAUSAL SLIDING WINDOW MHA Fodd TELOPS
sliding window causal MHA flops
flash_mask
sdpa
35.0
(LOPAE FLDPHE))
300
200
50
Sequence Lengtth
(6) CAUSAL SLIDING WINDOW MHA Bod TFLOPS
Figure 11: TELOPS of FULL/CAUSAL SLIDING WINDOW
29

---

Published as a conference paper at ICLR 2026
prefix, Im GQA Ewd flogs
603
-220
fllc
Flash maask
<dpa
tarch
252
Speed (tenPLoPPP))
203
252
200
150
1.0
328
4.6
Sa Sequence Length
(6) PREFIX TM GQA Fedd TRLOIS
prefix Jmm GQA_bwd flops
2330
flex
Fash mask
sópa
tarch
252
Speed (teraf
200
150
Sequence Length
(h) PREPIX TM OQA Rod TELOPS
profix Im. MHA fwd flops
600
-110
flix
Flash mask
sópa
tarch
253
Speed (brrFLOPP)))
200
253
200
155
100
Sequence Length
(c) PREFIX LM MHA Pad TELOPS
prefix, Im, MHA www.bw/ddppp...... flops
1330
Flash_maskk maakk
edpa
tarch
300
Speed (braFLOPers)
200
150
300
Sequence Length
(d) PREFIX LM MEIA Bwd TELOPS
30

---

Published as a conference paper at ICLR 2026
prefix, cumentt GQA flops
253
sdpa
Speed (teral LOPAN)
200
103
Sequencc Length
(c) PREFIX LM DOCUMENT GQA Pwd TFLOPS
prefix GQA flops
edpa
200
Speed (taraLLPPII)
100
1k
4.k
(3 PREFIX LM DOCUMENT GQA Bwd TFLOPS
prefix flos
edpa
Speee (teratt LoPun
250
300
44.
5k
26k
32k
43k
45k
Sequence
(g) PREFIX DOCUMENT MIAA TELOPS
flops
250
Speed (tera PLOPAi))
203
1.0
5k
2dkk
43k
45k
Figure 12: TFLOPS of PREFIX LM and PREFIX LM DOCUMENT
31

---

Published as a conference paper at ICLR 2026
share question, GQA, fwd foops
-228
flex
Flash maakk
<dpa
tarch
253
Speed (twaPLoPPa)
200
32k
4.6
Se Sequence Length
(6) SHARF QUESTION GQA Fue TFI LPS
share, question, BtwplonnoOoowwdddaaa bwvd flo ps
2310
flex
Fash mask
sdpa
tarch
Speed (teraf
103
100
Sequence Length
(b) SHARF QUESTION GQA Bud TALOPS
share question MHA, fwd flops
3003
-010
flex
flash mask
sópa
tarch
Speed (ErAFLOPPN)))
200
150
Sequence Length
(G) SHARE QUES TION MHA Fad TELOPS
share wirrregueettonMii. bww. flops
1330
Flssh Flash_maskk maak
sdpa
tarch
Speed (braFLOPPes)
2.0
Sequence Length
(d) SHARE QUESTION MIA Bwd TELOPS
32

---

Published as a conference paper at ICLR 2026
GQA, fwd, flops
030
sopa
200
150
Sequence Length
(e) CAUSAL BLOCKWISE GQA Pod TFLOPS
blockwise GQA bwd. flops
edpa
Speed (SraFLOPP)))
202
200
3k
4.k
CAUSAL BLOCKWISE GQA Bwd TELOPS
MWW.eiiiie__MMAA_EE.__ MHA.fwd flops
edpa
250
Speed (teraPL LOPeN
203
150
44.
5k
26k
32k 43%
45k
Sequence
(g) CAUSAL BLOCKWISE MEIA TFLOPS
Speed (teraTLOPPi))
200
355
302
5k
2dkk
43k
45k
(1)
Figure 13: TFLOPS SHARE QUESTION and CAUSAL BLOCKWISE
33

---

Published as a conference paper at ICLR 2026
global sliding GQA flops
1030
sdpa
173
Speed (teraFLoPwa))
123
300
Sequencc Length
GLOBAL SLIDING GQA Pwd TELOPS
global sliding GQA fopp
363
edpa
Speed (teraLLPPu))
4.kk
(b) GLOBAL SLIDING GOA Bed TELOPS
global subbal_sliding.. fadd
edpa
172
Speed (teraP LoPed
153
44.
5k
20k
32k
43k
45k
Sequence
(c) GLOOAAL SLIDING MHA TFLOPS
sliding. flops
390
340
Spand (tera LPPP))))
322
5k
26k
52k
43k
45k
(d)
Figure 14: TFLOPS of GLOBAL SLIDING
34

---

Published as a conference paper at ICLR 2026
causal, ctmment GQA flops
sdpa
Speed (OwraF) LOPAN)))
253
200
193
100
Sequence Length
(a) BLOCK CAUSAL DOCUMENT GQA Pwd TELOPS
Mock, causal cumeett GQA flops
sdpa
200
Speed (traFLoPPN))
150
100
36
4.k
(b) BLOCK CAUSAL DOCUMENT GOA Bwd TFLOPS
causal, document MHA fwdddwwdd... Hops
sdpa
Speed
150
300
44.
5k
20k
32k
43k
4.k
Sequence
(e BLOCK CAUSAL DOCUMENT MHA Fod TFLOPS
block,
bwd
300
250
Speed (tera
1.0
300
5k
26k:
32k
43k
4.5k
(d) DOCUMEET
Figure 15: TFLOPS of BLOCK CAUSAL DOCUMENT
35

---

Published as a conference paper at ICLR 2026
A.5.2 PERFORMANCE METRIC: PEAK MEMORY USAGE
We report the peak memory usage of the kernel a reference. The memory plots (Figure 17 Figure T8p clearly illustrate the detrimental impact of quadratic storage complexity on training: Naive Attention and SDPA scale only up to around 16K in the forward pass and about 8K in the backward pass. We truncate the plots at certain points and annotate the corresponding values. Un- der the same mask setting, different kernels exhibit similar peak memory the forward pass, show noticeable differences in the backward pass. Overall, the trend of peak memory scaling with sequence length across different kernels aligns well with intuition. Since sequence length is ex- tended single GPU, model parameters remain fixed, and the growth in peak memory is entirely determined by activations. In the standard attention module, activation memory computed as 11bshd + 5bhs² + 2bshd, where b is the batch size, the sequence length, h the number of heads, and d the hidden dimension (see Figure 16.. We do not consider any gradient checkpointing (Chen) al.) 2016) offloading (Ren et al., 2021 techniques. Although different kernels may employ various strategies to optimize memory usage, the overall growth trend of activations still approxi- mately follows a quadratic curve, which confirmed by our experimental results. the same time, recording peak memory further highlights the performance bottlenecks of the attention mechanism when handling ultra-Iong contexts. Due to the presence of activations, other distributed strategies such as tensor parallelism and data parallelism are insufficient to alleviate peak memory usage. Only context parallelism, which balances the workload across devices along the sequence dimension, can effectively address this issue.
Figure 16: Full Activations in Attention Module
36

---

Published as a conference paper at ICLR 2026
mem
Thunc 1N2
10th
Peak Memory (GB)
2.0
1.5
Sequence Length
FULL Fud Perk Menoey
full GQA bwd.
ThNTc
CLdYY
Ra mork
stas
10
Peak Memory (GB)
3K
8K
2,6c
32k
Sequence Lengtth
(bi FULL GQA Bed Peak Memory
full MHA fwdd me
Peak Memory (GB)
1k
2k
4k
8k
16k
24k
32k
40k
48k
Sequence Length
(c) FULL MHA Pud Peak Memory
full. MHA, bwd, mem
Peak Memory (GB)
32k
Sequence Lengtth
Mensory
37

---

Published as a conference paper at ICLR 2026
causal GQA mem
Thunc IN2
Anc
Rid, mork
otp
1ith
Peak Memory (GB)
2.0
1.5
Sequence Length
CATISAL GOA
causal GQA, bwd.
Purc
CUdVY
11 mwnk
Peak Memory (GB)
1kk
2k
16k
24k
48k
Sequence Length
CAUISAL GQA Bed Park Memory
causal MHA fwd mem
Peak Memory (GB)
1k
2k
4k
8k
16k
24k
32k
40k
48k
Sequence Length
[g) CAUSAL MHA Fwd Peak Memory
causal MHA bwd, mem
Peak Memory (GB)
32k
Sequence Lengtth
(b)
Memory
38

---

Published as a conference paper at ICLR 2026
document GQA, mem
F1E30
Pure 53
R2
612a
Peak Memory (GB)
11
Sequence Length
FULL DOCUMENT
full document, GQA, bwd
TrTTC
c.d0
nex
Ri mork
10
Peak Memory (GB)
1K
16k
24k
Sequence Lengtth
@ FULL DOCUMENT GQA Bed Perk Memory
full document MHA fwd mem
12
Peak Memory (GB)
1k
2k
4k
8k
16k
2.5g
32k
40k
48k
Sequence Lengtth
(k) PULL DOCUMENT MHA Fud Peak Memory
full, document MHA bwd. mem
H2.00
Peak Memory (GB)
2
32k
Sequence Lengtth
(1) DOCUMENT Peals Mensory
39

---

Published as a conference paper at ICLR 2026
causal document, mem
F1E30
Pure
CUUYY
fam_mwk
Peak Memory (GB)
11
Sequence Length
CAUSAL DOCUMMENT
Manory
usal document GQA bwd.
TrTTc
c.d0
nex
Riddy mork
10
Peak Memory (GB)
1K
16k
24k
Sequence Lengtth
CAUSAL DOCUMENT GQA Red Pesk Memory
causal document MHA fwd fwd_mem mem
12
Peak Memory (GB)
1k
2k
4k
8k
16k
2.5g
32k
40k
48k
Sequence Lengtth
(0) CAUSAL DOCUMENT MHA Pwdd Peak Memory
causal, document MHA bwd, mem m
H2.00
Peak Memory (GB)
B
2
32k
Sequence Lengtth
(9) CAUSAL
Memory
40

---

Published as a conference paper at ICLR 2026
sliding window full, fwd.
1220 Ponc
130) prace
Peak Memory (GB)
Sequence Langth
SLIDING Peak
sliding wwl_jwiiddww, fQAl_GQAAAAAA bwd,
Tuns
Peak Memory (GB)
16k Sequence Length
(n) FULL SLIDING WINDOW GQA Rod Peak Memory
sllidingwwin window. full. MHA. fwol mem
4120
Peak Memory (GB)
1k
2k
4k
8k
16k
24kc
32k
40k
48k
Sequence Length
(s) FULJ St MING WINDOW MHA Fed Peak Memory
sliding window full MHA, bwd mem
11000
Tunc
Peak Momory (GB)
2k
4k
Sk:
16k
242
32k
40k
48t
(0 FULL SLIDING WINDOW MILA Bud Peak Memory
41

---

Published as a conference paper at ICLR 2026
sliding window wiinow__cusall__e GQA_
Baapl_oAA____avvaaee Lvd_mem
Tunn
larch
Peak Memory (GB)
11
IK
2k
4k
8k
16k
24k
32k
40k
48K
Sequence Lengtth
(0) CAUSAL SLIDING WINDOW GQA Fud Peak Memory
sliding window causal GQA bwd, mme
11:00
Peak Memory (GB)
1k
2k
16k
24k
32k
40kc
48k
Lengtth
(n) CAUSAL SLIDING WINDOW GOA Food Peak Memocy sliding_window wwtpwwiiddwwwwee causal MHA fwdd mem
B1300
Vurc
fu5
codi
faskumakk
barhh
Peak Peak Memory (GB)
Sequence Length
CAUSAL WINDOW MELA Memury
sliding
HLD0
1ru1
f.5
cidin
fix
Rasi makk
>bo
Peak Memory (GB)
1k
32k
Sequence Lengtth
CAUSAL SLIDING WINDOW MILA Boe Peak Memory
Figure 17: Peak Memory of Static Regular Masks
42

---

Published as a conference paper at ICLR 2026
1220
Tur
Rich, prace
Peak Memory (GB)
Sequence Langth
PREHX Peak
prefix GQA, bwd,
Peak Memory (GB)
8k
14k
Sequence Length
(b) PREFIX TM GQA Roo Pesk Memory
prefix Jm. MHA, www.f/dddeeeee_o memm
4120
Peak Memory (GB)
1k
2k
4k
8k
16k
24k
22k
43k
48k
Sequence Length
(6) PREFIX 1M MHA Foe Perk Memory
prefix Im. MHA bwd mem
[1000]
Tune
Peak Momory (GB)
1k
8k
16k
24k
32k
40k
45k
(d) PREFIX LM MITA Bed Peak Memory
43

---

Published as a conference paper at ICLR 2026
profix_ document,
1220
Tur
Richh pracc
Peak Peak Memory (GB)
Sequence Langth
DOCUMENT GQA Memory
prefix_im dwcument,cccommm BereeeetttGAAAA GQA_ bwd.
Palt
Peak Memory (GB)
8k
16k
Sequence Length
PREFIX M DOCUMENT GQA Bwd Peak Memory
prefix Jm, document MHA fwd_mem
1000
Peak Memory (GB)
2
1k
2k
8k
16k
24k
32k
40k
45k
Sequence Length
(g) PREFIX 1MM DOCUMENT MHA Fied Pesk Memory
prefix Im document MHA bwd. mem
11000
Peak Momory (GB)
2
1k
8k
16k
24k
32k
40k
45k
Ohi PREFIX LM DOCUMENT MILA Peak Memory
44

---

Published as a conference paper at ICLR 2026
quastion fwd
1220 Tur
Rich praCe
Peak Peak Memory (GB)
Sequence Langth
Pesk
share, qww.tieeennnno. GQAA bwd
Peak Memory (GB)
8k
16k
Sequence Length
(i) SHARE QUESTION GQA Roo Pesk Memory
share question. MHA fwd__em
4120
Peak Memory (GB)
1k
2k
4k
8k
16k
24%
22k
43k
48k
Sequence Length
(1) SHARF QUESTION MHA Foe Renk Memory
share question MHA bwd. mem
[1000]
Peak Momory (GB)
1k
16k
24k
32k
40k
45k
(1) SHARE QUESTION MILA Bod Peak Memory
45

---

Published as a conference paper at ICLR 2026
causal blockwise fwd
1220
Tur
Rich praCe
Peak Peak Memory (GB)
Sequence Langth
CAUSAL BLOCKWISE Peak
causal www.kwwweeweeecccomooooo SQA.eaieeeeeee__eee bwrd mem
Peak Memory (GB)
8k
16k
Sequence Length
CAUSAL BI BLOCKWISE GQA Bud Peak Mamory causal blockwise, causal_blockwiise
MWA_fwd, mmem
10000
Peak Memory (GB)
2
1k
2k
8k
16k
24k
32k
40k
45k
Sequence Length
(o) CAUSAAL BLOCKWISS MHA Foud Pesk Menory
causal blockwise MHA, bwd mem
11000
Peak Momory (GB)
2k
9k.
166
2
32k
40k
48k
(p) CAUSAL BLOCKWISE MILA Bwd Peak Memory
46

---

Published as a conference paper at ICLR 2026
1220
Tur
Rich mrace
Peak Memory (GB)
Sequence Langth
GLOBAL SLIDING Peak
global, wwidilsiiiiggg... GQA, bwd.
Peak Memory (GB)
8k
14k Sequence Length
GLOBAL SLIDING GQA Bud Paak Memory
global. s sliding MHA fwdd mem
4120
Peak Memory (GB)
1k
2k
4k
8k
16k
24%
22k
43k
48k
Sequence Length
(9 DBAL. SLIDING MHA Fud Peak Menory
global sliding MHA bwd. mem
[1000]
Peak Momory (GB)
2k
9k
166
24c
32k
40k
48k
(U GLOBAL SLIDING MIIA Bed Peak Memory
47

---

Published as a conference paper at ICLR 2026
block causal document, GQA, fwd, mem
4120
Peak Memory (GB)
0.
1k
2k
4k
16k
24kk
32k
40k
48k
Sequence Length
(0) BLOCK CAUSAL DOCUMENT GQA Fod Peak Memory
block causal, document GQA, bwd. mom
14003
Peak Memory (GB)
1k
2k
6K
8k
16k
24kc
3.kk
40k
48k
Length
(V) BLOCK CAUSAL DOCUMENT GQA Bwd Peak Memory
document MHA, mem
M302
Fash, macc
Peak Memory (GB)
Sequence Length
(w)
Memory
block, causall.ocumeett
H302
VLix
1ox
Fach_mask
Peak Memory (GB)
10
Sequence Length
DOCUMEEN Bud Peak Manory
Figure 18: Peak Memory of Static Heterogeneous Masks
48

---

Published as a conference paper at ICLR 2026
A.6 DETAILS OF SPARSE KERNEL PERFORMANCE
In this section, present the detailed comparisons for sparse kernels, including TFLOPS perfor- mance and peak memory usage. There are only partial baselines in each figure or table because current block sparse kernels have such limitations in functionality: VSA and Triton VSA do not support GQA and 128 block size, FA2 Sparse does not support 64 block size, FA2 Sparse and Flash- Infer do not support backward computation, and Flex Attentioo faces severe memory issues so we discuss separately.
A.6.1 PERFORMANCE METRIC: FLOPS
TFLOPS of MHA with 64 block size. We report the performance of block sparse attention kernels with MHA and 64 block size in Figure 19 Only VSA, Triton VSA and FlashInfer support this setting. The left side shows the foward FLOPS with different sparsity ratios while the right side shows the backward FLOPS. Our findings reveal that VSA performs stably across different sparsity ratios, showing their robustness with adaptive computation. FlashInfer also performs stably but suffers from OOM issues with higher sparsity ratios because there are more blocks to compute, causing memory overhead with its metadata. However, we find that the performance of VSA reduces with the increase of the context length, especially for backward computation. This indicates that there may exists optimization opportunitics for larger context lengths and backward computation for training scenarios.
TFLOPS of MHA with 128 block size. We report the performance of block sparse attention kernels with MHA and 128 block size with forward computation in Table[3] Only FlashInfer and FA2 Sparse support this setting. Comparing with the performance of 64 block size in Figure [9 the TFLOPS of FlashInfer in 128 block size increases about 2x, proportional the scalability of FlashInfer block sparse kemnlls. Its OOM issues also decrease compared with 64 block size, because larger block size means smaller number of blocks, leading to smaller metadata storage. For the TFLOPS of FA2 Sparse, it demonstrates robustness across different sparsity ratios and context lengths. Its average performance is about 300+ FLOPS because it is not optimized for NVIDIA Hopper GPUs. There exists opportunities for tailored optimizations for specific hardware platforms to unleash the hardware performance.
the block size increase, showing
Separate explanation of FlexAttention TFLOPS. FlexAttention is separately discussed because is hard to generate the dynamic block sparse block through compilation, causing severe memory overhead with O(S?) block mask representations. So only test its TFLOPS performance with 4 heads and 16K context length in Table 6 performs bad in 64 block size due to the lack of optimizations with small block size. While 128 block size, it is comparable TFLOPS of FA2 Sparse in T3 The reason behind that FlexAttention uses 128 as its default block size, so the kernels it generates demonstrate relative good TFLOPS. also supports backward computation with the similar performance.
TFLOPS of trainable block sparse attention.
For NSA (Yuan et al., 2025), we follow the same setting, with 64 query heads and 4 key, value heads. In this configuration, each query group (comprising 16 query heads) selects the same key, value indices. Each query selects 32 Key-blocks with a block size of 64, resulting in total of 2048 selected tokens. Since the selection attention module dominates the end-to-end overhead and aligns with our block sparse framework. our experiments focus on this module. This configuration maps to q block size 16 and k block size 64 within our block sparse attention benchmark frame- work. We conduct both NSA-Triton and NSA-Tilelang in the benchmark, as shown in Table 10) implemented by Triton and Tilelang, respectively.
For DSA (DeepSeek-Al 2025), since it utilizes token-level sparse attention and is designed based on DeepseckMLA model architecture, we set the configuration to 128 query heads and I key, value head, matching DeepseekMLA. Each query group therefore consists of 128 query heads sharing the same key, value indices. We set the top-k selection size to 2048 tokens. In our framework, this corresponds to q_block_size 128 and k_block_size 1. We benchmark this setup using the DSA- DSA-Tilelang and Flas as shown in Table implemented by Tilelang and CUDA, respectively.
49

---

Published as a conference paper at ICLR 2026
sparsity MHO, fees
spersity MHHA, WWW.AHAAA.Wl..... floss
122,3000
Spaad (nEatLOPPs))
930
3.0
230
230
1.4
574
364
64.
11.0
1788
144
57%
-004
614
544
11.2
1244
Sequence Length
Sequence Length
(0 MI1A Fed TFLOPS with 0.2 sporsity ratio
(b) MHA Bwd TFLOPS with 0.2 spursity ratio
sparsityo.s MHO Fand Peps
spanchyo.5 5 JHHA, basd Plops
893
N223300
Sasttrk Infer
Spoodd (B6r3FLOPPN)
230
330
1.4
324
Sequence Length
Sequence Length
(6) MHA Ped TFLOPS with 0.5 sparsity catio
(d) MHA Rwd TRICISS with 0.5 sparsity mnlo
sparsity MIA Raps
sparsityO.0 2 MHA bwd
893
590
990
4.00
(HIOPLOPANE
230
2.3
230
220
(6) MHA Ped THUUPS 0.8 spacsity ratio
(5) MHA Bud TELOPS with 0.8 sparsity
Figure 19: TFLOPS of block sparse attention kernels with 64 block size and MHA
A.6.2 PERFORMANCE METRIC: PRAK MEMORY USAGE
Peak memory of MHA with 64 block size. We report the peak memory of block sparse attention kernels with MHA and 64 block sizes in Figure 20 Only VSA, Triton VSA and FlashInfer support this setting. The left side window shows the forward pass with different sparsity ratio while the right side window shows the backward pass. VSA and VSA trition share the same memory usage while flashinfer exhibit higher GPU memory consumption which may because of its metadata representa- tion for block sparse. It is also clear that The GPU memory consumption of VSA and VSA. -ttrtiition shows almost no correlation with the sparsity matiio. In contrast, the memory footprint of FlashInfer grows in proportion to the increase in the sparsity_raticn shown in Table in GQA scenarios, FlashInfer's ' GPU memory usage is significantly lower than in MHA scenarios, for both block sizes of 64 and 128. This is because the metadata required repre- sent the block sparse structure is greatly reduced. As for FA2 Sparse, its' memory consumption is greatly lower than flashinfer, which indicates flashinfer's 's poor performance in terms of GPU mem- ory usage. Peak memory of MHA with 128 block size. We report the peak memory of block sparse attention kernels with MHA and 64 block sizes in Table compared with MHA with 64 block size, flashinfer consumes less GPU memory. However, FA2 Sparse utilizes more GPU memory com- pared with GQA scenerio. Separate explanation of FlexAttention peak memory, FlexAttention is separately discussed because it is hard to generate the dynamic block sparse block through com- pilation, causing severe memory overhead with ( block mask representations. So we only test its TFLOPS performance with heads and 16K context length in Table 6 most of Attention'
Peak memory in GQA scenarios. As
50

---

Published as a conference paper at ICLR 2026
Table 4: TFLOPs of Sparse Kernels for GQA (64:8) Forward
Note: Seqlen Sequence Length, SR = Sparsity Ratio, x Not Supported. H100 GPU,
BlockSize 64
BlockSize 128
Kernels
Seqlen
SR 0.2
SR 0.5
SR 0.8
SR 0.2
SR 0.5
SR 0.8
16k
211.78
236.30
247.83
387.63
449.79
472.52
32k
232.81
243.77
203.54
422.95
455.97
485,06
48k
211.50
218.30
215.47
413.11
434.61
432.16
64k
190.73
210.82
204.16
414.25
427.77
399.22
FlashInfer
80k
201.71
198.39
206.43
351.87
393.34
410.88
96k
202.03
197.48
201.42
393.16
387.89
398.73
112k
198.14
196.82
199.41
393.41
389.33
397.02
128k
196.47
195.61
197.93
389.55
396.90
393,80
16k
X
X
X
312.54
326.62
332.00
32k
X
X
X
327.46
333.45
339.16
48k
X
X
331.99
281.27
331.24
64k
X
X
X
328.84
329.88
329.99
FA2 Sparse
80k
X
X
X
293.43
330.62
319.52
96k
X
X
X
328.38
318.18
321.05
112k
X
X
298.22
316.92
322.47
128k
X
X
317.48
317.11
321.04
Table
Forward TFLOPs of Sparse MHA Kernels (64:64, Block Size = 128)
Note: Seqlen Sequence Length, SR Sparsity Ratio, X Not Supported, H100 GPU,
Kernels
Seqlen
SR 0.2
SR 0.5
SR 0.8
16k
414.93
447.76
484.21
32k
425.38
416.37
425.88
48k
407.67
401.60
402.69
64k
398.01
394.79
392.05
FlashInfer
80k
342.30
384.46
OOM
96k
388.09
OOM
0OM
112k
380.67
OOM
OOM
128k
384.05
OOM
OOM
16k
315.12
326.09
331.27
32k
327.63
333.32
337.07
48k
331.79
282.22
335,95
64k
325.86
330.73
329.45
FA2 Sparse
80k
295.20
328.64
318.44
96k
325.09
317.42
320.26
112k
301.61
318.13
322.56
128k
308.39
317.80
321.67
GPU memory consumed during mask creation, while its runtime memory usage high. In block sparse scenarios, an efficient mask representation is crucial for GPU memory consumption.
51

---

Published as a conference paper at ICLR 2026
Table 6: TFLOPs of Sparse FlexAttention (SeqLen 16K)
Note: Seglen Sequence Length, SR Sparsity Ratio, H100 GPU.
BlockSize 64
BlockSize 128
Mode
Forward
Backward
Forward
Backward
29.05
57.22
339.74
299.43
GQA (4:1)
48.21
95.33
352.09
312.76
0.8
99.80
170.40
358.55
325.52
0.2
28.75
63.14
341.55
305.37
MHA (4:4)
0.5
48.76
103.37
345.84
337.60
0.8
101,58
186.44
352.18
349.00
22 PrEm
speraPyy.2 mwn.
144
244
4.b
6.6C
80C
112
3284
244
204
4.2
6.CC
20.4
21.5
328c
Seiguenca Lanaghh
Serguercca Leneghh
(4) MHA perk memory 0.2 sporsity ratios
(h) Rwd park menicey sparsity calio
sparsityO.s MHA meam
sport MHA buad,
0311
1PA
6dc
99c
1155
776
124
64c
Rerporse Lerepth
Rergursee Leregh
(c) MILA Fwd peuk memory with 0.5 spunity rañoo
(d) MIL Bwd peak menory with 0.5 spursity ratio
spersityO.i MHA Fwel prom
spors PHA Heed men
6/c
OC4
1120
246
224
436
Sequencee Lanighh
Sequueica Lenigh
(c) MHA Fwd pruk memory with 0.8 spumity rañio
(0 MHAA Bwd peak memory with 0.8 spunity ratio
Figure 20: Peak memory of block sparse attention kernels with block size and MHA
52

---

Published as a conference paper at ICLR 2026
Table 7: Peak Memory (GB) of Sparse Kernels for GQA (64:8) Forward
Note: Seqlen Sequence Length, SR = Sparsity Ratio, x = Not Supported. H100 GPU.
BlockSize 64
BlockSize 128
Kernels
Seqlen
SR 0.2
SR 0.5
SR 0.8
SR 0.2
SR 0,5
SR 0.8
5.468
5.502
5.540
5.448
5.466
5.485
32k
6.620
6.770
6,920
6.552
6.627
6,702
48k
7.837
8.174
8.512
7,684
7.853
8.022
64k
9.115
9.715
10.317
8.845
9.145
9.444
FlashInfer
80k
10.455
11.393
12.330
10.034
10.502
10,970
96k
11.857
13.207
14.557
11.250
11.923
12,599
112k
13.322
15.159
16,997
12.493
13,413
14,331
128k
14.846
17.247
19.647
13.766
14.964
16.165
16k
X
X
X
0,810
0.810
0.810
32k
X
X
X
1.393
1.393
1.393
48k
X
X
1.986
1.986
1.986
X
X
2.590
2.590
2.590
FA2 Sparse
80k
X
X
X
3.205
3.205
3.205
96k
X
X
X
3,830
3.830
3,830
112k
X
X
4.466
4.466
4.466
128K
X
5.113
5.113
5.113
Table 8: Forward Peak Memory of Sparse MHA Kernels (64:64, Block Size 128)
Note: Seqlen
Sequence Length, SR Sparsity Ratio, X Not Supported, H100 GPU.
Kernels
Seqlen
SR 0.2
0.5
0.8
16k
5,477
5.627
5.776
32k
6.791
7.392
7,991
48k
8,316
9.666
11.015
64k
10.051
12.452
14.849
FlashInfer
80k
11.994
15.746
0OM
96k
14.148
OOM
112k
16.514
OOM
OOM
128k
19.085
OOM
OOM
16k
1.251
1.251
1,251
32k
2.281
2.281
2.281
48k
3.329
3.329
3.329
64k
4.395
4.395
4.395
FA2 Sparse
80k
5,478
5.478
5.478
96k
6,587
6.578
6,578
112k
7.696
7.696
7.696
128k
8.832
8.832
8.832
Table 9: Peak Memory (GB) Sparse FlexAttention (SeqLen = 16K)
Note: Seglen Sequence Length, SR Sparsity Ratio, H100 GPU.
BlockSize 64
BlockSize 128
Mode
SR
Forward
Backward
Forward
Backward
0.2
0.285
0.383
0.294
0.392
GQA (4:1)
0.5
0.287
0.384
0.296
0.394
0.8
0.288
0.386
0.297
0.395
0.2
0.304
0.449
0.313
0.458
MHA (4:4)
0.306
0.450
0.315
0,460
0.8
0.307
0.452
0.316
0,461
53

---

Published as a conference paper at ICLR 2026
Table 10: TFLOPs of NSA for GQA (64:8)
Note: Seqlen Sequence Length, H100 GPU.
NSA-Triton
NSA-Tilelang
Seqlen
Forward
Backward
Forward
Backward
32k
97.35
140.34
60.91
3.97
64k
90.5
132.99
49.19
1.08
96k
83.39
129.65
45.16
1.3
128k
76.74
120.62
39.5
0.55
Table 11: TFLOPs of DSA for GQA (64:8)
Note: Seqlen
Sequence Length, H100 GPU.
DSA-Tilelang
DSA-FlashMLA
Seqlen
Forward
Backward
Forward
Backward
32k
334.13
114.64
610.49
64k
326.89
103.17
607.13
96k
331.28
98.49
589.91
128k
331.99
95.95
555.27
X
54

---

Published as a conference paper at ICLR 2026
A.7 DETAILS OF CONTEXT PARALLEL ATTENTION PERFORMANCE
Performance of ring P2P. The main limitation of ring P2P is that its communication volume cannot be adjusted. In each iteration, while computing with the KV pairs for the current stage, each device simultancously sends its own KV to the next device in the topology and receives the KV needed for the next stage. In multi-node, multi-GPU setup, all GPUs are connected in a ring-based P2P topology. The communication bottleneck is determined by the slowest inter-node link, forcing intra- node communication synchronize with the inter-node transfers, which leads to substantial overall bandwidth underutilization.
Ultimately, the overall efficiency of Ring P2P is determined by the actual per-GPU computation, which manifests in whether communication can be overlapped with computation. The FULL sce- nario represents the optimal performance case for Ring P2P. As shown in the Figure 21 in this scenario, each GPU executes the largest per-kernel computation. When inter-node communication can also be effectively overlapped with computation, further scaling the distributed setup does not significantly change the balance between computation and communication efficiency, the overall computational efficiency remains essentially constant.
In the CAUSAL scenario, the per-GPU communication volume of Ring P2P remains the same as in the FULL scenario. Even after load balancing, the per-stage computation per GPU is roughly half of that in the FULLL scenario (except for the first stage, which is 3/4). The reduced computation may no longer fully overlap with communication, leading to performance drop in Ring P2P under CAUSAL, as shown in the Figure 22 During the forward pass, when scaling nodes, we observe further performance degradation. We attribute this mainly to machine instability: even if only a single GPU underperforming in the ring, for example from a sudden drop in computation efficiency, can stall the entire topology and greatly reduce overall efficiency. In large-scale distributed settings, such effects are inevitable, and we report the experimental results faithfully.
For the DOCUMENT scenario, as shown in the Figure 23 Ring P2P exhibits similar trends in both the CAUSAL DOCUMENT and FULL DOCUMENT settings. Unlike the FULLICAUSAL. scenarios, it does not maintain relatively constant trend, which is expected. First, when han- dling veriable-length data, each segment must be padded according to its specific scale, leading to significant variation in Per-GPU computation per iteration, while communication volume remains constant. Second, sampling of variable-length data introduces differences in computational sparsity across iterations, resulting in fluctuations in the overall trend.
Hybrid Design. USP and Loong share very similar overall architectures, using the Ulysses design intra-node and Ring P2P inter-node. Overall, they achieve significant and stable performance improvements in the FULL and CAUSAL scenarios, as shown in Figure 21 and [2 In the FULL DOCUMENT and CAUSAL DOCUMENT scenarios, performance gradually decreases due re- duced overall computation, which aligns with expectations. Additionally, in the DOCUMENT sce- nario, as shown in Figure 24 and 23 both architectures demonstrate improved stability, mitigating the limitations of using only Ulysses or Ring P2P.
Here, we primarily explain why LoongTrain generally outperforms USP during the forward pass, yet performs on par with or slightly worse than USP during the backward pass.
In our benchmark reproduction and optimization, USP and Ring P2P both use the same RingAttn class. In the ring-topology iterations, the forward and backward data flows are essentially opposite. For example, during the forward pass, GPU, receives data from GPU,-1 and sends data to GPU,+1; in the backward pass. GPU, receives from GPU and sends to GPU,-
This design is both necessary and reasonable. At the end of the forward pass, GPU, actually holds the initial-stage KV data of GPU,, and on, with G finally holding GPUp's 's initial data. If the backward pass rotated data in the same direction as the forward pass, it would require either an additional P2P communication storing the initial-stage KV on each GPU in advance, both of which incur extra overhead.
Instead, we exploit the time difference between KV and gradient generation: the backward pass directly continues from the forward-pas KV states in reverse rotation, while gradients computed in the current stage sent during the next stage. After completing the same number of rotations, each
55

---

Published as a conference paper at ICLR 2026
GPU receives exactly its corresponding gradient (e.g., GPUg receives the gradient for KVu, and so on).
For Loong Train's DoublcRingAttn, the forward pass is consistent with USP. In addition, Loong leverages a heterarchical P2P architecture to implement a two-level sliding window, decomposing the full ring topology into wira-w and inter- windww groups. The intra-window group is iden- tical to the RingAttn class, but in the first stage, an additional P2P communication performed for the inter-window group prefetch the initial data for each GPU after the inter-window rota- tion. This design fully utilizes inter-node bandwidth, resulting in superior forward-pas performance compared USP.
However, Loong 's backward pass cannot directly leverage the forward-pes -aads end states. While this poses no issue for the last inter- inter-window, alters the initial state for each subsequent inter- window, and the final states differ depending on the specific intra- and inter-window configuration. As result, LoongTrain performs forward and backward passes using the same rotation order. Each GPU additionally stores the initial KV data, and ensure correct gradient propagation, extra P2P communication and synchronization of gradients is performed at the end of each inter-window. This guarantees that in the next inter-window, each GPU receives the corresponding KV data and gradi- ents. Consequently, LoongTrain gains no significant backward-pate advantage from the heterarchi- cal architecture, yet overall maintains performance comparable to USP. The observed fluctuations in trends are similarly attributed to the instability introduced by the additional P2P communications.
Ulysess. For Ulysess, the results straightforward: different sampling patterns naturally lead to variations in computation, and we recommend evaluating performance based on the specific appli- cation scenario.
A.8 ANALYSIS OF CONTEXT PARALLEL ATTENTION COMMUNICATION
Ulysess.
Forward (FWD): 2 All-to-all for Q and O, and 2 All-to-all for K and V. The communication volume 2- N- N2 the 2. N- 1-1th,d.
N2 thyd.
Backward (BWD): Alll for dQ and dO, and 2 All-to-all for dK and dV. The communication volume is 2- th 2 thgt
Overlap: All communications are exposed outside the computation.
Ring All-Gather.
Forward (FWD): 2 All-Gather for K and V. The communication volume is ty-thhud d * 2.
Backward (BWD): All-Gather for K and V, and Reduce-Scatter for dK and dV. The commu- nication volume is Nh-th d 2 + N-1 - thw d * 2.
Overlap: All communications are exposed outside the computation.
Ring P2P.
Forward (FWD): N P2P for K and V. The communication volume is N-1theod 2.
P2P for K and V, and N -1 P2P dK and dV. The communication volume is N-1 thg d thith
Backward (BWD):
Overlap: In theory, the FULL mask scenario achieves communication-comppaaaiii overlap. Focus- ing solely mask-sparse scenarios, computation and communication start simultaneously each stage across different ranks, with the communication volume remaining unchanged. The portion of 2 e flops
the entire pipeline's communication that can be overlapped approximately t/banduridth k, where t denotes the single computation area, and k merely represents correction coefficient that exists as constant for specific settings.
USP. We denote N, the Ulysses size and N, the Ring-P2P size, such that N, N₂ N. Forward (FWD): 2 All-to-all for Q and O, 2 All-to-all for K and V, and Np - P2P for K and V. N tho 2 + No th 2.
The communication volume is No
~ theyd 2 -
NooNa
56

---

Published as a conference paper at ICLR 2026
Backward (BWD): All-to-all for dQ and dO, All-to-all for dK and dV, NNp P2P for K and V, and N2 P2P for dK and dV. The communication volume is N. No the d No*N? thg
Na-1
N,-1 thg + Ny*N.
Np-1 thk 2.
No*Na
Overlap: the analysis remains the same as Ulysses and Ring P2P.
LoongTrain. We denote as the intra-ring size and k as the window number, such that l Np- Forward (FWD): All-to-all for Q and O, 2 All-to-all for and V, (n (n-1) 1) P2P for intra and N.-1 No thy
V, and k - 1 for inter K and V. The communication volume N,-1 thkod + R-1 2 (m-1)rk theud * + No*N. th d 2. No*N,
NgaN2
Backward (BWD): 2 All-to-all for dQ and dO, All-to-all for dK and dV, (n K and V, k - 1 P2P for inter K and V, n k P2P for intra dK and dV, and k P2P for inter dK and dV. The communication volume is Na-1
P2P for intra
No the * 2 + Na-1 NyaNa th * + N,+N,
(n- v-=1)* the * 2 +
k-1 k-1 thand 2 N. thkoo d+ * N,*Na tha. d 2 N,*N,
Overlap: The forward is consistent with USP. For the backward, the intra and inter communications of K and V are the same as USP. However, regarding intra dk and dV, unlike USP, which places intra K, V. dK, and dV in single buffer, LoongTrain separates K, V from dK, dV into two buffers, requiring two separate P2P communications. Thus, in each intra stage of Loong Train, the intra dK and dV communications putations, overlapping with the correction of dK and dV in the next stage. That said, LoongTrain incurs additional intra communication for dK and dV, leaving approximately nication exposed. After the final intra dK and dV communication concludes, Loong Train initiates an inter dK and dV communication, which overlaps with the subsequent intra computations. How- ever, the dKK and dV communication for the last window remains exposed.
asynchronously initiated only after the completion of com-
P2P commu-
full_fwe folo
700 (N128]
ulysses
ring_p2p
ring_allgether
usp
loongtrain
Speed (traPLOPPPN)
400
300
100
Distributed Scale (eGPUs)
flops
500
H1:00
ulysses
ring_p2p
ring_allgether
usp
loongtrain
Speed (LaraPLDP%(N)
503
300
Distributed Scale (#GPUs)
Figure 21: Forward and Backward TFLOPs of Context Parallel Attention FULL
57

---

Published as a conference paper at ICLR 2026
causal_fvd. Sausal_wvd_fopp (oop
703
[0.20
ulysses
fing_p2p
ring_allgether
usp
loongtrain
500
Speed (taraPLOPAN)
500
200
Distributed Scale (eGPUs)
causal flops
M1.20
ulysess
ring ppp
ring_allgather
usp
loongtrain
Speed (teraFLOPsM)
400
300
Distributed Scale (#GPUs)
Figure 22: Forward and Backward TFLOPs of Context Parallel Attention on CAUSAL
causal, document_fwe plops
P1:20
ulysses
ring_p2t
ring_allgather
loongtrain
500
409
300
Speed
200
100
Distributed Scale (#GPUs)
causal document bwd. flops
500
ulysses
rig_p2p
ring_aligan ang_aalgather
usp
Inongtrain
409
Speed (teraPLOPPN)
300
Distributed Scale (#GPUs)
Figure 23: Forward and Backward TFLOPs of Context Parallel Attention CAUSAL DOCU- MENT
58

---

Published as a conference paper at ICLR 2026
full document bwd flops
-1:00
ulysses
ring_p2p
ring_allgsther
loongtrain
Speed (taraPOPPP))))
16
24
22
48
56
80
(#GPUs)
Figure 24: Backward TFLOPs of Context Parallel Attention FULL DOCUMENT
59

---

Published as a conference paper at ICLR 2026
A.9 DETAILS OF DENSE KERNEL PERFORMANCE ON A800
We focus on reporting the FLOPs performance of the dense kernels on A800, since the memory consumption exhibits only small variations.
First, document input layout for the cuDNN fused kernel is specifically designed for the Hopper architecture. Therefore, attempted expand the input into basic four-dimensional format to meet the kernel's layout requirements. However, in most scenarios, the kernel still could not provide support, and we ultimately excluded these experiments. Additionally, FA3 itself is highly optimized implementation targeting the Hopper architecture, so we excluded the FA3 kernel as well. Similarly, FA2 only supports homogeneous masks, yet it still achieves competitive, often optimal performance in homogeneous mask scenarios, as shown in Figure 25 and 26
worth noting that under the sliding window mask configuration, due the window size design, when the sequence length is Ik, the configuration degenerates into a FULL mask, as shown in Figure 27 In such cases, all optimized kernels perform extremely well, which outcome of the experimental setup rather than an inherent advantage of any particular kernel.
reasonable
FlashMask demonstrates comparable or even superior performance to FA2 across multiple experi- ments with homogeneous masks, and its main advantage is its native support for the vast majority of heterogeneous mask patterns, as shown in Figure 30 these irregular sparse scenarios, FlashMask achieved the best results in our experiments.
We observed that FlexAttention delivers strong performance during the forward pass, but experi- ences noticeable drop in backward performance, shown in Figure 28] This is likely attributable to combination compilation optimizations, hardware capabilities, and evaluation methodology. We report this phenomenon faithfully in the paper while maintaining all experimental methods and measurement criteria consistent across all kernels.
Finally, although the Naive Torch implementation and PyTorch SDPA handle arbitrary mask types, both inevitably face significant computational and memory overhead, as shown in Figure [3 and 29 As result, their performance far inferior highly optimized kernels in most scenar- ios. This further suggests that for heterogeneous sparse attention, the fundamental bottleneck lies in kernel-level scheduling strategies and compute pipeline optimization, rather than merely mask support capability.
60

---

Published as a conference paper at ICLR 2026
full, BwW.G/AAAwwwaaoopoooooo
X030 fa2
cusii
flox
flash, mask
tondh
Spead (tersPLOPPs)
100
1.c
24
5k.
26k:
24k
32k
40k
4Bc
Sequence Length
NVIDIA
full AOI_GGA bwr, flops
200
A0000
1s2
cudrn
flex
flash_mask
sopa
torch
175
Speed (teraFLoP (PPA)
103
125
3k
18k
24k
32K
40kc
48c
Sequencce Length
TELOPS NVIIIA
full, MHA fwd flops
A030
1:2
flex
Flash_maskk mask
sdpa
tarch
200
Speed (traP(oPPe))
150
24c
Sequencc Length
(e) FULL MIJA Fed TELOPS NVIDIA A800 GPU
full Bwt_//AAAddd..d.a bwd_flops
200
X390
Tixe
flash, mask
sdpa
tarch
375
Speed (teraLoPPa))
150
225
102
4k
24cc
Sequence Length
(4) FULLI MHA Rwd TELOPS NVIDIA A800 GPII
61

---

Published as a conference paper at ICLR 2026
causal GQA, flops
A030
sdpa
Speed (teraf LOre)))) LOPAN)
100
Sequence Length
(e) CAUSAL GQA Fed TELOPS on NVIDIA A800 GPU
causal GQA, bwd, Nops
fashhmask
sdpa
175
Speed (traFLoPPe))
100
125
300
14kk
2.4k
40c
40c
Sequence
(5 CAUSAL GQA Bud TELOPS on NVIDIA A800 GPU MHA_L_MAA_fCC flops
fashhmaak
sdpa
177
Speed (terat LOPaN
159
202
26k
24k
32k
40c
48c
Sequence
(e) CAUSAL MHA Fed TFLOPS on NVIDIA A800 GPU
flops
flsh
Speed (tera PLOPAid
26k:
24k
32k
40x
48c
CAUSAL on
Figure 25: TFLOPS of FULL and CAUSAL on NVIDIA A800 GPU
62

---

Published as a conference paper at ICLR 2026
full document, GQA Fwd Πорs
Fle
flash, mask
sdpa
tarch
175
Speed (twaPLoPan)
150
323
100
24k
See Sequence Length
DOCUMMNTT GQA Fud TFI OPS NVIIA A800 GPU
dolumennt bow. fops
3.03
A030
fa2
fixx
flash, mask
sdpa
tarch
Speed (teraf
100
Sequence Length
FULL DOCUMENT GQA Bed TELOPS NVIDIA A800 GPU
full, document MHA fwd flops
fixx
flash, mask
sdpa
tarch
Speed (braFLOPP)))
Sequence Length
(0) PULL DOCUMENT MHA Fed TELOPS on NVIDIA A800 GPU
full, dellmmcttent, MHA bwd_
160
A030
Flash_mask
sdpa
tarch
1.03
Speed (traFlOPee)
2.29
Sequence Length
(8 FULL DOCUMENT MILA Bud TELLPS on NVIDIA A800 GPU
63

---

Published as a conference paper at ICLR 2026
causal, document GQA, Fwdd Roos
A030
Speed (OeraFLOPwN)
193
223
200
Sequencc Length
(e) CAUSAL DOCUMENT GQA Fwd TELOPS on NVIDIA A800 GPU
GQA, bwd flops
edpa
Speed (trrFLoPPg))
4k
8c
CAUSAL DOCUMENT GOA Bwd TELOPS on NVIIIA A800 GPU
MHA_fwd wwdd flops
175
sdpa
150
Speed (teraP LOPpN
100
1k
24
4k
8c
26k
24c
Sequence
(2) CAUSAL DOCUMENT Fwid TELOPS NVIDIA GPU
3.2
Speed PLOPPN)
1.23
1k:
24
4k
84
26k
24
CAUSAL DUCUMENT
Figure 26: TFLOPS of FULL/CAUSAL DOCUMENT on NVIDIA A800 GPU
64

---

Published as a conference paper at ICLR 2026
sliding window full GQA, fwd. flops
Flec
flash, mask
<dpa
tarch
173
Speed (twaPLoPan
153
123
100
Se Sequence Length
FULL SLIDING WINDOW GQA Fodd TII OPS NVIDIA A800 GPP
sliding www.wwwwwwww... full, GQA, bwd_ffops
340
A030
ficx
Flash mask
sdpa
tarch
1.23
Speed (teraf (tPraf LSPAN)
Sequence Length
(b) FULL SLIDING WINDOW GQA Bod TFL
OPLOPS OPS NVIDIA A800 GPU
slidiing wiiddwwwwf full. MHA Ewdd flope
200
fixx
flash, mask
sdpa
tarch
175
Speed (brrFLOPPe))
150
123
Sequence Length
(ULL SLIDING WINDDW MHA Pwd TFLOPS on N VIDEA A800 GPU
sliding, www.://idwwww.... full, MHA Mww.bwdddddd_a_ flops
X030
Besh_mask
5dpa
tarch
320
Speed (traFLOPee))
Cequence Length
(6) FULL SLIDING WINDOW MIA Bed TFLOPS on NVIDIA A800 GPU
65

---

Published as a conference paper at ICLR 2026
sliding causal, GQA flos
A030
flssh
sopa
253
(Owra))
159
200
Sequencc Length
(e) CAUSAL SLIDING WINDOW GQA Fwd IFLOPS on NVIDIA A800 GPU
sliding window causal, GQA, flops
160
flshhmaak
scpa
140
Speed (trrFLLoPec)
1.2
103
2.4k
40c
48c
(O CAUSAL SLIDING WINDOW GQA Bid IFLOPS on NVIDIA A800 GPU
sliding, www.wwwwwwwweeeeee Beaedd/_aausll_MBe. fweddddwd_rrr. fwd_flops
250
fashhmask
sdpa
200
Speed (terat LOPuN
150
300
26k
24k
32k
40k
48k
Sequence
(g) CAUSAL SLIDING WINDOW MHA Fud TFLOPS NVIDIA A800 GPU sliding
342
flsh,
120
Speed (tera PLOPAN)
302
2dkk
24kk
32k
40k
48c
(h)
Figure 27: TFLOPS FULL/CAUSAL SLIDING WINDOW on NVIDIA A800 GPU
66

---

Published as a conference paper at ICLR 2026
prefix, Im. GQA Ewd flogs
203
A093
fllx
Flash mmask
sdpa
tarch
173
Speed (twaPLoPP))
150
123
1.0
728
4.6k
Sie Sequence Length
PREFIX M GQA Fivd TRLOPS NVIDIA A8D0 GPU
prefix, Jmm bww.dddAwwwwwoooooo flops
A030
flex
Fash mask
sópa
tarch
300
Speed (teraF
325
200
Sequence Length
PREFIX LM GQA Rud TRI OPS NVIDIA A800 A800.IPU GPP
prefix Im. MHA fwd flops
200
A420
flox
Flahh mask
sópa
tarch
175
Speed (braFLOPPS))
153
125
300
Sequence Length
PREFIX LM MHA Fwd TELOPSS on NVIDIA A800 GPU
prefix, Jm, MHA www.dw/ddppp.....o..o flops
(030
Flashumat maakk
sdpa
tarch
Speed (swaFLOPee))
223
300
Sequence Length
(di PREFIX LM MITA Bud TELOPS on NVIDIA A800 GPU
67

---

Published as a conference paper at ICLR 2026
prefix, GQA flops
A030
sdpa
302
Speed LOPAN))))
Sequencc Length
(c) PREFIX LM DOCUMENT GQA Pwd TELOPS NVIDIA A800 GPU
prefix, ent GQA, flops
edpa
122
Speed (trrFFLPPM)))
3k
4.5k
(D PREFIX LM DOCUMENT GQA Bwd TELOPS on NVIDIA A800 GPU
prefix fops
edpa
100
Speed (teratt LOPMN)
4k.
5k
20k
32k
43k
45k
Sequence
(g) PREFIX LM DOCUMENT MHA Fwd TELOPS NVIDIA A800 GPU
flops
320
Speed (teraFLOPP
103
5k
2.kk
43k
45k
Figure 28: TFLOPS of PREFIX LM and PREFIX LM DOCUMENT on NVIDIA A800 GPU
68

---

Published as a conference paper at ICLR 2026
share question, GQA, fwd foops
flex
Flash, maakk
<dpa
tarch
Speed (twraPoPPxs
328
4.6k
See Sequence Length
SHARF QUESTIOO GQA Fwd TRI OPS on NVIDIA A800 GIPI
share, qwwstuuee GQA BaDDD__dd__pppp flops
140
A030
flex
Fash mask
sópa
tarch
Speed (teraf (EPrAf LDPAN)
303
Sequence Length
(b) SHARF QUESTION GQA Bud TELOPS NVIDIA A800 GPU
share question MHA fwd flops
380
flox
flash
mask
sópa
tarch
340
Speed (brrFLOPPS)))
122
300
Sequence Length
(ct SHARE QUESTION MHA FWd TELOPS NVIDIA A8IO GPU
share catarequeettinn.c MHA MHA www.k//dd_fgggo
A030
Flash, mmakk
sdpa
tarch
123
Speed (traaLOPMe))
Caquence Length
(6) SHARE QUESTION MIA Bed TELOPS NVIDIA A800 GPU
69

---

Published as a conference paper at ICLR 2026
GQA, fwd, flops
142
A030
sopa
Speed (OPra))
122
100
Sequence Length
(e) CAUSAL BLOCKWISE GQA Pwd TELOPS NVIDIA A800 GPU
blockwise GQA bwd. flops
160
edpa
142
Speed (traFLoPPu))
3k
4.5k
(E) CAUSAL BLOCKWISE GOA Bwd IPLOPS on NVIDIA A800 GPU
Mwweee_MAAAAW/ flops
sdpa
Speed (teraP) LoPun
5k
20k
32k
43k
45k
Sequence
(g) CAUSAL BLOCKWISE MHA Fwd TELOPS NVIDIA GPOT
blockwise bwd
342
Spand (tera PLOP))N)
322
3.0
5k
2dk:
43k
45k
Figure 29: TELOPS of SHARE QUESTION and CAUSAL BLOCKWISE on NVIDIA A800 GPU
70

---

Published as a conference paper at ICLR 2026
global slidingg GQA, fwvd flops
46:00
flash, mask
sopa
torch
Speed (teraf LOPAN)
20
Sequence Length
(a) GLOBAL SLIDING GQA Fwd THLOPS on NVIDIA A800 GPU
global slidingg GQA bwd floss
flash_mask
Speed (tHaFLOPPN))
00
20
Sequence Length
(b) GLOBAL SLIDING GQA Bwd TELOPS on NVIDIA A800 GPU
global BaiiblssiiigggHHAA WHA_wwd_foe flops
flash_maskk
00
Speed (teraP LOPPN)
20
24
1.5k
24c
40k
48k
Sequence Length
GLOBAL SLIDING MITA TELOPS on NVIDIA A800
global wwwiiiggdniiiie
flash
Speed (teraTLOPaN))
60
2<
1.kk
24c
45k
48k
GLOBAL
Figure 30: TFLOPS of GLOBAL SLIDING on NVIDIA A800 GPU
71

---

Published as a conference paper at ICLR 2026
causal, cutmentt flops
375
A030
sopa
Speed (Owra))
153
1.25
100
Sequence Length
(a) BLOCK CAUSAL DOCUMENT GOA Fed TFLOPS on NVIDIA A800 GPU
Mock causal documntt GQA flops
300
sdpa
140
Speed (traFLoPPe))
322
4.5k
(b) BLOCK CAUSAL DOCUMENT GQA Bwd TELOPS on NVIDIA A800 GPU
MHA MWA_twd Hops
edpa
375
Speed (teraPL LPPad
113
1.25
102
5k
20k
32k
43%
45k
Sequence
BLOCK CAUSAL DOCUMENT MFA TELOPS on NVIDIA A800
block,
bwd
3.0
Speed (tera PLOPaN)
1.2
2.0
5k
2dk
32k
43%
45k
Figure 31: TFLOPS of BLOCK CAUSAL DOCUMENT on NVIDIA A800 GPU
72


Reviews:
[{"review_id": "DJHjOmgWXT", "reviewer": "Reviewer_nrHk", "summary": "This paper proposes LongCA-bench, a unified benchmark for evaluating long-context attention mechanisms in large language model (LLM) training, covering both single-device kernels (dense and sparse) and distributed context parallel strategies. The benchmark addresses critical gaps in existing evaluations—such as incomplete operator comparisons and framework-specific context parallel designs—by integrating 12 dense/sparse kernels, 5 distributed attention mechanisms, and 14 attention mask patterns. It conducts large-scale experiments on up to 96 NVIDIA H100 GPUs, evaluating performance across sequence lengths (up to 512K tokens) and distributed scales. Key findings include hardware-optimized kernels (e.g., FlashAttention-3) outperforming general ones in regular masks, sparse kernels facing limitations in backward computation and flexibility, and hybrid distributed designs (USP, LoongTrain) balancing scalability and efficiency. The work provides actionable guidance for selecting attention mechanisms in ultra-long context training.", "questions": "see Weaknesses", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. LongCA-bench unifies diverse attention kernels (dense/sparse) and distributed mechanisms under a modular interface, enabling fair cross-method comparisons—addressing the fragmentation of existing evaluations.\n2. It systematically explores two understudied but impactful dimensions: 14 attention mask patterns (static/dynamic, regular/heterogeneous) and extreme long sequences (up to 512K) with large-scale distributed training (up to 96 GPUs), filling gaps in prior work.\n3. The benchmark uses real-world datasets (Pile, ProLong64K/512K) and realistic sampling strategies, ensuring results reflect actual LLM training scenarios. It also provides open-source code for reproducibility.\n4. Beyond performance metrics (TFLOPs, peak memory), the paper reveals trade-offs (e.g., hardware optimization vs. mask compatibility, computation-communication overlap in distributed systems) and identifies bottlenecks (e.g., sparse kernel backward pass inefficiency).", "weaknesses": "0. Why no linear attention kernels?\n1. The optimized distributed attention mechanisms only support 4 mask patterns (FULL, CAUSAL, FULL/CAUSAL DOCUMENT), excluding heterogeneous and dynamic masks—restricting its applicability to complex long-context tasks.\n2. The benchmark excludes FlexAttention from full evaluations due to severe out-of-memory issues, and most sparse kernels lack backward computation support or flexibility (e.g., fixed block sizes), limiting insights into trainable sparse attention.\n3. Key findings (e.g., FlashAttention-3’s superiority) are tailored to NVIDIA H100 GPUs, reducing generalizability to other hardware architectures (e.g., AMD GPUs, TPUs).\n4. Only 5 distributed strategies are evaluated, with Ring All-Gather’s results omitted due to resource constraints—missing opportunities to compare with emerging context parallel designs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qSuIuiC9HS", "reviewer": "Reviewer_Ehzx", "summary": "This paper proposes a unified benchmark framework for evaluating various dense/sparse attention kernels and context-parallel mechanisms in long-sequence scenarios (up to 512K tokens). The framework provides standardized assessments of computational efficiency (in terms of throughput and peak memory usage), scalability, and usability. The evaluation focuses on two main dimensions: (1) attention mask patterns, and (2) sequence length and scale of distributed computation. Firstly, the study performs unified data preprocessing for different attention mask patterns to ensure fair comparisons. Next, it integrates over a dozen attention kernels and incorporates three distributed mechanisms: All-to-All, Ring P2P, and Hybrid. Through these evaluations, the paper draws several insightful conclusions about the computational efficiency pros and cons of various dense and sparse attention kernel implementations. This lays a solid foundation for future research to weigh different backend implementations, explore directions for kernel optimization, further improve kernel implementations, and perform fair comparisons with existing methods.", "questions": "Please provide at lease some convincing explanations for the two points mentioned in the Weaknesses section. I am willing to raise my score if my concerns are addressed.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "(1) Extensive Method Integration: This work uses a unified interface to integrate 12 representative attention kernels and 5 distributed mechanisms.\n(2) Good Scalability: The evaluation is conducted on scenarios with sequence lengths up to 512K and across 96 GPUs.\n(3) Practical Insights: Through experimental evaluation, the authors obtain insightful conclusions regarding the impact of mask patterns, the trade-offs between kernel efficiency and usability, and the scalability characteristics of different distributed mechanisms.", "weaknesses": "(1) Architectural Limitation: The study is limited to the Hopper architecture and does not discuss the generalization of experimental conclusions to other architectures.  \n\n(2) Performance Metric Limitation: The research only focuses on throughput and memory usage as performance metrics. It does not analyze how metrics such as memory bandwidth utilization and inter-node communication load vary over time across different kernels and distributed mechanisms.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qho1PNbvpW", "reviewer": "Reviewer_1Et8", "summary": "The paper presents LongCA-Bench, a unified benchmark framework designed to evaluate long-context attention mechanisms across both single-device and distributed settings. The author integrates existing attention kernels (dense and sparse) and distributed context-parallel strategies into a modular, extensible evaluation platform. Through extensive experiments on up to 96 GPUs, the benchmark compares efficiency, scalability, and memory usage across diverse mask patterns, sequence lengths, and distributed configurations of existing methods. The work provides a valuable empirical reference for practitioners and researchers working on long-context training but does not seem to introduce novel attention algorithms or distributed mechanisms beyond benchmarking existing approaches.", "questions": "1. How can the proposed framework offer insight on sparse attention in distributed settings? Typically, how should we overlap communication with computation, and address the load balancing problem of arbitrary block sparse pattern? The author does not need to provide a complete solution to this problem, and what I would like to see is how this work could contribute to future research on these challenges.\n2. Could the authors provide more analysis to understand the trends observed in kernel or distributed scaling performance?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. **Comprehensive benchmarking.** The paper systematically benchmarks a wide range of attention implementations, including dense, sparse, and distributed mechanisms, under a unified framework. The experimental coverage (up to 96 GPUs and multiple mask types) is extensive and provides a clear view of current attention efficiency trends.\n2. **Sound experimental methodology.** The experiments are well-organized, use realistic settings (e.g., long context lengths, different mask patterns), and report consistent metrics (throughput, memory). The methodology is reproducible and technically competent.\n3. **Useful empirical reference.** The benchmark results could serve as a useful reference for practitioners and researchers seeking guidance on the performance trade-offs of existing attention kernels and distributed strategies.", "weaknesses": "1. **Limited Analysis.** The paper reports extensive throughput and memory results, but offers limited discussion on the underlying causes of observed performance trends of the benchmarked methods.\n2. **Incomplete coverage of the most critical setting — distributed sparse attention.** The integration of sparse attention (particularly dynamic block-sparse attention with TopK/TopP selection criterion) into distributed contexts remains an unexplored and practically important challenge. The paper does not explore this aspect (does not touch block sparse distributed attention, only discusses full/causal/document), which somewhat limits its impact given the stated motivation of benchmarking “long-context attention.”", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)