Hardware-Efficient Attention for Fast Decoding
Authors: Ted Zadouri
Hubert Strauss
Tri Dao
Venue: COLM 2025 (2025)
Keywords: Inference
Engineering for large LMs
Compute efficient LM
TL;DR: Incremental decoding slows attention, we propose new hardware-efficient variants that reorganize attention to preserve parallelization and model quality while boosting speed, GPU utilization, and throughput, all with minimal cache.

Abstract: The combination of excessive data movement, an expanding key-value cache, and the limited parallelism inherent in incremental decoding severely bottleneck attention. We explore the design of hardware-efficient attention optimized for LLM decoding. We examine how arithmetic intensity, parallelization, and model quality interact and assess whether the current architecture fully capitalizes on modern hardware. To maximize hardware-effiency, we first propose Group Tied Attention (GTA), a simple attention variant that combines and reuses key and value states to reduce memory transfers during incremental decoding while preserving model quality. We then introduce Group Latent Attention (GLA), a parallel-friendly latent attention combined with low-level optimization designed for fast decoding while showing high model quality. We empirically demonstrate the efficacy of these inference-aware variants in language modeling experiments, showing that GTA matches grouped query attention (GQA) quality with roughly 2x smaller KV cache, and GLA matches multi-head latent attention (MLA) but is easier to shard. Our optimized attention kernel for GLA is up to 2x faster than FlashMLA.

Published as a conference paper at COLM 2025
Hardware--fficient Attention for Fast Decoding
Ted Zadouri
Hubert Strauss'
Tri Dao
Princeton Language and Intelligence, Princeton University Together AI
(tz6037, hs6702, td8762)eprincetonn
Abstract
LLM decoding is bottlenecked for large batches and long contexts by loading the key-value (KV) cache from high-bandwidth memory, which inflates per- token latency, while the sequential nature of decoding limits parallelism. We analyze the interplay among arithmetic intensity, parallelization, and model quality and question whether current architectures fully exploit modern hardware. This work redesigns attention to perform more com computation per byte loaded from memory to maximize hardware efficiency without trading off parallel scalability. We first propose Grouped-Ticn -ied Attention (GTA), a simple variant that combines and reuses key and value states, reducing memory transfers without compromising model quality. We then introduce Grouped Latent Latent Attention (GLA), ferindll latent attention paired with loth-loweeee low-level optimizations for fast decoding while maintaining high
model quality. Experiments show that GTA matches Grouped-Query Attention (GQA) quality while using roughly half the KV cache and that GLA matches Multi-head Latent Attention (MLA) and is easier to shard. Our optimized GLA kernel up 2x faster than FlashMLA, for example, in a speculative decoding setting when the query length exceeds one. Further- more, by fetching smaller KV cache per device, GLA reduces end-to-end latency and increases throughput in online serving benchmarks by up to 2x.
1 Introduction
In light of test-time compute (OpenAT, 2024), inference efficiency now drives progress in AI, demanding greater emphasis on inference-aware architectures. The sequential nature of token-by-token decoding limits opportunities for parallelization. During decoding, Multi-Head Attention (MHA) (Vaswani et al., 2017) caches the key-value (KV) states of all prior tokens. These cached states scale linearly with batch size and sequence length and quickly exhaust high-bandwidth memory (HBM). Moreover, fetching this large KV cache from off-chip memory dominates execution time, significantly outweighing the relatively small computation performed by the matrix-vector workload at each decoding step. Memory fetches increase latency as the KV cache grows, resultnn in prolonged cycles of low GPU utilization (He & Zhai, 2024). This bottleneck hinders wide range of use cases: (i) latency-sensitive interactive applications; (ii) large batch LLM agents for multi-step reasoning (Yu et al., 2024b); (iii) test-time compute scaling (OpenAI, 2024; Snell et al., 2024); (iv) high-thrroughut batch inference; and (v) long-context video modeling (Wu et al., 2024). C
Collectively, these issues highlight the critical need for a
redesign of atten- tion. An ideal attention mechanism should (1) achieve high model quality, (2) scale efficiently across multiple devices, and (3) utilize modern hardware effectively at inference time.
Inference-aware variants of attention accelerate decoding. Multi-Query Attention (MQA) (Shazeer, 2019) caches only a single KV head, significantly reducing memory usage. Grouped-Query Attention (GQA) (Ainslie et al., 2023) offers a compromise by sharing KV heads among smaller groups of query heads for better quality. Multi-head Latent Attention (MLA), recently introduced by DeepSeek (DeepSeek-Al, 2024; 2025), compresses the hidden state into a joint latent vector through low-rank factorized projections and caches the single latent head with large dimension, and then up-projects it before computing attention. These variants, as mentioned, reduce the KV cache, thereby easing MHA's memory-bound performance during decoding by lowering data movement (Ivanov et al., 2021; Ootomo &

---

Published as a conference paper at COLM 2025
Yokota, 2023; Gholami et al., 2024). Regrettably, the computational capabilities of modern accelerators have outpaced the growth of memory bandwidth (Ghose et al., 2018). Arithmetic intensity (Williams et al., 2009), the ratio of arithmetic operations to bytes of memory access (FLOPs per byte), is commonly used to analyze whether workload is memory-bound or compute-bound. MQA increases arithmetic intensity by reusing a single KV head across all query heads, reducing the footprint of KV cache memory and therefore shortening its reload time from high-bandwidth memory while keeping FLOPs constant, but sacrificing quality and parallelism (Pope et al., 2022). GQA slightly increases the arithmetic intensity (proportionately to the group size) and scales efficiently during inference. However, with moderate tensor-parallel degree, each GPU still stores a sizable KV cache, and model quality diminishes for large group sizes. Alternatively, DeepSeek absorbs its low-rank MLA projection matrices during decoding and directly uses the single latent head to compute attention, effectively doubling the arithmetic intensity relative to MQA. However, MLA inherently replicates the latent across all devices, limiting parallel inference.
In this work, we redesign hardware-efficient attention variants through the lens of arithmetic intensity, focusing on scaling during the decoding stage while preserving model quality. We define the group size 8y as the number of query heads per distinct KV head; this group size largely determines the arithmetic intensity. Raising &q boosts the arithmetic intensity and proportionally shrinks the KV cache. However, past a threshold, further increases in &4 begin to trade off higher operations per byte of memory access, effectively GPU utilization, against parallelization, forcing duplication of projection weights and KV cache across devices, diminishing parallel scalability. Guided by these design principles, we propose two attention variants that combine high arithmetic intensity with efficient scaling across devices, comple- mented by low-level optimizations. To translate these design choices into practice, our kernels overlap compute with memory through asynchronous software pipelining and warp spe- cialization, employ cooperative offset calculator for paged KV, and in doing so, keep tensor cores fully loaded, pushing the kernels from being memory-bound towards compute-bound
We initially explore Grouped- -Ted Attention (GTA), which ties the key and value representations into one shared state that is used by small groups of query heads. GTA can reduce the KV cache size and improve the arithmetic intensity by up to a factor of 2 relative to its GQA counterpart with the same group size while preserving quality and parallelism.
We propose Grouped Latent Attention (GLA), parallel-friendly extension of latent attention that benefits from low-level optimizations. GLA achieves a similar quality to MLA and can be up to 2 faster; for example, in speculative decoding, when the query sequence length is two or greater. In online serving benchmarks, GLA reduces end-to-end latency and increases token throughput by up to 2x.
We demonstrate the efficacy of these variants in moderate-scale language modeling experiments trained on FineWeb-Edu. In an XL model (1,47B), GTA achieves perplexity of 10.12 (versus 10.20 for GQA), and GLA reaches a 60.0% average downstream accuracy with 10.21 perplexity (vs. 59.1% and 10.25 for MLA). In large model (876M), GTA achieves 11.2 perplexity and 57.6% average downstream accuracy (improved on GQA's 11.3 and 56. 56).%) For a medium model (433M), GLA yields an average downstream accuracy of 55.4%, slightly above MLA's accuracy of 54.9%.
largemodel
combine their algorithmic design with a collection of wworrrooooooo-lleeellpiiiiittttt leaiiiiiaatoosssiiinnnnn to attention-dccoding kernels 1.2-2x faster than FlashMLA (Li, 2025). We release these opti- mized kernels under a permissive ers. code is available at: https: //githu. com/Dao- con//Da-AILabb la..ab//rrrupdddee -atten--atent--- latent- attention
license to benefit researchers and practition-
2 Preliminaries
2.1 Inference-Aware Attention
MQA (Shazeer, 2019) reduces the KV cache size in MHA, thus accelerating decoding. It caches one KV head shared by all query heads. However, this aggressive sharing sacrifices model quality, and during distributed inference that partitions work at the head level, each device must replicate the single KV head to keep it accessible, thereby negating memory
Benchmarks were performed using the FlashMLA kernel version dated 28 March 2025.
2

---

Published as a conference paper at COLM 2025
Deceding step
beceting
sorks laghs
Lo 3
pary celpet
h: query sirput vict
(L,1,4ed)
(L,1,200)
[1,1,200]
30
Toint lobent
@ Gia (.4d
Di: (ag,
D,.
O,,
D
hg
On.
Og
0
0,
Figure 1: Hemory-loading schematics during decoding of MLA (Left) and GLA-2 (Right) illustrate reduced data movement and higher arithmetic intensity, achieving more FLOPs per byte accessed and easing the memory-bound bottleneck. In MLA, single latent head cVV with d, = 4d is loaded once from HBM to SRAM and reused as K and V for every query head (QK V. In GLA-2, two latent heads, each with de = 2dl, are likewise loaded once and reused as K and V for every query in their groups, eliminating mitigating cache duplication when queries are sharded across devices.
savings (Pope et al., 2022). GQA (Ainslie et al., 2023) addresses these issues by grouping query heads to share a distinct KV head. This approach eliminates most duplication across GPUs at no extra memory cost in distributed inference. As result, it reduces the KV cache size relative to the standard MHA while improving the model quality over the MQA. To preserve quality, it needs moderate number of groups; as result the per-GPU KV cache for long sequences or large batches can quickly exceed HBM capacity when using only few devices, e.g., with a tensor parallelism degree 2. The KV cache per device only shrinks substantially when using many devices, e.g. an eight-way split.
Multi-head Latent Attention (MLA), initially introduced in DeepSeek V2 (DeepSeek-Al, 2024) and repopularized through its use in DeepSeek R1 (DeepSeek-AIl et al., 2025) and FlashMLA (Li, 2025), compresses the hidden state of each token into a low-rank latent vector cVV caches only this vector, then projects it back to full-head keys and values to preserve distinct per-head features. For further reduction, its design only caches a single-head low-rank joint representation for both keys and values instead of separate ones. During decoding, the up-projection of the key wulk is absorbed in the query matrix Wº, and the up-projection of the value WuV in the output matrix WO, As a result, explicit keys or values never materialize; instead, each query attends directly to the latent cVV. MLA keeps positional information outside the compression path by concatenating a small decoupled Rotary Position Encoding (RoPE) (Su et al., 2023) with the latent, allowing the weight-absorption trick. The resulting prefill attention is or (QK + Qropee
2.2 Distributed Inference
Sequential decoding limits parallelization opportunities to only head axis for the attention component. Tensor parallelism (TP) partitions attention layers across devices with frequent synchronization (e.g,, all-gather) of activations. This overhead is mitigated by fast GPU interconnects such as NVLink (NVIDIA Corporation, 2024). TP avoids replicating the entire model per GPU, unlike Data Parallelism (DP) (Li et al., 2020), which becomes impractical for large models. It also circumvents idle GPUs during token-by-token decoding, a límitation seen in Pipeline Parallelism (PP) (Narayanan et al., 2021), which splits the model into sequential stages. Tensor parallelism is preferred during decoding because it distributes weights across GPUs, reducing the bottleneck of loading the KV cache when processing one token at time (Su et al., 2025).
3

---

Published as a conference paper at COLM 2025
3 Methodology
We describe our perspective on designing hardware-efficient attention variants by focusing on the arithmetic intensity. We then demonstrate how maximize the arithmetic intensity while efficiently parallelizing across devices by tying the key and value states and sharding the cached latent latent representatiion
3.1 Arithmetic Intensity in Decoding: A Hardware-Efficieen Perspective
During sequential decoding, the large GEMM workload (General Matrix Multiplication) used during training or prefill shifts to smaller batched GEMV (General Matrix- Vector Multiplication). Every loaded BF16 (2 bytes) element of the cached key matrix performs one MAC (2 FLOPs) with the single-token query element already in registers, yielding a 1:1 FLOP-to-byte ratio. This arithmetic intensity is far below the dense BF16 roofline of an Nvidia TPs
Hopper H100 SXM GPU (NVIDIA, 2022), 299 FLOPs per byte
9889 TFI OPs) leaving the T,35
tensor cores severely underutilized. In essence, the overall latency in distributed inference is limited by whichever is slower: the time to complete all FLOPs at the GPU's peak compute, the time to transfer the necessary data at the GPU's peak memory bandwidth, or the delay from inter-device communication at the available inter-connect bandwidth (Pope et al., 2022; Austin et al., 2025). In practice, with modest TP degree, e.g., eight-way shard, repeated loading of the large KV cache dominates decoding latency (Recasens et al., 2025). The GPU utilization of standard MHA, where the arithmetic intensity is ~~1 1 as shown in Table can drop to as low 7% during decoding (Recasens et al., 2025). In theory, it could accommodate around two and half orders of magnitude more FLOPs without increasing latency.
Attention Variant
GLA-2
GLA
MLA MQA
GQA
GTA
MHA
General Formulation
Db M4 +
Lhg
21hy
T
2.L 24 Mi
Arithmetic Intensity
1+
1+ 3ig
1+ 30g
By 84 24
26p R1
R 288
2 2hg
2 hg
R2 8g
288
x1
281
Table 1: Let be the KV sequence length, hg the number of query heads, hi the number KV heads, and define the group size 8q = iy (queries per KV head). The KV multiplicity is mio E {1,2), with "ko = 1 for shared KV states (K=V) and MRo =2 for distinct KV states (K* V). We assume L > hq.
her
3.2 Design Strategies for Maximizing Arithmetic Intensity
GQA is a general attention formulation that reuses one loaded KV head across each group of query heads. Since sharing the KV head does not reduce the number of operations (each query head still computes its attention scores) but does reduce memory reads, the arithmetic intensity increases proportionally to the group size gy or queries per distinct KV head (see Table 1). For example, in the case of MQA, which reuses a single KV head across all query heads, the arithmetic intensity is approximately the number of query heads, hqy
DeepSeek introduced MLA with an inference-focused design that absorbs low-rank matrices that avoid materializing KV states, thereby achieving high arithmetic intensity through three factors: caching single head latent, reusing the latent for both key and value, and employing a large number of query heads. Initially, MLA loads a single latent head and reuses it across all query heads, similar to MQA. The exact latent representatio tion loaded into on-chip memory now serves both the key and the value states when computing attention, effectively doubling the arithmetic intensity compared to the most aggressive MQA designs (see Figure 1, left, for an illustration of the MLA memory loading schematic). Although the DeepSeek paper (DeepSeek-AI, 2024) does not discuss this point explicitly, using low-rank projections (which reduce the parameter count) to shrink the latent cache also opens the possibility of increasing the column dimension of the up-projection matrices to recover those lost parameters. In other words, the model could reallocate the freed parameter budget to add more query heads, potentially preserving model capacity and increasing arithmetic intensity during decoding by sharing a single latent head across larger number of query heads. As shown in Table 1,
4

---

Published as conference paper at COLM 2025
KV Cache
Grouped-Query Attention
Multi Multi-Query Attention
Grouped-Tied Actention
Keys
Broodccat
Queri lea
byoea:byte Pd
9
9. hg
2.9
Figure 2: Overview of Grouped-Ties -ied Attention (GTA). A single projection produces a fied KV state that serves both key and value. The full tied KV dimension used the value. For the keys, half key dimension comes from the fied vector (no positional encoding applied), and the other half comes from a separate single-head projection (where RoPE is applied); this separate half is broadcast to all heads in the group and concatenated with the tied KV half. GTA roughly doubles the arithmetic intensity and halves the KV cache size relative to GQA with the same number of groups.
the arithmetic intensity depends on the number of query heads. In general, increasing &q reduces the KV cache size and increases the arithmetic intensity, effectively maximizing GPU utilization. Let mm { {1,2} denote the multiplicity of KV, where Mk = when K= (shared states) and Mgo = 2 when K + V (distinct states). In contrast to increasing the group size, &q. raising My from 1 to 2 increases the KV cache and reduces the arithmetic intensity.
hg
KVBytes mk B-L d, x sizeof( (axeof(dtype) dtype)
8q Arithmetic Intensity ~ 2 L.
2
2.89
~
2-hq-
Mike :2g
L
kk
84
Here, B is the batch size and d, is the dimension per head. Moreover, increasing the arithmetic intensity through a larger 84 comes with a trade-off in parallelization capability. This constraint can be quantified by the bounds on 84- Although increasing 8y increases the arithmetic intensity, 84 grows too large relative to the number of devices, it duplicates the parameters and diminishes parallelization gains. Once the duplication factor D equals N, where N is the TP shard count, each machine has a complete copy of the parameters and the KV cache; at that point, the model parallelism does not benefit. For zero redundancy parallelism, the number of KV heads hnge = ht/ 8q should be least N and at most hy, that is, 8q S hy N.
3.3 Hardware-Efficien & Parallelizable Attention
3.3.1 Grouped- Attention (GTA)
Singular-value plots reveal a steep decay in the key cache, where almost all variance is captured by few principal directions, so the keys reside in a low-rank subspace and are highly redundant (Saxena al., 2024). This effect is even more substantial before applying RoPE, where the keys collapse into an even smaller subspace (Yu et al., 2024a; Sun et al., 2024). Meanwhile, multiple studies show that applying RoPE a partial slice of head dimension preserves accuracy, so rotating the entire width yields little additional quality (Black et al., 2022; Barbero et al., 2025). If keys are intrinsically low rank and only slice of each head needs rotation for positional encoding, then rotating every channel and caching the full-rank key tensor wastes memory. Instead, we can rotate just slice of the head dimension required for positional information; the remaining unrotated channels, which tend to be in low-rank subspace and redundant
be shared tied with the value states. GQA already reduces KV cache and memory transfer by letting multiple query heads share one distinct KV head, and it scales efficiently across multiple devices. Building on GQA grouping design, combined

---

Published as a conference paper at COLM 2025
with the above low-rank and partial RoPE insights, we propose Grouped-Tied Attention (GTA), which unifies grouping, ties KV to a single state, and partially applies RoPE, all for the purpose of cutting the KV cache size while retaining quality.
In GTA, similar to GQA, each group of query heads shares one distinct KV head. GTA goes further by tying the key and value projection parameters to yield single state, called the tied KV, whose shape matches that of an individual key or value vector (see Figure 2, for an illustration of the differences between the architecture of GTA and GQA). The value path consumes the full dimensionality of the tied KV state, whereas the key path reuses only its first half as the unrotated half. The remaining RoPE component of the key comes from a separate one-head projection of the hidden state, broadcast across all groups, and concatenated with the unrotated half to form the full key vector. Empirical ablations show that applying RoPE to the shared half degrades quality even when the rotation is later inverted before reusing this half for the value path, so the tied portion is never rotated. After these steps, the query, key, and value states are defined as follows.
R BxLxhxxdd Xdy, KROPE MRREERRXXXXE IRBLXXIX da KNOPE R B x de 2
QE R BxLx KV,K, E
KNOPE
41 KV SAnn
V KV [AnA),
<= Krodeabroa
By tying the KV states, we load a single state into on-chip memory, reuse it for both keys and values, and share it across small set of query heads. This reuse reduces memory transfers, roughly doubles the arithmetic intensity, and halves the KV cache footprint relative to its GQA counterpart with the same number of groups. GQA-4 denotes four distinct key and value heads, whereas GTA-4 denotes four tied KV heads. Experiments show that perplexity (see 5.1.1) and performance in downstream tasks (see 5.1.2) remain comparable to its GQA counterpart.
3.3.2 Grouped Latent Attention (GLA)
MLA's low-rank KV joint compression caches a single latent head of dimension d, = 4 d, per token. Because tensor parallelism partitions the key and value up-projections Wuk wuv E R( d, 4y) in a column-paralla fashion across each device must retain the entire latent to reconstruct the keys and values for its heads. Consequently, the latent KV cache is duplicated in every tensor parallel rank, scaling the aggregate KV cache footprint in proportion to the number of ranks. In contrast, GQA stores 2 2/yy d, elements per token, with a larger KV cache than MLA when de < 2 kk d, and liv > However, for four-degree tensor parallel configuration, MLA and a GQA-8 occupy roughly the same KV cache size per device, although GQA-8 allocates 4x more bytes to its KV cache size from the base model. The MLA design achieves a high arithmetic intensity, but increasing the arithmetic intensity by sharing latent across all heads prevents shard-wise partitioning (see Section 3.2 for details). Hence, the expected per-device memory reduction from head level parallelism diminishes, limiting parallel efficiency in distributed inference.
We propose Grouped Latent Attention (GLA), which compresses tokens into ht latent heads, each with dimension de =2 d, (half of MLA's 4 dz)). During training, every latent head and its up-projection matrices reconstruct distinct key and value features for the query heads in its group. Consequently, the up-projection matrix for one latent head has column dimension 8y du rather than MLA's h2 d,, where &,, = h2y hhe is the group size or queries per distinct latent head. After weight absorption in decoding, each latent head attends only to the query heads in its group. Sharding the latent heads acToss tensor parallel ranks provides head-level parallelism without duplicating the latent KV cache when h2 = TP or reducing otherwise when h, < TP, thereby enabling efficient distributed scaling.
For concrete exposition, we demonstrate GLA with h 2 latent heads, which retains the KV cache size of MLA (4 dæ) but half the KV cache per device when TP2 2. Setting 2, splits the latent vector into two heads, KV and KV and partitions the query heads into two groups. During decoding, each TP rank computes local attention using its assigned latent head and query group, applies its slice of the output projection, and then participates in an AllReduce to sum the partial outputs into the final result. Formally:
6

---

Published as a conference paper at COLM 2025
Qo R 8x1x hy x(24)),
R 24))
KV KV c . E R BxLx24,,
W w7 IR
Oo = softmax (Ro(c))) %
O1 = softmax (RI(CTV)T) )
O OoWy 0 0 AllReduce (oo+0,).
When MLA runs in a TP plus DP hybrid setup to reduce KV cache duplication, uneven sequence lengths can create load imbalance since GPUs that finish short sequences idle until long ones finish processing, hurting latency-sensitive workloads. Small batches below the DP rank also leave many compute units idle. The larger latent head dimension of MLA, 4 d can also exhaust the KV cacheper device for long sequences or large batches, limiting the batch size or context length relative to GLA. For example, GLA-4 shards the latent into four heads, It =4 with 2 d, per head dimension; with the same configuration (TP=4, DP=2), it halves the cache per device, fetches a smaller cache per s het twice the KV cache size. our experiments, we set the lt. = for GLA, the same KV cache size as MLA (d. = 4 d',,, where matches the quality of MLA up to a 1.471 B model (see 5.1 for empirical results) while halving the cache footprint per device when TP> 2. The smaller cache per device also allows GLA to decrease the DP rank while increasing the TP degree, improving tolerance to workload imbalance.
During decoding, the arithmetic intensity of GLA is ~2 8q (double of GQA) as shown in Table 1; GLA-2 reaches ~ Ht FLOP's per byte of memory access, similar to MQA in its most aggressive sharing design, but GLA-2 has better quality. Refer to Figure 1 (right), which shows the memory loading schematic of G when computing attention. GLA maintains high arithmetic intensity, see Figure 14 (roofline analysis), parallelizes efficiently (see Figure 3 (right)) and provides high model quality (see 5.1). We benchmark the token throughput of GLA in Section 5.2; detailed latency throughput benchmarks are given in Appendix C.6
4 System Optimization: Asynchrony & Distributed Offset Calculation
We describe the system optimization to achieve peak performance for MLA, GTA, and GLA modern hardware such as H100 GPUs Thanks to very fast specialized matrix multiplication units such as Tensor Cores, we need careful software pipelining of memory loading and tensor core instructions to always keep the tensor cores busy. This is achieved through the warp specialization technique to exploit asynchrony on modern hardware.
4.1 Asynchrony with software pipelining and warp specialization
We use two techniques overlap compute and memory loading:
Software pipelining: we load the next KV block while the current KV block is being used in computation. This classical technique (Lam, 1988) avoids having the tensor cores waiting for memory loading.
2. Warp specialization: we have separate warps performing memory loading with either TMA (tensor memory accelerator) or asynchronous copy async instruction), and separate warps performing matiix-muliply-acumkke copy
(MMA). The former act as producer warps, while the latter act consumer warps (Bauer et al., 2014). This is commonly used in matrix multiplication (Thakkar et al., 2023) and attention (Shah et al., 2024a). This decoupling simple simplifies the software pipelining, allowing the warp scheduler to overlap the memory loading and compute.
4.2 Distributed offset calculation for paged KV
As new attention variants such as MLA, GTA, and GLA stress both the compute and the memory subsystems, one would have to perform memory loading as quickly as possible. Paged KV (Kwon et al., 2023) has become standard way of storing the KV cache. However, paged KV makes it difficult to use the TMA, a hardware unit that performs address calculation and bound checking to load contiguous blocks of memory. Instead, use the asynchronous copy instruction (cp. async) where each thread separately issues individual load instructions. The challenge is that address computation is surprisingly expensive,
7

---

Published as a conference paper at COLM 2025
as it requires 64-bit integer indexing, which translates to multiple instructions per integer multiplication. We instead have multiple threads in the same warp that cooperate to calculate the addresses. As an example, for head dimension 128, to load a block of size 128 128 from global memory to shared memory, we use 128 threads, with 16 threads loading per row:
1. Group 128 threads into 8 groups, each consisting of 16 consecutive threads. Each group g for g 0,,,,...., will be assigned to load rows gg -
For each thread which belongs to the group 8= [t/16], read the page index from the page table in row g + mod 16) 8. Using the page index, compute the global memory address of the paged KV corresponding to this row and store registers.
For row r in 8.g + 8,..... + 120, all threads f the same group use warp shuffle to get the global memory address from thread index g 16+ (r (r-8)/8-. Use this address to load the KV cache elements corresponding to this row.
We see that each thread only needs to store the address offset of 1 row (instead of 16 rows), as the address offsets of 16 rows assigned to each group are spread across 16 threads.
This enables high efficiency for arbitrary page size (such as page size 1), where page size 1 suffers no slowdown compared to page size 64, despite much larger number of address computations. This unlocks use cases such as prefix caching (Kwon et al., 2023) with RadixAttention (Zheng et al., 2024b) which requires page size 1. We benchmark the speed of paged KV with GLA in appendix C.5, showing a 1.2-1.5x speed up.
Experiments and Results
5.1 Model Quality
5.1.1
Validation Perplexity
In addition to reporting validation perplexity on 100M tokens from the FineWeb-Edu val- idation set, we evaluated perplexity on four additional datasets: Wikipedia and C4 (Colossal Clean Crawled Corpus) portions of RedPaajama v1 (Weber et al., 2024), Cosmopedia (Ben Allal et al., 2024) and Pile (Gao al., 2020), each evaluated using 100M tokens. We also report the average perplexity across these five datasets in Tables and 4. Further validation perplexity results for each dataset are in Appendix C.2.1.
Small (183M)
Medium (433M)
Large (876M)
Method
FineWeb-Edu
Avg.
FineWeb-Edu
Avg.
FineWeb-Edu
Avg.
MLA
16.318
40.290
12.561
28.230
11.363
24.929
GLA-2
16.371
40.604
12.456
27.586
11.293
24.492
GTA-4
16.607
42.680
12.785
29.952
11.232
24.994
GQA-4
16.578
42.520
12.922
30.144
11.340
25.286
MHA
16.715
41.826
12.979
29,990
11.501
25.837
MQA
16.972
43.907
13.068
30.524
11.413
25.206
Table 2: Validation perplexities (lower is better) Wenuee across three model sizes (small, medium, large), along with the average perplexity across five datasets (FineWeb-Edu validation set, Cosmopedia, RPV1 C4, RPV1 Wikipedia, and Pile). The lowest perplexity is in bold, and the second lowest underlined.
GLA-2 tends to improve upon MLA at medium and large scales, achieving lower validation perplexities both the FineWeb-Edu set and the five-dataset average while performing on par withith MLA in small model. Notably, GTA addresses key limitations of GQA and further reduces perplexity medium and large model scales. GLA-2 and 2 bly in the large model, achieving the lowest average perplexities of all variants (24.49-24.5) vs. the next-best MLA (24.93). GTA addresses the key limitation of GQA, namely, the inability to reduce the KV cache per device in lower TP' degree settings. For the large model, GLA-2 and - perform similarly on average, demonstrating strong performance compared to the other variants. This trend holds on the scale of the XL. (1.47B) model, GTA-4 slightly outperforms GQA-4 in perplexity (10.12 vs. 10.20 on FineWeb-
perform compara-
and GLA maintains advantage in perplexity over MIA (e.g., 10.21 vs. 10.25 on FineWeb-Fad see Table 4).
8

---

Published as conference paper at COLM 2025
5.1.2 Downstream Evaluation
We evaluated zero-shot performance on standard benchmarks: SciQ (Welbl et al., 2017),
OpenBookQA (Mihaylov al., 2018), ARC-Easy subset (Yadav al., 2019), HellaSwag (Zellers et al., 2019), PIQA (Bisk et al., 2020), WinoGrande (Sakaguchi et al., 2020), and MMLU
(Hendrycks et al., 2021).
Method
Winogrande
SciQ
PiQA
OpenBookQA
MMLU
HellaSwag
Arc Easy
Avg.
MQA
57.1
$6.0
71.3
37.8
25.7
52.2
64.7
56.4
GQA-4
54.9
87.3
71.2
38.4
25.7
52.2
68.6
56.9
MHA
53.8
86.0
71.4
37.4
25.2
51.5
66.7
56.0
GTA-4
57.3
88.1
72.6
39.8
25.2
52.6
67.5
57.6
GLA-2
55.0
88.1
72.0
40.6
25.6
52.4
69.1
57.5
MLA
54.1
86.8
71.5
40.4
26.0
51.6
66.5
56.7
Table 3: Downstream evaluation for the 876M models (higher is better). dg denotes the dimension of the RoPE. not specified for GTA, GLA, and MLA, then the RoPE dimension dR equals 32.
Method
Fine PPL
Avg. PPL
Avg. Downstream
KV cache (bytes/token)
TP=1
TP=2
MHA
10.311
21.206
60.1
8192
4096
GQA-
10.202
21.073
60.2
2048
1024
GTA-4
10.129
20.823
60.2
1152
640
GLA-2
10.218
21.163
60.0
1152
640
MLA
10.256
21.199
59.1
1152
1152
Table 4: Validation perplexity (lower is better) for the 1.471B model on FineWeb-Edu along with the average perplexity across five datasets. The lowest perplexity in bold, and the second lowest is underlined. Average downstream evaluation (higher is better) across seven datasets, where the highest accuracy is in bold, and the second highest is underlined. TP refers to the tensor parallelism degree. We report the token size bytes per device (lower better) for single layer. The lowest KV cache size is in bold, and the second lowest underlined.
Across downstream benchmarks, our proposed attention variants retain or exceed baseline accuracy. For the scale model 876M (listed in Table 3), GLA-2 produces an average accuracy of 57.5 57.5%, only 0.1 points below GTA-4 and effectively matches GQA-4, indicating that grouping or tying does not degrade quality. This behavior persists on the 1.471B XLallae where GLA-2 reaches an average accuracy of 60.0% compared to 59.1% for MLA, while GTA-4 and GQA- each record 60.2% (listed in Table 4). These results confirm that our hardware-efficient variants preserve improve downstream task performance from medium XL sizes. For more details on downstream evaluation, see Appendix C.2.2. Also, for ablation on small and medium model scales, see Appendix C.3.
5.2 Parallelization
GLA scales across GPUs by partitioning latent heads, reducing memory traffic for faster decoding. MI.A with tensor parallelism (TP) instead duplicates its single latent head with larger dimension on every GPU; to curb this cost, earlier systems fall back on hybrid tensor and data parallelism (DP) MLA, assigning different batch sequences to separate GPUs, benefit seen with large batches or high concurrent requests. We validate GLA's scalability on DeepSeek Coder V2 Base (236B parameters, 21B active) quantized to FP8 and served with our FlashAttention 3 kernels using SGLang framework (Zheng al., 2024b) to benchmark. Experiments compare pure TP on eight H100 GPUs with hybrid TP plus two-way or four-way DP. In the hybrid setup, only the attention submodule is replicated across DP groups; its outputs are all gathered before the MoE feedforward layer and redistributed to mitigate MIAA KV cache duplication. GLA configurations with zero redundancy sharding (latent heads evenly distributed across TP ranks) outperform MLA under an equivalent parallel configuration because of the fetching of smaller latent KV cache per device.

---

Published as conference paper at COLM 2025
Decoding speed, batch 128, query heads 128, query length
DeepSeek Coder V2 Base (2368) FP8
Throughos at Concurent Requests (ligher Better
Output Throughput (tokens//)
1400
1200
2000
GB/s
800
1000
600
200
0-
Sequence length
GLA-8
MLA
GLA-4
MLA
GLA-Z
MLA
Figure 3: Left: Decoding speed of MLA and GLA H100 80GB SMX5 GPU (theoretical BF16 compute 989 IFLOPS/s and memory 3350 GB/s), query length 1 where MLA close being bottle- necked by compute (reaching 610 TFi (PS/s) while GLA has not yet saturated compute (360 TFI TFL.OPS/S Right: Output throughput (higher better) for 64 concurrent requests for live server benchmark where GLA outperforms MLA under identical parallelism scheme. Also, GLA-8 with pure TP==8 outperforms MLA with a hybrid of TP and DP. The prefill/ /ecode sequence length is 8192 4096 respectively.
Figure 3 (right) shows that with 64 concurrent requests, GLA-8 (h, 8, d. 256) on eight GPUs delivers up to 2x the throughput of the single head latent MLA with de 512, as it loads smaller KV cache per device. The advantage persists under hybrid parallelism: GLA-8 with TP=8 outperforms MLL with (TP=2, DP=4), and even with an equivalent hybrid setup of TP and DP, GLA remains ahead when latent heads are sharded without duplication. At the same time, GLA-8 with pure TP=8 can still surpass hybrid TP and DP MLA. The GLA parallelization-frienddl design can lift peak throughput and remain resilient to adverse serving loads. Real-world workloads often feature sequence-length imbalance or small batches that trigger straggler effects and idle GPUs.
5.3
Speed
We benchmark the decoding kernels for MLA (1 latent head of dimension 512, RoPE dim 64) as shown in Figure left and GLA (2 latent heads of dimension 256 each, RoPE dim 64) as shown in Appendix C.7, Figure 13 (Left), with paged KV with page size 64. We compare our GLA implementation with the most optimized MLA implementation we can find, FlashMLA from DeepSeek (Li, 2025). Our GLA kernel is about 20% faster than FlashMLA in the standard decoding setup (query length 1) and more than faster the speculative decoding setup (query length 2). Our GLA kernel reaches up to 93% of the maximum memory bandwidth and 70% of the maximum TFLOPS on the H100 GPU, maximizing the use of both the memory and the compute subsystems.
Conclusion
We demonstrate that focusing on high arithmetic intensity and effective parallelization yields hardware-efficien attention mechanisms while decoding. We introduce Grouped- Tied halving Attention (GTA), which ties key-value (KV) states, halving the KV cache requirement and doubling the arithmetic intensity relative to GQA with the same number of groups. GTA achieves lower perplexity and better average downstream accuracy than the GQA baseline. We then propose Grouped Latent Attention (GLA), a parallelizable inference-award attention variant that achieves high arithmetic intensity and demonstrates decoding speeds up to 2x faster than the baseline DeepSeek FlashMLA (Li, 2025). GLA consistently matches exceeds MLA accuracy on all model scales and benchmarks evaluated. Moreover, sharding the latent heads across devices, with a head dimension smaller than MLA per device, as result, GLA fetches a smaller KV cache during decoding, reducing end-to-end latency and increasing throughput in online serving benchmarks by up to xx. Thus, GTA efficiently replaces GQA (requiring only half the KV cache memory). GLA is a practical replacement for MLA due to its scalable partitioning across GPUs and faster decoding speed.
10

---

Published as a conference paper at COLM 2025
References
Amey Agrawal, Ashish Panwar, Jayashree Mohan, Nipun Kwatra, Bhargav S. Gulavani, and Ramachandran Ramjee. Sarathi: Efficient Ilm inference by piggybacking decodes with chunked prefills, 2023. URL https: /arxiv org/abs//300/ 16369.
Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from 13245.
multi-head checkpoints, 2023.
Reza Yazdani Aminabadi, Samyam Rajbhandari, Minjia Zhang, Ammar Ahmad Awan, Cheng Li, Du Li, Elton Zheng, Jeff Rasley, Shaden Smith, Olatunji Ruwase, and Yuxiong He. Deepspeed inference: Enabling efficient inference of transformer models unprecedented scale, 2022. URL https: arviiv...v.vi org/abs/ /t0/0/002229.000... 00032.
Jacob Austin, Sholto Douglas, Roy Frostig, Anselm Levskaya, Charlie Chen, Sharad Vikram, Federico Lebron, Peter Choy, Vinay Ramasesh, Albert Webson, and Reiner Pope. How to scale your model. Online, 2025. Retrieved from https:/ /jjx-/jl.ggituubbio.. wtw.aiiioosaaaabooooo///
Federico Barbero, Alex Vitvitskyi, Christos Perivolaropoulos, Razvan Pascanu, and Petar Veličković. Round and round we go! what makes rotary positional encodings useful?, 2025. URL.https://arke
Michael Bauer, Sean Treichler, and Alex Aiken. Singe: Leveraging warp specialization for high performance on gpus. In Proceedings of the 19th ACM SIGPLAN symposssu on Principles and practice of parallel programming, pp. 119-130, 2014.
Loubna Ben Allal, Anton Lozhkov, Guilherme Penedo, Thomas Wolf, and Leandro von Werra. Cosmopedia, February 2024. URL https: //ugg///gag/ggggceeeeeee co/
TB/
cosmopedia.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. PIQA: reasoning about physical commonsense in natural language. In The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, New York, NY, USA, February 7-12, 2020, pp. 7432-7439. AAAI AAAl Press, 2020.
Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang, Michael Pieler, USVSN Sai Prashanth, Shivanshu Purohit, Laria Reynolds, Jonathan Tow, Ben Wang, and Samuel Weinbach. Gpt-neox-20b: An open-source autoregressive language model, 2022. URL 06745.
htts://aaxivvorr/abbstto
William Brandon, Mayank Mishra, Aniruddha Nrusimha, Rameswar Panda, and Jonathan Ragan Kelly. Reducing transformer key-value cache size with cross-layer attention, 2024. URL //axs: wttp://orgg///////////// 24005. 19911.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Gastrr, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Rerssrt-Voss,, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, llya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020. URL https::/aarke www.vorgg/aaeee orggabs//2055 14165.
Chi-Chih Chang, Wei-Cheng Lin, Chien-Yu Lin, Yan Chen, Yu-Fang Hu, Pei-Shuo Wang, Ning-Chi Huang, Luis Ceze, Mohamed S. Abdelfattah, and Kai- Chiang Wu. Palu: Compressing kv-cache with low-rank projection, 2024. URL
https://arxiv.or//ab/////e wwg.////////////////s 21118.
Yiting Chen and Junchi Yan. What rotary position embedding can tell us: Identifying query and key weights corresponding to basic syntactic high-level semantic information. In The Thirty-eighthh Annual Conference on Neural Information Processing Systems, 2024. URL
https://openrreviw.. http://rwweeeeeeeefoooooo//////e
11

---

Nils Graef and Andrew Wasielewski. Slim attention: cut your context memory in half without loss of accuracy k-cache is all you need for mha, 2025. URL https://aaxxi..or/abb/////// Published as a conference paper at COLM 2025
Zhuoming Chen, Ranajoy Sadhukhan, Zihao Ye, Yang Zhou, Jianyu Zhang, Niklas Nolte, Yuandong Tian, Matthijs Douze, Leon Bottou, Zhihao Jia, and Beidi Chen. Magicpig: Lsh sampling for efficient I1m generation, 2024. URL https: /topps://arge http://aaxxxvv...bb...sss/
16179.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning, 2023. URL
httpaaaiiii..orbb////br///////// 08691.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithuss 21060
through structured state space duality, 2024. URL
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashatten- tion: Fast and memory-efficient exact attention with io-awareness, 2022. URL hhpp:///axxi..orgaable
14135.
DeepSeek-AL. Deepscek-v2: A www.sntgecooooiicciii. and efficient oitpr langgggs language model, 2024. URL.ntts:.//axxiv.orr////////// 04434.
URL http://///..a///a/rrrrr////////0
DeepSeek-Al, Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, Xiaokang Zhang, Xingkai Yu, Yu Wu, Z. F. Zhibin Gou, Zhihong Shao, Zhuoshu Li, Ziyi Gao, Aixin Liu, Bing Xue, Bingxuan Wang, Bochao Wu, Bei Feng, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, Damai Dai, Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, and et. al. Deepseek-rl: Incentivizing reasoning capability in llms via reinforcement learning, 2025. URL https: /ttp://.///axx.oo////////// abs/2501 12948.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The pile: An 800gb dataset of diverse text for language modeling, 2020. URL wwwpaaiwww..ooggaba.e/aaea
000277.
Amir Gholami, Zhewei Yao, Sehoon Kim, Coleman Coleman Hooper,, Michael W. Mahoney, and Kurt Keutzer. Ai and memory wall, 2024. URL httpss: http:///og//bb///////// 14123.
Saugata Ghose, Abdullah Giray Yağlıkıı, Raghav Gupta, Donghyuk Lee, Kais Kudrolli, WilliamX. X. Liu, Hasan Hassan, K. Chang, Niladrish Chatterjee, Agrawal, Mike O'Connor, and Onur Mutlu. What your dram power models are not telling you: Lessons from a detailed experimental study, 2018. URL //tps//s/////xii..ooo
100118070000 05002.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek
Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, Amy Yang, Angela Fan, and et. al. The llama 3 herd of models, 2024. URL 21783.
https://arxxw.or//bbb////
Inc. Groq. whatt processinn.rooonnnnonoo unit?, 2024. URL https: / /gr.///g.o..oooo.o/ooo com/wp pdf. Groq white paper.
2024. URL upllaass82224/0///rqggoggggggg.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, 00752.
Jiaao He and Jidong Zhai. Fastdecode: High-throughput gpu-efficient I1m serving using heterogeneous pipelines, 2024. URL
abs//033. 11421.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn DawnSong, and Ja- cob Steinhardt. Measuring massive multitask language understanding, In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021, 2021.
Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Michael W. Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami. Kvquant: Towards 10 million context length llm 18079.
inference with kv cache quantization, 2024.
12

---

Published as a conference paper at COLM 2025
Jingcheng Hu, Houyi Li, Yinmin Zhang, Zili Wang, Shuigeng Zhou, Xiangyu Zhang, Beeung-Yunng Shum, and Daxin Jiang. Multi-matrix factorization attention, 2025. URL https://arxiv. https////xxiv.oo//abbbab
arg/ 19255.
Andrei Ivanov, Nikoli Dryden, Tal Ben-Nun, Shigang Li, and Torsten Hoefler. Data movement is all you need: A case study optimizing transformers, 2021. URL
https://arxxi.oo//bbb///// org/abs/2007. .00072.
https://arxiv,
Tao Ji, Bin Guo, Yuanbin Wu, Qipeng Guo, Lixing Shen, Zhan Chen, Xipeng Qiu, Qi Zhang, and Tao Gui. Towards economical inference: Enabling deepseek's multi-head latent attention in any transformer-based Ilms, 2025. URL httttp::/aaxxv..rrrgg awwpbbb......e/eee 14837.
Seijin Kobayashi, Yassir Akram, and Johannes Von Oswald. Weight decay induces low-rank attention layers, 2024. URL htLLhttss///xxiitn www.oogg////...ooooooooo aaw///10...../////// 23819.
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory man- agement for large language model serving with pagedattention, 2023. URL
https://aarii..or/ab////2333 wwtp:/ra////..o0000// 06180.
https://arxiv.r
Monica Lam. Software pipelining: An effective scheduling technique for vliw machines. In Proceedings of the ACM SIGPLAN 1988 conference on Progranming Language design and Implementation, pp. 518-328, 1988.
Jiashi Li. Flashmla: Efficient mla decoding kernels. https://github.. http:/u..ooom/eeegeeeeeeeee www.//eeeeeeeeeeeeeeeeeeeeee.. ai/ 2ashMLA,,025..
Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, and Soumith Chintala. Py- torch distributed: Experiences on accelerating data parallel training, 2020. URL https://arxiv.org/bbb/2266 15704.
Zhenghao Lin, Zihao Tang, Xiao Liu, Yeyun Gong, Yi Cheng, Qi Chen, Hang Li, Ying Xin, Ziyue Yang, Kailai Yang, Yu Yan, Xiao Liang, Shuai Lu, Yiming Huang, Zheheng Luo, Lei Qu, Xuan Feng, Yaoxiang Wang, Yuqing Xia, Feiyang Chen, Yuting Jiang, Yasen Hu, Hao Ni, Binyang Li, Guoshuai Zhao, Jui-Hao Chiang, Zhongxin Guo, Chen Lin, Kun Kuang, Wenjie Li, Yelong Shen, Jian Jiao, Peng Cheng, and Mao Yang. Sigma: Differential rescaling of query, key and value for efficient language models, 2025. URL https://arxiv. absg//bs//aa 2501. 13629.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, 2019. URL
Anton Lozhkov, Loubna Ben Allal, Leandro von Werra, and Thomas Wolf. Fineweb-edu: The finest collection of educational content, 2024. URL https:
//huggingface. fetatttseet//tgggggggeee///////e
Accessed: [date].
Fanxu Meng, Zengwei Yao, and Muhan Zhang, Transmla: Multi-head latent attention is all you need, 2025. URL htttpp://arxxv.orr/abs///// 07864
Meta Al. The Llama 4 Herd: The Beginning of New Era of Natively Multimodal Al Innovation, 2025, URL https: //ap: http:/aaaaaaaaeeeeeeeeeeeee Accessed 29 Apr 2025.
www.m www.mmttt......m wwtp./////oololmaaaa-ao///// -tw.aammmmmmmmmmmmmimm----eeeeee -mmodal-iitelllgennee/.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can suit of armor conduct conductelectricity? new dataset for open book question answering. In EMNLP, 2018.
Subhabrata Mukherjee, Arindam Mitra, Ganesh Jawahar, Sahaj Agarwal, Hamid Palangi, and Ahmed Awadallah. Orca: Progressive learning from complex explanation traces of gpt-4,2022 URL htttps://arxxv..rrraabs///
Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Anand Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phanishayee, and Matei Zaharia. Efficient large- scale language model training on gpu clusters using megatron-lm, 2021. URL 04473.
13

---

Published as a conference paper at COLM 2025
NVIDIA. NVIDIA H100 tensor core 8pu architecture. wrr.g-teenorrrrerrttt2e2--eeeeep
https://resouress..vdia.ooo 2022.
NVIDIA. NVIDIA Blackwell Blacweell architecture technical brief. https://resourcess.vddi.... emrus-blaakell-aacchhttttuuuuu 2024.
NVIDIA Corporation. Nvlink, wtt.tttpp.//ww...aaaaaaaaaaee/rr
nvlink/.
Hiroyuki Ootomo and Rio Yokota. Reducing shared memory footprint to leverage high throughput tensor cores and its flexible api extension library, 2023. URL https://arxiv, wrgg/bs//3/0..... https://arxx..orrr/////88..88 15152.
OpenAI. Openai oll system card, 2024.
Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Anselm Levskaya, Jonathan Heek, Kefan Shivani Agrawal, and Jeff Dean. Efficiently scaling transformer inference, 2022. U
05102.
Pol G. Recasens, Ferran Agullo, Yue Zhu, Chen Wang, Eun Kyung Lee, Olivier Tardieu, Jordi Torres, and Josep LI. Berral. Mind the memory gap: Unveiling gpu bottlenecks in 08311.
large-batch Ilm inference, 2025, URL htttss:/arrxvv.org//aa
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge scale. In The Thirty-Fourth AAAI Conference on Ar- fificial Intelligence, AAAT 2020, The Thirfy-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelli- gence, EAAI 2020, New York, NY, USA, February 7-12, 2020, pp. 8732-8740. AAAI Press, 2020.
Utkarsh Saxena, Gobinda Saha, Sakshi Choudhary, and Kaushik Roy. Eigen at- tention: Attention in low-rank space for kv cache compression, 2024. URL
https://arxiv.cm wttpr/iiw.or//bb/24000. 20550/0620.0..66....
Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention-3: Fast and accurate attention with asynchrony and low- now-precision. Advances in Neural Information Processing Systems, 37:68658-6866 2024a.
Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision, 2024b. URL https:////aa /ttp////iiii.rr///////0.0.0000. 08608.
Noam Shazeer. Fast transformer decoding: One write-head is all you need, 2019. URL
Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Daniel Y. Fu, Zhiqiang Xie, Beidi Chen, Clark Barrett, Joseph F. Gonzalez, Percy Liang, Christopher Ré, lon Stoica, and Ce Zhang. Flexgen: High-throughput generative inference of large language models with single gpu, 2023. URL URL.tttss:/aaxxv.orrr/bb////
06865.
PrajwalSinghaninn Siddharth Singh, Shwai He, Soheil Feizi, and Abhinav Bhatele. Loki: Low- rank keys for efficient sparse attention, 2024. URI.htts:://rriv.rrgg.
2406. 2426202542.......
Charlie Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling 1lm test-time com- pute optimally can be more effective than scaling model parameters, 2024. URL wwt.://g/////......// http:://arri..or//abb///...
https://arxiv..
03314.
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, 2023. URL
https://arxiv.cm wttp:///rgg/bb///ee/ee/ 00..1000000 09864.
Qidong Su, Wei Zhao, Xin Li, Muralidhar Andoorveedu, Chenhao Jiang, Zhanda Zhu, Kevin Song, Christina Giannoula, and Gennady Pekhimenko. Seesaw: tigh-throughput 11m inference via model re-sharding,
https:://axxiv.orr//att 06433.
14

---

Published as a conference paper at COLM 2025
Hanshi Sun, Li-Wen Chang, Wenlei Bao, Size Zheng, Ningxin Zheng, Xin Liu, Harry Dong, Yuejie Chi, and Beidi Chen. Shadowkv: Kv cache in shadows for high-throughput long-context Ilm inference, 2024. URL /tL. wttp://..wog///////////// 21465.
Vijay Thakkar, Pradeep Ramani, Cris Cecka, Aniket Shivam, Honghao Lu, Ethan Yan, Jack Kosaian, Mark Hoemmen, Haicheng Wu, Andrew Kerr, Matt Nicely, Duane Merrill, Dustyn Blasig, Fengqi Qiao, Piotr Majcher, Paul Springer, Markus Hohnerbach, Jin Wang, and Manish Gupta. CUTLASS, January 2023. URİ https: / /gi thub cow com.NWIIDA/cutlass.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, and et. al. Llama 2: Open foundation and fine-tuned chat models, 2023. URL 09288.
http:///axiiv.or/bbb//////
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https://proceedinngg http:///prreeeddnnn.nerrii.....e 3f5ee243547dee91 (bdb5s44 caa
Ben Wang and Aran Komatsuzaki. Gpt-j-6b: a 6 billion parameter autoregressive language moddel,22211
Maurice Weber, Daniel Fu, Quentin Anthony, Yonatan Oren, Shane Adams, Anton Alexan- drov, Xiaozhong Lyu, Huu Nguyen, Xiaozhe Yao, Virginia Adams, Ben Athiwaratkun, Rahul Chalamala, Kezhen Chen, Max Ryabinin, Tri Dao, Percy Liang, Christopher Ré, Irina Rish, and Ce Zhang. Redpajama: open dataset for training large language models, 2024. 12372.
Johannes Welbl, Nelson F. Liu, and Matt Gardner. Crowdsourcing multiple choice science questions, 2017 URL
Samuel Williams, Andrew Waterman, and David Patterson. Roofline: an insight- ful visual performance model for multicore architectures. Commun. ACM, 52 (4):65-76, April 2009. ISSN 0001-0782. doi: 10.1145/14988665149955 https://doi.co http:///dooorr/1/00/ /ttrg//...1////..... 1145/1498765 1498785.
URL
Haoning Wu, Dongxu Li, Bei Chen, and Junnan Li. Longvideobench: A bench- mark for long-context interleaved video-language understanding, 2024. URL
https://arxiv.orgg wttp:vvv..or//b///0/...// 15754.
Heming Xia, Tao Ge, Peiyi Wang, Si-Qing Chen, Furu Wei, and Zhifang Sui. Speculative ietttiinnspeeeaatttii speculative execution for accelerating seq2seq generation, 2023. URL https://arxiv. wttp://.org/bb//////0/ 16487.
decoding: Exploiting
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Chen, Han, and Mike Lewis. stfccinnt language models with attention sinks, 2024. URL https: //axxiv orggabs/2399...../ 17453.
Vikas Yadav, Steven Bethard, and Mihai Surdeanu. Quick and (not dirty: Unsupervised selection of justification sentences for multi-hop question answering. Kentaro Inui, Jing Jiang, Vincent Ng, and Xiaojun Wan (eds.), Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing,
2019, Hong Kong, China, November 3-7, 2019, pp. Association for Computational Linguistics, 2019, doi: 10.18653/v1 v//119-1260. URL htttps:://oe //oi org.
Bowen Yang, Bharat Venkitesh, Dwarak Talupuru, Hangyu Lin, David Cairuz, Phil Blunsom, and Acyr Locatelli. Rope to nope and back again: A new hybrid attention strategy, 2025.
URL htlLtttss///axiv...r www.xivorrg/bbb/551111 18795.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training, 2024.
URL
https://ariiv.og/aas//222.20
15

---

Published as conference paper at COLM 2025
Hao Yu, Zelan Yang, Shen Li, Yong Li, and Jianxin Wu. Effectively compress kv heads for Ilm, 2024a. UR..Kttss//arrav..rrrbb///// 07056.
Yangbin Yu, Qin 2024b. URL httpss//arrxiv.or//ab////// 022023
Qiang Fu, and Deheng Ye. Affordable generative agents,
Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, Y. X. Wei, Lean Wang, Zhiping Xiao, Yuqing Wang, Chong Ruan, Ming Zhang, Wenfeng Liang, and Wangding Zeng. Native sparse attention: Hardware-alignee and natively trainable sparse attention, 2025. URL
org/ /rg//br//500 abs/2502. 11089.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can machine really finish your sentence? In Anna Korhonen, David R. Traum, and Lluís Märquez (eds.)) Proceedings of the 57th Conference the Association for Computational Linguistics, ACI 2019, Florence, Italy, July 28-August 2, 2019, Volume 1: Long Papers, pp.
4791-4800. Association for Computational Linguistics, 2019. doi: 00.18653/v1/P19-114 URL https://doi org/10 10. 20653/v1/P19-14777
Yifan Zhang, Yifeng Liu, Huizhuo Yuan, Zhen Qin, Yang Yuan, Quanquan Gu, and Andrew Chi-Chih Yao. Tensor product attention is all you need, 2025. URL
https://arxiv.orr/abb/2......
Zhenyu Zhang, Ying
Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, Zhangyang Wang, and Beidi Chen. of H20: Heavy-hitter oracle for efficient generative inference of large language models, 2023. 14048.
URL http://arxii..orra//2332332
Lianmin Zheng, Liangsheng Yin, Zhiqiang Xie, Chuyue Sun, Jeff Huang, Cody Hao Yu, Shiyi Cao, Christos Kozyrakis, Ion Stoica, Joseph E. Gonzalez, Clark Barrett, and Ying Sheng. Sglang: Efficient execution of structured language model programs, 2024a. URL
Lianmin Zheng, Liangsheng Yin, Zhiqiang Xie, Chuyue Livia Sun, Jeff Huang, Cody Hao Yu, Shiyi Cao, Christos Kozyrakis, Ion Stoica, Joseph E Gonzalez, al. Sglang: Efficient execution of structured language model programs. Advances in Neural Information Processing Systems, 37:62557-6258. 2024b.
16

---

Published as a conference paper at COLM 2025
A Discussion
We discuss related work, limitations, and some future directions. For related work, Appendix reviews the attention mechanisms that reduce the KV cache during pretraining, surveys broader hardware efficient designs, and extends the related work to cover post hoc strategies for faster LLM decoding, studies on applying partial RoPE, and low-rank projections, insights that shaped our methods.
GTA reduces the unsharded KV cache size by roughly half relative to GQA with the same number of KV heads, mainly by caching one tied state per head. With TP=8 and eight KV
heads, GTA-8 stores 1.5 du per token, where " it comes from separate RoPE, while GQA-8 stores 2 d,, a minimal savings of 25%. However, for eight KV heads and with 4, the gap widens because GTA-8 allocates 2.5 d, per token, whereas GQA-8 uses 4 d, so a moderate TP degree provides greater memory relief. Our GLA experiment utilizes two latent heads to match MLA's unsharded KV cache size, where GLA halves the KV cache per device when TP2 2 while MLA duplicates the KV cache across devices. Up to the 1.471B scale, model quality remains comparable. As next step, it is worth evaluating larger-scale models with more latent heads while keeping the head dimension fixed at 2 d to further validate GLA's performance. For example, the family of Llama 4 models (up to 400B parameters) employ GQA-8 with hg = 40 40 hkp =8, where the KV cache size per token on each device for TP=8 degree is 2 dg, and beyond this degree it would duplicate the cache. A matching GLA-8 with eight latent heads would cache slightly more, 2.5 d per token on each device where d
I comes from decoupled RoPE; whether GLA-8 improves quality over GQA-8 with roughly similar cache budget is an open question. For half of the KV cache, based on our 1.471B scale experiment, GLA-2 achieves 10.218 validation perplexity versus 10.202 for GQA-4 on
FineWeb-Edd Furthermore, the additional 4 per token KV cache footprint of decoupled RoPE can be mitigated by applying RoPE only in partial layers, as shown by Cohere (Yang et al. 2025) and Llama 4 (Meta AI, 2025). We leave these scaling studies for future work.
Section 3.2 shows that the arithmetic intensity scales with the number of query heads. Ap- pendix C.3 presents an ablation that replaces full query and output projections with low-rank versions to cut the parameters. The reduced capacity offset by widening the column dimen- sion to add additional query heads per group in GQA and GTA. The ablation aimed to improve quality (or preserve quality when the KV cache is further reduced) while boosting the arith- metic intensity, effectively improving GPU utilization. The validation perplexity was slightly worse by 0.1 to 0.2 relative to the baselines when ablating across few ranks and query heads. detailed study of this trade-off is left for future work. Finally, exploring how parallelization and arithmetic intensity interact in other architectures, such as Mamba (Gu & Dao, 2024; Dao & Gu, 2024) and Linear Attention (Yang et al., 2024), could further exploit modern hardware.
B Related Work
B.1
Attention Variants
B.1.1
Algorithmic
Follow-up works to DeepSeek's MLA include Multi-matrix Factorization Attention (MFA) (Hu al,, 2025), which resembles MQA but uses a larger head dimension and factorized query projections, where it shares similar limitations with MQA, particularly regarding inefficient KV cache utilization per device due to duplication and lack of compatibility with tensor parallelism. Tensor Product Attention (TPA) (Zhang et al., 2025) factorizes queries, keys, and values with two rank R projection matrices per state, so the per token cache is (Rx+Rv) dh+ elements. We report validation perplexity in our ablations for TPA in the Appendix C.3. The low-rank structure of TPA supports straightforward head-level TP sharding, but its KV cache scales linearly with R and already exceeds MLA once R> 4; ranks four and above are important for quality, especially at larger scales, so can quickly lose the manory-saving advantage.
Pre-training:
Post-hoc adaptation: In (Lin et al., 2025) they propose SIGMA, which utilizes novel differentially scaled QKV module specifically optimized and applied during the fine-tuning stage to improve inference efficiency by compressing K while lightly compressing V and
17

---

Published as a conference paper at COLM 2025
increasing Q Slim Attention (Graef & Wasielewski, 2025), post-training approach that keeps only the key vectors and recreates the values on the fly, cutting the context m for any multi-head attention in half. Our proposed method cuts the KV cache further than both of these approaches. differs from the techniques above, as it attempts to restructure the attention architecture during pre-training, as opposed to changing the existing attention of the already pre-training model. However, similar distillation methods that have been applied to MLA (Meng et al., 2025; Ji et al., 2025) can be adopted for GLA in the post-training stage to realize the benefits of the low KV cache footprint and easy parallelization.
B.1.2 Systems
FlashAttention (Dao et al., 2022) reorders attention computation with an I/O-aware tiling strategy that keeps data in high-speed memory, avoiding the need to materialize the entire attention matrix. It drastically reduces memory overhead and yields significant speedups, particularly during decoding with large sequence lengths. FltshAttentionn22 (Dao, 2023) further refines the attention kernel by reducing non-matrix multiplication operations and improving parallel work partitioning, delivering additional gains in hardware utilization. Its system-level improvements provide notable increase in throughput over the original PlashAttention Then FlashAttention-3 (Shah et al., 2024b) leverages geott-egneattion GPU features, such as asynchronous memory operations and low-precision computation, pushing attention efficiency closer to hardware limits. Natively trainable sparse attention (NSA) (Yuan et al., 2025) introduces hardware-alignen sparse attention with dynamic hierarchical pattern. This design reduces computational complexity while maintaining near-full attention fidelity, enabling efficient decoding over extremely long sequences. Our proposed methods are orthogonal to these system-level attention optimizations.
B.2 Additional Approaches to Accelerating Decoding
The following post hoc adaptation design features work in conjunction with the GLA and GTA design architecture, enhancing its performance capabilities.
Algorithmic: There have been many algorithmic efforts such as token eviction (Zhang et al., 2023; Xiao et al., 2024) or sharing KV cache between adjacent layers (Brandon et al., 2024), batching to improve GPU utilization (Mukherjee et al., 2023), and speculative decoding (Xia et al., 2023). Systems: On the system side, there has been work on quantization (Hooper et al., 2024), CPU offloading (Aminabadi et al., 2022; Sheng et al., 2023; He & Zhai, 2024), and memory management using PagedAtter Attention (Kwon et al., 2023) to mitigate memory fragmentation problems. Hardware: In addition, there have been efforts on the hardware side that benefit inference, such FP4 support (NVIDIA, 2024) or NVLink (NVIDIA Corporation, 2024), to hardware chips designed solely for fast inference (Groq, 2024).
B.3 Low-Rank Projections
Empirical findings indicate that, before applying RoPE, key activations have sharply decaying singular-value spectrum (Yu et al., 2024a; Chen et al., 2024), implying that many dimensions contribute minimally. Furthermore, (Singhania et al., 2024) show that keys exhibit substantially reduced intrínsic dimensionality across models, suggesting an inherently low-rank space. For example, (Kobayashi et al., 2024) finds that the regularization of the weight decay drives the combined key-query mapping to an even lower rank. In contrast, value activations exhibit mixed low-rank tendencies. Some studies have shown that cached values do not compress well without a severe accuracy penalty (Chang et al., 2024; Singhania et al., 2024). However, additional findings demonstrate that partial compression can be achieved with acceptable performance degradation (Saxena et al., 2024; Sun et al., 2024). These inconsistencies suggest that the effective rank of values is model- and method-deependent. GTA utilizes the insights from these aforementioned works for its design of tying the key and value within each query group, thereby reducing the KV cache size.
B.4 Rotary Position Encoding (RoPE)
Per head dimension (Barbero et al., 2025) suggests that it may not be necessary to apply RoPE (Su et al., 2023) to every head dimension since the highest frequencies already provide positional discrimination. Other studies (DeepSeek-AL AI, 2024; Black et al., 2022; Wang &
18

---

Published as a conference paper at COLM 2025
Komatsuzaki, 2021) find that they can preserve the quality of the model by applying RoPE to only half of the dimension. Inspired by the same insight, GTA rotates only partial slice of the head dimension and ties the rest to the portion with half of the value head dimension, reducing the KV cache size while preserving accuracy.
Per layer (Chen & Yan, 2024) shows that RoPE contributes the most in the early transformer layers, where attention is focused on local syntactic relations. In contrast, the deeper layers shift toward semantic cues, implying diminishing returns from positional rotation later in the stack. (Yang et al., 2025) present RNoPE, design that interleaves RoPE layers with NoPE layers and limits RoPE to a sliding window, achieving markedly better retrieval at very long context lengths and influencing the architecture choices of Llama 4 (Meta AI, 2025). Overall, applying RoPE to partial layers is beneficial for GLA, considering that the decoupled RoPE, with single head that is half the head dimension, can be eliminated.
C Full Experimental Results
C.1 Experimental Setup
We empirically validate that our simple methods, GTA and GLA, (1) achieve quality comparable to GQA and MLA, (2) are easily parallelizable, and (3) run efficiently on modern hardware such as the H100 GPU. For example, GLA achieves an upstream and downstream quality similar to that of MLA, yet is easier to shard, and our GLÀ kernel is up to 2x faster than DeepSeek FlashMLA in a speculative decoding setup. GLA,, denotes the configuration in which the query latent is sharded, removing duplication of its down projection across devices and cutting the parameter count per device; we include this variant mainly as an ablation since the query latent is not cached.
We train models on four scales: small (183M), medium (433M), large (876M) and XL (1.471B) parameters on the FineWeb-Edu-DDe ration of the GPT-3 model (Brown et al., 2020) with the Llama architecture (Grattafiori et al., 2024). The small model is trained on 25 billion tokens, whereas the medium, large, and XL models are each trained on 50 billion tokens. We use the Llama 3 tokenizer (Grattafiori et al., 2024) with vocabulary size of 128K tokens. We use the Adam (Loshchilov & Hutter, 2019) optimizer with ($1,82) (0.9, 0.95), a weight decay of 0.1, and gradient clipping at 1.0. We follow the training recipe from Gu & Dao (2024), using learning rate scaled by x relative to GPT-3 for model of the same size, with decay of cosine to 1% of the maximum learning rate. We use the configuration of the GPT-3mode (Brown et al., 2020). We first use the configuration of the GPT-3 model for a given parameter size for our MHA baseline, which has the largest parameter budget. Then, we widen the MLPs of every other attention variant until each model matches the MHA parameter count. Essentially, MHA's parameter size is the anchor point.
dataset (Lozhkov et al., 2024), following the configu-
In GQA-4 & GTA-4, the 4 represents the number of groups or the number of KV heads, Rg. X1
GLA, refers to the GLA version in which the latent query is also sharded.
Model Size
#Param
Micro-batch Size
Batch Size
Learning Rate
#Layer
dmodel
ht
dp
Small
183.65M
16
512
10-4
12
768
12
64
Medium
433.77M
16
512
1.45 10-4
24
1024
16
64
Large
876.55M
8
512
10-
1.2x
24
1536
16
96
XL
1471.12M
8
256
1.0
24
2048
16
128
Table 5: Model configuration for the four model sizes in our experiments. We GPT- 3 model configuration, with the Llama 3 architecture as the backbone and its tokenizer as well.
19

---

Published as conference paper at COLM 2025
Method
Model Param
Intermediate size
MLA
183.65M
2128
GLA-2
183.51M
2208
MHA
183.45M
2048
MQA
183.53M
2520
GIA-4
183.40M
2462
GQA-4
183.53M
2392
Table 6: Model parameters and FFN intermediate size for small model.
Method
Model Param
Intermediate size
GLA-2
433.5 M
3152
GL
433.89M
3280
MLA
433.55M
3062
GTA-4
433.57M
3320
GQA-4
433.77M
3248
MHA
433.77M
2736
MQA
433.77M
3376
Table 7: Model parameters and FFN intermediate size for medium model.
Method
Model Param
Intermediate size
MHA
876.55M
4096
GQA-
876.55M
4864
MQA
876.55M
5056
GTA-4
876.55M
4976
MLA
876.73M
4640
MLA (dR:48)
876.74M
4592
GLA-2 (dg:48)
876.96M
4914
GLA-2
876.73M
4768
G
876.44M
4936
Table 8: Model parameters and FFN intermediate size for large model. dg denotes the RoPE dimension and the default is 32 for this model size
Method
Model Param
Intermediate
MLA
1470.58M
6120
GLA-2
1470.78M
6292
MHA
1471.12M
5464
GTA-4
1471.22M
6638
GQA-4
1470.83M
6486
Table 9: Model parameters and FFN intermediate size for XL. model.
C.2 Quality
C.2.1
Validation Perplexity
Method FineWeb-Edu Cosmopedia RPV1
Pile
RPV1 Wikipedia
Avg
C4
MHA
16.715
20.542
31.628
40.444
99.800
41.826
GQA-4
16.578
20.599
32.059
43.841
99.525
42.520
MQA
16.972
22.094
32.245
44.308
103.915
43.907
GTA-4
16.607
20.768
32.911
42.181
100.932
42.680
GLA-2
16.371
20.542
31.628
40.444
94.037
40.604
GLAy-2
16.333
20.110
31.517
38.725
92.820
39.901
MLA
16.318
20.063
31.484
39.528
94.056
40,290
Table 10: Validation perplexity for the small model (lower is better). The lowest perplexity in bold, and the second lowest underlined. RPV1 refers to RedPajama vl.
20

---

Published as conference paper at COLM 2025
Method
FineWeb-Edu Cosmopedia
RPV1
Pile
RPV1
Avg
C4
Wikipedia
GLA-2
12.456
13.722
24.308
27.676
59.766
27.586
-
12.433
13.917
24.263
28.224
60.359
27.840
MLA
12.561
14.039
24.507
28.602
61.438
28.230
GQA-4
12.845
14.532
25.159
30.401
65.871
29.761
GTA-4
12.785
14.812
25.009
30.447
66.708
29.952
MHA
12.979
14.666
25.331
30.772
66.201
29.990
GQA-4 :4,4:4-d4;h4;44))
12.922
15.024
25.282
31.510
65.980
30.144
MQA
13.068
15.163
25.585
31.504
67.302
30.524
Table 11: Validation perplexity for the medium model (lower is better). The lowest perplexity is bold, and the second lowest underlined. goR refers the rank the low-rank query and output projections, resulting in reduced model parameters. To offset these lost parameters for fair comparison with the baselines, we increase the query heads h2 to 48. It's beneficial since arithmetic intensity depends the number of query heads. RPVI refers to RedPajama vl.
Method
FineWeb-Edu
Cosmopedia
RPV1
Pile
RPV1
Avg
C4
Wikipedia
MHA
11.501
12.605
22.496
27.651
54.933
25.837
GQA-
11.340
12.358
22.219
26.635
53.878
25.286
MQA
11.413
12.437
22.383
26.521
53.274
25.206
GTA-4
11.232
12.159
22.059
26.136
53.383
24.994
MLA
11.363
12.468
22.294
25.685
52.837
24.929
MLA (dg:48)
11.245
12.021
22.053
25.246
52.212
24.555
GLA,-2 48):48)
11.337
12.144
22.234
24.620
52.612
24.589
GLA,-2
11.276
12.100
22.126
24.681
52.371
24.511
GLA-2
11.293
12.106
22.130
24.698
52.233
24.492
Table 12: Validation perplexity for the large model (lower is better). dR refers to the RoPE dimension and the default is 32 for this model size. RPVI refers to Redl'ajama vl.
Method FineWeb-Edu Cosmopedia RPV1
Pile
RPV1
Avg
C4
Wikipedia
MHA
10.311
10.540
20.117
22.432
42.628
21.206
GQA-4
10.202
10.418
19,986
22.642
42.119
21.073
GTA-4
10.129
10.399
19.849
22.184
41.551
20.823
GLA-2
10.218
10.482
20.020
22.298
42.796
21.163
MLA
10.256
10.561
20.041
22.516
42.624
21.199
Table 13: Validation perplexity for the XL model (lower is better). Bold indicates the lowest score in each column; underlined indicates the second lowest. RPV1 refers to RedPajama vl.
21

---

Published as conference paper at COLM 2025
Method
FineWeb-Edr PPL
Avg
Avg Downstream
KV cache (bytes/token)
PPL
TP=1
TP-2
TP=4
MHA
10.311
21.206
60.1
8192
4096
2048
GQA-4
10.202
21.073
60.2
2048
1024
GTA-4
10.129
20.823
60.2
1152
GLA-2
10.218
21.163
60.0
1152
640
640
MLA
10.256
21.199
59.1
1152
1152
1152
Table 14: Validation perplexities (lower is better) for the 1.471B model on FineWeb-Edu along with the average perplexity across five datasets (FineWeb-Edt validation, Cosmopedia, RedPajama v1 C4, RedPajama v1 Wikipedia, and Pile). The lowest perplexity is in bold, and the second lowest underlined. Average Downstream evaluation (higher is better), where the highest accuracy is in bold, the second highest is underlined. TP refers to the tensor par- allelism, and we report the KV cache of token in bytes per device across various TP degrees.
C.2.2
Downstream Evaluation
Method
Winogrande
SciQ
PiQA
OpenBookQA
MMLU
HellaSwag
Are-Easy
Avg
G -
55.2
84.9
70.5
35.6
25.2
47.9
66.3
55.1
GQA-4
52.4
83.6
69.7
36.0
25.5
45.7
64.9
54.0
GQA-4
53.8
85.7
36.2
25.4
46.3
64.6
54.5
GTA-4
54.2
85.5
69.0
34.0
25.9
46.8
64.2
54.2
MQA
55.5
84.6
69.5
37.0
26.2
45,9
60.5
54.2
GLA-2
56.7
84.1
70.3
37.2
26.2
48.2
65.3
55.4
MLA
54.5
86.1
70.2
36.8
25.1
47.2
64.2
54.9
MHA
55.2
84.8
69.3
35.0
25.5
46.2
63.0
54.1
Table 15: Downstream evaluation for the medium model (higher better). Bold indicates the highest score in each column;; underlined indicates the second highest. goR denotes the rank of the low rank query and output projections, set to 4d,. To compensate for the reduced parameter count and ensure a fair comparison with the baselines, we increase the number of query heads hg to 48, which also benefits arithmetic intensity since scales with the number of query heads.
Method Winogrande
SciQ
PiQA
OpenBookQA
MMLU
HellaSwag Arc-Easy
Avg
GLA-2
57.4
91.8
73.9
40.4
26.1
72.1
GQA-
59.0
91.5
74.1
41.6
25.2
58.5
71.6
60.2
GTA-4
58.2
91.0
75.1
40.8
25.3
58.6
72.5
60.2
MLA
56.4
89.5
73.5
39.4
25.3
58.1
71.8
59.1
MHA
60.5
90.7
73.1
41.0
25.9
57.6
71.9
60.1
Table 16: Downstream evaluation for the XL model (higher is better). highest score in each column; underlined indicates the second highest.
Bold indicates the
C.3 Ablations
Different attention variants reduce the learned parameters and the representational capacity per layer; therefore, these saved parameters need to be redistributed elsewhere. For example, Llama 2 (Touvron et al., 2023) increases the width of the FFNs for MQA and GQA in their ablation to make a fair comparison to MHA. In addition, MQA initially proposed to increase the width to match the parameters to MHA (Shazeer, 2019). Meanwhile, DeepSeek-Al (2024) adjusts the depth of the model, increasing the number of layers for a fair comparison. Altering make head-to-head comparisons,
the depth is less common because it is challenging there less flexibility in moderately scaled models to match parameters. (Pope et al., 2022) shrink the head dimension of MHA to match the parameters of MQA and TPA, while (Zhang
22

---

Published as conference paper at COLM 2025
et al., 2025) increases the number of query heads to align the parameter count, essentially distributing the saved parameters into the query projections. For instance, in the case of GQA and GTA, the KV heads need to be divisible by the query heads, so there is less flexibility in terms of altering the number of query heads to match as closely as possible to the baseline for fair comparison. In the case of MIA and GLA, increasing the query heads is beneficial since, during decoding, we do not materialize KV. Essentially, single latent head is shared across all or groups of query heads; therefore, increasing the query heads trivially improves GPU utilization while decoding, as we demonstrated earlier Table where the arithmetic intensity for MLA and GLA boils down to the number of query heads.
C.3.1 Ablations: Small Model 188M Parameters
Attention Type
Model Param
FineWeb-Edu
RPV1
Cosmopedia
RPVI Wikipedia
Pile
Avg
C4
MHA
183.45M
16.71
32.24
20.54
99.79
40.44
41.94
GQA-
174.01M
17.07
32.93
21.62
104.1
45.09
44.16
MQA
170.47M
17.39
33.70
22.53
108.0
46.79
45.68
GTA-
171.95M
17.04
32.95
21.82
103.5
44.67
44.00
TPA (r=2)
172.10M
17.06
32.92
21.81
99.58
44.17
43.11
TPA (r=4)
174.90M
16.90
32.63
21.59
99.43
43.15
42.74
MLA
181.44M
16.40
31.61
95.46
40.04
40.80
-
175.54M
16.56
31.91
20.66
94.49
40.13
40.75
Table 17: We ablate by keeping the width the FFNs (2048) and number of query heads (hg 12)
constant across the attention variants. Validation perplexity for the small model (lower is better). Bold marks the lowest value in each column, and underlined marks the second-lowest. We include the benchmark for Tensor P'roduct Attention (TPA) (Zhang et al., 2025) with ranks 2 and for the low-rank projection matrices of keys and values. Bold marks the lowest value in each column, and underlined marks the second-lowest. RPVI refers to RedPajanm vl.
Method
FineWeb-Edu
RPV1
Cosmopedia
RPV1
Pile
Avg
C4
Wikipedia
GLA,-2 (h,:20))
16.450
31.706
20.849
95.273
40.116
40.879
MLA
16.338
31.516
20.111
94.273
39.168
40.281
GLA,-2
16.337
31.517
20.110
92.820
38.726
39,902
GTA-4 (dg:16)
16.870
32.732
21.089
103.508
43.103
43.460
GTA-4 (gox: :4 24:
16.517
31.946
20.727
99.962
41.462
42.123
GTA-4 (90g 36)
16.496
32.048
20.537
101.030
43.787
42,780
GQA-4 (gog: :3 36)
16.546
32.048
20.786
99.684
43.787
42.570
GQA-4 (40g:4 h4: 24)
16.405
31.754
20.530
97.979
43.302
41.994
Table 18: The ablations are for different query head counts and projection ranks. Given that arithmetic intensity during decoding depends on the number of query heads, we increase the number of query heads, Hg. and the query and output projections are low-rank, denoted by qog, to compensate for the added parameters. dg denotes the RoPE dimension. For instance, dR = 16 for GTA-4, we apply RoPE to only 25% of the head dimensions instead of 50% in our proposed approach. Validation perplexity for the small model (lower is better). Bold marks the lowest value in each column, and underlined marks the second-lowes. RPV1 refers to RedPajama vl.
23

---

Published as conference paper at COLM 2025
Method
Model Param
FineWeb-Edu
RPV1
Cosmopedia
RPV1 Wikipedia
Pile
Avg
C4
MHA
183.45M
12
16.71
32.24
20.54
99.79
40.44
41.94
MQA
183.45M
17.03
32.97
21.79
104.30
44.51
44.12
GQA-4
183.45M
20
16.79
32.44
21.19
101.50
43.05
42.99
GTA-
181.40M
20
16.83
32.58
21.22
104.70
44.46
43.96
GTA-
186.11M
24
16.55
32.07
20.49
99.13
42.78
42.20
TPA (r=2)
183.05M
21
16.74
32.33
22.32
103.60
43.32
43.66
TPA (r=4)
183.68M
19
16.75
32.32
21.63
99.78
41.81
42.46
MLA
183.02M
16.33
31.70
20.84
95.27
39.16
40.66
G
183.51M
13
16.44
31.51
20.11
94.27
40.11
40.49
Table 19: We ablate by keeping the width of the FFNs (2048) constant across different variants, but increasing the query heads, Ha. to match the parameters for fair comparison. Recall that the arithmetic intensity of attention during decoding depends the number of query heads. Validation perplexity for the small model (lower is better) across different numbers of ht and identical FFN width. Bold marks the lowest value each column, and underlined marks the second-lowest.
C.3.2 Ablations: Medium Model 433M Parameters
In the primary experiment, the medium model (433 M) trained on 50B tokens, whereas the ablation studies and baseline within this section are trained on 25B tokens.
Method
Model Param
FineWeb-Edu
RPV1
Cosmopedia
RPV1 Wikipedia
Pile
Avg
GTA-4
433.57M
13.250
25,804
15.051
70.733
31.091
31.19
MLA
433.55M
13.066
25.367
14.515
66.667
30.027
29.93
GLA-2
433.60M
12.985
25.216
14.422
65.730
29.515
29.57
G
433,89M
12,957
25.108
14.434
67.182
29,909
29.92
MLA
433.55M
13.087
25.411
14.582
66.847
29.725
29,93
GQA-4
433.77M
13.395
25.999
15.211
70.403
31.650
31.33
MHA
433.77M
13.552
26.286
15.330
72.488
32.124
31,96
MQA
433.77M
13.574
26.436
15.615
72.789
33.513
32.39
TPA (r=2)
433.77M
13.186
25.612
15.186
68.709
31.269
30.79
TPA (r=4)
433.96M
13.143
25.538
14.672
66.877
30.396
30.12
Table 20: The width of the FFN is modified to match parameters closely possible across variants. They are all trained on 25B tokens. Validation perplexity for the medium model (lower is better). Bold marks the lowest value in each column, and underlined marks the second-lowest. We benchmark TPA using low rank key and value projection matrices at ranks 2 and 4. RPV1 refers to RedPajama vl.
24

---

Published as conference paper at COLM 2025
Method
Model Param
Winogrande
SciQ
PiQA
OpenBook
MMLU
HellaSwag
Are
Avg
QA
Easy
TPA (r=2)
433.77M
53.8
84.1
68.3
36.3
25.8
45.4
63.5
TPA (r=4)
433.96M
51.8
83.7
68.6
35.4
25.2
45.5
65.7
53.6
GTA-4
433.57M
55.2
84.3
69.2
34.9
26.0
45.3
63.1
54.0
MLA
433.55M
54.0
82.9
69.3
39.9
25.4
46.1
65.4
54.7
GLA-2
433.60M
55.5
83.8
70.0
35.2
25.5
46.2
66.6
54.6
-
433.89M
54.3
85.9
37.2
46.3
63.8
54.5
MLA
433.55M
53.2
83,9
69.3
39,9
25.4
45.9
64.3
54.5
GQA-4
433.77M
55.0
82.8
69.2
34.5
25.3
45.0
63.6
53.6
MHA
433.77M
51.7
85.5
69.3
35.6
25.4
44.2
62.8
53.5
MQA
433.77M
51.5
83.7
68.3
37.4
25.7
44.4
62.6
53.3
Table 21: The width of the FFN is modified to match parameters closely possible across variants. All models are trained on 25 B tokens. Downstream evaluation for the medium model (higher is better). Bold indicates the highest score in each column; underlined indicates the second highest. benchmark TPA using low-rank key and value projection matrices at ranks 2 and 4. RPVI refers to RedPajama v1.
Method
Model
hg
FineWeb-Edu PPL
Avg
Avg Downstream
KV Cache (bytes/token)
Param
PPL
TP=1
TP=2
GLA-2
434.73M
26
13.236
30.358
53,9
576
320
MQA
433.77M
31
13.703
33.022
53.4
256
256
GQA-4
433.77M
28
13.567
32.019
52.6
1024
512
GTA-4
428.26M
28
13.401
31.475
53.6
576
320
-
434.76M
32
13.321
30.909
53.4
576
320
MLA
434.32M
23
13.249
30.875
54.2
576
576
MHA
433.77M
16
13.552
31.956
53.5
4096
2048
TPA(r=2)
433.47M
29
13.367
31.171
53.2
744
624
TPA(r=4)
432.59M
26
13.404
31.030
54.0
1440
1232
Table 22: We run ablation by keeping the width of the FFNs (2736) constant across different variants but increasing the query heads, hy, to match the parameters for fair comparison. Recall that the arithmetic intensity of attention during decoding depends on the number of query heads. We report the validation perplexity (lower is better) for Fine WeneWeb-Eu, along with the average perplexity across five datasets: FineWeb-Edu validation set, Cosmopedia, RedPajama v1 C4, KejPajama v1 Wikipedia, and Pile. The lowest perplexity is in bold, and the second lowest is underlined. We report the average downstream evaluation (higher is better), where the highest accuracy is in bold, and the second-highes is underlined. TP refers to the tensor parallelism, and we report the KV cache of token in bytes per device across various TP degrees. We benchmark TPA using low-rank key and value projection matrices at ranks 2 and 4. RPVI refers to RedPajama v1.
25

---

Published as conference paper at COLM 2025
Method
Model Param
hy
FineWeb-Edu
RPV1
Cosmopedia
RPV1 Wikipedia
Pile
Avg
C4
GLA-2
434.73M
26
13.236
25.710
14.665
67.764
30.416
30.358
MQA
433.77M
31
13.703
26.721
15.732
75.651
33.305
33.022
GQA-4
433.77M
28
13.567
26.341
15.514
72.141
32.534
32.019
GTA-4
428.26M
28
13.401
26.036
15.217
71.095
31.628
31.475
-
434.76M
32
13.321
25.859
15.085
69.604
30.677
30.909
MLA
434.32M
23
13.249
25.734
15.089
69.569
30.735
30.875
MHA
433.77M
16
13.552
26.286
15.330
72.488
32.124
31.956
TPA(r=2) 433.47M
29
13.367
25.936
15.322
70.091
31.138
31.171
TPA(r=4)
432.59M
26
13.404
26.059
15.409
69.805
30.473
31.030
Table 23: We ablate by keeping the width of the FFNs (2736) constant across different variants but increasing the number of query heads, hg, to match the parameters for fair comparison. Recall that the arithmetic intensity of attention during decoding depends on the number of query heads. Validation perplexity for the medium model (lower is better). Bold marks the lowest value in each column, and underlined marks the second-lowest. There is less flexibility for GQA and GTA to match parameters since the number of KV heads kkp needs to be divisible by the hg. We benchmark TPA using low-rank key and value projection matrices at ranks 2 and 4. RPV1 refers to RedPajama v1.
Method
Model
Winogrande
SciQ
PiQA OpenBook
MMLU
HellaSwag
Are
Avg
Param
QA
Easy
GLA
434.73M
52.5
84.1
69.2
36.3
25.6
45.4
64.3
53.9
MQA
433.77M
54.6
83.9
67.9
34.7
25.8
43.7
63.6
53.4
GQA-4
433.77M
52.2
82.3
68.0
34.9
24.9
43.9
62.6
52.6
GTA-4
428.26M
53.4
85.1
68.2
34.9
25.9
44.8
63.6
53.6
GLA,-2
434.76M
53.2
83.7
68.4
35.2
25.1
44.8
63.6
53.4
MLA
434.32M
69.5
35.4
24.4
44.9
65.9
54.2
MHA
433.77M
51.7
85.5
69.3
35.6
25.4
44.2
62.8
53.5
TPA(r=2) 433.47M
51.9
84.3
68.9
35.0
25.1
44.7
62.1
53.2
TPA(r=4)
432.59M
52.9
83.3
68.4
38.2
25.8
44.7
65.1
54.0
Table 24: We run ablations by keeping the FFN width (2736) constant across variants while increasing the number of query heads h, to match parameters for fair comparison. The arithmetic intensity of attention during decoding depends on h2. Downstream evaluation for the medium model (higher is better). Bold indicates the highest score each column; underlined indicates the second-highest. GQA and GTA have less flexibilitty because hipp must divide hg. TPA is benchmarked with low-rank key and value projections at ranks and 4. RPV1 RedPajama vl.
26

---

Published as conference paper at COLM 2025
Method
Model Param
Winogrande
SciQ
PiQA
OpenBook
MMLU
HellaSwag
Arc
Avg
QA
Easy
GLA
434.73M
52.5
84.1
69.2
36.3
25.6
45.4
64.3
53.9
MQA
433.77M
54.6
83.9
67.9
34.7
25.8
43.7
63.6
53.4
GQA-4
433.77M
52.2
82.3
68.0
34.9
24.9
43.9
62.6
52.6
GTA-4
428.20 M28
53.4
85.1
68.2
34.9
25.9
44.8
63.6
53.6
-
434.76M
53.2
83.7
68.4
35.2
25.1
44.8
63.6
53,4
MLA
434.32M
53.9
85.6
69.5
35,4
24.4
44.9
65.9
54.2
MHA
433.77M
51.7
85.5
69.3
35.6
25.4
44.2
62.8
53.5
TPA(r=2) 433.47M
51.9
84.3
68.9
35.0
25.1
44.7
62.1
53.2
TPA(T=4)
432.59M
52.9
83.3
68.4
38.2
25.8
44.7
65.1
54.0
Table 25: We run ablations by keeping the FFN width (2736) constant across variants while increasing the number of query heads hg match parameters for fair comparison. The arithmetic intensity of attention during decoding depends on Mg. Downstream evaluation for the medium model (higher is better). Bold indicates the highest score in each column; underlined indicates the second-highest. GQA and GTA have less flexibility because lite must divide hq- TPA is benchmarked with low-rank key and value projections at ranks 2 and 4. RPV1 = Redl'ajama v1.
C.4
Per Token KV Cache Size per Device
Method
KV cache per Token
KV cache per token per Device (2 GPUs)
KV cache per token per Device (4 GPUs)
KV cache per token per Device (8 GPUs)
MHA
64dy
32d;
16d
8d₃
GQA-4
16dn
8d;
4d
2d
MQA
2dp
2d
2d
2d,
MLA
4.5d
4.5d
4.5d
4.5d
GLA-2
4.5d
2.5d
2,5d,
2.5d
GTA-4
8.54;
4.5d
2.5d
1.5dg
Table 26: An example of KV cache per token for llama 8B model configuration with :2 and hkp across various TP degrees. d,, denotes head dimension.
C.5 Speed
We show here that our technique (distributed offset calculation) significantly speeds up the attention kernel when using paged KV. Typically, the attention kernel speed slows down when the page size is small since there is more overhead of address calculation (Kwon et al., 2023). However, a smaller page size reduces fragmentation and unlocks new use cases, such as prefix caching (Zheng et al., 2024b). We benchmark the speed of the decoding kernels for GLA (2 latent heads of dimension 256 each, RoPE dimension 64) with paged KV, as shown in Figure 4. We compare page size 1 and page size 64, with without distributed offset calculation. With distributed offset calculation, page size does not suffer from the slowdown, matching the speed of page size 64. On the other hand, without distributed offset calculation, page size is 1.3x slower than page size 64. We see that the distributed offset calculation gives speedup of 1.2x for page size 64 and speedup of 1.5x for page size 1.
C.6 End-to-End Latency and Throughput
Live server setup. We use SGLANG, production-orienter service framework, and run every experiment in its live server mode, ensuring that HTTP parsing, dynamic queueing, and GPU kernel invocation are timed together (Zheng al., 2024a). Evaluating this mode exposes the queueing overhead and network latency absent from offline testing, thus revealing how GLA and MLA behave under real deployment constraints. The load generator sends 1280 prompts with a chosen concurrency limit, which controls the number of active requests at once. The server combines these active requests into small batches so the limit affects load pressure rather than the fixed batch size. We used the pre-trained weights of the Coooseee-CCoer-V2 Base (236B parameters with 21B active parameters),
the fly,
27

---

Published as conference paper at COLM 2025
Decoding speed, batch 128, query heads 128, query length 2
Page size 1, no dist. offset calc
3000
Page size 64, dist. offset calc
Page size 1, w/ dist. offset
Page size 64, w/ dist. offset calc
2442
2460390
244
212
1023203
2000
GB/s
37V
1433
1000
1K
2K
4K
8K
16K
32K
64K
Sequence length
Figure 4: Decoding speed of GLA on H100 80GB SMX5 GPU (theoretical max BF16 compute 989 TFLOPS/s and memory 3350 GB/s), for query length 2, with BF16 format. Distributed off- set calculation gives 1.2-1.5x speedup, allowing page size 1 to match the speed of page size 64.
quantized to FP8, and served with FlashAttention3 kernels. For the benchmarks, set the page size 64. To simulate GLA, we restructure the MLA latent dimension to GLA with randomly initializing weights since we benchmark performance, not accuracy, in this phase. We also employ chunked -preiiiissss (Agrawal et al., 2023), with tile length of 8192 tokens, and run the prefill kernel one block at time. Decode batches are formed independently, prefill tokens never mix with decode tokens by default.
Parallelism and metrics. Every transformer block sharded across eight GPUs with tensor parallel, while the MoEs feedforward layers are further partitioned by expert parallel. We also benchmark mix of data parallelism and tensor parallelism, and whenever data parallelism is enabled, only the attention submodule is replicated across data parallel groups. Its outpuls are all gathered before the MoEs ferd-forward layer, then redistributed to mitigate the KV cache duplication of MLA. We benchmark broad spectrum inference workloads to assess GLA and MLA under both identical parallelism configurations and in cases where GLA employs only tensor parallelism, while MLA combines tensor and data parallelism. We report four service-level metrics: end-to-end (E2E) latency, time-to-first token (TTFT), inter-token latency (ITL), and output throughput. All values the figures are sunmarized by their median, which is less sensitive to heavy-tail behavior in large-scale interactive systems, Additionally, we provide the mean values the tables.
C.6.1 Tensor Parallelism: GLA MLA
In this configuration with TP degree across x8 H100 GPUs, GLA-8 employs eight latent heads, where each token has to cache latent dimension of 256, whereas MLA maintains a 512-dimensional latent cache duplicated across devices. Both methods have decoupled the RoPE dimension of 64. Figures 5 and Table 27 reveal consistent gains for GLA-8 at every load level. With 16 concurrent requests, GLA-8 red redu the median end-to-end latency from 136 to 117 seconds, à reduction of approximately 15%, while increasing token throughput by approximately 17%. When the concurrency limit rises to 64, GLA-8 completes in 179 seconds compared to 381 seconds for MLA, cutting latency by 53% percent; the first token now arrives after 12 seconds rather than about minutes, and throughput grows by about 70% to 1461
28

---

Published as conference paper at COLM 2025
DeepSeek Coder (236B)
DeepSeek Coder Base (236B) FP8
Median E2E Latency (lower better)
Token Throughput (higher better)
Output Throughput (tokens/))
1400
Latency (seconds)
500
1200
400
1000
300
800
600
100
128
128
Concurrent Requests
Concurrent Requests
Figure 5: Median end-to-end latency (left), lower is better, and output throughput (right), higher is better, of MLA and GLA-8 under pure TP' service eight GPUs. Prefill/Decoud are fixed at 8 K/4 K tokens as concurrency is swept over 16, 64, 128. Because GLA-8 stores roughly half the KV-cache per token under TP degree of 8, it fetches less data during decoding and consistently outperforms MLAA.
tokens per second. Even with 128 concurrent requests, GLA-8 still reduces latency by around 24% and maintains a throughput lead of nearly 60%. These advantages stem from the smaller KV cache footprint of GLA-8 per device, which reduces memory traffic, allows more active requests to fit on the GPUs, and shortens the waiting time before computation can begin.
Median E2E Latency
Median
Median ITL
Output
Prefill/Decode length
Max conc.
Method
TTFT
Throughput (token/s)
/#Prompts
(s)
(s)
(ms)
GLA-8 (TPS)
8K/4K
16/1280
116.83
3.74
27.10
561.06
MLA (TP8)
8K/4K
16/1280
136.23
3.98
31.77
481.09
GLA-8
8K/4K
64/1280
179.32
11.96
38.16
1460.61
MLA (TP8)
8K/4K
64/1280
381.13
192.70
43.03
858.95
GLA-8 (TP8)
8K/4K
128/1280
432.54
223.09
45.99
1362.84
MLA (TP8)
8K/4K
128/1280
572.20
392.07
43.04
858.69
Table 27: Median service-level metrics for MLA and GLA on x8 GPU TP server; the table reports end-to-end latency, time to first token, inter-token latency, and output throughput at concurrency limits of 16, 64, and 128. GLA surpasses MLA on every measure, cutting latency by more than half and lifting throughput by roughly 70% at the mid-load point of 64 concurrent requests.
Mean E2F Latency
Mean
Mean
Output
Prefil//Decode length
Max conc.
Method
TTFT
ITL
Throughput
/#Prompts
(s)
(s)
(ms)
(token/s) /s)
GLA-8 (TP8)
SK/4K
16/1280
116.80
3.59
27.64
561.06
MLA (TP8)
8K/4K 4K
16/1280
136.21
3.83
32.32
481.09
GLA-8 (TP8)
8K/4K
64/1280
179.45
11.94
40.90
1460.61
MLAA (TP8)
SK/
64/1280
301.57
118.99
44.58
858.95
GLA-8 (TP8)
8K/4K
128/1280
370.62
168.84
49.27
1362.84
MLA (TP8)
8K/4K
128/1280
589.07
406.52
44.58
858.69
Table 28:
Mean service-level metrics for MIA and GLA-8 on x8 GPU TP server; the table reports end-to-end latency, time to the first token, inter-token latency, and output throughput at concurrency limits of 16, 64, and 128
29

---

Published as a conference paper at COLM 2025
DeepSeek Coder (236B) Median E2E Latency (lower better)
DeepSeek Coder Base (236B) FP8
Token Throughput (higher better)
TP2
GLA-2
Output Throughput (tokens))
250
2500
Latency (Seconds)
225
2000
200
1500
175
150
1000
125
128
128
Concurrent Requests
Concurrent Requests
Figure 6: Median end-to-end latency (left), lower is better, and output throughput (right), higher is better, of MLA and GLA under expert parallelism across 8 GPUs with 4 DP groups solely for attention. Prefill/eeoode Decode are fixed at 8K/4K tokens as concurrency is swept over 16, 64, 128. GLA outperforms MLA consistently under various concurrencies.
C.6.2 Data Parallelism + Tensor Parallelism: GLA vs. MLA
Figures 6 and Table 29 demonstrate how the balance between compute capacity and memory traffic changes when parallel data attention is introduced. Under mixed TP 2 + DP 4, as shown in Figures 6 and Table 29, GLA-2 with two latent heads each with dimension 256, shortens the median end-to-end latency from 137 to 120 seconds and increases throughput from 477 tokens per second to 544 tokens per second at light load of 16 concurrent requests. At 64 concurrency, the advantage grows to roughly 16% lower latency (196 s relative to 166 s) and 19% higher throughput of 1334 tokens per second vs. 1584 tokens per second. Similarly, in mixed TP + DP 2, as shown in Figures 7 and Table 31, GLA-4 consistently outperforms MLA under various concurrent requests.
However, when the limit reaches 128 requests, as shown in Figures 8 and Figures 9, MLA with the hybrid of TP with degree 2 and DP with degree 4, overtakes GLA-8 in pure TP with degree 8 by using the extra replicas to spread the batch and saturate all compute units; MLA now delivers about 56% more tokens per second (2122 tokens per second vs. 1363 tokens per second) and finishes roughly 43% earlier (247 seconds relative 433 seconds). The cross-over occurs because the added compute lanes of data parallelism offset its cache duplication overhead once the server is heavily loaded. In contrast, GLA-8 in pure TP has already reached the memory bandwidth ceiling and cannot scale further, demonstrating that data parallelism is useful only at large concurrency.
Median F2F Latency
Median
Median
Output Throughput (token/s)
Prefil/Decoode length
Max conc.
Method
TTFT
/#Prompts
(s)
(s)
(ms)
GLA-2 (TP2, DP4)
SK/
16/1280
120.43
5.56
27.95
543.77
MLA (TP2, DP4)
8K/4K
16/1280
137.33
5.92
31.97
477.30
GLA-2 (TP2, DP4)
8K/ [K/4K
64/1280
165.86
14.12
35.01
1583,51
MLA (TP2, DP4)
8K/4K
64/1280
196.47
14.78
42.35
1334.18
GLA-2 (TP2, DP4)
8K/4K
128 128/1280
211.98
25.32
40.90
2474.20
MLA (TP2, DP4)
8K/4K
128 1280
246.81
26.93
49.12
2121.88
Table 29: Median service-level results for GLA-2 and MLA when both run with eight-way tensor parallelism and four-way data parallel attention. GLA-2 shows lower latency and higher throughput at the two lighter loads; at the heaviest load, MLA narrows the gap, but GLA-2 still leads by about 14% on latency (lower is better) and throughput (higher is better).
30

---

Published as conference paper at COLM 2025
Mean
Mean
Mean
Output
Prefill/Decode length
Max conc.
Method
E2E Lateney
TTFT
ITL
Throughput
/#Prompts
(s)
(s)
(ms)
(token/s)
GLA-2 DP4)
8K/4K
16/1280
120.51
5.18
28.16
543.77
MLA (TP2, DP4)
8K/4K
16/1280
137.29
5.50
32.18
477.30
GLA-2 (TP2, DP4)
8K/4K
64/1280
165.52
14.20
36.95
1583.51
MLA (TP2, DP4)
8K/4K
64/1280
196.46
14.76
44.37
1334.18
GLA-22 DP4)
8K/4K
128/1280
211.86
25.39
45.53
2474.20
MLA (TP2, DP4)
8K/4K
128/ 1280
247.04
26.57
53.84
2121.88
Table 30: Mean service-level results for GLA-2 and MLA when both run with eight-way tensor parallelism and four-way data parallel attention. GLA-2 shows lower latency under various metrics.
Median E2E Latency
Median
Median
Output
Prefill/Decode length
Max conc.
Method
TTFT
ITL
Throughput (token/s)
/#Prompts
(s)
(s)
(ms)
GLA-((TP4 DP2)
SK/4K
16/1280
118.34
4.51
27.48
553.29
MLA (TP4, DP2)
8K/4K
16/1280
135.86
4.71
31.66
482,42
GLA-4 (TP4, DP2)
8K/ /K
64/1280
170.66
12.80
36.07
1542.96
MLA (TP4, DP2)
8K/4K
64/1280
205.39
13.36
44.51
1276.25
GLA-4 (TP4, DP2)
8K/4K
128/ 1280
222.36
23.87
43.36
2357.85
MLA (TP4, DP2)
8K/4K
128/1280
462.03
237.35
49.66
1341.89
Table 31: Median service-level results for GLA-4 and MLA when both run with eight-way tensor parallelism and four-way data parallel attention. GLA-4 shows slightly lower latency (lower is better) and higher throughput (higher is better) at the two lighter concurrent requests; at the heaviest load GLA-4 performs significantly better than MLA.
Mean
Mean
Mean
Output
Method
PrefillDeccode
conc. /#Prompts
E2E Latency
TIFT
ITL
Throughput
length
(s)
(s)
(ms)
(token/s)
GLA-4 (TP4, DP2)
8K/4K
16/1280
118.44
4.21
27.89
553.29
MIA (TP4, DP2)
8K/4K
16/1280
135.84
4.44
32.09
482.42
GLA-4 (TP4, DP2)
8K/4K
64/1280
169.87
12.77
38.36
1542.96
MLA (TP4, DP2)
8K/4K
64/ 64/1280
205.37
13.38
46.88
1276.25
GLA-4 (TP4, DP2)
8K/4K
222.30
23.87
48.46
2357.85
MLA (TP4, DP2)
8K/4K
128/1280
380.39
165.82
52.40
1341.89
Table 32: Mean service-level results for GLA-4 and MLA when both run with eight-way tensor parallelism and four-way data parallel attention.
31

---

Published as conference paper at COLM 2025
DeepSeek Coder (236B)
DeepSeek Coder Base (236B)
Median E2E Latency (lower better)
Token Throughput (higher better)
Output Throughput (takemsks)
Latency (seconds)
2000
300
1500
200
1000
128
128
Concurrent Requests
Concurrent Requests
Figure 7: Mean service-level metrics for MLA and GLA on x8 GPU TP server, the table reports end-to-end latency, time to first token, inter-token latency, and output throughput at concurrency limits of 16, 64, and 128 with fixed prefill/ decode length 8K/4K
DeepSeek Coder V2 Base (2368) FP8
DeepSeek Coder V2 Base (236B) FP8
Throughppt Concumert Roquists HHighar Bettir)
Throughout at 128 Concuriont Requeiss (H glle Bettar
Output Throughout (tokens/S)
Output Throughout hhuu (tokens/S)
1400
1200
600
1000
500
CLA-8 1P8
MLA
CLA-4
MLA
GLA-2
MLA
GLA-8 TP8
MLA
GLA-4
MLA
GLA-2
MLA
TP4 DP2
TP2
TP4 DP2
TPM
TP2 DP4
Figure 8: Token throughput 64 concurrent requests (left) and 128 concurrent requests (right), where higher is better. The prefill/ decode sequence length is 8192 $192/4096. 4096. GLA outperforms MLA under equivalent parallelism configurations, but MLA with hybrid of TP and DP at the 128 concurrent requests has higher throughput than GLA under pure TP.
Median E2E Latency
Median
Median
Output
Method
Prefill/Decode length
Max conc.
TTFT
Throughput (token/s)
/#Prompts
(s)
(s)
(ms)
GLA-2
32K/48 /KK
16/1280
166.18
18.11
32.47
394.76
MLA (TP2, DP4)
32K/ 4K
16/1280
188.37
36.13
34.02
347.88
GLA-2 (TP8)
64K/4K
16/1280
219.90
61.94
35.70
224.29
MLA (TP2, DP4)
64K/4K
16/1280
313.68
118.37
37.00
208.59
Table 33: Median service-level results for GLA-2 only with eight-way tensor parallelism and MLAA under mix parallelism scheme with eight-way tensor parallelism and four-way data parallel attention. GLA-2 has 14% higher throughput (higher is better) relative to MLA for prefill length of 32K while 7% higher throughput for 64K prefill length.
32

---

Published as conference paper at COLM 2025
DeepSeek Coder Base (2368)
DeepSeek Base (236B)
Median EZE Larency Concument Noquests [Lower Bettor)
Median EZE Lasency at 128 Concurint Requests faaee Setter
600
500
Latency (seconds)
300
Latency (seconds)
400
250-
200
300
150
200
100
100
GLA-8
MLA
TP8
TPS
TP4 DP2
TP4 DP2
TP2 DP4
TP2 DP4
TP8
TPS
TP4 DP2 TP4 DP2
TP2 DP4 TP2 DP4
Figure 9: Median E2E latency at 64 concurrent request (left) and 128 concurrent request (right), where higher is better. The prefill decode sequence length is 8192/4096 GLA outperforms MIA under equivalent parallelism configurations, but MLA with hybrid of TP and DP at the 128 concurrent requests has higher throughput than GLA-8 under pure TP.
DeepSeek Coder (2368) Median Latency (Lower Better)
DeepSeek Coder (2368) FP8 Throughput (Higher Better)
MLA
Output Throughput [tokens/s]
502
Latency (seconds)
150
100
1.5
350
100
Prefill Length
Prefill Length
Figure 10: Median end-to-end latency (left), lower is better, and output throughput (right),
higher is better, under TP and DP groups solely for attention and expert parallelism for GLA and MLA respectively, across 8 GPUs. Demonstrating long-context with moderately high number of concurrent requests.
Mean
Mean
Mean
Output Throughput (token/s)
Prefill/Decode length
Max conc.
Method
E2E
Latency
TTFT
ITL
/#Prompts
(s)
(ms)
GLA-2
32K/4K 4K
16/1280
166.00
18.09
36.12
394.76
MIA (TP2, DP4)
32K/40 4K
16/1280
188.37
31.25
38.36
347.88
GLA-2 (TP8)
64K/4K 4K
16/1280
291.90
112.39
43.63
224.29
MLA (TP2, DP4)
64K/4K
16/1280
314.16
102.16
51.77
208.59
Table 34: Mean service-level results for GLA-2 only with eight-way tensor parallelism and MLA under mix parallelism scheme with eight-way tensor parallelism and four-way data parallel attention.
C.6.3 Data Parallelism: Workload Imbalance
The random ratio parameter is a fraction of the minimum length that the benchmarks' random-request generator may assign to any individual prefill or decode sequence. For example, with random ratio of 0.125 and sequence length of 4096 tokens, each request created with lengths drawn uniformly from the integer range of 512 4096 tokens, giving every batch a consistent lower bound while retaining a realistic spread of sequence sizes. The random ratio is applied to prefill sequence lengths (131K) and decode (4K). The experiments
33

---

Published as a conference paper at COLM 2025
in this section demonstrate workload imbalance with varying sequence lengths across the batch, which can leave GPUs idle.
In Figure 11 and Table 35, we demonstrate where for long prefill of 131K and relatively long decode of 4K, where the sequence length uniformly sampled within the batch, GLA-8 with TP degree 8 has about 2.7x higher throughput than MLA in hybrid TP with degree 2 in four data parallel ranks. Because every NCCL collective in a data-parallel group must be entered by all ranks, one replica that is still busy with very long sequence forces every other replica, and its tensor-parallel shards, of that single straggler, with pure TP-8, there is no extra data- data-parallel barrier, so only the eight shards that hold the weights pause for one another; long sequence slows that shard group, but leaves the rest of the cluster working, keeping GPU utilization much higher.
wait, so throughput collapses the speed
Median E2E Latency
Median
Median
Output
Method
Prefill/eccode
Rand. ratio
Max conc. /#irrompts
TTFT
ITL
Throughput
length
(s)
(s)
(ms)
(token/s)
GLA-8 (TP8)
131K/4K
0
4/1280
80.21
6.42
25.10
101.59
MLA (TP2, DP4)
131K/4K
4/ 1280
203.04
32.03
28.64
37.50
GLA-8(TP8)
131K/4K
0.125
4/1280
89.57
7.58
25.45
100.68
MLA(TP)) DP4)
131K/4K
0.125
4/1280
233.69
38.54
28.66
37.20
GLA-8 (TP8)
32K/4K
0.125
4/1280
55.97
1.14
22.32
165.78
MLA (TP2 DP4)
32K/4K
0.125
4/1280
73.82
3.29
25.73
125.31
Table 35: With a random ratio 0, each request chooses its prefill and decode lengths uniformly from a single token up to the maximum lengths. With random ratio of 0.125, the lengths are sampled uniformly, but now the range starts 12.5% of the maximum specified length. GLA-8 with pure TP has higher throughput and lower median end-to-end latency than the hybrid TP DP MIL configuration across both long and moderate context settings.
Mean E2E Latency
Mean
Mean
Output
Method
PrefillDecoode length
Rand.
Max conc.
TTFT
ITL
Throughput
ratio
/#Prompts
(s)
(s)
(ms)
(token/s)
GLA-8 (TP8)
131K/4K /
4/1280
80.93
7.78
35.54
101.59
MIA (TP2, DP4)
131K/4K
0
4/1280
219.43
41.82
86.31
37.50
GLA-8 (TP8)
131K/4K
0.125
4/1280
91.73
8.72
35.93
100.68
MIAA (TP2, DP4)
131K/4K
0.125
4/1280
248.35
46.89
87.21
37.20
GLA-8 (TP8)
32K/4K
0.125
4/1280
55.66
1.19
23.59
165.78
MLA (TP2, DP4)
32K/4K
0.125
4/1280
73.65
3.79
30.26
125.31
Table 36: With random ratio 0, each request draws its prefill and decode lengths uniformly from one token up to the maximum lengths. With a random ratio of 00.15,,, the range begins at 12.5% of the maximum length. Across both long and moderate context settings, GLA-8 in pure tensor parallel form sustains higher throughput and lower mean end-to-end latency than MLA that combines tensor and data parallelism.
34

---

Published as conference paper at COLM 2025
DeepSeek (2368)
Latency
DeepSeek Coder V2 Base (2368) FP8: Taken Throughput
Wandoad Imbalance from Variablle Longth (ower Setten
Lengh
GLA-B (TPO)
MLA (TP2 DP4)
GLA-8 (TP8)
MLA (TP2 DP4)
200
Latency (seconds)
Output Throughput
150
30-
Prefil Length UE1. Length U[16K, 131K]
Prefill Length- Ut3, 131K]
Prefill Length- U[10K, 131K]
Length
Figure 11: Median end-to-end latency (left), lower is better, and output throughput (right), higher is better, where the sequence length can vary and is sampled from uniform distribution. GLA using pure TP outperforms MLA with hybrid TP and DP.
p99
p99
p95
Output
Prefill/Decode length
Rand. ratio
Max conc.
Method
E2E
Latency (s)
TTFT
ITL
Throughput
/#Prompts
(s)
(ms)
(token/s)
GLA-8 (TP8)
131K/4K
0
4/1280
175.47
19.91
26.78
101.59
MLA (TP2, DP4)
131K/4K
0
4/1280
566.48
117.44
30.80
37.50
GLA-8 (TP8)
131K/4K
0.125
4/ 1280
182.06
19.98
26.98
100.68
MLA (TP2, DP4)
131K/4K
0.125
4/1280
572.05
119.69
30.77
37.20
GLA-8 (TP8)
32K/4K
0.125
4/1280
99.08
2.49
23.49
125.31
MLA (TP2, DP4)
32K/4K
0.125
4/
135.87
8.61
27.48
165.78
Table 37: For ninety-ninth percentile values of latency, TTFT, and ITL (lower is better), GLA-8 with pure tensor parallel remains faster for the extreme long-context workload, while MLA with hybrid parallelism shows higher output throughput in the moderate context run.
C.6.4 Latency Sensitive Workloads
In latency-sensitive workloads, the predominant objective is minimize end-to-end response time, particularly time to first token, to meet strict service level objectives rather than to maximize aggregate throughput. Because larger batch can increase the queueing and prefill delay, latency-sensitee of throughput to deliver faster responses. In Table 38, GLA-8 with pure TP at eight degrees manages to reduce latency by x2 and cut the time to first token by almost x4 relative to MLA with hybrid of TP and DP, where it is necessary to mitigate the duplication of the KV cache.
serving keeps the batch size very small, at the expense
Median E2E Latency
Median TTFT
Median ITL
Output Throughput (token/s)
Prefill/Decode length
Max conc.
Method
/#Prompts
(s)
(s)
(ms)
GLA-8 (TP8)
64K/256
3/1280
24.60
12.96
24.54
31.17
MLA (TP8, DP4)
64K/256
3/1280
54.25
46.76
28.14
14.14
Table 38: Under latency-sensitee scenarios, GLA with only tensor parallelism outperforms MLAA with mix of TP and DP solely for attention for long context short decode scenarios by over 50% for both end-to-end median latency (lower is better) and output throughput (higher is better).
35

---

Published as conference paper at COLM 2025
DeepSeek (2368)
DeepSeek Coder Base (236B) FP8
Median E2E Latency (lower better)
Token Throughput (higher better)
Output Throughput (tokenshs)
Latency (seconds)
30
1000
800
20
10
400
16K
32K
16K Decode Sequence Length
32K
Decode Sequence Length
Figure 12: Demonstration of MLA and GLA for TP for degree of on long decode tasks. With 256 number of prompts and 32 concurrent requests across various decode sequence lengths with fixed 2K prefill sequence length.
Mean
Mean
Mean
Output
Prefill/Decode length
Max conc. /#Prompts
Method
E2E Latency
TIFT
ITL
Throughput
(s)
(s)
(ms)
(token/s)
GLA-8 (TP8)
64K/256
3/1280 3/1280
24.62 54.26
12.76
46.51
31.17
MI.A (TP2, DP4)
64K
46.47
30.55
14.14
Table 39: Mean service-level results for GLA and MLA. GLA shows lower latency (lower is better) than MLA under various metrics.
C.6.5 Decode Heavy Workloads
In decode-heavy workloads, the generated continuation is so long that the sequential decode phase dominates wall-clock time, resulting in latency and memory bandwidth for the KV cache being the primary bottleneck. Since the model will be performing sequential decoding most of the time, batching offers minimal benefit. In Figure 12, where there is short prefill of 256 and long decoding of up to 32K, with GLA-8 and MLA across eight-degree parallelism, GLA-8 can generate up to 2.5x higher throughput.
C.6.6 Small Context and Short Chat
Small Context describes inference requests in which both the prompt and the generated continuation are very short relative the model context window. For example, a voice assistant answers brief query in single response. In Table 40, GLA-8 with eight latent heads in eight-degree parallelism has a lower latency relative MLA with hybrid TP with degree 2 and DP with four data-parallel ranks since it is single batch setting, the GPUs in the three out of four DP rank parallel groups remain idle, and GLA-8 has fetch half the KV cache per layer; therefore, it outperforms MLA.
Median E2E Latency (s)
Median
Median ITL
Output Throughput (token/s) /s)
Prefill/Decode length
Max cone.
Method
TTFT
/#Prompts
(s)
(ms)
GLA-8 (TP8)
256/128
1/1280
2.49
0.11
18.72
51.45
MLA (TP2, DP4)
256/128
1/1280
2.91
0.12
21.94
43.96
Table 40: Under short chat scenario where there is usually one concurrent request, GLA with only tensor parallelism has 17% higher throughput (higher better) than MLA with mix of tensor parallelism and data parallelism
36

---

Published as conference paper at COLM 2025
Mean
Mean
Mean
Output
Prefill/Decode length
Max conc.
Method
E2E Latency
TIFT
ITL
Throughput
/#Prompts
(s)
(s)
(ms)
(token/s)
GLA-8 (TP8)
256/128
1/1280
2.49
0.11
18.73
51.45
MLA (TP2, DP4)
256/128
1/1280
2.91
0.12
21.95
43.96
Table 41: Mean service-level results for GLA and MLA. GLA shows lower latency (lower better) than MLA under various metrics
Median B2E Latency
Median
Median ITL
Output Throughput
Method
Prefill/Decode length
conc.
TTFT
/#Prompts
(s)
(s)
(ms)
(token/s)
GLA-8 (TP8)
2K/2K
8/1280
47.18
0.86
22.54
346.92
MLA (TP2, DP4)
2K/2K
8/1280
56.35
0.82
27.04
290.91
Table 42: For moderate size prefill and decode sequence lengths with moderate number of concurrent requests, GLA with only tensor parallelism has roughly 19% higher throughput (higher better) than MLA with mix of tensor parallelism and data parallelism.
Mean
Mean
Mean
Output
Method
Prefil/Deccode length
conc. /#Prompts
E2E Latency
TIFT
ITL
Throughput
(s)
(ms)
(token/s)
MLA (TP2, DP4)
2K/2K 2K/2K
8/1280
56.37
0.81
27.12
290.91 346.92
GLA-8 (TP8)
8 8/1280
47.22
0.82
22.67
Table 43: Mean service-level results for GLA and MLA. GLA shows lower latency (lower is better) than MLA under various metrics
C.7 Kernel Execution Time
We benchmark the latency of the attention kernels in these two settings on H100 GPUs (ignoring communication overhead) in Tables and 45. GLA with = can be 1.3-1.5 times faster than MLA with DP in these settings.
Seqlen
MLA (DP)
GLA (TP=2)
Seqlens in batch
MLA (DP)
GLA (TP=2)
2048
15.0 js
16.1 μs
1024 +15+ [192]
23.8 μs
25.4 us
8192
20.8 us
19.1 μus
1024 +15+ 16384
29.8 jts
26.2 jts
32768
35.9 us
27.6 μs
1024 +15 32768
41.1 Jts
30.6
131072
81.0 js
55.0 μs
1024 15 66536
56.0 jts
42.6 uts
Table 44: Attention kernel latency (us) for
Table 45:
Attention kernel latency (μs) with H100 GPUs (8B model), imbalanced workload
MLA vs. GLA on two GPUs with batch=1
37

---

Published as conference paper at COLM 2025
NVIDIA GPU Hardware Evolution
Decoding speed, batch query heads 128, query longth
Peak BF16 Compute Arithmetic Intensity
3000
CLA (ours)
Blockwell $209
Performance (TFLOPSIS)
2000
Blackwell 8100
GB/S
2000
1500
1000
1000
500
Ampere A100
Volta
1K
2K
16K
32K
64K
100
150
200
250
300
Sequence length
FLOP.byte
Figure 13: Left: Decoding speed of MLA and GLA on H100 80GB SMX5 GPU (theoretical max BF16 compute 989 TFLOPS/s and memory 3350 GB/s), for query length 2. At query length 2, GLA saturates compute (700 TFLOPS/S s) and memory (3030 GB/s). Right: Peak BF16 theoretical peak FLOPs (TFLOPS/s) versus the arithmetic intensity for successive NVIDIA GPUs (Volta V100 with FP16). Performance computing has historically grown faster than bandwidth, with the H100 architecture (NVIDIA, 2022), which has the most drastic FLOPs-to-byte ratio increase relative its predecessor. The decoding workload lies far left, so every device, even Blackwell B200, stays memory-bound bound and reaches only few percent of its nominal TFLOP rate.
H1.0 GPU (BF16)
H100 GPU (BF18)
Peak Compute: 989 TFLOPs. Peak Bancleidth: 3.35 TB/s query secten:
Peak Compute: 989 TELOPs, Peak Bandwidth: 3.35 TB/s quary seeten:
Attanable Campo Conpiaatioaall Perfurmance (IFLOPOI))
Gol imeads 123 weaed, 20-26)
Computational reof
Attalnale Computational Perfermanse (TFLO
Gon 120 biheds, meaes c--30l
Computational reof
M34-321 meaadt reads jond)
MOG-328 ro peras medd, brnd
pr122, shead
rehoeds 125 reeck
Nok-z inbeadi.edd rhead, Jam2)
Mukid wwwede.geeaee ge16 ie cheads (wil
140-10
MuA-2 irhendt. 16. mesddihowl
NO6-22 Inboadieedd choud, Juni)
1: bandwilth roof
to bandwiath rof
G04.134
144-122
104
Arithmetic Intensity (FLOPs: Byta)
Arithmetic Intensity (FLOPs Byte)
Figure 14: Roofline analysis of BF16 decoding on single H100 80GB SXM5. In this figure only, the numeric suffix (e-g., GQA-128) indicates the number of query heads hGp elsewhere in the paper, denotes ht.. Left, Lg =1: With 128, MLA attains an arithmetic intensity of ~ 2- 2-hh-hhh200 226,, near the compute roof of ~ 295 FLOPs/byte of H100, whereas GLA-128 with two latent heads remains on the I/O roof with arithmetic intensity of ~ ft 1128 similar to MQA. Right, 2 e.g. in speculative decoding setting when query length 2, for hy : 128 climbs beyond roof and becomes compute bound, while GLA with two latent heads, sits at the inflection point, and can run up to x faster.
38


Reviews:
[{"review_id": "M0Av5p5aP5", "reviewer": "Reviewer_R9T9", "summary": "This paper propose novel attention variant GTA(Group Tied Attention) and GLA(Group Latent Attention) for decoding efficiency. Multi-Head Attention (MHA) is known for low arithmetic intensity during decode stage, GQA (Grouped Query Attention) and MLA (Multi-Latent Attention) increase arithmetic intensity. However, MLA do not support tensor parallelism when used in distributed inference because the number of key/value head is only one and can not be further sharded on head dimension. The GLA proposed in this paper addresses this issue by combining GQA and MLA, splitting head dimension into groups and breaks the data dependency of different groups to make GLA compatible with TP. GLA designs efficient CUDA kernel pipeline for Hopper and achieve significant speedup in compute-bound setting (MTP enabled). Evaluations on downstream tasks show competitive performance compared to alternatives.", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "* A simple and effective extension to MLA, address an important issue that MLA is not compatible with TP.\n* Train Llama-3 architecture (maximum scale at 876M) with GLA from scratch on 100B scale datasets, experiments on down-stream tasks show competitive performance.\n* Kernel optimization for GLA, higher TFLOPs/s than MLA. Distributed offset calculation to alleviate pointer arithmetic issue for LDGSTS style data-loading.", "reasons_to_reject": "* The advantage over GQA is not significant, lack study of comparing GLA with GQA with higher group ratio.", "questions_to_authors": "* Can we adapt existing models to GLA without training or with post-training, like in TransMLA(https://arxiv.org/abs/2502.07864)\n* How does GLA compare to GQA with higher group size (such as 32)?\n* Where does kernel performance over MLA coming from? I can imagine that compared to MLA, GLA don't need to broadcast QK results to all warps participated in PV computation, thus requiring less synchronization, but not sure if it will be the same case for Blackwell, where   mma output already stored in TMEM.\n* page_size = 1 and MLA pipeline design is also explored in flashinfer, which is also worth mentioning in related work.", "extra_scores": {}}, {"review_id": "XaznEXRYzz", "reviewer": "Reviewer_CoaW", "summary": "This paper introduces two hardware-efficient attention mechanisms, Group Tied Attention (GTA) and Group Latent Attention (GLA), designed to accelerate Large Language Model (LLM) decoding by addressing bottlenecks related to data movement, expanding Key-Value (KV) caches, and limited parallelism.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "1. The paper tackles significant and well-recognized challenges in LLM inference: excessive data movement, the expanding KV cache, and limited parallelism in decoding.\n2. The authors ground their designs in an analysis of arithmetic intensity, providing a clear rationale for how their methods improve hardware efficiency.\n3. The authors details crucial low-level optimizations like asynchrony through software pipelining/warp specialization and a novel distributed offset calculation for paged KV. The latter significantly improves paged KV performance, especially for small page sizes.\n4. The experiments cover multiple model scales, datasets, and evaluation metrics (perplexity, downstream tasks, latency, hardware utilization), providing a robust assessment of the proposed methods.", "reasons_to_reject": "1. The largest scale of the experiment is around 900M parameters with 50B tokens, which is relatively small. The downstream task performances comparing MLA to GLA is inconsistant across scales (i.e. GLA is better at 433M but worse in 876M ). I would suggest having a few 1.3B/100B setting experiments to see a clearer scaling trend.", "questions_to_authors": "1. How does GTA/GLA affects retrieval intensive tasks performance? What about long contecxt?", "extra_scores": {}}, {"review_id": "0LioRIXddc", "reviewer": "Reviewer_J2Hz", "summary": "This paper tackles the hardware inefficiency of attention mechanisms during LLM decoding, specifically the bottleneck caused by large KV cache transfers and low arithmetic intensity. The authors propose two variants: Group-Tied Attention (GTA), which ties key and value states within groups to reduce KV cache size and improve arithmetic intensity over GQA; and Group Latent Attention (GLA), a parallel-friendly modification of latent attention (MLA) that shards the latent state to enable efficient tensor parallelism. The paper provides theoretical motivation based on arithmetic intensity, details system-level optimizations like asynchronous loading and distributed offset calculation for paged KV, and presents empirical results showing comparable/better model quality (perplexity, downstream tasks) for GTA/GLA against GQA/MLA respectively, along with significant kernel speedups (up to 2x for GLA vs. FlashMLA).", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "- Addresses a highly relevant and critical problem in LLM inference efficiency.\n- Provides clear motivation grounded in hardware characteristics (arithmetic intensity).\n- Proposes novel attention variants (GTA, GLA) with specific design goals (efficiency, parallelizability).\n- Includes non-trivial system-level optimizations (distributed offset calculation) demonstrating practical implementation awareness.\n- Strong empirical validation across multiple model scales, datasets, and metrics (quality, speed, parallelization scenarios) supporting the claims.\n- GLA's design directly addresses a known limitation (parallelizability) of the recent MLA approach.", "reasons_to_reject": "- The conceptual novelty of GTA could be seen as incremental over GQA (tying K/V).\n- The speedup results focus on the attention kernel; end-to-end inference improvements are not directly reported (though the kernel is the bottleneck).\n- Communication overhead associated with GLA's tensor parallelism (TP) is benchmarked separately but not deeply analyzed in the main evaluation sections.", "questions_to_authors": "- Could you elaborate on the practical implementation complexity of sharding the latent state and projections in GLA, especially for g > 2 shards?\n- How does the communication overhead (e.g., AllReduce) in GLA with TP scale with the number of shards (g) and model size in practice, and how does it impact the observed kernel speedups in an end-to-end scenario?\n- Table 1 shows arithmetic intensity. What specific operations and data types (e.g., BF16) are assumed in these calculations?", "extra_scores": {}}, {"review_id": "BJGzeJgWo9", "reviewer": "Reviewer_TDTP", "summary": "In this work, the authors analyze the attention design from the perspective of arithmetic intensity and parallelization, and propose two novel attention variants named GTA and GLA, respectively. GTA ties the KV states onto a single projection, and GLA brings parallelization-aware modification to MLA. The experiments demonstrate the efficacy of the two variants on both quality and efficiency.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "1. A timely topic concerning LLM inference efficiency. \n2. A new perspective. The authors theoretically analyze the impact of various attention variants on computational intensity and parallelization patterns. \n3. Sufficient quality experiments. The authors conduct comprehensive experiments on multiple datasets to demonstrate that the generation quality is well preserved.", "reasons_to_reject": "1. The main logic is not so clear. This paper attempts to analyze attention variant design from a hardware-efficiency perspective. However,  while Section 3 demonstrates certain trade-offs, it fails to establish clear design principles. Notably, increased computational intensity or modified parallelization schemes do not inherently guarantee acceleration.\n2. Lack of algorithm-level insights. The paper's primary contribution lies in its reformulation of the attention computation's mathematical expression, which necessitates deeper algorithmic insights. The design of GTA references certain data characteristics,  but most are derived from prior works. In designing GLA, the authors directly presents the modified computation scheme without justifying how the model quality is preserved.\n3. The acceleration of GLA is not sufficiently general, as it is constrained to specific scenarios and inputs.\n4. Writing issues. The variable definitions appear too late in the manuscript, and under-sized fonts in Figure 2 affect reading.", "questions_to_authors": "1. In GTA, How is $K_{rope}$ paralleled across multiple GPUs?\n2. What is the relationship between the design principles between GTA and GLA?  It seems they are totally independent.\n3. In Section 5.2, it is better to also consider attention-specialized parallelization methods like sequence parallel (SP). \n4. Does the observed speedup in Figure 3 originate from low-level optimizations,  considering the uniform sequence lengths?", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper presents a timely and valuable contribution to the field of efficient Large Language Model inference. It introduces two novel attention variants, Group-Tied Attention (GTA) and Group Latent Attention (GLA), designed to alleviate critical bottlenecks in LLM decoding, namely excessive data movement, expanding KV caches, and limited parallelism. The work is commendably grounded in a hardware-aware perspective, focusing on optimizing arithmetic intensity and enabling effective parallelization. The authors motivate their architectural modifications through a clear analysis of arithmetic intensity and hardware characteristics. This provides a solid foundation for understanding why these variants are expected to be more efficient. GTA ties key and value states to reduce KV cache and improve arithmetic intensity over Grouped Query Attention, while GLA modifies Multi-head Latent Attention to be more parallel-friendly by sharding the latent state, a key limitation of prior latent attention methods.
Decision: Accept