Organizations: The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China · Beijing Academy of Artificial Intelligence, Beijing, China · Université Paris Cité, Paris, France
FullAttn often assigns negligible normalized mass to much of the causal score space, yet dense kernels execute the complete post-score path after forming each QK tile. We introduce MALA, a fused attention primitive that preserves score access to every legal causal interaction and uses normalized contribution to allocate post-score computation. Forward uses its evolving online-softmax normalizer, while backward reuses the finalized normalizer to derive nested retained support using only standard attention state. A common tolerance governs training and inference, allowing for adaptive retention of the work. MALA reduces low-contribution post-score computation. A matched-work study at 8K isolates the benefit of distribution-adaptive allocation: under exactly matched total post-score work, MALA approaches a per-instance reference-mass oracle, with mean omitted mass of 0.0188% versus 0.0182%. Across context lengths from 1K to 32K tokens, the same tolerance maintains low output and gradient errors relative to the reference. Across a broader controlled associative-recall comparison, MALA closely tracks FullAttn as context grows, reaching 89.67% accuracy at 8K compared with 89.97% for FullAttn. In an attention-operator benchmark at 128K tokens with tensor parallelism, MALA reduces forward and backward latency during training by 2.2x and 3.0x and decoding latency during inference by 1.6x relative to FullAttn. Across scaling-law training from 0.6B to 14B parameters, MALA closely tracks FullAttn in perplexity while reducing total training FLOPs. The resulting 14B models and 32B models from separate continued training achieve comparable knowledge, reasoning, and long-context retrieval scores to FullAttn. These results indicate that allocating post-score computation according to normalized attention contributions can retain the evaluated capabilities of FullAttn while reducing attention computation.
Figures & tables
Figure 1: MALA lets attention allocate its own post-score compute. Every legal causal tile undergoes QK score discovery. Forward uses its evolving softmax normalizer to allocate forward post-score execution; backward reuses the finalized normalizer to derive nested retained support while storing only standard attention state.
Allocation Scope
Selection Signal
Mean Slots
Omitted Mass (%)
Output Error (%)
Mean ↓
P95 ↓
Mean ↓
P95 ↓
Position
Final reference mass
≈ 1,024
0.6012
0.9923
0.6231
4.2370
Layer/Head/Position
Final reference mass
≈ 1,024
0.1721
0.6247
0.2817
0.9237
Input/Layer/Head/Position
Final reference mass
≈ 1,024
0.0182
0.0664
0.0164
0.1013
Input/Layer/Head/Position
MALA online mass bound
≈ 1,024
0.0188
0.0693
0.0174
0.1017
Table 1: Matched-work allocation. Every policy executes exactly the same total post-score work, with a shared average of approximately 1,024 slots per query. The first three rows use finalized per-instance reference mass for region selection and differ only in how work is allocated across decisions. Omitted probability mass and relative output L2 error are reported as percentages; statistics pool input-layer-head-query-block allocation decisions.
Figure 2: Operator fidelity. Forward and Backward allocated work is reported as mean post-score key slots per query. Actual omitted reference probability mass, relative output error, and relative gradient errors are measured over every query position against the same operator with screening disabled. Query positions are first aggregated within each sequence-layer-head state; reported means and P95 values pool these states, and marker halos reflect confidence intervals from sequence bootstrap.
Figure 3: Associative Recall . Accuracy with 256 key-value pairs across sequence lengths and model dimensions. Fixed-budget sparse methods use a common 1,024-slot ceiling, with realized raw slots determined after causal clipping; MALA obtains its work from the canonical tolerance. The task directly tests the 256 bindings, with no injected distractors. MALA closely tracks FullAttn as the context grows, while the other sparse mechanisms lose a substantial fraction of the associations.
Figure 4: Operator Latency and Memory . End-to-end attention-operator latency (bars, left axes) and per-rank peak memory (markers, right axes) for the training forward and backward passes and autoregressive decoding with tensor parallelism TP=8 on 8 H100 GPUs. Measurements include each method’s complete score-discovery and selection path. MALA reduces latency while retaining FullAttn-level memory through in-kernel allocation.
Figure 5: Scaling Laws . Perplexity as a function of total training FLOPs for models from 0.6B to 14B parameters during 4K pre-training (left) and 32K long-context training (right). The FLOP accounting includes complete score discovery and method-specific allocation or selection costs. MALA closely follows FullAttn in Perplexity while shifting the frontier toward lower training FLOPs.
Model Scale
Method
Knowledge
Reasoning
RULER 32K
RULER 128K
14B
FullAttn
72.32
64.46
89.42
65.84
14B
MoBA
70.01
58.07
86.85
60.17
14B
DSA
70.11
62.04
87.59
63.05
14B
MALA (ours)
72.48
64.66
89.45
65.75
32B
FullAttn
75.62
75.67
92.70
82.03
32B
MALA (ours)
75.75
76.10
92.71
82.56
Table 2: Model-level evaluation summary. Average knowledge and reasoning scores, together with RULER scores at the native 32K context length and after YaRN extrapolation to 128K. Detailed per-task results and uncertainty estimates are reported in Appendix D .
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Method
Allocation Rule
dv/dv,ref
Nominal Slot Ceiling
Raw-Slot Reference
FullAttn
Complete causal history
1×
Full
Full
SWA / Seer / InfLLMv2 / NSA
Native support or selector
1×
1,024
after causal clipping
MoBA
Native block routing
1×
1,024
after causal clipping
DSA
Token-level Top- K
4×
1,024
after causal clipping
MALA
Normalized-mass tolerance τ=1
1×
-
measured near reference
Appendix
Table 3: Associative-recall allocation configurations. Fixed-budget sparse methods use a common 1,024-slot ceiling, with realized work computed after causal clipping and native block expansion. The raw-slot comparison aligns retained interactions; DSA therefore uses the same slot ceiling despite its larger value dimension. MALA uses τ=1 , with work determined by normalized-mass allocation.
Method
nh
nhkv
dqk
dv
Allocation
Target Eq. Slots
Precision
FullAttn 1 1 1 The FullAttn implementation is available at github.com/Dao-AILab/flash-attention .
64
8
128
128
Full causal context
Full
BF16
MoBA 2 2 2 The MoBA implementation is available at github.com/MoonshotAI/MoBA .
64
64
128
128
B=128 , native routing
≈ 1,024
BF16
DSA 3 3 3 The DSA indexer implementation is available at github.com/deepseek-ai/DeepGEMM . , 4 4 4 The DSA implementation is available at github.com/Dao-AILab/flash-attention/flash_attn/cute .
64
1
576
512
Top- K=256 , FP8 indexer
≈ 1,024
BF16/FP8
MALA
64
8
128
128
τ=1 , fused allocation
measured ≈ 1,024
BF16
Appendix
Table 4: Operator benchmark configurations. The sparse methods target comparable forward post-score FLOP-equivalent work, while reported latency and memory include their complete score-discovery and selection paths.
Algos
Params
PT/LCT Tok
Batch
PT/LCT LR
nlayers
dmodel
nh
nhkv
dqk
dv
PS FLOPs/Slot
PT Alloc
LCT Alloc
FullAttn
≈ 0.6B
12B,1.5B
0.256M
1e-3,1e-4
28
1024
16
8
128
128
4096
4096
32768
MoBA
≈ 0.6B
12B,1.5B
0.256M
1e-3,1e-4
28
1024
16
16
128
128
4096
K=1024
K=1024
DSA
≈ 0.6B
12B,1.5B
0.256M
1e-3,1e-4
28
1024
16
1
576
512
16384
K=256
K=256
MALA
≈ 0.6B
12B,1.5B
0.256M
1e-3,1e-4
28
1024
16
8
128
128
4096
τ=1
τ=1
FullAttn
≈ 1.7B
34B,4B
0.512M
8e-4,8e-5
28
2048
16
8
128
128
4096
4096
32768
MoBA
≈ 1.7B
34B,4B
0.512M
8e-4,8e-5
28
2048
16
16
128
128
4096
K=1024
K=1024
Appendix
Table 5: Self-Attention Variants Scaling Laws Configurations. Complete model, training, and attention configurations used in the scaling-law experiments. PS FLOPs/Slot describes allocated forward post-score work; the FLOP-axis additionally includes complete score discovery and method-specific selection.
Table 6: Knowledge Benchmark Results . Knowledge evaluation of the scaling-law models at 14B and the continued-training models at 32B. MALA produces averages comparable to FullAttn at both scales and remains close across individual tasks.
Table 7: Reasoning Benchmark Results . Reasoning evaluation of the scaling-law models at 14B and the continued-training models at 32B. MALA matches FullAttn on average at 14B and 32B without a fixed retained-work budget.
Table 8: Retrieval Benchmark Results . RULER evaluation at the native 32K context length and after YaRN extrapolation to 128K. MALA remains comparable to FullAttn across both model scales and context lengths.
Attention mechanisms let neural networks weight information from other tokens or regions according to relevance, enabling context-dependent processing.
Multi-Head Latent Attention (MLA) enables expressive multi-head attention with compact caches for its content and decoupled RoPE paths, yet cache memory still scales linearly with context length and batch size. In this work, we establish a systematic model of MLA's dual-path quantization errors, characterizing their distinct effects on attention-output distortion and explaining the pronounced amplification of RoPE-path errors. Guided by this analysis, we introduce QuantMLA, a function-aligned framework for low-bit dual-path quantization. We derive path-specific transformation spaces that preserve full-precision computation while remaining fully fusible into model parameters offline, eliminating online transformation overhead. Within these spaces, QuantMLA learns path-specific transformations with function-aligned objectives: attention-output reconstruction captures the content path's coupled matching and aggregation errors, while positional QK reconstruction preserves the RoPE-induced component of the attention logits and admits a theoretical bound on output distortion. Across four MLA model families, QuantMLA enables, to our knowledge, the first reported joint INT4 caching of the content and RoPE caches with minimal accuracy degradation. Further compressing the content cache to INT2 while retaining the RoPE key cache at INT4 maintains competitive performance on challenging reasoning and code benchmarks. We develop a native low-bit MLA attention kernel that integrates unpacking and dequantization directly into attention computation. The physical cache layout provides 3.59x compression at 128K context, while a cache-pressure serving workload achieves 5.168x higher whole-job output throughput than BF16. The code will be released upon acceptance.
Zunhai Su, Yuxuan Sun, Jianchao Tan +5
The University of Hong Kong · Meituan LongCat Team · South China University of Technology +1
FullAttn repeatedly exposes the complete causal history to every attention head, creating substantial redundant computation and memory traffic even with IO-efficient dense kernels. We introduce CoWA, a structured attention architecture that distributes access to the causal history across KV heads. All heads share near-diagonal and prefix-sink windows, while complementary long-range windows partition the remaining history. Their union provides full causal coverage although each head attends sparsely to distant tokens. This position-defined attention pattern requires no learned router or indexer, is used consistently during training and inference, and aligns with KV-head tensor parallelism. A window-matched ablation at 8K isolates the effect of complementary long-range allocation: CoWA with 100% collective coverage reaches 89.73% accuracy, compared with 89.97% for FullAttn, while duplicated long-range windows perform substantially worse. Across a broader controlled associative-recall comparison with matched token budgets, CoWA closely tracks FullAttn as the context grows, whereas other sparse patterns lose a substantial fraction of the associations. In an attention-operator benchmark at 128K tokens with tensor parallelism, CoWA reduces forward and backward latency during training by 7.4x and 8.6x and decoding latency during inference by 3.0x over FullAttn. Its per-rank peak operator memory matches FullAttn during training and is 7.6x lower during decoding. Across scaling-law training from 0.6B to 14B parameters, CoWA closely tracks FullAttn in perplexity while reducing total training FLOPs. The resulting 14B models and 32B models from separate continued training achieve comparable knowledge, reasoning, and long-context retrieval scores to FullAttn. These results show that full causal coverage can be a collective property of the head ensemble rather than a duplicated property of every head.
Jingze Shi, Zhangyang Peng, Xianduo Li +6
The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China. · Beijing Academy of Artificial Intelligence, Beijing, China. · Université Paris Cité, Paris, France.
Multi-Head Latent Attention (MLA) compresses the KV cache into a latent representation, reducing inference memory and bandwidth. Existing MHA2MLA/TransMLA methods do not support QKNorm and require costly recovery on modern LLMs, especially MoE. We present MLAfication, a QKNorm-compatible, parameter-efficient extension of MHA2MLA. It combines calibrated KNorm, parameter-efficient distillation, contribution-aware Partial-RoPE, and activation-aware joint KV initialization. Across MHA and GQA models, including QKNorm- and MoE-equipped variants, MLAfication preserves performance under substantial KV-cache compression, surpasses the current SOTA on commonsense and long-context benchmarks, and extends MLA conversion to modern architectures. On Qwen3-30B-A3B, MLAfication delivers up to 5.35× the serving throughput of BF16 GQA and, combined with 4-bit KV quantization, reduces the KV-cache footprint by 94.73%. Our code is available at https://github.com/JT-Ushio/MLAfication.