FlexLinearAttention: Compiling a Unified Abstraction into Scalable Kernels for Linear Attention
Authors: Haojie Duanmu
Size Zheng
Ningxin Zheng
Jianqiao Lu
Xuegui Zheng
Xingcheng Zhang
Li-Wen Chang
Xin Liu
Dahua Lin
Venue: ICLR 2026 Poster (2026)
Keywords: linear attention
efficiency
compiler
kernels
Primary area: infrastructure, software libraries, hardware, systems, etc.
TL;DR: FlexLA is a domain-specific compiler that generates scalable kernels for diverse linear attention variants, from a simple unified abstraction, providing comparable or even better performance than expert-tuned libraries.

Abstract: The quadratic complexity of softmax attention poses a major bottleneck for long-context modeling, motivating a surge of linear attention variants with linear complexity. Unlike softmax attention, which benefits from optimized kernels, linear attention lacks general-purpose, hardware-efficient support and scalable distributed implementations. We introduce **Flex**ible **L**inear **A**ttention (FlexLA), a domain-specific compiler that automates the generation of high-performance, scalable kernels for a wide range of linear attention models directly from high-level PyTorch code. At its core, FlexLA employs an intuitive programming abstraction that decomposes any linear attention algorithm into three canonical phases: intra-chunk computation, inter-chunk state propagation, and output merging. This unified abstraction enables FlexLA to perform domain-specific optimizations, automatically generating kernels that fuse computation and communication at a fine-grained tile level and eliminating host synchronization. Our evaluation demonstrates that FlexLA combines programmability with performance: a wide range of linear attention variants can be implemented in just a few dozen lines of code, while the generated kernels deliver 1.01x-4.9x the performance of sate-of-the-art expert-optimized library and scale with near-linear efficiency on scalar gated linear attention to 16 million tokens on 128 GPUs, surpassing the state-of-the-art distributed baseline by up to 7.2x.

Published as a conference paper at ICLR 2026
FLEX LINEAR ATTENTION: COMPILING A UNIFIED ABSTRACTION INTO SCALABLE KERNELS FOR LIN-
EAR ATTENTION
Haojie Duanmu" * 123 Size Zheng! 1, Ningxin Zheng', Jianqiao Lu', Xuegui Zheng* Xingcheng Zhang", Li-Wen Chang', Xin Liu', Dahua Lin
'ByteDance Seed 2Shanghai Jiao Tong University SShanghai Al Laboratory "The Chinese University of Hong Kong
ABSTRACT
The quadratic complexity of softmax attention poses major bottleneck for long- context modeling, motivating surge of linear attention variants with linear com- plexity. Unlike softmax attention, which benefits from optimized kemnlls, lin- ear attention lacks general-purpose, hardware-efficient support and scalable dis- tributed implementations. We introduce Flexible Linear Attention (FlexLA), a domain--peccffic compiler that automates the generation of higg-performance. scalable kernels for wide range of linear attention models directly from high- level PyTorch code. At its core, FlexLA employs an intuitive programming abstraction that decomposes any linear attention algorithm into three canonical phases: intra-chunk computation, inter-chunk state propagation, and output merg- ing. This unified abstraction enables FlexLA to perform domain-specific opti- mizations, automatically generating kernels that fuse computation and communi- cation at umee-gainedd tile level and climinating host synchronization. Our eval- uation demonstrates that FlexLA combines programmability with performance: a wide range of linear attention variants can be implemented just few dozen lines of code, while the generated kernels deliver 1.01x-4.9x the performance of sate-of-the-art expert-optimized library and scale with near-linear efficiency on scalar gated linear attention to 16 million tokens 128 GPUs, surpassing the state-of- saate-of-the-at ff-the-art art distributed baseline by up to 7.2x.
INTRODUCTION
Transformer models rely on self-attention, which has quadratic time and memory complexity with respect to sequence length. As models handle increasingly long contexts, this quadratic bottleneck severely limits scalability. In response, many efficient-attention mechanisms have been proposed. In particular, linear attention methods remove the softmax nonlinearity and reorder computations to achieve linear computation complexity and constant-memory inference. This has led to prolifera- tion of innovative architectures, such as Mamba (Gu & Dao 2023 Dao & Gu 2024). RetNet (Sun)
ctt al., 2023), RWKV (Peng et al. 2002). GLA Yang et al. 2023). ). HGRN (Qin et al. b) and Gated DeltaNet Yang et al. 2024a). These architectures demonstrate capabilities competitive with, or even superior to, standard transformers.
Unlike softmax attention, which has benefited from highly optimized and staddaanddd kemnlss like Flash-Attenn Attention (Dao et al., 2022) and Ring- Ging-Attention (Liu et al. (203) that efficiently map its computation and communication modern Al infrastructure, the landscape for linear attention is far more fragmented. Flash-Linear- Linear-Attention (FLA) (Yang & Zhang 2024) provide a valuable collection of triton kernels for a series of linear attention variants. But essentially relies on ex- pert developers to provide manual implementation for each variants. The rapid evolution of linear attention variants means that a one-size-fits-ll solution does not exist. This forces researchers into costly and inefficient cycle of manual kernel development for each new variant, a process fraught with two major challenges.
Work done ByteDance Seed. Corresponding Author.

---

Published as a conference paper at ICLR 2026
PyTorch Implementation
Our DSL
Performance
scaiar_pinco.
4//ko: [a-T0]- [SIT]- [800]
def durrkesdeek,,
Lomsu, 6 g-corsunt)) 2.5w(())
Scalar GLA: 256K Sequence Length
for in in reneece x -0:
K.deryy (E_5u E.cresunthiee, (10espp))
b, h // h.
for chure,de
slb,
:1-s2oneO
reture K_decay
Cal c_l, c.v, C.6. C_ chon()
def desay eddell 5, 0):
a compute intre-chink decay
e.cumman, 0,s c_g (oss C_E-SUNc))
k.x.m)0)..p0))
Gaa-ggima 10t
der norge ecdela, H, v, 5, realej:
decay local comsan).enst (None,
decay, vec 0.rumtum00) (0)
a 2. corpute intre-curk state
deladesyy (c.k decay_k)
decay_P qt.
decay vedl Note] decay Me:
tr1111)))
deray seely Mote] -trey, veeisons, n (tetunk 8 (.kT)))r)))0)) decay P) curr_state,
currsstate
1. corpute
c.))
chink
returr (o 9
[ok decay_p)
vocay mone) scale
Lleeoesttontionn SP SP,GROM, chunc, mode, decay, sode,
decay lenn]
curr_state_) scale
a'4, update state state state Suc 1 cure_stace
toath Eager
Figure 1: Comparison between PyTorch and our proposed DSL in writing scalar gated linear atten- tion (data type conversion is omitted for simplicity). And the performance comparison of different approaches. Test shape: 32 heads with 128 head dimension. Torch-Eager fails to parallelize the computation, while Torch-Compile also performs poorly. FLA provides expert optimized triton ker-
nels, while program generated by FlexLA offers comparable performance to handwritten kernels with the scalability to distributed environments.
First, the implementation of expertise. While many linear attention models share a conceptual similarity, their specific state update rules and memory access patterns can differ substantially. Achieving hardware efficiency necessitates not only fusing the state update rule into single kernel but also manually tuning hardware-specific parameters like pipeline schedules and tile sizes. Even with high-level DSLs like Triton (Tillet & Cox. 2019), developers must often delve into low-level hardware details, such as managing barriers or dealing with shared memory capacity limitation, to extract maximum perfor- mance for each variant. This creates high barrier to entry and slows down the pace of innovation.
kernels is an arduous task requiring deep hardware
Second, existing solutions lack robust support for distributed execution, which is non-negotiable for scaling to contexts of hundreds of thousands or millions of tokens. FLA offers a bunch of single-device kernels but do not address the distributed scaling problem. When sequence lengths exceed the memory capacity of single accelerator, distributed sequence parallelism becomes es- sential. However, enabling sequence-parale custom communication schedules tailored to each variant's state update rule. Existing sequence parallel schemes, such LASP and LASP-2 (Sun et al. 2024a] 2025) are designed for specific architectures and employ generic communication primitives (c.g., All-Gather from NCCL). The mismatch of existing communication primitive and dataflow of distributed linear attention leading to significant network bandwidth underutilization(Chou et al. 2023).
linear attention is non-trivial: it typically requires
Can we provide solution to bridge the gap between the rapid evolution of linear attention algorithms and the difficulty of developing scalable kernels? We observe that many of these difficulties stem from not exploiting the common structure underlying linear-attention variants. Our central insight is that most linear-attention variants share a small set of canonical operations and data exchanges. Based on this, we introduce Flexible Linear Attention (FlexLA), a compiler-driven framework that allows implementing the majority of linear attention variants in few lines of id- iomatic PyTorch code and scaling them to distributed system. As shown in Figure 1 our DSL expresses linear attention in three modular functions. The compiler translates the DSL. into high performance kernels: reducing latency from 34.6 seconds (PyTorch eager) to 9.2 ms, even better than SOTA written kernel from FLA. More importantly, the latency was further reduced to 2.7 ms when scaled to 4 GPUs distributed system.
The frontend of FlexLA ingests user-defined computation logic for linear attention variant. Its backend then intelligently maps this logic, along with potential communication operations, onto hardware accelerators and network interfaces, applying domain-specific optimizations to generate high-performance, distributed-aware kernels. FlexLA is built around three key principles: A Linear-A A Ateetion-Speccffic Programming Abstraction. We formulate our programming abstrac- tion based on the canonical chunk- parallll representation of linear attention. This model decom- poses the computation into three intuitive phases: a compute phase that processes local chunks of the sequence in parallel, an update phase that communicates and updates the inter-chunk states, and
2

---

Published as a conference paper at ICLR 2026
merge phase that combines the global state with local chunk results. This abstraction aligns di- rectly with the mathematical structure of chunk-wise parallel form, allowing researchers to translate their algorithm's formulation into our framework with minimal effort by providing simple Py PyTorch Fusion. At compilation time, FlexLA lowers the three user-provided callables into Triton code. We leverage Triton-Distributed (Zheng et al. 2025a) compiler backend, which extends Triton with native communication primitives. This allows our compiler to generate fine-grained, tile-level communication instructions that are fused directly with computation. By creating custom communication patterns tailored the algorithm, we by- pass the overhead and limitations of standard libraries like NCCL, enabling a more efficient use of the underlying network fabric and dramatically improving hardware utilization. Optimization of System Bottlenecks. Beyond kernel fusion, we identify and optimize other crit- ical system-level bottlenecks that affect real-world system performance. FlexLA employs suite of techniques, including Ahead-of-Time (AOT) compilation with static kernel dispatcher built on top of Triton to reduce runtime overhead and an adaptive parallelism scheduler that dynamically explores the optimal configuration of compute resources, further boosting the end-to-end efficiency of linear attention execution.
callable. Native Computte-Commnnieee
Targeted
To demonstrate flexibility, we implement broad range of linear attention variants using FlexLA, each requiring only dozens of lines of code. This programmability does not come at the cost of performance. On a single GPU, our generated kernels achieve 1.01x to 4.9x the performance of the state-of-the-art FLA library of expert-tuned kernels. Furthermore, in distributed settings, FlexLA demonstrates near- linerr scalability on up to 128 GPUs, outperforming the leading open-source baseline by up to 7.2x.
2 PRELIMINARY
2.1 LINEAR ATTENTION ARCHITECTURE
Given a sequence X2, Xt. E RLxd the input of attention block: qú 1i,k, V; Wax, WKX, Wvx, where x,, qí, k, V., y; e Rd and the weights W Wk. W. E pd Trans- formers employ softmax attention as a token mixer (Vaswani et al. 2017
O Σ exp(a kj)
V
(1)
O softmax @ @M)V
p=1 exp(q] kp)
Equation 2 is the matrix form of Equation where Q [q qr. K [k....., kt] VI. E Lxd and M [ 1 is a causal mask.
Such matrix form is well suited to modern accelerators, which excel at large matrix multiplications, but incurs (lL2d) complexity. If we remove the softmax operation, the computation becomes associative: o = q: (k,v,) which reduces the complexity to O(Ld)). The recurrence form is
expressed as:
S. = S.-
S..
Here S E Rdxd is the state (or memory) updated in each time step. Equation 3 highlights the key idea of linear attention: replacing the exponential kernel in softmax attention with a linear recurrence. Although this formulation resembles RNNs (Hochreiter & Schmidhuber, [997). The critical difference is that dependencies across time steps remain linear, which makes parallel training possible. Indeed, linear attention can be written in fully parallel form:
O (QK O M)V
Modern linear attention often augments the recurrence with a decay or gating mechanism, e.g., S, G @ S k (Gu & Dao 2023). or with more sophisticated update rules such as the delta rule (Yang et al. 2024b): S, S B,kkk,) + B.v,k,.. which enhance memory utilization.
3

---

Published as a conference paper at ICLR 2026
2.2 CHUNK- WISE PARALLEL FORM OF LINEAR ATTENTIONS
The fully parallel form in [qqaaiin 4 achieves maximum hardware utilization but retains quadratic complexity. Conversely, the recurrent form in Equation 3 has linear complexity but is inherently sequential and cerdwarr--ieeffcieee In practice, linear attention strikes a balance by adopting chunk-
wise parallelization (Hua et al., 2022, Sun al.
2023 Yang et al., 2023).
Specifically, the sequence of length L is partitioned into & chunks of size C. Let Q K [i) Vii
RCxd denote the query, key, and value matrices of the i-th chunk, and let S Rd be the state after processing chunk i. The chunk-wise formulation separates computation into intra-chunk and inter-chunk parts, and then merge both together to get the final output shown in Figure 2
iC
Intra Chunk
VI
Inter Chunk
Sy Sk-1)
Σ
KJVv
(5)
V2
V3
Matrix: K(j T V14
K1
K2
K3
K4
V4
VI
V2
( =
Sk-1 +
Qua o M)
V
Q1
inter
intra
Q3
P3
03
K3
(6)
Q4
P4 (14
K4
Figure Chunk-wise parallel form demonstration for linear attention.
The inter chunk item be viewed readout memory from start of the sequence to the start of current chunk while the intra chunk item is processing the information in current chunk. When chunk size C is sequence length L,, it becomes the fully parallel form in Equation The chunk size is a tradeoff between parallelism and FLOOSs
2.3 DSLS AND DOMAIN SPECIFIC COMPILER
A domain-specific language (DSL) trades generality for performance by restricting program ex- pressivity particular domain, creating opportunities for targeted optimization. Deep learning compilers like TVM (Chen et al.] 2018). ThunderKitten (Spector et al.) 2024). TileLang (Wang al. 2025) and torch. comp le (Ansel al.) 2024) exemplify this approach: by operating constrained set of primitives (i.e. tir in TVM and aten operators in Porohh,,, they can systematically explore focused design space to apply transformations like operator fusion and loop tiling, thereby automatically generating high-performanee code from high-level descriptions.
This paradigm is particularly effective for operations that possess rich computational structure but high implementation complexity. By exposing domain-relevant abstractions, a DSL. allows develop- ers to specify what to compute, while delegating the complex details of how to execute efficiently to the compiler. This separation of concerns is key to enabling automated, domain-specifi opti- mizations without burdening the user with low-level hardware details.
Applied to the domain of linear attention, effective DSL must therefore provide abstractions that are expressive enough to capture the diverse patterns found in various state update rules and parallel scan formulations. Simultaneously. its compiler must be able to recognize these common patterns and systematically generate optimized kernels for them, bridging the gap between high- level algorithmic design and performance, hardware-aware code. Because chunked parallel forms are complex to implement and offer opportunities to fully exploit hardware, this work focuses on the prefill phase of linear attention (used in inference or the forward pass of training). The backward pass can be implemented in similar manner (Qin et al.
3 FLEXLA
In this section, we propose unified abstraction of diverse linear attention variants. This abstraction enables programmers to easily express linear attention semantics without worrying about implemen- tation details and kernel performance.
4

---

Published as a conference paper at ICLR 2026
Table 1: comparison of representative linear attention variants that can be easily mapped to our
three-phase abstraction, including HGRN (Qin et al. 2023), RetNet (Sun et al. 2023). Mamba2 (Dao
Gu (024), GLA (Yang et al., 2023). and GDN (Yang al. (024a). Despite diverse state
types (vector vs. matrix) and decay mechanisms (element-wise product vs. matrix multiplication).
Ve. kr; qr. are value, key and query projections; ax,, Br, Pt, d are gates; @ the Hadamard product.
Model
Update rule
Read-out
State + Decay type
HGRN
h, a4 14-1 (1 at.) U.
O h4 q:
vector data-dependent vector
RetNet
S: = 7 S v,k,
O, S.q.
matrix + data-independent scalar
Mamba2
S: ?/s S viki
O: S.q.
matrix + data-dependent scalar
GLA
S: = -S-1 (10,
/ vzke
0D = 5:q:
matrix + data--dpenddnn vector
GDN
S. ar Sc Bekckh + B.nkk
O2 S.q.
matrix data-depenndnnt matrix
3.1
PROGRAMMING ABSTRACTION
Despite the numerous linear attention variants designed by researchers, we unify these variants into three commonly shared phases based on chunk-wise parallel form introduced in ssussection 2.2 Intra-Chunk Computation. The first phase computes local state within each chunk of the input sequence. In this stage, computation across different chunks is embarrassingl parallel, as there are no data dependencies between them. Each chunk is processed independently, transforming its se- quence of inputs into relative state summary. Inter-Chunk State Propagation. The second phase addresses the dependencies between chunks. To compute the correct global state at the beginning of each chunk, the state summaries from all preceding chunks must be accumulated. For instance, in the case of vanilla linear attention, this propagation corresponds to a prefix sum (scan) operation, as shown in Equation 5| This phase is inherently sequential due the temporal dependencies between chunk states. And cross-device communication happen in this phase. Merging and Output Gen- eration. The final phase merges the results of the intra-chunk and inter-chunk computations. Here, operations are once again parallel across chunks. Each chunk utilizes the global state propagated from Phase and its local inputs to compute its final output sequence.
Based these insights, designed our programming abstraction around three corresponding callable: chunk_mode, decay_ __ode and merge_mode correspond to three phases. This ab- straction empowers users to implement new linear attention variant by simply defining its chunk- wise parallel logic in idiomatic PyTorch code, decoupling algorithmic expression from system op- timization. Furthermore, this three-phase decomposition also helps us optimize the program: we separate the parts of the entire program that can be executed in parallel and the parts that must be executed serially and may involve cross-device communication. With this information, FlexLA can perform more aggressive and accurate optimizations.
As shown in Table I prominent linear attention variants employ vastly different state representa- tions and decay mechanisms at the token level. To map them to our abstraction, these token-level updates are reformulated into chunk-level matrix operations. For example, denoting the intra-chunk queries, keys, and values as matrices Q, K, and V. the chunk-wise parallel form for Mamba2 can be specified as:
S = II a S V
0 -=
+ +SQ
(7)
where Stg is the state summary for the 1-th chunk, a is the chunk decay, and M and G represent specific masking and gating matrices.
Under framework, a user only needs to implement these equations within our three-phase pro- gramming abstraction in native Py Torch code. FlexLA then automatically handles all subsequent code generation, performance tuning, and scaling to distributed systems.
3.2 COMPILATION AND CODE GENERATION
Given the chunk-wise parallel form description in our DSL, FlexLA performs a series of graph- level and system-level transformations to produce optimized program as illustrated in Figure 3 The user-defined function is first captured Torch.fx graph (Reed et al. 2022) via tracing. We choose fx graph as our intermediate representation (IR) because is the most powerful tool in the PyTorch ecosystem, most torch operators can be directly captured, which makes our DSL
5

---

Published as a conference paper at ICLR 2026
Chunk-wise Parallel Form Description
Torch.fx Graph
Compute-comee compute-communication Fused Program
Graph Rewrite with Domain Knowledge
Distributed
1. Try to Fuse Subgraph
ENV
def chunk mode(...):
2. Hand-Written OP Substitution
def decay mode def merge, mode
Common Subexpression Elimination Transpose Elimination
GPU:
NIC:
Without overiap
Tilee level overlap (ours)
Transformed Graph
Code Generation
def chunk chunk_decay :.
AOT Compilation
Compute Capability
lower_ triangular Satpguulaaannnerseeeeee inverse() ()
Runtime Overhead Elimination
Network Topology
Tillng Plan, Profile--asee Auto-Ture Tune
copy to comn buffer()
Hand-written operators
def fused_ communication(....)
Compute Triton
def merge_node((...
Communication = OpenSHMEM Primitive
Figure 3: An overview of the FlexLA compitation pipeline. Our compiler ingests a high-level de- scription of linear attention algorithm, specified using our DSL (chunk_mode, decay_mode, merge merge_mode). This description is traced into a graph structure, which then undergoes a series of domain-specific optimization passes. The optimized graph is then compiled Triton with native communication support, enabling fine-grained tile-level overlap between computation and commu- nication for better utilization of GPU and Network Interface Card (NIC).
expressive enough implement most of linear attention variants. We first replace special op code in fx graph as custom instructions (e.g. placeholder is replaced with load instruction). Domain-specifie optimization passes then applied. To enable fusion of non-trivial operators (e.g. lower_triange ar_vrrerr__eeeeeeeeee in GDN), we provide a bunch of custom triton kernels com- monly used in linear attention domain and mapping them to corresponding torch operators. Once the compiler see these operators, for instance, torch. inverse, they will be marked and subse- quently substituted with custom triton source code in code generation phase. Common optimization passes like transpose elimination are also applied in this phase. The IR is further rewritten with system-resource awareness. FlexLA take the hardware information to generate hardware-soe speciffc instructions. For instance, Tensor Memory Accelerator (TMA) availability is marked as attribute of load instruction.
Recent studies have demonstrated that fine-grained compute-communtin fectively hide latency in distributed settings (Chang et al., 2024 Zheng et al. 2025b). We adopt this technique in FlexLA by automatically fusing computation and communication at the tile level, thereby reducing data dependency scope and eliminating the frequent GPU-host synchronizations inherent to traditional overlap strategies (Jangda al. 2022). Within our programming abstraction, all cross-device communication is confined to the second phase, i.e., inter-chunk state propagation. Consequently, FlexLA first analyzes the data dependencies in this phase, then determines the corre- sponding computational tiling and communication tiling strategies based on the network topology, selects the appropriate communication mode, and generates computation and on-device communi- cation instructions.
fusion can more ef-
Finally, the IR is lowered into triton source code. Different with torch inductor (Ansel et al. 2024) targeting on official triton, FlexLA targets on Triton-Distibuued (Zheng et al., 2025a). who addtion- aly provide Ene-grained communication control. On device computation is translated into compu- tation primitives provided by Triton while the communication logic is mapped to the OpenShmem- style communication primitives provided exclusively by Triton-distributed, which are ultimately translated into GPU-initiated communication operations.
3.3 PERFORMANCE OPTIMIZATIONS
With domain-specific knowledge of linear attention, FlexLA can explore compact yet effective op- timization space. In particular, the choice of whether to fuse different phases introduces an important trade-off. For example, fusing chunk mmde with decay_meee _moee avoids materializing intermediate states in global memory, thereby reducing memory traffic. However, such fusion also limits available parallelism, since computations can no longer be scheduled independently at the chunk level. There are many such trade-offs and FlexLA will handle all of these to get better performance. FlexLA employs parallelism scheduling algorithm that dynamically chooses optimal parallelization

---

Published as a conference paper at ICLR 2026
Torcs-Cample pile
FLA
Our
MORN
Vanilla- LA
Scaler GLA
Voctem
CON
80
Figure 4: Latency comparison of different linear attention variants under varying sequence lengths on a single H100 GPU. Each subplot corrssonnds to one model, with the top row showing short to medium sequences (1K-16K) and the bottom row showing long sequences (32K-256K5
strategy based on input shapes and hardware information internally. The specifics of this parallelism scheduler are detailed in Appendix B
Beyond the optimization for target program. we further incorporate set of optimizations tailored to the compilation system itself. In practice, the Triton runtime introduces overheads on the order of hundreds of microseconds, often exceeding the actual most of linear attention kernel execution time at short to medium sequence length (e.g. 2K or even 4K). While approaches like CUDA- mitigate launch overhead for workloads with static input shapes, they often incur significant memory costs and are unsuitable for the dynamic workloads in inference scenarios. To address these limitations natively, we extend Triton compiler with custom Ahead-of- (AOT) compilation module. Specifically, our module compiles Triton source code into pre-linked dynamic libraries ahead of execution. At runtime, FlexLA employs a profile-guided static dispatcher that bypasses the Triton runtime entirely, invoking the optimal pre-compiled binary directly through the CUDA Driver API. The static dispatcher automatically generated by FlexLA from an offline performance database, ensuring that the empirically best- best-performing kernel is selected for any given workload without incurring runtime overhead from hash lookups or dynamic compilation logic.
Another problem is redundant compilation. Since our system is specialized for linear attention, we observe that certain input tensor dimensions (e.g., head dimension and number of heads) remain relatively static across runs, while sequence length is typically dynamic. This property facilitates efficient AOT compilation: FlexLA allows users to specify constant dimensions and their admissible ranges via input metadata. FlexLA enumerates the Cartesian product of these ranges and generates all potentially required kernels in advance. Furthermore, by leveraging Py Torch's symbolic tracing, our system supports tensors with symbolic shapes, ensuring that recompilation is unnecessary unless static dim change.
4 EXPERIMENTS
We implement several high-performance linear attention kernels using FlexLA, including HGRN, vanilla linear attention, scalar GLA, vector GLA, and Gated DeltaNet (Qin et al.) 2023) Dao Gul 2024 Yang et al. 2023 2024a). While many linear attention models differ in their parameterization, the computational patterns we implement cover over ten existing model designs (Appendix C).
4.1 SINGLE-DEVICE EVALUATION
Figure 4rreoorts the latency of kernels generated by FlexLA across sequence lengths ranging from IK to 256K on single H100 GPU. We compare with two baselines: Torch-Compile, repre- senting a general-purpose compiler without domain-speciff knowledge, and FLA (commit hash: 027 211.. the sfaee-of-hh-as library of providing expert-tunee tundd Triton kernels for linear attention. For all variants except HGRN, we fix BatchSize = 1, NumHeeds = 32, and HeadDim = 128, and vary sequence length. For HGRN, follow its original single-head configuration. A batch size of
7

---

Published as a conference paper at ICLR 2026
Sequence Length per GPA2.
Sequence Langth per GPU: 1.k
Sequence Langth per GPU: 324
Orripsstioo on
Ovehhed (Dun)
Ovelhad
Latunny (mne)
32
64
128
16
64
Number of GRUs
Number of GPUs
Number of
(a) Gated Delta Network
Sequence Length per
Sequence Longth per
Sequenee Length per
12
Latono Uel
Compulation (LASP2)
Mimber of GR
Bumber of GPC
Mimeer of GPIls
(b) Gated (scalar) Linear Attention
Figure 5: Weak scaling performance comparison for GDN and GLA models. Both figures show latency breakdown across varying number of GPUs with a fixed sequence length per GPU.
one is standard and reasonable choice, as a single long sequence in linear attention is computation- ally equivalent to a batch of packed shorter sequences.
Corph-Compile consistently exhibits poor performance, as it fails to apply the domain-specific fusion strategies required for linear attention. FLA achieves strong performance by carefully hand-tuning 10-aware tiling. Our work, FlexLA, automating this optimization process, consistently matches outperforms FLA. On HGRN, FlexLA achieves 1.64-2.02x speedup, highlighting its ability to discover optimization opportunities beyond expert tuning (We check the generated code and find FlexLA use a more efficient threads allocation). On scalar and vector GLA as well vanilla linear attention, FlexLA provides stable improvements (1.1-2.0x) while for GDN, performance converges FLA at longer sequences. The speedup is most pronounced on short medium sequence length. In this regime. system-level overheads and the choice of parallelism strategy are the dominant fac- tors, and FlexLA effectively eliminate these bottlenecks, as discussed in subsection 3.3
4.2 SEQUENCE PARALLEL EVALUATION
We further evaluate weak scaling behavior under distributed training, comparing against LASP2 (Sun et al.] 2025). the strongest open-source baseline at the time of writing. ZeCO (Chou al. 2025) reports improved scaling via pipelined communication, but its implementation is not public. We benchmark two representative workloads: GDN and scalar GLA, which feature matrix- multiplication and element-wise decay mechanisms, respectively. Experiments were conducted on a cluster of up to 128 NVIDIA H20 GPUs, interconnected with NVSwitch within nodes and Infini- Band between nodes. We fix BatchSize 4, Numleeads = 32, HeadDim 128, and maintain constant workload per GPU while scaling the total sequence length from 128K (on 4 GPUs) to 4 million tokens (on 128 GPUs). The ideal outcome for weak scaling is constant execution time.
Figure 5 show that FlexLA exhibits near-ideal weak scaling for both workloads: latency remains flat as GPU count increases. This attributable to two core features of our compiler. First, generates communication patterns that avoid the data redundancy incurred by the All-Gather primitive used in LASP2. Second, its ability to fuse computation and communication effectively hides the latency of the local state update. This is particularly impactful for GDN, whose matrix-based update is time consuming. In contrast, the communication and computation redundancy of LASP2 is amplified as the number of nodes increases, causing its performance degrade significantly (e.g., from 49 2.2ms on 4 GPUs to 345ms on 128 GPUs for GDN). This confirms that FlexLA eliminates redundant communication and achieves scalable performance large GPU clusters.
8

---

Published as a conference paper at ICLR 2026
4.3 ABLATION STUDY
AOT compilation with static dispatcher. We measure the end-to-end latency of the execution of Scalar GLA kernel including both kernel ex- ecution and to demonstrate Triton runtime overhead and the efficiency of our solution as shown in Figure 6 At a sequence length of 1024, this overhead (207μs) is over 4.4 times the actual kernel execution time (47us). Our static dispatcher mitigates this issue by reducing the over- these latency-
Triee Oigattc
Latency (us
200
head by 46%. yielding 1.6x end-to-end speedup sensitive inputs. As the sequence length grows and execution time be- comes the dominant factor, our dispatcher consistently maintains a neg- ligible overhead that is effectively eliminated at 8192 tokens (reducing from 101us to just I
1024
2048
4096
8152
Sequence Length
Figure 6: Execution time decomposition.
Tile Level Compute- SoapuuttCComunicationn Overlapping. We con- duct a targeted study to isolate and quantify the benefits of our tile-level compute-communication fusion. We set two baselines for the inter-rank state propagation: where each rank i waits to receive all state data from rank before performing its local update and sending to rank + 1.
Table 2: State communication time on 8xH800 GPUs.
State
Serial communication,
size is 67MB.
Method
Time (us)
2 Pipelined baseline, which chunks the state and uses standard NCCL send/recv operations to overlap the computation of one chunk with the communication of the next. Our method, in con- trast, generates a single kernel that fuses the state update com- putation with communication primitives at the tile level. The re-
Serial Torch-Pipeline
873
902
Ours
560
sults, presented in Table 2 show the standard Pipelined (PyTorch) baseline is slightly slower than the naive Serial implementation, demonstrating that host-managed pipelining can be counterproductive due to the significant overhead of launching numerous small operations and the required host-device synchronization. Conversely, our fused kernel substantially reduces the total time, achieving 1.56x speedup over the serial baseline.
5 RELATED WORK
There are a wide range of approaches to address the quadratic complexity of softmax attention. Sparse attention mechanisms (Zaheer et al., 2020 Xiao al. 2023 Yuan al. 2025) leverage structured or un- un-structured sparsity in attention to skip computation. Quantized attention (Shah al.)
Zhang et al. 2024) use exploit low-precision arithmetic unit in modern hardware to get higher throughput. There are also various techniques to reduce key-value (KV) cache overhead. KIVI (Liu ct al. grouped query attention (GQA) (Ainslie et al. 2023). and multi-heat latent attention (MLA) (Liu 2024a) alter the attention architecture to reduce memory overhead.
2024
2024b) and SKVQ (Duanmu et al. 2024) directly compress KV cache using quantization while
ct al.
Another line of work proposes architectural alternatives with lower complexity, including linear tention variants (Katharopoulos et al., 2020 Dao & Gu 2024 Peng et af. 2023 Yang et al. (004a)
as well test-time-training approaches (Sun al. 2024b Behrouz et al. 2024). For linear
tention sequence parallelism, LASP (Sun et al. 20224) first extend linear attention to distributed
environments with serial send-receive primitive. LASP2 (Sun al. 2025) improves on this by
leveraging collective communication primitives, yet both approaches still incur significant band- width under-utilization. ZeCO (Chou et al. (025) introduces a pipelined send-receive scheme to hide send/receive latency but relies on manual chunk-size tuning and does not detail its implemen- tation. Therefore, it was not included in our comparison.
AI compilers have been developed to optimize torch. compile (Ansel et al. 2024), TVM (Chen et al. 2018), and TASO [Jia et al.) 2019) are effective for common operators but their optimization spaces do not cover linear attention. Operator-leved compilers such as Triton (Tillet & Cox 2019), ThunderKitter Spector et al. (2024), (Zheng al., 2025a) provide expressive
broad range of workloads.
TileLang (Wang et al. 2025), and abstractions for modern accelerators, but supporting the large and growing family of linear attention variants still requires substantial manual development effort. The most closely related

---

Published as a conference paper at ICLR 2026
work, FlexAttention (Dong al. 2024). targets block-sparse softmax attention and is limited to single-device settings.
CONCLUSION
We presented FlexLA, Nomain--specion compiler for linear attention that unifies diverse algorith- mic variants under a common three-phase abstraction. By generating hardware-efficient kernels with native distributed execution support, FlexI.A bridges the gap between rapidly evolving linear attention research and the complexity of hand-tuned implementations. Our evaluation demonstrates that FlexLA achieves both high performance and broad applicability across modern linear attention models. We hope this work will accelerate the development of new architectures and inspire further research at the intersection of deep learning algorithms and domain-specific compilation.
ACKNOWLEDGEMENTS
The authors would like thank the diligent anonymous reviewers for their constructive feedback. We also extend our gratitude to the fla-org open-source community for their invaluable contributions the linear attention infrastructure. This work supported by the Shanghai Municipal Science and Technology Major Project
REPRODUCIBILITY STATEMENT
We are committed to the reproducibility of our work. Most of algorithms, implementation details, and experimental setups descrihed in the paper. Due to organizational policies requiring internal review prior to relese,, release, the source code is not included with the submission. However, we are committed to open-sourcing the code and will provide it reviewers upon request for the purpose of evaluation.
REFERENCES
Joshua Ainslie, James Lee-Thorp, Michiel De Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head check- points. arXiv preprint arXiv: 13245, 2023.
Jason Ansel, Edward Yang, Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky, Bin Bao, Peter Bell, David Berard, Evgeni Burovski, al. Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation. In Proceedings of the 29th ACM International Conference Architectural Support for Programming Languages and Operating Systems, Volume 2, pp. 929-947, 2024.
Behrouz, Peilin Zhong, and Vahab Mirrokni. Titans: Learning to memorize at test time, arXiv preprint 1XXi:22501.00666 2024.
Win Chang, Wenlei Bao, Qi Hou, Chengquan Jiang, Ningxin Zheng, Yinmin Zhong, Xuanrun Zhang, Zuquan Song, Chengji Yao, Ziheng Jiang, al. Flux: Fast software- based communication overlap on gpus through kernel fusion. arXiv preprint arXiv:2406.06883 068588 2024.
Feiyang Chen, Yu Cheng, Lei Wang, Yuqing Xia, Ziming Miao, Lingxiao Ma. Fan Yang, Jilong Xue, Zhi Yang, Mao Yang, et al. Attentionengine: A versatile framework for efficient attention mechanisms on diverse hardware platforms. arXiv preprint arXiv:2502.15344 2025.
Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyun Wang, Yuwei Hu, Luis Ceze, al. (TVM): An automated (End-to-End) op- timizing compiler for deep learning. 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18), pp. 578-594, 2018.
Yuhong Chou, Zehao Liu, Ruijie Zhu, Xinyi Wan, Tianjian Li. Congying Chu, Qian Liu, Jibin Wu, and Zejun Ma. Zeco: Zero communication overhead sequence parallelism for linear attention. arXiv preprint arXiv:2507.0000 2025.
10

---

Published as a conference paper at ICLR 2026
Tri Dao and Albert Gu. Transformers SSMs: Generalized models and efficient algorithms through structured state space duality. In International Conference on Machine Learning (ICML), 2024.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io to-awareness. Advances in neural information processing systems, 15:16344-6635, 2022.
Soham De, Samuel L Smith, Anushan Fernando. Aleksandar Botev, George Cristian-Murant Al- bert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mix- ing gated linear recurrences with local attention for efficient language models. arXiv preprini arXiv:2402.1 19427, 2024.
Juechu Dong, Boyuan Feng, Driss Guessous, Yanbo Liang, and Horace He. Flex attention: A programming model for generating optimized attention kernels. arXiv preprint ar..v:24120000 2024.
Haojie Duanmu, Zhihang Yuan, Xiuhong Li, Jiangfei Duan, Xingcheng Zhang, and Dahua Lin. Skvq: Sliding-window key and value cache quantization for large language models. arXiv preprint arXiv:2405.062/9 2024.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:23/2.00755 2023.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):
1735-1780, 1997.
William Hu, Drew Wadsworth, Sean Siddens, Stanley Winata, Daniel Y. Fu, Ryann Swann, Muham- mad Osama, Christopher Ré, and Simran Arora. Hipkittens: Fast and furious amd kernels, 2025. URL httpt:///a//xii...oo. wwtp:iviorrggaooooo abs 251 08083
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc Le. Transformer quality in linear time. In Inter- national conference machine learning, pp. 9099-9117. PMLR, 2022.
Abhinav Jangda, Jun Huang, Guodong Liu, Amir Hossein Nodchi Sabet, Saced Malcki, Youshan Miao, Madanlal Musuvathi, Todd Mytkowicz, and Olli Saarikivi. Breaking the computation and communication abstraction barrier in distributed machine learning workloads. In Proceedings of the 27th ACM International Conference Architectural Support for Programming Languages and Operating Systems, pp. 402-416, 2022.
Zhihao Jia, Oded Padon, James Thomas, Todd Warszawski, Matei Zaharia, and Alex Aiken. Taso: optimizing deep learning computation with automatic generation of graph substitutions. In Pro- ceedings of the 27th ACM Symposium on Operating Systems Principles, pp. 47-62, 2019.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In International conference ma- chine learning, pp. 5156-5165. PMLR, 2020.
Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi Dengr, Chong Ruan, Damai Dai, Daya Guo, et al. Deepseek-v2: A strong, economical, and efficient mixture- of-experts language model. arXiv preprini arXiv:2405.0444 044444 2024a.
Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near- infinite context. arXiv preprint arXiv:2310.0188. 2023.
Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. Kivi: A tuning-fin free asymmetric 2bit quantization for kv cache. arXiy preprint arXiv:2402.022500 2024b.
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng. Michael Chung, Matteo Grella, et al. Rwkv: Reinventing mns for the transformer arXiv preprint arXiv:2305. 13048, 2023.

---

Published as a conference paper at ICLR 2026
Zhen Qin, Xuyang Shen, Dong Li, and Yiran Zhong. Accelerating linear attention design by uni- fying forward & backward propagation. In ES-FoMo III: 3rd Workshop on Efficient Systems for Foundation Models,
Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, and Yiran Zhong. Hgrn2: Gated linear rass with state expansion. In First Conference on Language Modeling, b.
Zhen Qin, Songlin Yang, and Yiran Zhong. Hierarchically gated recurrent neural network for se- quence modeling. Advances in Neural Information Processing Systems, 36:33202-3322 -32222, 2023.
Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixuan Sun, and Yiran Zhong. Lightning attention-2: A free lunch for handling unlimited sequence lengths in large language models. arXiv preprint 4rvi:22401. 400.04400..4668,, 2024.
James Reed, Zachary Vito,o, Horace He, Ansley Ussery, and Jason Ansel. torch. fx: Practical pro- gram capture and transformation for deep learning in python. Proceedings of Machine Learning and Systems, 4:638-651, 2022.
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, (166586-66888 2024.
Julien Siems, Timur Carstensen, Arber Zela, Frank Hutter, Massimiliano Pontil, and Riccardo Grazzi. Deltaproduct: Improving state-tracking linear rmns via householder products. arXiv preprint arXiv:2502.10025 2025.
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.
Weigao Sun, Zhen Qin, Dong Li, Xuyang Shen, Yu Qiao, and Yiran Zhong. Linear attention se- quence parallelism. arXiv preprint arXiv:2404.02882, 2024a.
Weigao Sun, Disen Lan, Yiran Zhong, Xiaoye Qu, and Yu Cheng. Lasp-2: Rethinking sequence parallelism for linear attention and its hybrid. arXiv preprint 2rXiv:2502.07563, 2025.
Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang, Yann Dubois, Xinlei Chen, Xiaolong Wang, Sanmi Koyejo, et al. Learning to (learn at test time): Rnns with expressive hidden states. arXiv preprint arXiv:2407.046200 2024b.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.08622 2023.
Philippe Tillet and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, pp. 10-19, 2019.
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.
Lei Wang, Yu Cheng, Yining Shi, Zhengju Tang, Zhiwen Mo, Wenhao Xie, Lingxiao Ma, Yuqing Xia, Jilong Xue, Fan Yang, et al. Tilelang: A composable tiled programming model for ai systems. arXiv preprint arXiv:2504. 175577 2025.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17455 2023.
Songlin Yang and Yu Zhang. Fla: A triton-based library for hardware-efficient implementtonnnnnnnnnnnttt of lincar attention mechanism, January 2024. URL http::///ittub.ccom//Lt
Flash
Songlin Yang, Bailin Wang. Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with harfiwte-efficient training. arXiv preprint 1rXiv:2312.06635, 2023.
12

---

Published as a conference paper at ICLR 2026
Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. arXiv preprint arXiv:2412.06664, 2024a.
Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transform- with the delta rule over sequence length. Advances in neural information processing systems, 37:115491-11552, 2024b.
Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, YX Wei, Lean Wang, Zhiping Xiao, et al. Native sparse attention: Hardware-aligned and natively trainable sparse attention. arXiv preprint arXiv:2502.11089 2025.
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.
Jintao Zhang, Jia Wei, Haofeng Huang, Pengle Zhang, Jun Zhu, and Jianfei Chen. Sageat- tention: Accurate 8-bit attention for plug-and-play inference acceleration. arXiv preprint 1rXiv:2410.02366, 2024.
Size Zheng, Wenlei Bao, Qi Hou, Xuegui Zheng, Jin Fang, Chenhui Huang, Tianqi Li, Haojie Duanmu, Renze Chen, Ruifan Xu, et al. Triton-distributed: Programming overlapping kernels on distributed ai systems with the triton compiler. arXiv preprint arXiv:2504. 19442, 2025a.
Size Zheng, Jin Fang, Xuegui Zheng, Qi Hou, Wenlei Bao, Ningxin Zheng, Ziheng Jiang, Dongyang Wang, Jianxi Ye, Haibin Lin, al. Tilelink: Generating efficient compute-communicatin over- lapping kernels using tile-centric primitives. arXiv preprint arXiv:2503.20313, 2025b.
13

---

Published as a conference paper at ICLR 2026
A USAGE OF LLMS
For the preparation this manuscript, we utilized LLM, as writing assistant to enhance the quality of the prose. Our process interactive: the authors provided initial drafts and specifie sentences to the LLM and used its suggestions to refine the text. Majority of prompts aimed at refining sentences improve conciseness, addition correcting grammar and improving overall clarity. The core scientific ideas, methodology, and experimental results were developed exclusively by the human authors, who bear full responsibility for all claims and content within this paper.
B PARALLISIM SCHEDULER
FlexLA automatically explore different parallel schemes for linear attention during compilation. With the domain knowledge of linear attention, we can shrink the search space into two mainly used schemes, which is partially explored in FLA (Yang& & Zhang 2024).
In this section, we provide a thcoretical analysis of two parallel execution strategies for linear at- tention using vanilla linear attention GPU example. Then derive heuristic scheduling algorithm [ that FlexLA used to build a parallelism scheduler, selecting the optimal strategy based on the input tensor shapes and hardware characteristics.
B. NOTATION
The vanilla linear attention is defined by the recurrence:
S. +
(8)
S;
(9)
where kr, q E RD and V, E RDD are the key, query, and value vectors at timestep I, and St E D De is the state matrix. We consider batched inputs Q,K IR BxTxxXDD and V E RBXT*HxDs The sequence of length T divided into Nc chunks of size C, such that T Ne C. Let NsM be the number of Streaming Multiprocessors (SMs) the target GPU, representing its parallel execution capacity.
B.2 STRATEGY 1: DECOUPLED THREE-PHASE PHASE EXECUTION
This strategy directly maps the chunk-wise parallel algorithm three-phase programming model. Each phase is separate kernel launch.
Parallelism Analysis. The degree of parallelism varies significantly between phases.
Phase 1 (Intra-Chunk State Computation): Computations for each chunk are indepen- dent. The total number of parallel tasks B H Nc. This phase exhibits the highest degree of parallelism.
Phase (Inter-Chunk State Propagation): The prefix sum (scan) is sequential across the Nc dimension. Parallelism is limited to the batch (B) and head (H) dimensions. Further parallelism, which we denote as Pstates can be achieved by partitioning the state matrix S E RD and processing its partitions on different thread blocks. The total number of parallel tasks is B x H x Petate-
Phase 3 (Merge): Similar Phase 1, the merge operation is independent across chunks, offering high degree of parallelism with B x H x Nc tasks.
Memory Access Analysis. The defining characteristic of this strategy is the materialization of intermediate states in global memory (GMEM) between phases.
Phase Reads K and from GMEM. Writes the intermediate, chunk-local states R Bx Dxxx back to GMEM.
GMEM Traffic
BHT D.) BHNoDDD,
(10)
Reads
14

---

Published as a conference paper at ICLR 2026
Phase 2: Reads the B x H x Ne chunk-local states from GMEM, performs the scan, and writes the updated global states S globbl E Bx D2 xD back to GMEM.
GMEM Traffic = BHNgDaD, + BHNoDaD, - Nc (11)
Reads
Writes
The total GMEM traffic introduced between Phase 1 and Phase is B H X Ne x DR
B.3 STRATEGY FUSED INTRA- AND INTER
INTER-CHUNE CCHNNK EXECUTION
This strategy fuses Phase and Phase 2, computing the local state chunk and immediately uses it to update the global state, all within on-chip memory.
Parallelism Analysis. By fusing the phases, the execution is constrained by the most sequential part, which is the inter-chunk scan. Therefore, the maximum number of parallel tasks is identical to that of Phase 2 in the decoupled strategy: B x H x Pstate- The high parallelism across the chunk dimension (Nc) sacrificed.
Memory Access Analysis. The primary advantage of this strategy is the elimination the inter- mediate GMEM traffic.
Fused Phase (1+2): Each of the B x H Pstate parallel tasks loads its corresponding slice of K and V from GMEM chunk by chunk. The intermediate, chunk-local states are generated and accumulated entirely within SMEM. The only state written to GMEM is the final updated global state for each chunk, which required by the Merge phase.
GMEM Traffic = BHT(D, + Do) + -BHNDDDDy
(12)
Writes
Compared to Strategy 1, this approach saves 2 x H No D&D, worth of GMEM traffic, which is the cost of one full read and one full write of the intermediate state tensor.
B.4 HEURISTIC SCHEDULING ALGORITHM
The choice between the Decoupled and Fused strategies presents classic trade-off between paral- lelism and memory locality.
Strategy (Decoupled) is favored when the GPU has high degree of parallelism (NsM is large) that is not saturated by the task parallelism of Strategy 2 (B x H x Patate). The performance gain from launching more parallel tasks in Phase and 3 must outweigh the latency incurred by the extra GMEM I/O.
Strategy 2 (Fused) is favored when the task parallelism of the scan (B X H x Pstate) is already sufficient to saturate the GPU's SMs, or when the cost of reading/writing the intermediate states is the dominant performance bottleneck.
We can formulate a simple heuristic to guide this choice as shown in Algo
In practice we set Tchnkk to to reduce the runtime overhead and omit the micro-benchmark effort. The insight is that most linear attention variants is memory bound instead of compute- intensive.
C COMPUTATIONAL PATTERNS AND MODEL COVERAGE
This appendix elaborates on the claim that our implemented kernels for few representative models support much broader range of linear attention variants. The mapping from representative implementations to the computational patterns and the models they cover is detailed below and summarized in Table 3
15

---

Published as a conference paper at ICLR 2026
Algorithm Parallelism Scheduler)
Require: Shapes B, I, H, dk. de, chunk size C, tile sizes tg. to
Require: Device params: memory bandwidth BW (bytes/s), Pmax (hardware concurrency cap)
Require: Cost param: chomp_chunk (measured per-chunk compute time in seconds)
Require: Thresholds: Pmin (saturation threshold, tasks), element size (bytes)
Ensure: Return strategy (FUSED, DECOUPLED]
1: N₂ + [T/C]
2: Pstate [dg/ 1. [du/to]
3: Paco < B No Pataee
4: Prused B H Patate
5: GMEMa & (2C(dz + dw) + 5dgdd)
6: GMEM fused <- (2C(ds + dp) 4dgdu) .8
b Fast path: fused already provides enough parallelism, prefer fused to save GMEM IO 7: if Prused 2 Pmin then
8:
return FUSED
9: end if
Otherwise estimate per-chunk runtime under each mapping
10: Pdec <- min( Pdecs Pinax)
11: Prused
min( Prused P Pmax)
GMEMdee
T Toomp_ebunk
Taee
BW Poec
Paec
GMEM
Teomp chunk
13: Tissed +
+
BW PPussd
Prused
14: if T Pusedd I Taee then
15: return FUSED
16: else
17:
return DECOUPLED
18: end if
Vector-State Linear RNNs (represented by HGRN) Our HGRN kernel embodies the compu- tational pattern of linear recurrent networks where the state is vector, updated via element-wise operations with gated inputs. This is common pattern for models aiming for high efficiency with compact state. Models sharing this fundamental structure include the original HGRN (Qin al. 2023) and Hawk (RG-LRU) (De et al. 2024
Matrix-State with Scalar Decay (represented by scalar-GLA) The scalar-GLA kernel repre- sents the widely adopted matrix-state linear attention pattern. In this formulation, the state is matrix updated via an outer product of key and value vectors, combined with a simple data-dependent or data-independent scalar decay. This pattern is foundational to many prominent and powerful models such as RetNet (Sun et al. 2023), Mamba-2 (Dao & Gu 2024), and Lightning Attention (Qin et al. 2024
Matrix-State with Vector Decay (represented by vector-GLA) Our vector-GLA implementa- tion captures the pattern of matrix-state models that employ more expressive, data-dependent vector decay. This allows for per-feature state transition dynamics, offering richer representa- tion than single scalar decay. Models in this category include the original Gated Linear Attention (GLA) (Yang et al. 2023). HGRN-2 (Qin et al. b). and RWKV-6 (Peng et al. 2023).
Delta-Rule Updates (represented by GDN) Finally, our Gated DeltaNet (GDN) kernel is repre- sentative of family of models based on the delta rule for state updates. This update mechanism can be interpreted as applying a series of Householder transformations to the state matrix, enabling more complex state transitions. This pattern is central to the family of DeltaNets, including the original DeltaNet (Yang et al., 2024b). GatedDeltaNet (Yang et al., 2024a), and DeltaProduct (Siems et al. 2025).
16

---

Published as a conference paper at ICLR 2026
Table 3: Mapping of representative kernels implemented subsection 4.1 models they cover. This demonstrates the generality of our compiler.
the broader set of
Representative Kernel
Covered Models (Examples)
HGRN
HGRN Qin et al.
2023).
Hawk (RG-LRU) (De al
2024
Scalar-GL.
RetNet (Sun et al.
2023).
Mamba-2
Dao & Gu
2024 Lightning
Attention (Qin et al.
2024)
Vector-GLA
GLA Yang et al. et
2023).
HGRN-2
Qin et al. b). RWKV-6 Peng
2023)
GDN
DeltaNet (Yang et al., 2024b
GatedDeltaNet Yang et al.
2024a)
DeltaProduct (Siems al. 2025)
D DISCUSSION ON EXTENSIBILITY AND SUSTAINABILITY
The longevity and utility of a Domain Speciffc Language (DSL) depend heavily on whether its ab- straction captures the invariant properties of the target domain rather than transient heuristics. In this section, we discuss the future-proofing of FlexI.A from two perspectives: algorithmic expres- siveness and hardware sustainability.
D.1 ALGORITHMIC EXPRESSIVENESS GROUNDED IN ASSOCIATIVITY
The universality of FlexLA's three-phase abstraction Caura-Chunk, Inter- Chter-Chunk, and Merge) derives from the mathematical foundation of efficient sequence modeling: the associative property.
Mathematical Invariance. Virtually all Linear Attention and State Space Duality (SSD) models aim achieve (N) computational complexity by formulating the attention mechanism rence or parallel prefix scan. Mathematically, any algorithm that can be decomposed into a chunk parallel form fits the FlexLA abstraction. This an ad-hoc design choice but direct mapping of the underlying algebraic structure required for parallelization.
Handling Complex Dependencies. FlexLA is expressive enough to handle complex state updates found in modern architectures, provided they satisfy chunk-wise associativity. For instance:
The Delta Rule: Despite involving data-dependent updates (e.g., h2 h + Be 1)) the Delta Rule preserves the associative structure within local chunks, allowing FlexLA to effectively fuse operations.
Element-wise Decay: Varying decay rates (as seen in Vector-Gated GLA) are fully sup- ported, as the clement-wise multiplication distributes over addition, maintaining the scan property.
Theoretical Limitations and Edge Cases. The boundary of FlexLA's applicability strictly de- fined by non-associative recurrences. If an architecture introduces a dependency where the state update h, = ff,hhhh,,,, ve) involves non-linear function that prevents parallel(e.g. passing the hidden state through a complex MLP at every step before the next update), it cannot be expressed in Flexl.A
A notable example is the Test-Time Training (TTT) layer Sun et al., 2024b). While TTT can be viewed through the lens of linear attention or fast weights, its gradient- updates involve non- linearities that break associativity. Consequently, TTT cannot be accelerated via the parallel prefix scans used in FlexLA. However, worth noting that this limitation mutual: by abandoning associativity, such models fundamentally forego the massively parallel efficiency that characterizes the linear attention regime targeted by FlexLA.
D.2 HARDWARE SUSTAINABILITY HIERARCHICAL DECOUPLING
To ensure sustainability amidst rapid hardware evolution, FlexLA employs a strict hierarchical de- coupling between the algorithmic description and sarewfce-specific instructions.
17

---

Published as a conference paper at ICLR 2026
Layered Compilation. FlexLA operates a high-level graph compiler rather than low-level sembler. It does not directly emit hardware- -ereeiii-- assembly (ISA). Instead, it lowers the Py PyTorch- based algorithmic description into intermediate kernel DSL. In current implementation, target Tri fritc ton, effectively delegating low-level complexities such as register allocation, instruction scheduling, and tensor core management (e.g., WGMMA on Hopper) the Triton compiler.
Backend Agnosticism. This layered design makes FlexLA inherently adaptable. While the cur- rent backend is Triton, the architecture allows for retargeting the mapping layer to other emerging DSLs (e.g.. ThunderKitten (Spector et al. 2024) or TileLang (Wang al. 2025)). These DSLs provide better performance with the cost of hardware overfit (e.g. ThunderKittens target on NVidia GPU, HipKittens (Hu al. 2025) target solely on AMD GPU while Triton support multiple hard- ware backend). Triton allows us minimize the use of inline assembly (e.g. PTX on NVGPU), reserving it only for specific extensions not yet exposed by the intermediate representation.
Forward Compatibility. Consequently, as hardware architectures evolve (e.g., the transition to NVIDIA Blackwell), FlexLA benefits automatically from updates made by the community to the intermediate compiler stack. This ensures that user-level code remains stable and performant with- requiring modification, solving the maintenance bottleneck often associated with hand-written kernels.
E DISCUSSION ON GENERALIZABILITY TO SOFTMAX ATTENTION
While Forge explicitly designed for Linear Attention, the underlying mathematical associativity, which enables our optimization is shared by Softmax attention (exploited via the online softmax trick (Dao al. 2022)). In this section, we analyze the theoretical feasibility of extending Forge to support Sofumax attention, and the design rationale behind our decision to specialize in the Linear Attention domain.
Extending the Forge abstraction to support Softmax-based mechanisms is theoretically feasible but would necessitate two primary modifications to the current three-phase abstraction:
Generalizing the Inter-Chunk State. In Linear Attention, the state propagated between chunks is typically feature map (e.g., S E Rd governed by linear recurrence. con- trast, Softmax attention requires the propagation of normalization statistics, specifically the running maximum m and the running sum / to stabilize the computation of exponen- tials. Extending FlexLA would require modifying the decay_mode phase to support the propagation these scalar or vector statistics alongside, or instead of, the matrix state.
Introducing a Rescaling Primitive. The merge_modk mode in FlexLA is currently designed for linear combinations (accumulation). Softmax attention, however, mandates renormal- ization step when merging partial results. Specifically, the output of a preceding block O must be scaled down based on the difference between its local maximum m, and the new global maximum mnewi
O Onea O X
+ O2 X ma- mnew
(13)
Supporting this would require introducing native rescale (output, old stats, new stats) primitive into the FlexL. DSL.I
Despite the feasibility of these extensions, we deliberately narrowed the scope of FlexLA to Linear Attention to maximize expressiveness.
The Tension between Universality and Specificity. Constructing "universal" DSL often neces- sitates a rigid abstraction that compromises the ability to model complex, domain-specific patterns. A recent framework, AttentionEngine (Chen et al. 2025), attempts unify both Softmax and Linear Attention. However, to maintain this broad generality, its abstraction struggles to express advanced Linear Attention variants that feature intricate dependencies, such the Delta Rule or Vector- Gated GLA. By specializing in Linear Attention, Forge avoids these constraints, supporting these complex patterns effortlessly.
18

---

Published as a conference paper at ICLR 2026
From ecosystem perspective, Softmax attention is already well-served by highly optimized li- braries and compilers (c.g., FlexAttention (Dong et al. 2024)). In contrast, the Linear Attention variants proposed
landscape is characterized by rapid algorithmic fragmentation. where frequently but lack efficient kernel support. FlexL. addresses this specific "N-to-1" compiler chal- lenge, solving critical bottleneck that is currently further optimizing Softmax kernels.
pressing for the research community than
Finally, while the abstraction of FlexLA is specialized, the system-level optimizations are broadly applicable. For example, our Ahead-of-Time (AOT) compilation pipeline and static dispatcher, designed eliminate runtime Python overheads, are universal optimizations. These techniques can be directly adopted by Softmax-focuse DSLs (such as FlexAttention) to significantly improve performance, particularly in short-sequence regimes where dispatch latency dominant factor.
F EXPERIMENTS: DISTRIBUTED EVALUATION ON H100 CLUSTER
To verify the robustness of distributed scaling strategy, conducted additional experiments an 8x H100 node connected via NVLink.
Table presents the latency comparison between Ring-Attention standard distributed softmax attention baseline) (Liu et al. 2023) with FlashAttention-3 which exploit advanced hardware feature on Hopper GPU, 2 and FlexLA (both LASP2 and FlexLA using the Scalar GLA variant). The experiments were conducted with batch size of 1, NumHeads = 32, and HeadDim = 128. The SeqLen in table means sequence length per GPU(global sequence length is equal to SeqLen x NumberOfGPUs),
The result shows FlexLA consistently outperforms the LASP-2 baseline across all sequence lengths.
For instance, at sequence length of 16k on 4 GPUs, FlexLA achieves a 1.17x speedup over LASP- (0.96ms vs. 1.m).. Even with the high bandwidth of NVLink on H100s, FlexLA maintains its efficiency advantage. This confirms that our fine-grained compute-communication overlap strategy effective not only on wanddwidhh-oosstrinne hardware (like the H20 used in the main result) but also nigh-performance flagship clusters. As expected, Linear Attention methods (both LASP- 2 and FlexIA)) are orders of magnitude faster than Ring-Attention especially at longer sequences (e.g., at 512k global sequence length, FlexLA is over 160x faster than Ring-Attention).
These results validate that the performance gains reported in the main paper are not artifacts of the H20 hardware or larger batch sizes, but rather stem from the fundamental efficiency of Forge's generated kernels and scheduling logic.
Table 4: Distributed Latency Comparison on H100 GPUs (Batch Size 1). We compare the end-to-end latency of Forge (Scalar GLA) against Ring- Ring-Attention (Softmax) and LASP-2 (State- of-the-art Linear Attention Sequence Parallelism). Forge consistently achieves the lowest latency across all sequence lengths and GPU configurations.
Latency (ms)
# GPUs
SeqLen Per GPU
Ring-Attn
LASP-2
Forge (Ours)
8192
8.40
1.02
0.79
16384
17.77
1.13
0.96
4
32768
61.65
1.66
1.57
65536
229.43
2.90
2.76
8192
19.32
1.00
0.77
16384
38.09
1.15
0.97
8
32768
133.14
1.67
1.62
65536
462.30
2.92
2.84
G CODE EXAMPLE
In this section we provide code examples to implement different linear attention variants and the code generate by FlexLA. Firt we show the Sealar GLA in Listing and the generated code in List-
19

---

Published as a conference paper at ICLR 2026
ing FlexLA enable pre-compiled kemnel (AOT) and static dispatch at runtime. In this example, we tune the kernel on H=[1.4,8, 16,32] and D=[64, 128], so there are a lot of auto-generated code encoded pre-tuned information for dispatcher (our AOT optimization). Using FlexLA implement this kernel only involves 50+ lines of code (I.OC), which the generated triton code and host launcher around 1200 LOC. Even without including the static dispatcher and host launcher in count, the generated Triton kerel alone has over 400 LOC. And this is only partially generated content, be- cause FlexLA also generates code for other parallel strategies.
Then we demonstrate the implementation of DeltaNet using FlexLA in Listing 3
Note that DeltaNet involves complex data dependencies where intermediate results computed in the chunk phase (specifically U and W) are required in the merge phase. FlexLA introduces a primitive forge cache result explicit mark these tensors. 'he compiler then automatically handles the memory layout and data movement to ensure these values are efficiently reused across phases without redundant re-computation manual memory management.
G. IMPLEMENTATION AND CODE GENERATION FOR SCALAR GLA
def chunk __c____ede________ seeee_scllrrg____ (k: Tensor, v: Tensor, g: Tensor) Tensor:
"**k: [c, K], v: [C, V], g: [C]
g_cumsum = g, cumsum (dim=0)
g_ last g. sum (dim=0)
9_cumsum (g_cumsum_1a) last g_cumsum) unsqueezee (() exp
k_fp32 peperuute(1) ([1, tt (t.dtype)
k_decay - (k_ g_cumsum) to (v.dtype)
state = k decay v
chunk state
10
11
def decay_mode_scalarr (pree
(prev_s: Tensor, chunk_state: Tensor, g: Tensor)
-> Tensor: """g: [C] """
13
g_sum g.sum (dim=0) exx ()
14
return prev_s sum chunk_state
15
16
def seagge__od__scaaarrooo
g: Tensor,
18
k: Tensor,
19
v: Tensor,
20
g: Tensor,
chunk_state: Tensor, scale: Tensor,
25
-> Tensor:
24
"*"q: [C, K], k: [C, K], v: [C, V), chunk_state [K, V] "**
25
g_cumsum = cumsum (0)
chunk_state chunk_state to dtype)
27
p (q (.T))tril(0)
28
p (P (g_cumsum[.... None] g_mumm [None, ...1).expp .to (v.
dtype)
return (p C q chunk_state g_cumsum.exp [..., None]) scale
31 CONST CenstExpr ("H", H)
32
CONST K ConstExpr ("K", K)
33
CONST_V ConstExpr ("V", V)
meta
35
"c": SymbTensor ( ["T", CONST_H, CONST_K), dtype-dtype),
36
"k": SymbTensor ["T", CONST_H, CONST_K), dtype-dtype),
37
"v": SymbTensor(( ["T", CONST_H, CONST_V] dtype=dtype),
"g": SymbTensor (["T,, CONST_H), dtype=torch. float32),
"prev_s". SymbTensor (["NS", CONST_H, CONST CONST_V], dtype=torch. float32), cchunk_state!, SymbTensor (["NC", CONST_H, H, CONST_X, CONST_V), dtype=
40
dtype),
sscale": 1 math. sgrt (K),
20

---

Published as a conference paper at ICLR 2026
LinearAttentien/
46
input_meta=metta
45
sp_group-pg if args. clse Nene,
46
enablc_@ot=args.
code_dir code_dir-args.dir,
49
50
Listing 1: Implementation of Scalar GLA Forge
(use_chunk wetneleeeee___eeeee sintue
"*bf16:16,
"*bf16:16,
*.fp32,
#.bf16:16,
*.bf16:16,
"*fp32,
*+i32,
"132,
10
"132,
"+132,
12
"SUSE_ IN INITIALLSTATT,,
13
"SH,
14
"SK,
"SV,
"SCHUNK, "
17
"SBLK K,
18
BBLK_V, "QUSE_TMM TMA"
19
22
23
def get
int, T: int, H: int, K; int, V: """Static dispatcher for chunkchunk, decay kernel. Auto-generatedd."
int):
25
IK, D DIK,, VI
26
key (B, T, H, 1))
27
if key in (
(I, 4096, 4, [64, 641), (2, 1024, 4, 164, 641),
30
):
31
BLK_K K = 32
32
BLK_V = 32
num_warps num_ctas
34
35
num_stages 3
36
maxnreg None elif key in
(1, 1024, 1, 164, 641),
39
(1, 1024, 1, 1128, 1281),
40
(1, 1024, 4, 164, 641),
41
(1, 1024, 4, [128, 1281),
(1, 1024, 8, [64, 641),
43
(1, 1024, 8, [128, 1281),
44
(1, 1024, 16, [64, 641),
45
(1, 1024, 32, [64, 64)),
(1, 2048, 1, [64, 64]),
(1, 2048, 1, [128, 1281),
18
(1, 2048, 4, 164, 641),
d9
(1, 2048, 4, [128, 1281),
50
(1, 2048, 8, (64, 641).
(1, 2048, 8, [128, 128]),
(1, 2048, 16, [64, 64]),
21

---

Published as a conference paper at ICLR 2026
(1, 2048, 32, 164, 641),
(1, 4096, 1, [64, 641),
55
(1, 4096, 1, [128, 1281),
56
(1, 4096, 4, [128, 1281),
(1, 4096, 8, [64, 641),
(1, 4096, 8, [128, 1281),
59
(1, 4096, 16, [64, 641),
60
(1, 4096, 32, 164, 641),
(1, 8192, 1, 164, 641),
(1, 8192, 1, 1128, 1281),
(1, 8192, 4, [64, 641),
64
(1, 8192, 4, [128, 1281),
65
(1, 8192, 8, 164, 641),
(1, 8192, 8, [128, 1281),
(1, 8192, 16, [64, 64]),
68
(1, 8192, 32, [64, 64]),
89
(1, 16384, 1, [64, 641),
70
(1, 16384, 1, [128, 1281),
(1, 16384, 4, [64, 641),
72
(1, 16384, 4, [128, 128)),
73
(1, 16384, 8, [64, 641),
74
(1, 16384, 8, [128, 128]),
15
(1, 16384, 16, 164, 641),
(1, 16384, 32, 164, 641),
77
(1, 32768, 1, 164, 641),
78
(1, 32768, 1, [128, 1281),
79
4, [64, 641), (1, 32768,
(1, 32768,
4, [128, 1281),
(1, 32768,
8, 164, 641),
82
(1, 32768, 8, [128, 128)),
83
(1, 32768, 16, [64, 64]),
84
(1, 32768, 32, [64, 641),
(1, 65536, 1, 164, 641),
86
(1, 65536, 1, [128, 1281),
87
(1, 65536, 4, [64, 641),
88
(1, 65536, 4, [128, 128]),
(1, 65536, 8, 164, 641),
(1, 65536, 8, [128, 1281),
91
(1, 65536, 16, 164, 641),
92
(1, 65536, 32, 164, 641),
(1, 131072, 1, 164, 641),
(1, 131072,
[128, 1281),
95
(1, 131072, 4, [64, 641),
96
(1, 131072, 4, [128, 1281),
97
(1, 131072, 8, [64, 641),
(1, 131072, 8, [128, 128]),
(1, 131072, 16, [64, 64)),
(1, 131072, 32, 164, 641).
101
(1, 262144, 1, 164, 641),
(1, 262144, 1, 1128, 1281),
(1, 262144, 4, [64, 64]),
(1, 262144, 4, [128, 1281),
105
(1, 262144, 8, [64, 641),
106
(1, 262144, 8, [128, 1281),
(1, 262144, 16, [64, 641),
(1, 262144, 32, [64, 641),
(2, 1024, 1, [64, 641),
110
(2, 1024, 1, [128, 1281),
(2, 1024, 4, [128, 1281),
(2, 1024, 8, [64, 64]),
(2, 1024, 16, 164, 641),
114
(2,, 2048, 1, [64, 641),
115
(2, 2048, 1, 1128, 1281),
(2, 2048, 4, 164, 641),
(2, 2048, 4, 1128, 1281),
22

---

Published as a conference paper at ICLR 2026
(2, 2048, 8, 164, 641),
(2, 2048, 16, 164, 641),
120
(2, 4096, 1, [64, 641),
(2, 4096, 1, [128, 1281),
(2, 4096, 4, [64, 64]),
(2, 4096, 4, [128, 1281),
124
(2, 4096, 8, [64, 641),
125
(2, 4096, 16, 164, 641),
(2, 8192, 1, 164, 641),
(2, 8192, 1, 1128, 1281),
(2, 8192, 4, [64, 641),
129
(2, 8192, 4, [128, 1281),
130
(2, 8192, 8, 164, 641),
(2, 8192, 16, 164, 641),
(2, 16384, 1, [64, 64]),
133
(2, 16384, 1, [128, 128)),
134
(2, 16384, 4, [64, 641),
135
(2, 16384, 4, [128, 1281),
(2, 16384, 8, [64, 641),
(2, 16384,
16, [64, 64]),
138
(2, 32768,
1, [64, 641),
139
(2, 32768,
1, [128, 128]),
(2, 32768,
4, 164, 641),
(2, 32768,
4, 1128, 1281),
142
(2, 32768,
8, 164, 641),
143
(2, 32768,
16, [64, 641),
144
(2, 65536,
1, [64, 641),
(2, 65536, 1, [128, 1281),
(2, 65536, 4, [64, 64]),
147
(2, 65536, 4, [128, 128)),
148
(2, 65536, 8, [64, 64]),
(2, 65536, 16, [64, 641),
(2, 131072,
164, 641),
151
(2, 131072, 1, [128, 1281),
152
(2, 131072, 4, [64, 641),
153
(2, 131072, 4, [128, 128]),
(2, 131072, 8, [64, 64]),
(2, 131072, 16, [64, 641),
156
):
157
BLK K 32
BLK V 32
num_warps num_ctas num_stages 4
161
162
maxnreg None
elif key in
(1, 1024, 32, [128, 1281)
(1, 2048, 32, 1128, 1281).
166
(1, 4096, 32, [128, 1281),
(1, 8192, 32, 1128, 1281),
(1, 16384, 32, [128, 128]),
(1, 32768, 32, [128, 1281),
170
(1, 65536, 32, [128, 1281),
171
(1, 131072, 32, [128, 1281),
(1, 262144, 32, [128, 1281),
(2, 1024, 32, [128, 1281),
(2, 2048, 16, [128, 1281),
175
(2, 2048, 32, [128, 1281).
(2, 4096, 16, [128, 1281),
(2, 4096, 32, [128, 128]),
(2, 8192, 32, [128, 1281),
179
(2, 16384, 16, [128, 1281),
180
(2, 16384, 32, 1128, 1281),
(2, 32768, 32, 1128, 1281),
(2, 131072, 16, 1128, 1281),
23

---

Published as a conference paper at ICLR 2026
):
BLK_K 32
185
B 64
186
num warps num_ctas = num_stages
189
maxnreg None elif key in (2, 1024, 16, 1128, 1281),
190
BLK_K 32
194
BLK_V 64
195
num_warps 8 num ctas = stages maxnreg None elif key in
198
199
200
(1, 1024, 16, [128, 1281),
(1, 2048, 16, [128, 1281),
(1, 4096, 16, [128, 1281),
203
(1, 8192, 16, [128, 1281),
204
(1, 16384, 16, [128, 128]),
[1, 32768, 16, 1128, 1281),
(1, 65536, 16, 1128, 1281),
207
(1, 131012, 16, 1128, 1281),
208
(1, 262144, 16, [128, 1281),
209
(2, 1024, 8, [128, 1281),
(2, 1024, 32, [64, 641),
(2, 2048, 8, [128, 1281),
212
(2, 2048, 32, [64, 641),
213
(2, 4096, 8, [128, 1281),
214
(2, 4096, 32, [64, 641),
(2, 8192, 8, [128, 1281),
216
(2, 8192, 16, [128, 1281),
217
(2,, 8192, 32, [64, 641),
218
(2,, 16384, 8, [128, 128]),
(2,, 16384, 32, [64, 64]),
(2, 32768, 8, [128, 1281),
221
(2, 32768, 16, 1128, 1281),
222
(2, 32768, 32, 164, 641),
(2, 65536, 8, 1128, 1281),
(2, 65536, 16, [128, 1281),
(2, 65536, 32, 164, 641),
226
(2, 131072, 8, [128, 1281),
227
(2, 131072, 32, [64, 64)),
):
BLK_K 32
230
V 64
231
num_warps 8
num_ctas- =
num_stages 4
maxnreg None
235
elif key in (2, 65536, 32, [128, 1281), (2, 131072, 32, [128, 1281).
236
):
BLK_K 64
240
BLK_V 128
241
num_warps num_ctas =
num_stages
244
maxnreg None
245
else:
raise ValusError (f unuuuporetee config for Fuse chunk decayy_keree 3THD=(key|")
24

---

Published as a conference paper at ICLR 2026
248
return
249
"CHUNK"; 64,
250
"USE_ INS__INIIIALLSSAAAS STATEE: Truc, "H": H,
"K": D[O),
253
"V": D[1], K": KLK_,, :BLK_V": BLK_,, "num_warps" warps": warps, "nun_stages"; nun_stages,
254
255
258
259
262
263
Baot_ paot__ompill__ppcce
264
{
"signature"; fuse fuse_ehunle serneelsatensssgguuu
"grid": ("(SK $BLK_K BLX_V", "H_MUI_NS"),
SBLK_K", " (8V &BLK_V /
267
"triton_algo_infoss [
268
get fut__cuke_c httpe__hhunn_dccyyyeeee decay_ (B, 7, a, 3, V)
for H, (K, v) inl
270
(1, 262144, 32, 164, 6411,
271
(1, 262144, 8, [128, 128]),
272
(2, 1024, 16, [128, 1281),
(2, 1024, 4, [64, 641),
(2, 131072,
[128, 1281),
275
(2, 131072, 1, [64, 641),
276
(2, 131072, 16, [128, 128]),
277
(2, 131072, 16, [64, 641),
(2, 131072, 32, [128, 1281),
279
12, 131072, 32, (64, 641),
280
(2, 131072, 4, [128, 128]),
281
(2, 131072, 4, [64, 64]),
(2, 131072,
[128, 128]),
(2, 131072, 8, [64, 64]),
284
12, 32768, 32, 1128, 1281),
285
(2, 65536, 16, 1128, 1281]
288
289
}
290
configs-[
triten. (rrttnn.ooniigg(BBIX..... BLK_K, "BLK_V": BLK_V, " USE_TMA sum_stages-rummsttggee for num_warps in 14, 81 for nun_stages in 13, 41 for BıJ_K in [128, 64, 32] for BLK_V [128, 64, 32] for USE_ in [True, False]
1, L___wapps-num_warpp,
293
297
298
1,
key=[],
301
8triton. jit
302
def fuse_chhun__docaykkeeaa
303
k,
g.
306
s,
307
out_0, out_ 1,
308
seqlens,
25

---

Published as a conference paper at ICLR 2026
NS,
311
H_MUI_NS,
312
Bhunk__ofaetts_wit___ee
313
USE_ LIL_IINIIILASTTET tl. cl.constexpt,
H: tl.constexpr,
K: :ll.constxpr,
316
V: t:l.censtepp,
317
CHUNK: t1.constexpt,
BLK_K: ti.constexpr,
BLK V: tl. conntteee,,,, USK_TMA: constexpr
321
322
NUM_BLK_K (K BLKKK 1) // BLK_K
NUM BLK V = (V BLK V 1) // BLK_
NUM BLK KV NUM BLK NUM BLK
325
326
k i_,, 1_sh a1.programmidvon tl.program t1.program_idd(2)
327
1_,, i_h = i_sh // H, iss $ H
bos,
ts).to(t1.int33)) (L.
Fis (nt.3))).in)33
329
T ecs bos
330
NC tl.cdiv(T, CHUNK)
boh = tl. Load (chunk www.oooffttttteoo.ooooeee
tttaddotttsststsssakeee.. with ini tt (t1.
333
out out 0 (boh
h) (tt(ELiin)))) prev_s prev_s + (i_s H i_h) *K- -
334
335
initial_decay - 1.0
out_l: [NC + NS, H]. NOTE: AADODAM in log space ptr_out_l out
boh H i
338
# store initial decay
339
(t.streepptr_out
caatial_decayy
340
ptr_out_1 += H [BK, BV]
342
prev_ te.zeros K, BLK dtype=tl. flcat if USF_INITIAI_STATR:
343
344
(K, v), (v, 1), (i_k BiK_,, i_v HI.K_V, (BIK_K, BIK_V), (1, 0))
ptr_prev_s t1 ww..aaeelllleccoommmmm (tke_bbocc_ptrrprrrreeeee
bik_prev_s tl lodd (ptr_prev__s boundary_checkk (0, 1)) (t] float32)
346
out ptr_out_0 ma.eake block por 0, (K, V), (V, 1), k*LK K, V* BLX V), (BLK K, BLK (1, 0)) # fat: skip tl.store (ptr_out_0, sk<prov_..topptroouut..... elementtt.l out_0 AAKAY
349
350
if USE_TMP
k_ desc tl baake__tesor_ddecree k bos shape= K), strides= [ K, 11, block_ (ock_shape=/CHHNN, B),
i_h
353
354
355
desc www.ake www_mee_tensor_desrriiiee v bos * * i_h v,
358
359
shape= VI, strides= v, 11,
Block_shape=[Cee
BLK_V),
362
363
out_0_desc tenmaakk_ene out_0, shape- [NC, K, v], strides- [H V, v, 1], block_shae (1, BL,_K, BLK_V],
364
367
368
369
for i_c in range (NC)
26

---

Published as a conference paper at ICLR 2026
load (trans) 'k': (T, K,) (BLK CHUNK, if not USE_TMA:
372
373
cur_k k bos + H + K + 1 K $ fmt: skip (K, T,), (1, B K,), (1_k
374
ptr_k_o =
BLX K i_ CHUNK,), (BLK_K, K, CHUNK, (0, 1,1) # fat: skip
0 tl.load (ptr_k_0, Soundary_checke
1,)) fnt:
skip
376
else:
377
0 = k. Less..Loadl1... * CHUNK, 1_k K]) .trans
load ' (T, H, v,) (CHUNK, BIKVV,) if not USE_TMA:
380
381
cur_y v bos * H V i V 4 fmt: skip ptr_v_ = tl.make block ptroccc__preur_______ v, (T, v,), (a v, 1,), (i_c CHUNK, i BLK EIK_V,), (CHUNK, BLK_V,), (1, 0,1) fat: skip
385
vik t1.load (ptr_v_ boundary_cheece
0,)) fmt:
skip
384
clsc:
bik_v_1 v_desc.load CHUNK, L BLK
387
388
@all_external func:
local www.a_ooaaaaaurreaeeeeaaeeeauaae to pre-compute g
391
Load 'g': (T, H,) => (CHUNK,)
392
cur_g - g bos * + i_h 1 fmt: skip
393
ptr_g_2 = d1.make_block_pa (cur_g, (T,), (H,), (i_c CHUNK,), ( CHUNK,), (0,))
fmt skip
bik_g_2 t1.1oad 1attrdrdpttgg
Saundary_checkee (0,)) fmt: skip
395
load lad__aas_gg:11 => ()
396
# load 'g' [(CHJN 1)]: :: (T, 8,) => ()) cur_g
397
bos * H i_h 1 + (CHUNK - 1) H fmt: skip
ptr_last_g_3 cur i_ CHUNK
$ fmt: skip
399
last_g_3 (ptr last_g_3)
400
# sub: torch.Size ('CHUNK', )) => CHCHUNK ,) sub_4 - last_g_3 blk blk_g_2 unsqueeze: ('CHUNK', -> (1, (CHUNK))
401
unsqueeze_5 sub_ [None, :]
404
exp: (1, CCUNK') => (1, CHUNK')
405
exp_6 6 = ex. (unsqueeze 5)
# to: ('BLK K', CHUNK') ) => ('BLK K', 'CHUNK')
blk_k_0_float32_ blk to ttl. float32)
mul: ('BLK_'' CHUNK' (1, (CHUNK') ) ('BLK CCUNKK'
409
mul_8
exp_6 $ to: ('BLK_K' 'CHUNK') => ('BIK_K' K', "CHUNK")
410
mul_8 8___flaat166999 matmul: ('BLK_K, "CUNK") ('CHUNK', BLK_V')) ("BLK_K', V
(t.. (tl.bfloat16)
413
matmul 10 ti.dot (mul 8 bfloaa166_,, v to (nul_ bfloati6__ ddypp)
$ exp: () ->
exp_1_11 tl.exp plast_g33)
416
mul: ('BLK_K' BLL_V'), O) -> ("BLK_K,, (BLK_V")
417
mul_ 12 = bik_pree prev_re exp 11
# add: ('BLK K', BLK V), ('BLK K', (BLKVV) E> ('BLK K', 'BLK
add_13 mmu_ll111 1_12 matmul_10 mul: ('s3/, ,), () => ('s3',)
420
421
tul_tensor_14 initial_decay # store => ('BLK_K K', "BLK_V"
cxp_1_11
assume output layout: [NC_ [IC_WITH_ING INI, H, K, V] out_ out put_O.dttppe.e eleeent_myeeyy
424
425
it not USE TMA:
ptr_out_0 ma.make black ptr (out 0, (3, V.), (V, 1,), ki
BLK X, iv BIK V,), (BLK BLK V,), (2, 0)) $ imt: skip
27

---

Published as a conference paper at ICLR 2026
tt.store (pt add 13.to ty) houndary_cheee (2,
)
428
out += H
429
clsc:
out_o_desc___ store ([i_e, i_k BIK_K, i_v BLK_V), edd_13.to to (_t [None, ;, :1)
431
bik_prev__ add_13.to (bd_ $ store => ('s3',)
432
435
# output layout:
WITH INI, H, s3]
out_l_ty eet_
i_c H, wul_tensor_14.to initial_decay mul_tenscr_14
436
437
def Baancchffuee_cun__deeee
g: torch. Tensor, prev_s: tenso,
440
441
v: torch. Tensor,
442
k; torch. torch.Tenson Tensor,
cu_seqlens: torch. IntTensor, cached_results, dict, dist_scan: Dist.ScanContxt. Nene, lazy_updater bool True, aictlstr, torch. Tensor]: ...
445
446
447
449
perform "decayed scan" on 'chunked 'chunked_states states'
450
451
state_dee _eate_dyppe_aeeeeeeeeeeee_l_ee chunk_state.de
chunk_decay None
prev_ see__ran_sstat_e_ee sum None
454
455
TODO: now fix chunk_siae to 64 CHUNK 64
456
chunk_indices pregare_chunk__nddceeiccei (ouselllee CHUNK)
chunk_offsets thunk_offaets_withe htnporrecccnnccffsstttttoeeeeeee CHUNK) CHUNK)
459
seeparr_ccuuk_offfetsstttta (cu_seqlens,
NS, NC 1en (cu_ senlens) 1, chunk_indices.
NC_WITH_INI aetiron.get.io FORGE USE AOT', (0')
462
use act os ensiroi.o..ooo
463
FORGE USE ACT True if use ase_aatt.ovwe lower () in ['1', 'true', "yes'] else None
H, K, V = *k. shape, v.shape(-1) spcateed_sttts k.nex_empty ( INC_WITILINI, X, V])
466
467
def alloc_fn(size) int, alignment: int, stream: int):
return forch.empte (size, device=' dtype=torch. int8)
470
(alloc_
471
472
def grid (meta):
BL._K meta ['BIK_K[]
BLK_V
[1 BLBKK__V
475
WUM_BI__ (K BLK_K 1) // BLK_K
476
NUM BLK V IV + BLK V 1) // BLK V
NUM BLK KV = NUM BLK K NUM BLK_ V
MUL NS
NS
return BLK NUMLBLK_V, NS,
480
481
chunk_dccay = torch, empty (INC_NITH_INI, H), dtype-torch. float if FORGE_USE ______UE______ AOT is None:
juse_chunk_decayy kernel [grid]
kernel
484
g-g,
485
prev_s-prev_s,
v=v, k=k,
28

---

"chunk_stated updated_states, "chunk_decay": None, None, Published as a conference paper at ICLR 2026
out 1=chunk decay, cuLesggeensruuueegllene, Н-Н,,
490
491
NS-NS,
K-K,
V-V,
494
CHUNK-CHUNK,
495
H_MUL_NS- NS,
out jut_OOuupdate__state
Sauunkofffeet__eeeeeeeeeeeeeeeee with wni-chunn_offsettee sesettswwtttt_ts ini,
USK_INITIAI
499
500
else:
503
from Coge.aot_ute import forge_aot_ope
304
305
alge_info forge_aot_ops. fuse_chunk chean_ cercelkeeelllcoo for _K, _ jet_fise_ccun___dccc__eeeee
tot.eey_eeiiiitttotoooooooooooee trgoootaagoooooooee in
(NS, T, K, V)
i)
507
setattr (algotifaalg_if.),) _,, _v)
torge_aot_ops fuse waase_chunk._e chunk_decay_caaaaeeeeeee
0, # torch. outa. ceaplaacuueet_eeeee strram .euda strear, k. (trata_prr(), v.data_ptr(), v
510
511
512
g.data_ptr() $
prev_s.data_str() A prev_s
uptated_states... dat.daaaaaaaaa (tt() # out_
315
Baank__dcaa.data___o cu_sellens.dat__tzr)) # cu_seqlens NS, # NS
out_l
516
517
NS, # NS
ohunk_offsets__with ini (t caunk_offsets.wthhh.. alge_info,
520
521
524
final_chunk chunk indices chunk offsets chunk_offsets with [1:]
325
final state_ttalloloe local updated_states shank_ indices, ...]
final_ Lanal_decayy_eeee chunk dele_deecaylll final chunk_ indices ...]
if cist_scan.pg.
529
prov_rank_state_sum = waiisssscan.frrreeeee
530
decsy_opee=eccyyyyee.aa.. Lazy_update-Trus,
533
534
535
cached_results. (()) if prev_rank_state_suu return
not
538
539
"chunk_state state": updated_statee
"chunk "chunk_decay"; chunk decay,
"prev_rand "aaat__smm,, prev_rank_stare
542
543
clsc:
544
return
547
(prev_rank_state___
548
549
merge merge_moce saenel_signattin = (
*=bi16:16,
29

---

Published as a conference paper at ICLR 2026
**bf16:16,
**bf16:16,
553
*+fp32,
554
*.bf16:16,
"fp32,
**bf16:16,
557
"132:16,
558
**bE16:16,
559
"132:16,
**bi16:16,
"+132,
362
"132,
363
"132,
*.132,
"132,
366
"SFUSE SP EST_STATE_UPDATE "SH, "SK,
367
368
"&V, "SCHUNK, "
371
"SBLK_K, K,
572
"SBIK_V V,
"SUSE TMA"
515
576
577
def jet_mmerg__nod__eeee__ntooo (B: int, т: int, H: int, K: int, V: int): "**Static dispatcher for prrge_mode_kernes Auto-generated,
[K, key (B, T, H, D)
380
581
if key in
582
(1, 4096, 4, [128, 1281),
(1, 131072, 32, [128, 1281),
584
(1, 262144, 32, [128, 1281),
385
(2,, 1024, 8, [128, 1281),
586
(2,, 4096, 1, [128, 1281),
(2, 8192, 1, [128, 1281),
):
589
BLK K 128
390
BLK v = 128 warps num_ctas = num_stages
594
maxnreg None elif key in
595
(1, 4096, 32, [128, 128]),
(1, 16384, 8, [128, 1281),
598
(1, 16384, 32, 1128, 1281),
599
(1, 32768, 8, [128, 1281).
(1, 32768, 16, 1128, 1281),
(1, 32768, 32, [128, 1281),
(1, 65536, 8, [128, 1281),
603
(1, 65536, 32, [128, 1281),
604
(1, 131072, 4, [128, 1281),
(1, 131072, 8, [128, 1281),
(1, 131072, 16, [128, 1281),
(1, 262144, 4, [128, 1281).
608
(1, 262144, 8, [128, 1281).
609
(1, 262144, 16, [128, 1281),
(2, 4096, 16, [128, 1281),
(2, 4096, 32, [128, 1281),
612
(2,, 16384, 8, [128, 1281),
613
(2, 16384, 32, 1128, 1281),
(2, 32768, 16, 1128, 1281),
(2, 32768, 32, 1128, 1281),
30

---

Published as a conference paper at ICLR 2026
(2, 65536, 8, [128, 1281),
(2, 65536, 16, [128, 1281)
618
(2, 131072, 4, [128, 1281),
619
(2, 131072, 8, [128, 1281),
(2, 131072, 16, [128, 1281).
622
BLK_K 128
623
BLK_V 128
warps ctas num_stages maxnreg None
627
628
elif key in
(1, 1024, 16, [128, 1281),
(1, 1024, 32, [128, 1281),
631
(1, 2048, 8, [128, 1281),
632
(1, 2048, 16, [128, 1281),
633
(1, 16384, 1, [128, 1281),
(1, 32768, 1, [128, 1281),
(2, 2048, 4, [128, 1281),
636
(2, 65536, 32, [128, 1281),
637
(2, 131072, 32, [128, 1281),
638
):
BLK K 128
640
BLK V 128 num_warps num_ctas - num_stages maxnreg None elif key in
641
642
645
646
(1, 2048, 32, [128, 128]),
647
(1, 4096, 16, [128, 1281),
(1, 8192, 16, [128, 1281),
(1, 8192, 32, [128, 1281).
650
(1, 16384, 16, [128, 1281),
651
(1, 65536, 4, [128, 1281),
(1, 65536, 16, [128, 128]),
(2, 1024, 32, [128, 1281),
654
(2, 2048, 32, 1128, 1281),
655
(2, 8192, 8, [128, 1281),
[2, 8192, 16, 1128, 1281),
(2, 8192, 32, [128, 1281),
(2, 16384, 16, [128, 1281),
659
(2, 32768, 4, [128, 1281),
660
(2, 32768, 8, [128, 1281),
(2, 65536, 4, [128, 128]),
):
BLK_K 128
664
BLK_V = 128
num_ warps num_ctas -
nun_stages
668
maxnreg None elif key in
669
(1, 1024, 1, 164, 641).
(1, 4096, 1, 164, 64]).
672
):
673
BLK_K 128
674
BLK_V = 32
num_warps num_ctas num_stages 3
677
678
maxnreg = None
elif key in
(1, 1024, 1, 1128, 1281),
31

---

Published as a conference paper at ICLR 2026
(1, 1024, 4, 164, 641),
(1, 1024, 8, [64, 641),
683
(1, 2048, 1, [64, 641),
684
(1, 2048, 1, [128, 1281),
(1, 2048, 4, [64, 64]),
(2, 1024, 1, [64, 64]),
687
(2, 2048, 1, [64, 641),
688
):
689
BLK_K = 128
BLK 32
num_warps num_ctas -
692
693
num_stages 4
maxnreg = None
elif key in
696
(1, 1024, 8, [128, 1281),
697
(1, 1024, 32, [64, 641),
698
(1, 2048, 4, [128, 1281),
(1, 4096, 4, [64, 641),
(1, 4096, 8, [64, 64]),
701
(1, 4096, 8, [128, 1281),
702
(1, 8192, 1, [128, 1281),
703
[1, 8192, 8, 164, 641),
(1, 8192, 8, 1128, 1281),
705
(1, 8192, 32, 164, 641),
706
(1, 16384, 1, [64, 641),
707
(1, 16384, 32, [64, 641),
(1, 32768, 1, [64, 641),
(1, 32768, 4, [64, 64]),
710
(1, 32768, 8, [64, 64]),
711
(1, 32768, 16, [64, 641),
712
(1, 65536, 1, [128, 1281),
(1, 131072,
[128, 1281),
714
(1, 131072, 32, [64, 641),
715
(1, 262144, 1, [64, 641),
716
(1, 262144, 32, [64, 64]),
(2, 1024, 4, [64, 641),
(2, 1024, 4, [128, 1281),
719
(2, 1024, 8, 164, 641),
720
(2, 1024, 16, 164, 641),
[2, 1024, 32, 164, 641),
(2, 2048, 4, [64, 641),
723
(2, 2048, 8, [64, 641),
724
(2, 2048, 8, [128, 1281),
725
(2, 2048, 16, [64, 64]),
(2, 2048, 16, [128, 128]),
(2, 2048, 32, 164, 64]),
728
(2, 4096, 4, 164, 641),
729
(2, 4096, 4, [128, 1281),
(2, 4096, 8, 164, 641).
[2, 4096, 8, [128, 1281),
732
(2, 4096, 16, [64, 641),
733
(2, 4096, 32, 164, 641),
734
(2, 8192, 1, 164, 641),
(2, 8192, 4, [128, 1281),
(2, 16384, 1, [64, 64]),
737
(2, 16384, 1, [128, 1281),
738
(2, 16384, 4, [128, 1281),
739
(2, 16384, 8, [64, 64)),
(2, 16384, 32, [64, 64]),
(2, 32768, 16, [64, 641),
742
(2, 32768, 32, [64, 641),
743
(2, 65536, 1, 164, 641),
744
(2, 65536, 1, 1128, 1281),
(2, 65536, 16, 164, 641),
32

---

Published as a conference paper at ICLR 2026
):
BLK_K 128
748
B 64
749
num warps num_ctas = num_stages
752
maxnreg None elif key in
753
754
(1, 4096, 1, 1128, 1281),
):
BLK_K 128
757
BLK_V 64
758
num_warps 8 num ctas = stages maxnreg None clif key in
761
762
763
(1, 1024, 4, [128, 1281),
(1, 1024, 16, [64, 641),
(1, 2048, 8, [64, 64]),
766
(1, 2048, 16, [64, 641),
767
(1, 2048, 32, [64, 64]),
(1, 4096, 16, 164, 641),
(1, 4096, 32, 164, 641),
770
(1, 8192, 1, 164, 641),
771
(1, 8192, 4, [64, 641),
772
(1, 8192, 4, [128, 1281),
(1, 8192, 16, [64, 641),
(1, 16384, 4, [64, 64]),
775
(1, 16384, 4, [128, 128)),
776
(1, 16384, 8, [64, 64]),
777
(1, 16384, 16, [64, 641),
(1, 32768, 4, [128, 1281),
779
(1, 32768, 32, [64, 641),
780
(1, 65536, 1, [64, 641),
781
(1, 65536, 4, [64, 64]),
(1, 65536, 8, 164, 64]),
(1, 65536, 16, [64, 641),
784
(1, 65536, 32, 164, 641),
785
(1, 131072, 1, 164, 641),
(1, 131072, 4, 164, 641),
(1, 131072, 8, [64, 64]),
788
(1, 131072, 16, [64, 641),
789
(1, 262144, 1, [128, 1281),
790
(1, 262144, 4, [64, 641),
(1, 262144, 8, [64, 64]),
(1, 262144, 16, [64, 64)),
793
(2, 1024, 16, [128, 1281),
794
(2, 8192, 4, 164, 641),
795
(2, 8192, 8, 164, 641).
(2,, 8192, 16, [64, 64]),
(2, 8192, 32, 164, 64]),
798
(2, 16384, 4, 164, 64]),
799
(2, 16384, 16, [64, 641),
(2, 32768, 1, 164, 641),
(2, 32768,
1, [128, 1281),
(2, 32768,
4, 164, 641),
803
(2, 32768,
8, [64, 641),
804
(2, 65536, 4, [64, 64)),
(2, 65536, 8, [64, 64]),
(2, 65536, 32, [64, 64]),
807
(2, 131072, 1, [64, 641),
808
(2, 131072, 1, 1128, 1281),
(2, 131072, 4, 164, 641),
(2, 131072,
164, 641),
33

---

Published as a conference paper at ICLR 2026
(2, 131072, 16, 164, 641), (2, 131072, 32, [64, 641),
812
813
):
814
BLK_K = 128
BLK_V V 54
num_warps num_ctas
817
818
suaes 4
maxnreg = None
elir key in (2, 1024, 1, [128, 1281),
822
(2, 2048, 1, [128, 1281),
823
(2, 4096, 1, 164, 641),
):
BLK K 128
826
BLK_ V 64 num_warps
827
828
num_ctas =
num_stages
maxnreg None
831
else:
832
raise ValueError (f*Unsupported config for me merge_mode kereeed 8THD = =)ky!")
834
return
835
"CHUNK"; 64, "FUSE_SP_S _PP_SSTATE__PAATE" True,
836
"H": Н, "K": DI0I,
839
"y": D[1], "BLK_K": BLK_K, "BLK_V"; BLK_V, "num_warps": warps "nun_stages"; num_stages,
840
841
843
844
845
848
849
Paot _ite__complle____ wtw.cccleeppceeeseeeeeeeeeeeeeee spaces
"merge mode kernel":
"signature"; merge_ serts__mmd__eeeee_ssggege
"grid": (SK $BLK_K 1) SBLK_K) ((8V $BLK_V 1)
BLX_V) "$8", "NC"),
853
"triton_algo_info"" all_aooooooo:
get_merge_mod__eeree for H, in[
T, H, 8, V)
856
(1, 1024, 8, 164, 641),
857
(1, 2048, 1, [128, 1281),
858
(1, 2048, 16, 1128, 1281),
(1, 2048, 4, [64, 64]),
(1, 2048, 8, [128, 1281),
861
(1, 262144, 32, [128, 128]), (1, 32768, 1, [128, 1281), (1, 32768, 4, [128, 128]),
862
(1, 4096, 1, [128, 1281),
(1, 4096, 1, [64, 64]),
866
(1, 4096, 4, [128, 1281),
867
(2, 1024, 16, [128, 128]),
(2, 1024, 8, [128, 128]),
(2, 131072,
[128, 128]),
870
(2, 131072, 1, [64, 64]),
871
(2, 1310/2, 16, 1128, 1281),
(2, 131072, 16, 164, 641),
(2, 131072, 32, 1128, 1281),
34

---

Published as a conference paper at ICLR 2026
(2, 131072, 32, 164, 6411,
875
(2, 131072, 4, [128, 1281),
876
(2, 131072, 4, [64, 641),
877
(2, 131072, 8, [128, 128]),
(2, 131072, 8, [64, 64]),
(2, 16384, 16, [128, 128]),
880
(2, 16384, 4, [128, 128]),
881
(2, 16384, 8, 164, 641),
882
(2, 2048, 1, [128, 1281),
(2, 2048, 1, 164, 641),
(2, 2048, 16, [128, 128]),
885
(2, 2048, 4, [128, 1281),
886
(2, 32768, 32, [128, 128]),
(2, 32768, 32, [64, 641),
(2, 32768, 8, [128, 1281),
889
(2, 4096, 1, [64, 641), (2, 4096,
890
4, [64, 641),
891
(2, 4096,
8, [128, 1281),
(2, 65536, 1, [128, 1281),
(2, 65536, 1, [64, 64]),
894
(2, 65536, 16, [64, 641),
895
(2, 65536, 4, [128, 128]),
(2, 8192, 1, 1128, 1281),
(2, 8192, 32, 1128, 1281)
898
899
1,
900
})
Striton. autotune(
903
configs=[
904
triton. ConniogggiiggiiB.FF.. BLK_K, "BLK_V": BLX_V, :ISE_TMA:: USE_TMA 1, num_warps=nun_warps, mun_stagps-nunnetagee)
for num_warps in 14, for num_stages in 13, 4] for BLK_K [128]
906
907
908
for BLK_V in [128, 64, 32] for USF_IMA in [True, False]
},
911
key=ll,
912
Striton. jit
def merge_mode__ernel
9,
916
k,
917
v,
g,
chunk_state,
920
scale,
921
chunk_decay,
935
out_0, cu cu_seqlens, NS,
926
MUL NS, chunk_indices, NC,
929
930
931
WUSE_SP_STTTT__UDDTE tl.censtexpr,
H: tl.constexpr,
K: tl.constexpr,
934
V: constexxpr,
935
CHUNK: tl.constexpr,
BLK K: tl.constexpr,
BLK V: tl.constexpr,
35

---

Published as a conference paper at ICLR 2026
USE TMA: tl. constexpr,
940
NUM_BLK (K BLK 1) // BLK_K
941
BUMBLKKY = (V BLKV 1) // BLK_V NUM BLK KV NUM BLK K B
944
i_v, i_h, i_ tl. www.rrorrammddde ss, i = ( indices i_go * 2) (t t1.load
(1.program_id(1) cl.program_id (2)
945
chank
chank_indices indices + i_g 2 t)t.o((..i)))3)
Lou_adde__eeeeeeeeeeee i (t.to(t1.ntt3))
(t...
ti. (ou
947
T ecs bos
948
NC - tliv CHUNK)
FIXME: 'ikk temporarily k
951
952
out_0 = out_ (bos * H 1_h) *
953
if
desc t1. wwammaketenssce descriptor d + i_h K,
956
q bos * shape- K], strides=|E K, 11, shock_shape=
957
958
KI,
960
961
dedec tl.
962
k bos * H K + i_h K,
shape- K],
strides=
11,
965
Block_shape-(CHE (CHUNK, B],
966
967
V. dede t1. mal.ake_tensor. dccnsorrdsscrppoor v 1 V.
bos
shape= V), strides- [H v, 11, block_shape- [CHUKK BIK_V],
970
971
[CHUNK,
chunk_state_des
tl wwl.maakeeeeeee ct._nke_tensoo_deerrrbbb (
974
chunk_stee state + ge + to)tl.ttti))) -H K V
x
V,
shape= V], strides-[V. [V, 11, block_shape- [BLK_K, BLK_V],
977
978
979
if FUSE_SP SUSE_SP_STATE_UDD S___SSTTT_UPPAE_
prev_rank staetsssstteeeeee staat__sm__eeeeeeeeeeeeeeeeeeeee ses_
982
azev_rank_state_ee in K V
i_s stride dl prev rank state_sun
983
shape= [K, Vl, strides- [V, 1],
block_shape- [BIK_K, BLK_V],
986
987
# load 'q': (T, H, K,) => (CHUNK, K,) if not USE USE_TMP TMA:
cur_q q bos HI H+ ptr_g_0 ttl.mace_llockkptrrorr..ooo CHUNK, i_k BIK_K,), (CHUNK, BLK_K,), (1, 0,1) # fat: skip
fmt: skip (I, K,), (H + K, 1,), (i_c
991
blk_g 0 - t1. load (ptr_qgo boundary_cheaet (1, 0,) # fmt: skip
else:
994
0 d_eesc. load ([i_c CHUNK,
K])
995
Load (trans) k (T, H, K, => (BLK K, CHUNK,) it not USE TMA:
36

---

Published as a conference paper at ICLR 2026
cur k bos ptr_k_1 til.mak_bloock_srrtuu K fmt: skip (K, T,), (1, K,), (ikk K, i_c + CHUNK, ), (BLX_X, CHUNK, ), (C, 1,)) A fmt: skip
1000
b = tl. load (ptr_kl, boundary checkkrr 1,)) F fmt: skip
else:
1002
blk_k_1 k dedec load ([i_c CHUNK,
BLK 3LK_K]).trans
1003
1004
Load 'y (I, H, V,) => (CHUNK, BIKVV,
1005
it not USE_IMA
H+ V # imt: skip htl.makeblock_ptt ptrr (T, v,,, (H * v, 1,), (i_c CHUNK, i_v BiK_V,), (CHUNK, BiK_V,), (1, 0,) A fmt: skip
cur_y
1007
ptr_v_2
1008
b - tl. coal ( boundary_check- (1, 0,)) # fmt: skip
1009
else:
blk 2 v_desc. ([i_c CHUNK, iv VI)
1011
1012
call_oxternal_fure func: "chunk_local_cumsi to pro-conpute g
1013
Load 'g': (I, cur_g g bos make__boock_prrrcrrrr.o
(CHUNK,)
1015
i_h
fmt: skip
1016
ptr_q_3
(T,), (H,), (i_c CHUNK, (CHUNK
,), (0,)) # fmt:: skip
1017
b_3 ti. Lo.doad(ptr boundary (oundary_check=(0)))
(0,)) int: skip NITH INI, V,) (BLK BLK_V,)
Load "chunk_state it not USE TMA:
1019
1020
cur_chunk_state - chunk_state (i_go i_s tt int64) * * fnt: skip tl.make_block_pte (cur_ chan__cchkkktteeee V. 1,), (i_k BLK K, i_ BLK_V,,, (BLK K, BLK_V, V,), (1, 0,)) $ fmt skip
v + i_h * K ptr_chunk_state,
V,),
1022
blk buk_chunk_stater
t1.lcad httptttttchh__stteeeee Boundary_check-e
0,)) # fnt: skip else:
1023
1024
blk www.chhuunttteeeeeee
kii BLK i_v
B1)_v1)
1025
1026
comments not available for Cp 'if_beg' if SUSL_SP_STATR_JUDATES load 'prev rank state sum': (NS, H, K, V,) => (BLK K, V,) if not USE TMA:
1028
1029
1030
cur prev. rank_statee www.stttteeseeeeee sun = prev rank state sun + iss
Btr_a_dd___s__reeaaa__tttteee cur _ur_prov_rank_state__ee
i_h K A fmt: skip
1031
pur_pree_rankkstare__uu
21..ake_block_pe
(K, V,), (V, 1,), (1_k + BLK_K, i_ BLK
,, (BLK_K, sik_V,), (i, 0,)) # fmt: skip
1032
Bitprev_rank_ste sttst_stttuu_at
tl.
boundary_check- 0,)) fmt: skip
1033
else:
1034
paep_reea___ce_eeeeee rank state 5 prev_rank_state_eun seteaattessu___eeee_eee load k K, i B[)
prev
1035
1036
load 'chunk_decay"' (NC_NITH_INI, H,) (,) trr_chunk_deeas - chunk_decay (i_ge i_s) to (t1. (nt6))
1037
stride _e__edd__________e_______ ne with_int ini chwnkhhhhhhhhheee decay + i
# fmt: skip
1038
ptr berr_chunk_desee _ecay__eee______e
cur chunkh decay fmt: skip
[kk_chunk_decay 6 t1. chunk decay_6)
1040
to: ('BLK_K' BLK_V')
('BLKW,, (BLK_V")
1011
Bal__pro__rak_sttaee sum 5_foa_t33 = aak_prov_rankk_tttto_oauu
to (t.l.. f(lat())
1042
# mul: ( ('BLK_K)
BL) => ('BLK_K (BLK_V)
1043
mul_tensor_
b1k sappree_ran__sttte___ee
5 floatt322
1044
$ to: ('BIK K', BLKV') => (BBLK V
1045
mul tensor 8 bfloat16 = mul tensor 8. (t..btlttt16)
37

---

Published as a conference paper at ICLR 2026
add: ('BLK KY, "BLK ) ("BLK K', VO) ('BIK "BLK V
1047
add_tensor_10
blk_chunk_statt__ # comments not available for 'bind_var'
1048
1049
blk p1k__chunk_stat__ add_tensor_10
1050
comments not available for op 'end_if'
1051
# to: ("BLK_K', "BIK V() => ("BLK_K* 'BLK_VI
1052
Baauuk_stttt___________ # matmul: ( CHUNK', KL)_K(), (' ('CHU,K', "BLK => ( ('CHUNK', CCHNKK
bik_chunk whwnkhussaaaaaeeecee
1053
1054
matmul_ tl.dot (olk___, bik_k_l) (blk .dtppe)
1055
tril: ('CHUNK', 'CHUNK() => ('CHUNK', , CCHUN''
1056
tril_13 - t1. where atangrange (0, CHUNK) [:, Nonel >= tl. arrnea (0, CHUNK [None, ;,, matmul_12 0)
unsqueeze: ('CHUNK'
('CHUNK', 1)
1058
unsqueeze_
blk g 3 None]
1059
unsquecze: (' CHUNK' ,) => (1, 'CHUNK')
1060
unsquccze_1_15 b_k_g_3 [Nonc, :]
sub: (' CHCHU' 1), (1, CHUNK') => CHUNK', CHUNK')
sub_16 unsqueeze_14 unsqueeze__e 1_15
1063
exp: ('CHUNK' CCHNKY => ( ("CHUNK', 'CHUNK))
1064
exp_17 clubxp(sub_11)
1065
mul: ('CHUNK', CHUNK') ) C'CHU'K', CHUNKY) => (' CHUNK' CHUNKK
mul 18 tril 13 exp 17
1067
to: (' C'UNK'K' "CONK')
(' CHUNN' CHUNK")
1068
[u1_11_bflloat16_1 mul 10l_1 (t1.bfloat16)
1069
matmul: CHUNK', (CHUNK'), ('CHUNK', (BLK_V') => ("CHUNK', BLL__
matmul 10 t1.dot (mul 18 1f__fflcat66___1111 19, blk_v_2)
mul mul_18 @flcat16_19.dee dtype)
1071
# matmul: (' CHUNK', K'), ('BLK_K', VL) => ('CHUNK', "BLK_V
matmul_2_20 2_21 t1. ( bik_chunk_state bfloat1e__111 (
blkgg 0. 0.dtype)
1073
exp: CHUNK', ,) -> ('CHUNK', ,)
1074
exp_1_22 - tl tl.exp (b1k_g_3)
1075
unsqueeze: CHUNK" -> CHUNK' 1)
1076
unsqueezz____3 exp 2_2221 [;, None]
1077
mul: (' CHUNK'"' 'BLK V ('CHUNK' 1) C'UNK' 'BIK
1078 mul 2.24 = matmul 2 21 unsqueeze 23
1079
# add: (' CHUNKKK,, , 'BLK ) C'CHUNK' BBLK V => ('CHUNK', BBLK V')
1080
add_24 25 matmul_1_20 + mu1_1_24
1081
mul: (' CHUNK', BBLKV') ) () ('CHU,K', (BLX
1082
mul_2_26 = add_25 scale # store => (' CHUNK' (BLK_V)
1083
1084
# assume output layout: [T, H, V]
out_0_ty out_o www.ddtepeeeeeeeee element_ty
1086
ptr_out_0 ta.make plock__loockptt_ (cu 0, (I, v,), (*** 1,), c
CHUNK, i_v + BIK_V,), (CHUNX, BIK,V,,, (1, 0)) # imt: skip (ptr_oou mul 2 to (out ty),, boundary_check= 01)
1087
1088
1089
def
launch_merge_man
1090
chunk_decay: torch.Tensor.
1091
v: torch. Tensor,
1092
k: torch. torch.Tensar Tensor,
1093
9: Tensor,
1094
g: Tensor,
1095
prov_rank_statc_sumt torch. Tensor,
1096
chunk_stato: torch. scale,
1097
1098
cu_seqlens: torch. IntTer, cached_results, dist,
1099
1100
fuse_ htt_ssp_pppaeee bool,
1101
-> tensor
1102
т, H, K, V *k. shape, v.shapel
38

---

Published as a conference paper at ICLR 2026
1104
if scale is None:
1105
scale shspee ** -0.5
1106
1107
TODO: new fix chunk_size to 64
CHUNK 64
1109
chunk_indices
deepare_chunk_indicse indices (cu_seqlens, CHUNK)
1110
chunk_offseta chunk thunk_ofisets_aiee ini prepare_chunk thunk__ofseee settsetsswitthh_eeeeesetss ini (cu_seqiens, CHUNK)
(ou_seqlenn CHUNK)
1111
prepare, NS, NC (en Len(cu_seelenn)) seqiens) 1, chunk indices shapel0)
1112
1113
NG_WITH_ out torch. empty_like
1114
1115
use_act os. environ. gee ('FORGE_USE_AN)
(0r)
1116
ACT True if use_ act ave_act.lover. lower () in ['1', 'true', "yes'] else
None
1117
1118
def alloc_fn(size: int, alignment: int, stream: int):
1119
return torch. torch.cmpty (sizc, device=' atype=torch. int8)
triton. seittonsst_aalocccee (alloc_fn)
1122
1123
def grid (meta)
1124
BLK_K K meta BBBK K'
BLK V meta BLK V NUM BLK K (K BLK 1) // BLK K
1126
1127
NUM_BLK_V (V BLK 1) // BLK_V
1128
NUM_BLK_KV NUM WIM_BEK_K NUM JUM_LLK_ NS H_MUL_NS BLK BLK_KV, NC,)
1129
1131
1132
if FORGE_USE. AOT is None:
1135
chunk_decay=chunk decay, v-v,
1135
1136
1137
k-k,
1138
q=q,
g-g,
1140
prev_ rank state ststtte_sssrreeeeeee rank state sum, stride d0 ns _s_ns_pppe prev__e_ee__________eee_eeee_ee_
1141
prev_ rank_rank___ _aa__rrnn__________aa__eeeaee_ee sttsssttattttteeeeeeeeeeeeeestst sum=H+K-V, sannk_sstteechukk... scale=scale, cuLseglens=cuueeggennns !=:,,
1142
1143
1144
1145
1146
NS=NS, K=K, V=V, CHUNK-CHUNK,
1147
1148
11:9
1150
L_MUL_NSEE NS, out out_0=out,
1151
1152
thun_iindices-chu___idiiii
1153
NC=NC, 1UB0B//e_...YY...purrrrrsssssssL
1154
1155
)
1156
1157
else:
1158
1159
from forge. aotgu.attuuiinn import forge_aot_ops
1160
1161
ingo_info trrre_aot_oss.meree_nooom kernel wet.eee_tittttoo__eeeeeeeeeeeeee trittrrnalllt ciittonaalg_inne cero_inffttttt
1162
for k, _ setattr (algo_ino,, _K, _v)
get__merge_mdde_eeree____ine (N3, T, H, items
1163
1164
httpsseeeeeemeeeee_aeeeeeee forge_aot ops.t_ppp..eeeeeeee ceret_opssmerreeeee mode kernel
1165
0, # torch. cuca. stream () cuda strear,
q.data (tr(),
39

---

Published as a conference paper at ICLR 2026
d.data ptr
1168
v.data_ptr(), g. () $
1169
1170
chunk_state
1171
scale, # scale
if
else 0,
# prev prev_rank_state____ sun
1173
Н*K*V, # stride_d0_ne ns_ htt__peeee_ra__stteeeeeee lll.ecay.datappt//
1174
it chunk _ecak else 0, # chunk_decay
# stride dt_ne_wwitteeeeeeeeee_e_________ with ini aani_chunk__ae decay out out.data_ptra (), cut_
1176
1177
o__eeglens.data_ptr//
# cu_seqlens
1178
NS, # NS
1129
H NS, # H MULL chunk NC, # NC algo_info,
flankindices
1181
1182
1183
lechedresults.uppe ({)) return out
1186
1187
1188
def fusea_op():
cached_
1190
ACTO: precompute decay outside
1191
g_cumsum chunk_local_cumsum (gaaamgggNNooo [None, ...l, 64, cuLseqlens-ccu_eegeeee
(0) (0)
updated_states
lamch_fuse_chunkdeeay
g=g_cumsum,
1194
prev_s-prev_s,
1195
y=v,
1196
k-k,
1198
cachecLreelltsccchceLceeatt dist_acan-dist_eccan,
119
1200
lary_update-
1201
1202
chunk_state updated_states| chhnkksstttk)
1203
prev_rank__0 Btank__ttte__su__ sum updated states. (rprev rank state_sunt None)
1204
chunk deccy = updated_states. chunk
chunk chunk_decay', None)
1205
launch merncheeerrreeee mode
1206
chunk_dcay-chnn_ddeayy
1207
vev,
1208
k=k,
1209
q=q,
1210
g-g_cumsum,
1212
chunk_state=-chuk__tttee
1213
scale=scale,
1214
cu aol_seeeenn=ou_seeeenn
1215
1216
fuse_sp_updtte-ffsss______
1217
1218
return e, Stunk_state
Listing Generated Scalar GLA kernel by Forge
G.2 IMPLEMENTATION FOR DELTANET
def Chunk_mede Beank__mdd__dettnneee deltanet (k: Tensor, v: Tensor, b: Tensor) Tensor:
I = forge. identity (k, k. sizi (0)
# Note: forge handles the specific fp32 accumulation T (I pnrgeematmuloout_fp333. (k bl..., Nonel, k.T) .erl..rrl. (-1))
inverse () (k. dtype)
40

---

Published as a conference paper at ICLR 2026
bl..., None] }
bl..., Nonel)
Explicitly cacho results for reuse in other phases
forge. cache_resuht (J, "u")
forge.cache_result (9, "x")
k.T
11
return
12
13
def Secaaymecde cetee_ddltanettt (prev_s: Tensor, k: Tensor, Tensor, chunk_state
Tensor) -> Tensor:
***
15
Inter-Chunk State Propagation
16
# Calculate decay matrix based en cached
decay forge serge.identity (k, k.size (1)) k.T 3 return decay prev_s.tor (decay.
19
Chunk_state
20
def merge_modc_deltancc (q: Tenser, k; Tensor, ui Tensor, w: Tensor, chunk_state: Tensor, scale: Tensor) -> Tensor:
=**
23
Output Merging
24
Note: 'u' and 'w' are automaticallo injected from the cached results
new
chunk state return ((q k.T) (tril((1)) (0) new v q chunk state) scale
27
Listing 3: Forge Implementation of DeltaNet with Intermediate Result Caching
41


Reviews:
[{"review_id": "9BgdLrJAzc", "reviewer": "Reviewer_VUej", "summary": "The AI community has been rapidly innovating on AI architectures, but it it painstaking to obtain hardware-efficient implementations. The kernels need to respect the GPU memory hierarchy, support multi-GPU execution, and remain easy to implement. This work observes that there are a few rules underlying linear attentions and encodes these patterns into the Forge DSL.", "questions": "See above.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This is a very important problem space. Many sub-quadratic models have been fast in theory but not in practice, and the challenge of developing kernels has prevented the community from understanding how to obtain the best quality for a fixed wall clock time.\n- The paper considers a relevant baseline framework – FLA – and shows consistent speed ups.", "weaknesses": "It would be useful to understand how Forge kernels compare to highly optimized kernels, since Triton kernels are routinely slower than CUDA kernels:\n- How does Forge generalize to linear attentions that use very large state sizes (e.g., Taylor approximations like ReBased, Based, Learned feature maps with large feature dimension, Mamba-2 with large state size)?  - Does register management, careful use of wgmma/tcgen05, etc., become important? \n- How do Forge kernels compare to other DSLs (e.g., TileLang, ThunderKittens which have open-sourced efficient linear attention kernels) in the single-GPU setting? It would be useful to have head-to-head speed comparisons. \n\nThe writing would benefit from additional clarity and explanation in certain places:\n- Are the kernels for inference/forwards pass or backwards as well? Does Forge consider decoding kernels as well, or prefill-only? It would be useful to clarify and discuss the scope of Forge in the paper, and how the ideas could generalize. \n- Do the kernels remain numerically stable for training and inference, for instance as compared to the FLA kernels?\n- How do the multi-GPU implementations compare to popular multi-GPU softmax attention implementations? Do we see a speed up from linear attention? \n- It would be useful to provide clearer explanation on where the experimental gains arise from: L370-377 attribute it to Forge’s “system level overheads” and “parallelism strategy” but this is a vague explanation", "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": "C0OGyWkBDg", "reviewer": "Reviewer_PFK1", "summary": "What about the following summary, am I missing anything?\n- Forge introduces a domain-specific language (DSL) for linear attention kernels - this is the key novelty and the abstractions defined there\n- The DSL is compiled to triton kernels. The compiler can do optimizations such as compute communication fusion, and a range of targeted optimization of system bottlenecks\n- in single GPU, forge generated kernels can give an increase of 1.01-4.9x speedup compared to hand-tuned kernels\n- Forge can also generate distributed kernels accross multiple GPUs, which is demonstrated up to 128\n\nOverall, forge is a well thought through engineering solution that can help to develop linear attention kernels with good performance, without the effort of manually implementation", "questions": "I would love to understand what would be required to expand this framework to softmax based attention. The authors discuss associativity at  the beginning of the paper, but associativity also has been exploited for optimizations with softmax attention (see e.g. LeanAttention)", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- well-engineered compiler system that automates generation of linear attention kernels, both for single GPU and distributed kernels\n- integrates well into existing pipelines such as Triton, this will make it easier to adopt.\n- important optimizations for compute- communication fusion, and ahead-of-time (AOT) compilation\n- good empirical results, showing 1.01x-4.9x speedup over FlashLinearAttention on single GPUs and near-linear scaling up to 128 GPUs\n- practical significance for research on models that leverage new linear attention mechanism, reducing manual kernel engineering effort which is one of the limitations to scale more researchy models", "weaknesses": "- limited novelty (the DSL/abstraction for linear attention), but good integration of existing work\n- too narrow. linear attention is not widely used, and mainly for more researchy models. Not relevant for predominant model architectures. This limits the practical use of this framework to a very specific set of research explorations. Expanding this to more common forms of attention would be desirable.", "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": "OmfQs2iTr7", "reviewer": "Reviewer_RnXQ", "summary": "Forge is a DSL + compiler for linear attention. It abstracts linear attention variants to three definitions, and uses a compilation backend with a distributed version of triton-lang to optimize linear attention across computation and communication resources. Forge showcases its efficiency by comparing against Flash Linear Attention, a library of hand-designed kernels for linear attention, and performs on par or better than it.", "questions": "* Why did the authors find this particular abstraction more expressive than the rest?\n* Are there linear attention variants this abstraction does not support? Is there an example of this from recent publications in the community?\n* Why did the authors not include source code with the text?\n* Can you include the chunk, decay and merge sections for the five variants studied in the paper? This would be ideally be in a code repository, but in absence of that it should be added to the appendix.\n* Have the authors read AttentionEngine [1]? Seems like relevant work.\n\n[1] Chen, Feiyang, et al. \"Attentionengine: A versatile framework for efficient attention mechanisms on diverse hardware platforms.\" arXiv preprint arXiv:2502.15349 (2025).", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Thank you for submitting your work. This was a pleasant read and I am optimistic it will be valuable to the community. I found the presentation to be very well-done, and the explanations of how the DSL is translated to intra and inter-GPU parallelized code was interesting to read. The experiments are expansive and the results are promising.", "weaknesses": "There are two issues I am concerned about.\n* The central issue of abstractions is being future-proof. This boils down to 1) whether this abstraction is expressive enough for a wide range of designs, and 2) whether maintaining the abstraction (translation and compilation) is sustainable with hardware changes. These questions are challenging to answer, since to some extent they involve making educated guesses. But I still would have liked to see some discussion about these points in a paper of this sort.\n* I could not find an anonymized code repository in the text. It is odd not to share code for papers with DSL contributions. Please add an anonymized code repository link to the paper text.", "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": "2sDHd4U1RU", "reviewer": "Reviewer_r3XC", "summary": "This paper presents Forge, a domain-specific compiler designed specifically for linear attention mechanisms. Forge abstracts the computation into three canonical stages based on chunk-wise parallelization: intra-chunk computation, inter-chunk state propagation, and output merging. This unified abstraction both simplifies the development of new linear attention variants (such as Mamba, RetNet, RWKV, GLA, HGRN, and GDN) and enables systematic optimizations for both single-device and distributed scenarios. Forge leverages Triton-Distributed as its backend to generate high-performance kernels that fuse computation and communication, and incorporates an adaptive scheduling system for parallelism and a static dispatcher to reduce runtime overhead. Experimental results demonstrate that Forge matches or outperforms state-of-the-art hand-written kernels (e.g., Flash Linear Attention) in both performance and scalability, while also requiring much less manual effort to implement new variants.", "questions": "1. According to the experimental results, Flash Linear Attention already supports multiple linear attention variants and delivers strong performance across them. While Forge’s three-phase abstraction enables users to easily implement different variants, Flash Linear Attention is also Triton-based, which inherently provides customization capabilities. How does Forge compare to Flash Linear Attention in terms of flexibility when supporting diverse linear attention variants?\n2. How does the Parallelism Scheduler mentioned in Appendix B adapt to distributed environments? Are there differences in scheduling strategies between single-GPU, intra-node, and inter-node scenarios?\n3. Are there any limitations or edge cases where the three-phase abstraction might restrict the expression of certain linear attention updates, particularly for future architectures that may involve more complex dependencies?\n4. What is the overhead (if any) of using Forge's compilation pipeline compared to maintaining optimized hand-written kernels, especially as new hardware or Triton versions are released?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Proposes a unified, three-phase abstraction for linear attention that captures a wide range of state-of-the-art variants, enabling easy mapping from high-level algorithms to efficient implementations.\n2. Automates the generation of high-performance, hardware-aware kernels from simple PyTorch code, reducing the barrier for researchers to experiment with or deploy new linear attention models.\n3. Integrates native support for fine-grained, tile-level compute-communication fusion using Triton-Distributed, improving distributed execution and network bandwidth utilization compared to existing methods.\n4. Demonstrates strong empirical results, with Forge-generated kernels achieving up to 4.9x speedup over expert-tuned baselines and near-linear weak scaling up to 128 GPUs and 16 million tokens.\n5. Includes practical system-level optimizations such as ahead-of-time compilation and static dispatching to eliminate runtime overheads, particularly beneficial for short to medium sequence lengths.\n6. Provides evidence that the abstraction is expressive enough to cover a broad family of linear attention models with minimal code changes.", "weaknesses": "1. As a compiler-based framework, the proposed abstraction is tailored specifically for linear attention and lacks extensibility to other attention mechanisms, which limits its general applicability.\n2. Although Forge demonstrates competitive performance compared to Flash Linear Attention, the paper does not sufficiently emphasize its advantages, particularly in terms of implementation flexibility and extensibility relative to Triton-based Flash Linear Attention.\n3. While Forge automates optimization, the actual benefits in terms of code brevity or maintainability versus existing Triton-based manual approaches are not quantified in user studies or qualitative analysis.\n4. The distributed scaling results may be somewhat overstated: for example, in Figure 1, the single-GPU latency is 9.2ms, while 4 GPUs achieve 2.7ms. The real scaling efficiency is about 85%, not precisely \"near-linear\" as asserted.\n5. The distributed experiments (Section 4.2) use H20 GPUs (with slower communication) and BatchSize=4, whereas single-GPU experiments use H100 GPUs and BatchSize=1. This difference in hardware and batch size may exaggerate the benefits of compute-communication overlap and makes the experimental comparison less fair.", "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)