Scaling Parameter and Context in Attention: Native Sparse Attention from Mixture-of-Head

Zizhuo Fu; Runsheng Wang; Meng Li

cs.CLSep 30, 2026

Scaling Parameter and Context in Attention: Native Sparse Attention from Mixture-of-Head

Authors: Zizhuo Fu, Runsheng Wang, Meng Li

Organizations: Peking University

Abstract

Scaling attention parameters can improve language model quality, but retaining full token histories makes additional heads costly at long contexts. Furthermore, since attention retrieves and combines contextual information, parameter scaling should also support longer contexts. We therefore ask whether attention parameter scaling can directly enable efficient and effective context scaling. We introduce NAMOH, an architecture-native sparse attention mechanism that activates $K$ of $H$ heads per token. Each head retains only its assigned tokens and performs causal attention within this subsequence. Head selection thus jointly determines active parameters and available context without scanning the full history. Under balanced assignments, increasing $H$ at fixed $K$ shortens head histories and reduces per-token key-value (KV) access without increasing total KV storage. We further support head-relative rotary position embeddings to shorten positional spans within routed subsequences, aiming to mitigate position-induced attention noise. Experiments show that NAMOH can outperform fully activated models with the same total parameters, while enabling more efficient long-context inference than smaller dense models with matched active parameter counts. It remains compatible with GQA and existing sparse attention mechanisms. We hope this work offers a new path for scaling attention, with parameter scaling directly enabling context scaling.

Figures & tables

Figure 1: Parameter and context scaling in frontier language models. (a) Attention and FFN parameters in frontier models. Since the rise of MoE, parameter growth has been dominated by FFN experts, with comparatively little growth in attention. (b) Maximum context windows in frontier models. Sparse and linear-time mechanisms support further context expansion.

Figure 2: NAMOH with two active heads out of four. (a) Prefill routes tokens $0$ to $7$ before QKV projection. Each head attends within its ordered subsequence, and projected outputs are combined with routing weights. (b) During decoding, tokens $8$ and $9$ access and extend only their selected heads’ KV caches. Position indices follow each head’s local token order.

Cost	MHA	GQA	SA	MoH	NAMOH	GQA +SA	GQA + NAMOH	SA + NAMOH	GQA+SA + NAMOH
Total weights	$4Hd_{m}d_{h}$	$2Hd_{m}d_{h}\frac{g+1}{g}$	$4Hd_{m}d_{h}$	$4Hd_{m}d_{h}$	$4Hd_{m}d_{h}$	$2Hd_{m}d_{h}\frac{g+1}{g}$	$2Hd_{m}d_{h}\frac{g+1}{g}$	$4Hd_{m}d_{h}$	$2Hd_{m}d_{h}\frac{g+1}{g}$
Active weights	$4Hd_{m}d_{h}$	$2Hd_{m}d_{h}\frac{g+1}{g}$	$4Hd_{m}d_{h}$	$4Hd_{m}d_{h}$	$4Kd_{m}d_{h}$	$2Hd_{m}d_{h}\frac{g+1}{g}$	$2Kd_{m}d_{h}\frac{g+1}{g}$	$4Kd_{m}d_{h}$	$2Kd_{m}d_{h}\frac{g+1}{g}$
KV storage	$HT$	$HT/g$	$HT$	$HT$	$KT$	$HT/g$	$KT/g$	$KT$	$KT/g$
KV activation	$HT$	$HT/g$	$\alpha HT$	$KT$	$TK^{2}/H$	$\alpha HT/g$	$TK^{2}/(Hg)$	$\alpha TK^{2}/H$	$\alpha TK^{2}/(Hg)$

Table 1: Attention weight and KV cache costs. Total and active projection weights, KV storage after

T

tokens, and historical KV activation for single-token decoding.

Attention Mechanism	KV Stor.	KV Act.	MMLU	GSM8K	HEval	ARC-E	ARC-C	HellaSwag	PIQA	OBQA	BoolQ	WinoG.	Average
MHA-32	$1$	$1$	32.79	5.31	6.10	65.87	34.13	45.88	69.37	37.20	46.64	53.99	39.73
MHA-16	$1/2$	$1/2$	32.35	3.11	4.27	63.76	34.39	46.84	70.08	35.80	46.18	52.49	38.93
GQA-32KV16	$1/2$	$1/2$	27.48	4.62	7.32	65.61	35.58	47.56	69.48	36.00	48.40	52.01	39.41
SA-32 ( $1/4$ )	$1$	$1/4$	25.27	2.58	3.66	60.87	33.21	45.20	59.90	37.60	44.22	50.12	36.26
SA-16 ( $1/2$ )	$1/2$	$1/4$	24.36	2.88	3.05	63.22	34.74	45.00	59.36	35.60	53.39	49.49	37.11
NAMOH-32A16	$1/2$	$1/4$	32.74	5.84	9.76	66.04	36.43	48.11	70.67	36.60	49.33	55.81	41.13

Table 2: Evaluation results of models trained from scratch with different attention mechanisms and KV budgets. KV storage (KV Stor.) and activation (KV Act.) are normalized to MHA-32. For sparse attention (SA), the fraction in parentheses denotes the activated KV fraction per head.

Figure 5: Task-dependent head utilization in NAMOH-32A8. Activation frequencies of individual heads in one attention layer on MMLU, GSM8K, PIQA, and HumanEval.

Figure 6: Routing shapes attention connectivity in NAMOH. (a) Token-to-head assignments. (b) The number of shared active heads for each causal token pair. (c) Layer-wide attention weighted by query-specific routing gates. (d) Attention maps for heads 0, 8, 16, and 24.

Load Balancing Variants
Method	MMLU	GSM8K	HEval	ARC-E	ARC-C	HellaSwag	PIQA	OBQA	BoolQ	WinoG.	Average
CV ( Shazeer et al., 2017 )	31.31	4.62	7.93	64.86	36.26	47.78	69.26	35.80	55.38	51.54	40.47
fp ( Fedus et al., 2021 )	32.40	5.31	9.76	64.66	37.11	46.16	70.40	36.80	56.81	52.96	41.24
Loss-Free ( Wang et al., 2024 )	33.23	3.11	7.93	62.76	34.23	45.12	68.23	35.00	55.63	52.96	39.82
Shared-Head Variants
+ 2 Shared (Full)	30.99	5.38	8.64	64.10	37.64	47.12	70.97	37.00	58.31	53.28	41.34

Table 5: Further refinement of NAMOH-32A8 with load-balancing strategies and shared heads.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Mechanism	Prefill	Single-token decoding
MHA	$HT^{2}/2$	$HT$
GQA	$HT^{2}/2$	$HT$
SA	$\left(\alpha+\frac{1}{B}\right)\frac{HT^{2}}{2}$	$\left(\alpha+\frac{1}{B}\right)HT$
MoH	$KT^{2}/2$	$KT$
NAMOH	$T^{2}K^{2}/(2H)$	$TK^{2}/H$
GQA+SA	$\left(\alpha+\frac{\chi_{g}}{gB}\right)\frac{HT^{2}}{2}$	$\left(\alpha+\frac{\chi_{g}}{gB}\right)HT$

Appendix

Table 6: Leading attention and indexer terms. The prefill column reports

C_{\mathrm{attn}}+\chi_{g}C_{\mathrm{meta}}

; the decoding column reports

gA_{\mathrm{KV}}(T)+\chi_{g}D_{\mathrm{meta}}(T)

. Insert these terms into Equation 58 with the active weights in Table 1 to obtain total arithmetic costs. Routed histories are balanced, and finite-length and block-rounding terms are omitted here but retained in the derivations.

Attention Mechanisms

Attention mechanisms let neural networks weight information from other tokens or regions according to relevance, enabling context-dependent processing.

See all 517 papers

HyScale: Scaling Long-Context Training for Learned Sparse Attention
Oct 23, 2026Yangshen Deng, Xinjing Huang, Yulong Wang +8Efficient Language Model Training Sparse Attention
University of Edinburgh · Tencent Hunyuan
Hierarchical Sparse Attention Done Right: Toward Infinite Context Modeling
Jul 3, 2026Xiang Hu, Xinyu Wei, Hao Gu +10Long-Context Language Model Inference Long-Context Language Modeling
Tencent HY Team · ShanghaiTech University · The Hong Kong University of Science and Technology +1
DashAttention: Differentiable and Adaptive Sparse Hierarchical Attention
May 18, 2026Yuxiang Huang, Nuno M. T. Gonçalves, Federico Alvetreti +5Self-Attention Long-Context Modeling
1Tsinghua University · 2Instituto Superior Técnico, Universidade de Lisboa. · 3Instituto de Telecomunicações. +6

Oct 23, 2026J/K move · Enter open · S save

HyScale: Scaling Long-Context Training for Learned Sparse Attention

Yangshen Deng, Xinjing Huang, Yulong Wang, Yunzhuo Liu+7

University of Edinburgh · Tencent Hunyuan

Efficient Language Model Training Sparse Attention Long-Context Language Modeling

Open paper AskPDF arXiv

We present HyScale, a long-context training system for learned sparse attention. Learned sparse attention is a promising architecture for million-token contexts, widely adopted in production since DeepSeek-V3.2. Selection, however, grows into the training bottleneck. For each query, it scores every previous token whereas sparse attention reads only

k

; worse, current systems materialize a full score table and repeat selection during recomputation. Our key insight is that the output of selection is not an activation, but rather a set of indices. Building on it, HyScale introduces HyTopK and selection replay. HyTopK finds each query's exact top-

k

on the fly without a score table, while selection replay reuses forward-pass selection in recomputation. Together with holistic system optimizations tackling end-to-end training challenges, HyScale speeds up GLM-5.1 training on 1M-token sequences by 2.07

\times

over Megatron-LM on 512 H20 GPUs. It also speeds up the CSA layers of DeepSeek-V4-Flash and the QSA layers of Qwen3.8-Flash-Next by 2.01

\times

and 1.57

\times

at 1M tokens on 32 H800 GPUs.

Abstract

Attention Mechanisms

Explore similar work

HyScale: Scaling Long-Context Training for Learned Sparse Attention

Hierarchical Sparse Attention Done Right: Toward Infinite Context Modeling

DashAttention: Differentiable and Adaptive Sparse Hierarchical Attention