CoWindow Attention: Full Causal Coverage Is a Collective Property

Jingze Shi; Zhangyang Peng; Xianduo Li; Yanlin Qi; Xiaotian Lin; Haoxian Chen; Liangdong Wang; Guang Liu

cs.AISep 26, 2026

CoWindow Attention: Full Causal Coverage Is a Collective Property

Authors: Jingze Shi, Zhangyang Peng, Xianduo Li, Yanlin Qi, Xiaotian Lin, Haoxian Chen, Liangdong Wang, Guang Liu,

Organizations: The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China. · Beijing Academy of Artificial Intelligence, Beijing, China. · Université Paris Cité, Paris, France.

Abstract

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.

Figures & tables

Panel of Figure 1 — Figure 1: CoWindow Attention Overview. Left: All KV heads share the near window and receive complementary long-range windows whose union covers the full causal attention region, preserving direct access to every preceding key position without duplicating the full history in every head. The shared sink window is omitted for visual clarity. Right: As sequence length grows, CoWA reduces forward and backward latency during training and decoding latency during inference, while matching FullAttn’s peak operator memory during training and reducing it during decoding. Separate model-level evaluations show comparable performance on the RULER benchmark.

Layout	Window Assignment	Long-Range Coverage	1K	2K	4K	8K
FullAttn	Full history per head	100%	100	100	100	89.97
SWA	No long-range window	0%	100	17.48	14.12	5.34
1 unique window	$1$ window $\times$ $8$ heads	12.5%	100	53.53	32.99	21.32
2 unique windows	$2$ windows $\times$ $4$ heads	25%	100	62.74	40.37	32.92
4 unique windows	$4$ windows $\times$ $2$ heads	50%	100	83.91	61.75	52.32
CoWA	$8$ windows $\times$ $1$ head	100%	100	100	100	89.73

Table 1: Collective Coverage with Matched Per-Head Window Widths. The 1-, 2-, 4-, and 8-window layouts use matched per-head window widths and vary the number of distinct long-range windows assigned across eight KV heads. FullAttn and SWA are dense and local references and are not part of this window-matched comparison. We report associative-recall accuracy (%) at

d_{model}=512

Figure 2: Associative Recall . Accuracy with 256 key-value pairs across sequence lengths and model dimensions. At sequence lengths 1K, 2K, 4K, and 8K, all sparse methods use matched per-query token budgets of 1,024, 1,152, 1,408, and 1,920 tokens, respectively. We do not add random padding or additional noise tokens. CoWA closely tracks FullAttn as the context grows, whereas the other sparse patterns lose a substantial fraction of the associations.

Figure 3: Operator Latency and Memory . End-to-end attention-operator latency (bars, left axes) and maximum per-rank peak allocation (markers, right axes) for the training forward and backward passes and decoding with tensor parallelism

TP=8

on 8 H100 GPUs. CoWA has the lowest latency and decoding-operator peak allocation among the tested methods at long sequence lengths.

Figure 4: Scaling Laws . Perplexity as a function of training FLOPs for models from 0.6B to 14B parameters during 4K pre-training (left) and 32K long-context training (right). Each marker reports final perplexity; its outer ring denotes one standard deviation over the perplexity measurements collected during training. CoWA closely follows FullAttn in perplexity while requiring fewer FLOPs.

Model Scale	Method	Knowledge	Reasoning	RULER 32K	RULER 128K
14B	FullAttn	72.32	64.46	89.42	65.84
14B	MoBA	70.58	59.71	87.17	61.52
14B	DSA	70.23	63.28	87.67	64.05
14B	CoWA (ours)	72.70	64.87	89.13	66.60
32B	FullAttn	75.62	75.67	92.70	82.03
32B	CoWA (ours)	76.07	75.53	92.58	81.78

Table 2: Model-Level Evaluation Summary . Average knowledge, reasoning, and RULER scores at 14B and 32B. RULER is evaluated at the native 32K context length and under YaRN extrapolation to 128K. Detailed per-task results and uncertainty estimates are reported in Appendix B .

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Mechanism	Selected Attention	Selection Overhead	Auxiliary State
FullAttn	$O\!\left(n^{2}h(d_{qk}+d_{v})\right)$	-	-
MoBA	$O\!\left(nwh(d_{qk}+d_{v})\right)$	$O\!\left(n^{2}hd_{qk}/B\right)$	$O\!\left(n^{2}h/B\right)$
DSA	$O\!\left(nwh(d_{qk}+d_{v})\right)$	$O\!\left(n^{2}h_{I}d_{I}\right)$	$O\!\left(n^{2}h_{I}+nw\right)$
CoWA	$O\!\left(nwh(d_{qk}+d_{v})\right)$	$O(1)$ per QO block	$O(H_{k})$

Appendix

Table 3: Comparison of Different Attention Variants . The sequence-dependent attention and sparse-selection costs. The auxiliary state column describes routing or indexing state beyond the QKV tensors and attention outputs.

Algo	warmups	runs	batch	$n_{h}$	$n_{h_{kv}}$	$d_{qk}$	$d_{v}$	selector	precision
FullAttn 1 1 1 The FullAttn implementation is available at github.com/Dao-AILab/flash-attention .	25	100	1	64	8	128	128	-	BF16
MoBA 2 2 2 The MoBA implementation is available at https://github.com/MoonshotAI/MoBA .	25	100	1	64	64	128	128	$B=128$	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 .	25	100	1	64	1	576	512	FP8 indexer	BF16/FP8
CoWA	25	100	1	64	8	128	128	fused window	BF16

Appendix

Table 4: Latency and Memory Benchmark Configurations . The operator shapes and sparse-selection mechanisms used for the running-time curves. All measurements use tensor parallelism

TP=8

on 8 H100 GPUs.

Algo	PST FLOPs	1K	2K	4K	8K	16K	32K	64K	128K
FullAttn	32768	1024	2048	4096	8192	16384	32768	65536	131072
MoBA	32768	1024	1024	1536	2048	3072	5120	9216	17408
DSA	139264	256	288	352	480	736	1248	2272	4320
CoWA	32768	1024	1152	1408	1920	2944	4992	9088	17280

Appendix

Table 5: Per-Query Token Budgets . The token budgets used at every benchmark sequence length.

Algo	SeqLen	Fwd Lat.	Bwd Lat.	Dec Lat.	Fwd Mem.	Bwd Mem.	Dec Mem.
FullAttn	1,024	0.015	0.055	0.003	4.531	14.094	0.504
FullAttn	2,048	0.043	0.149	0.004	9.063	28.188	1.004
FullAttn	4,096	0.153	0.468	0.007	18.125	56.375	2.004
FullAttn	8,192	0.572	1.699	0.009	36.250	112.750	4.004
FullAttn	16,384	2.179	6.455	0.012	72.500	225.500	8.004
FullAttn	32,768	8.570	24.449	0.018	145.000	451.000	16.004

Appendix

Table 6: Latency and Memory Benchmark Results . End-to-end latency (ms) and peak memory (MiB) for the training forward and backward passes and inference-time decoding at every sequence length. Latency is the synchronized tensor-parallel group latency, and memory is the maximum per-rank peak allocation.

Algos	Params	PT/LCT Tok	Batch	PT/LCT LR	$n_{layers}$	$d_{model}$	$n_{h}$	$n_{h_{kv}}$	$d_{qk}$	$d_{v}$	PST FLOPs	PT Bgt	LCT Bgt
FullAttn	$\approx$ 0.6B	12B,1.5B	0.256M	1e-3,1e-4	28	1024	16	8	128	128	8192	4096	32768
MoBA	$\approx$ 0.6B	12B,1.5B	0.256M	1e-3,1e-4	28	1024	16	16	128	128	8192	1536	5120
DSA	$\approx$ 0.6B	12B,1.5B	0.256M	1e-3,1e-4	28	1024	16	1	576	512	34816	352	1248
CoWA	$\approx$ 0.6B	12B,1.5B	0.256M	1e-3,1e-4	28	1024	16	8	128	128	8192	1408	4992
FullAttn	$\approx$ 1.7B	34B,4B	0.512M	8e-4,8e-5	28	2048	16	8	128	128	8192	4096	32768
MoBA	$\approx$ 1.7B	34B,4B	0.512M	8e-4,8e-5	28	2048	16	16	128	128	8192	1536	5120

Appendix

Table 7: Self-Attention Variants Scaling Laws Configurations . The model and hyperparameter configurations used in our self-attention variants scaling laws experiments.

14B-Non-Thinking, temperature=0.6, top_k=20, top_p=0.95, max_gen_toks=32768, seed=[0, 42, 233, 666, 1234]
Model	MMLU	MMLU-Pro	BBH	HellaSwag	OBQA	WinoGrande	Avg
	Gen\|5-shot	5-shot	CoT\|3-shot	5-shot	5-shot	5-shot
FullAttn	78.21 $\pm$ 0.31	64.96 $\pm$ 0.46	82.59 $\pm$ 0.48	80.64 $\pm$ 0.38	49.4 $\pm 2.17$	78.16 $\pm$ 1.14	72.32
MoBA	76.34 $\pm$ 0.46	58.73 $\pm$ 0.83	82.3 $\pm$ 0.41	80.2 $\pm$ 0.34	49.7 $\pm$ 2.05	76.24 $\pm$ 1.45	70.58
DSA	72.72 $\pm$ 0.53	59.92 $\pm$ 0.62	81.62 $\pm$ 0.63	80.01 $\pm$ 0.33	49.53 $\pm$ 2.24	77.63 $\pm$ 1.17	70.23
CoWA (ours)	79.04 $\pm$ 0.28	64.89 $\pm$ 0.43	83.64 $\pm$ 0.51	80.3 $\pm$ 0.31	49.9 $\pm$ 2.01	78.48 $\pm$ 1.23	72.7

Appendix

Table 8: Knowledge Benchmark Results . Knowledge evaluation of the scaling-law models at 14B and the continued-trained models at 32B. CoWA produces averages comparable to FullAttn at both scales and remains close across individual tasks.

14B-Non-Thinking, temperature=0.6, top_k=20, top_p=0.95, max_gen_toks=32768, seed=[0, 42, 233, 666, 1234]
Model	PIQA	GSM8K	Hendrycks-Math	ARC-C	AGIEval	GPQA-Diamond	Avg
	5-shot	CoT\|4-shot	CoT\|4-shot	0-shot	0-shot	CoT\|3-shot
FullAttn	81.63 $\pm$ 0.86	90.32 $\pm$ 0.93	62.46 $\pm$ 0.65	61.95 $\pm$ 1.44	50.56 $\pm$ 0.71	39.87 $\pm$ 3.47	64.46
MoBA	77.73 $\pm$ 0.93	91.29 $\pm$ 0.97	39.0 $\pm$ 1.38	62.45 $\pm$ 1.7	48.96 $\pm$ 0.87	38.88 $\pm$ 3.81	59.71
DSA	80.42 $\pm$ 0.77	90.52 $\pm$ 0.99	55.99 $\pm$ 0.82	62.04 $\pm$ 1.29	47.6 $\pm$ 0.63	43.11 $\pm$ 4.52	63.28
CoWA (ours)	81.5 $\pm$ 0.72	90.58 $\pm$ 0.83	62.31 $\pm$ 0.62	62.1 $\pm$ 1.62	50.44 $\pm$ 0.76	42.33 $\pm$ 3.17	64.87

Appendix

Table 9: Reasoning Benchmark Results . Reasoning evaluation of the scaling-law models at 14B and the continued-trained models at 32B. CoWA achieves average scores comparable to FullAttn at both scales.

14B Native 32K Sequence Length, seed=[0, 42, 233, 666, 1234]
Model	NIAH-S	NIAH-MK	NIAH-MQ	NIAH-MV	RULER-VT	RULER-CWE	RULER-FWE	RULER-QA	Avg
FullAttn	100 $\pm$ 0	98.0 $\pm$ 0.65	98.8 $\pm$ 0.68	95.5 $\pm$ 0.72	99.44 $\pm$ 0.2	75.7 $\pm$ 0.69	91.67 $\pm$ 0.71	56.28 $\pm$ 1.2	89.42
MoBA	99.8 $\pm$ 0.004	94.4 $\pm$ 0.63	98.2 $\pm$ 0.73	95.1 $\pm$ 0.7	99.0 $\pm$ 0.32	67.35 $\pm$ 0.96	91.44 $\pm$ 0.69	52.12 $\pm$ 1.64	87.17
DSA	100 $\pm$ 0	95.6 $\pm$ 0.68	98.0 $\pm$ 0.77	96.1 $\pm$ 0.76	99.24 $\pm$ 0.26	69.25 $\pm$ 0.87	91.74 $\pm$ 0.67	51.42 $\pm$ 1.43	87.67
CoWA (ours)	100 $\pm$ 0	96.96 $\pm$ 0.71	98.7 $\pm$ 0.69	95.2 $\pm$ 0.73	99.28 $\pm$ 0.21	72.6 $\pm$ 0.8	92.53 $\pm$ 0.66	57.74 $\pm$ 1.14	89.13
32B Native 32K Sequence Length, seed=[0, 42, 233, 666, 1234]

Appendix

Table 10: Retrieval Benchmark Results . RULER evaluation at the native 32K context length and under YaRN extrapolation to 128K. CoWA achieves average scores close to FullAttn at both context lengths. NIAH-S averages NIAH-S-1, NIAH-S-2, and NIAH-S-3; NIAH-MK averages NIAH-MK-1, NIAH-MK-2, and NIAH-MK-3; and RULER-QA averages RULER-QA-SQuAD and RULER-QA-HotpotQA.

Efficient Attention

Efficient attention reduces the compute, memory traffic, and cache costs of Transformer attention while preserving useful long-range interactions.

See all 134 papers

Sep 26, 2026

MassAlloc Attention: Let Attention Allocate Its Own Compute

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.

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

Attention Mechanisms Efficient Attention LLM Inference Acceleration Sparse Attention

Ask PDF

Mar 6, 2026

Stem: Rethinking Causal Information Flow in Sparse Attention

The quadratic computational complexity of self-attention remains a fundamental bottleneck for scaling Large Language Models (LLMs) to long contexts, particularly during the pre-filling phase. In this paper, we rethink the causal attention mechanism from the perspective of information flow. Due to causal constraints, tokens at initial positions participate in the aggregation of every subsequent token. However, existing sparse methods typically apply a uniform top-k selection across all token positions within a layer, ignoring the cumulative dependency of token information inherent in causal architectures. To address this, we propose Stem, a novel, plug-and-play sparsity module aligned with information flow. First, Stem employs the Token Position-Decay strategy, applying position-dependent top-k within each layer to retain initial tokens for recursive dependencies. Second, to preserve information-rich tokens, Stem utilizes the Output-Aware Metric. It prioritizes high-impact tokens based on approximate output magnitude. Extensive evaluations demonstrate that Stem achieves superior accuracy with reduced computation and pre-filling latency.

Lin Niu, Xin Luo, Linchuan Xie +4

Tencent · School of Biomedical Engineering, Division of Life Sciences and Medicine, University of Science and Technology of China, Hefei, Anhui, 230026, China · Suzhou Institute for Advanced Research, University of Science and Technology of China, Suzhou, Jiangsu, 215123, China +5

Causal Attention Sparse Attention Self-Attention LLM Inference Acceleration

Ask PDF

Sep 28, 2026

SMat-Attention: Structured Long-Context Sequence Modeling

Long-context sequence models face a fundamental tradeoff: softmax attention uses flexible token-level interactions at quadratic cost, whereas linear attention obtains linear-time training and constant-time decoding by compressing history into a fixed-size state. In this work, we ask whether we can connect these regimes through a tunable notion of structure. To this end, we introduce Structured Matrix Attention (SMat-Attention) via a family of causal masks with structured long-range routing whose row supports have VC-dimension

d

. In our construction,

d=1

recovers the standard causal mask, and increasing

d

permits richer subset-routing patterns. We give chunkwise forward and backward algorithms to enable hardware-efficiency. For sequences of length

T

, the hard-routing construction takes

O(T^{2-3/d}+T)

work, despite the mask being dense, for our prescribed family. In fixed-horizon streaming, decoding after the distant prefix takes constant time per token using

O(T^{1-1/d})

cached states. SMat-Attention therefore makes VC-dimension an explicit knob governing access-pattern complexity, prefill cost, and decoding memory. Empirically, subset-routing and rule-assisted multi-key retrieval experiments illustrate the masks' routing expressiveness. Extensions to Mamba-2 and Gated DeltaNet using learned routing with top-

k

query reads retain subquadratic prefill, improve recall accuracy over the backbones in several settings, and achieve comparable small-scale language-modeling performance.

Emile Anand, Abdullah Ateyeh, Archer Wang +1

Massachusetts Institute of Technology

Long-Context Modeling Linear Attention Attention Mechanisms Attention Routing

Ask PDF

Abstract

Efficient Attention

Explore similar work

MassAlloc Attention: Let Attention Allocate Its Own Compute

Stem: Rethinking Causal Information Flow in Sparse Attention

SMat-Attention: Structured Long-Context Sequence Modeling