MOIRA: Mass-Oriented Indexing with Ragged Attention for Long-Context Decoding

Dich Nhat Minh Nguyen; Tran Dang Duong Nguyen

cs.AIOct 3, 2026

MOIRA: Mass-Oriented Indexing with Ragged Attention for Long-Context Decoding

Authors: Dich Nhat Minh Nguyen, Tran Dang Duong Nguyen

Abstract

Long-context decoding is limited by memory bandwidth, because every output token reads the KV cache of every layer. Sparse decoding reduces this cost by reading only part of the KV cache. We observe that the number of pages a query needs varies widely across KV heads, layers and steps. Fixed budgets are simple, but they are sized for demanding cases and tuned per workload; adaptive budgets follow this variation more flexibly, but existing designs pay for it with extra selection cost or training. At the kernel level, FlashAttention-3 (FA3) and FlashInfer are designed for rows of similar length: with page lists whose length differs per KV head, they either pad the lists (forfeiting much of the sparse saving), leave thread blocks unbalanced, or rely on a host-side plan that runs outside the CUDA graph. We propose MOIRA, a training-free sparse decode path in vLLM whose budget adapts per KV head and per layer. For every request, layer, KV head and step, a coverage rule keeps the smallest set of pages whose estimated attention mass reaches a fraction $γ$ . A new kernel, self-planning attention, lets each thread block derive its own share of the work from the list lengths, so the whole decode step stays inside the CUDA graph. On an H200, at RULER's 128k context, MOIRA with $γ=0.99$ matches dense accuracy while reading about 30% of the pages and reduces the time per output token (TPOT) by 2.2-2.5 $\times$ relative to dense FA3; with $γ=0.98$ it reduces TPOT by 2.7 $\times$ and stays within the noise of dense. Under high serving load it raises throughput by up to 51%. These results suggest that a budget adapted per head and layer, paired with a kernel that keeps such budgets inside the CUDA graph, makes sparse decoding both flexible and fast.

Figures & tables

Figure 1: The budget a row needs varies, and fixed budgets or coarse lengths forfeit the savings. Oracle coverage sets $\mathcal{C}^{\star}(0.95)$ over 60 LongBench requests, 64-token pages. (a) Mean size of the set, as a fraction of the context, for each (layer, KV head) of Llama-3.1-8B: the first layer needs most of its context, while most layers need only a tenth. (b) Tokens read by budget policies relative to the coverage set, when each policy must give 90% of the rows, over all layers and steps, at least 0.95 of the mass: every fixed or profiled budget reads 1.7–2.4 $\times$ as many tokens as the coverage set (dashed line). (c) Fraction of the dense KV cache read by batches of 32 requests (total over decode steps 16–256) when the per-(layer, KV head) lengths are coarsened: padding to the longest head doubles the reads.

Panel of Figure 2 — Figure 2: Existing decode kernels on per-head lists of different lengths. Each bar is one (request, KV head) row with its selected pages; each segment is the work of one thread block; red is padding. Times: attention of one layer of Llama-3.1-8B on lists from real decode steps, scaled to 32k context $\times$ 16 requests (H200). Only self-planning attention reads exactly the selected pages, gives every thread block a bounded, nearly equal amount of work, and runs inside a CUDA graph.

Panel of Figure 3 — Figure 3: One decode layer of Moira , for one (request, KV head) row. Step 1 approximates the attention mass of each page; Step 2 keeps the fewest pages whose total mass reaches $\gamma$ ; Step 3 reads those pages, with each thread block taking one chunk of the row by itself; Step 4 combines the chunks. Values are illustrative; Step 3 uses the row lengths of Figure 2 .

Llama-3.1-8B
	accuracy		latency (ms)
Method	RULER $\uparrow$	KL $\downarrow$	Attn	Host	TPOT	$\times$
dense FA3	71.51	–	15.59	–	19.0	1.00
Quest-36k a	69.63	0.022	21.98	2.11	–	–
Twilight-36k	67.94 †	0.053	14.73	6.96	–	–
FFD $\delta 7$ b	64.85 †	0.071	13.03	–	18.8	1.01

Table 1: Accuracy and latency at 128k. RULER-13 (13 tasks

\times

20 documents); † : significantly below dense (paired bootstrap by document, 95%). KL: mean token KL to dense on 65 teacher-forced RULER answers (0.0003 with every page kept). Attn: selection + attention of one decode step at 4

\times

130k on the GPU; Host: FlashInfer’s plan() call per step, with its index arrays already built. TPOT: measured end to end in vLLM with CUDA graphs on RULER prompts (batch 4); –: the method does not run in vLLM; dense FlashInfer: 18.8 / 20.0 ms.

\times

: dense FA3 TPOT / TPOT. Page reads of Moira (and of the c rows, same lists) at

\gamma{=}0.98

/ 0.99: 19.6 / 27.4% on Llama-3.1-8B, 24.2 / 32.5% on Qwen3-4B. Bold: best sparse value.

	Pages	Accuracy	Trun-	Closed
Budget	read (%)	(%)	cated	acc. (%)
dense	100	75.2	89	92.1
Fixed budget
proportional top- $k$	39.5	41.0 †	231	78.3
fixed $K{=}75$	33.1	69.8 †	131	95.7
fixed $K{=}90$	39.8	72.7	108	93.8

Table 2: Adaptive vs. fixed budgets on AIME (Qwen3-4B-Thinking, 480 solutions). All sparse arms use the same estimator and kernels and differ only in the rule that sets the page count. Truncated: solutions that hit the 32,768-token limit. Closed acc.: accuracy on the solutions that finish before the limit. Accuracy also counts the 0–2 truncated solutions per arm whose answer, extracted from the truncated text, is correct. † : significantly below Moira

\gamma{=}0.98

(paired bootstrap by problem, 95%).

Figure 4: Attention time per layer as per-head lengths spread (Llama-3.1-8B geometry, 22% of the pages kept on average; the x axis is the coefficient of variation (CV) of the per-(request, KV head) lengths). “real”: per-(request, layer, KV head) page counts from decode steps of Moira on LongBench prompts of about 8k tokens, applied as fractions of each context (0.33–0.41 of the pages; short contexts need a larger fraction than RULER at 128k). FlashInfer additionally needs the host plan given in each panel, per layer and step. The 64k panel has no FlashInfer arm.

	Llama-3.1-8B				Qwen3-4B
Shape	dense	Moira	$\times$	reads	dense	Moira	$\times$	reads
8k $\times$ 32	12.66	9.18	1.38	0.42	12.42	8.98	1.38	0.49
16k $\times$ 32	19.24	10.83	1.78	0.37	20.59	11.68	1.76	0.43
32k $\times$ 16	16.46	8.71	1.89	0.31	17.00	8.92	1.91	0.39
64k $\times$ 8	19.56	7.45	2.62	0.16	20.43	8.39	2.44	0.25
130k $\times$ 4	19.32	7.32	2.64	0.15	20.26	7.53	2.69	0.21

Table 3: TPOT (ms) of Moira (

\gamma{=}0.98

) and dense FA3 on prompts cut from LongBench documents (batch

\times

context as listed), mean of two interleaved pairs (whose speedups differ by at most 0.13), and page reads of Moira . Bold: lower TPOT.

	dense		Moira $\gamma{=}0.98$
Shape	step	FA3	step	att	est	sel	mrg	GEMM
8k $\times$ 32	12.72	8.16	9.38	4.02	0.38	0.23	0.18	3.96
32k $\times$ 16	17.00	12.48	9.45	4.20	0.40	0.25	0.13	3.86
64k $\times$ 8	20.34	15.89	8.44	3.04	0.44	0.41	0.12	3.83
130k $\times$ 4	20.16	15.76	8.18	2.78	0.43	0.42	0.13	3.83

Table 4: One decode step of Llama-3.1-8B by kernel (ms, Nsight Systems median of 128 steps). att, est, sel, mrg: Attend , Estimate , Select , Merge ; GEMM: matrix multiplications. Tracing at CUDA-graph node level lengthens the step by up to 1 ms relative to Table 3 . Bold: shorter step.

Model	set, $c$	FA3	FI	Moira	FA3	FI	Moira
		median TPOT (ms) $\downarrow$			output tok/s $\uparrow$
Llama-8B	A, 1	5.63	5.86	6.36	152	147	137
	A, 16	28.6	34.2	17.2	511	446	743
	B, 1	8.16	8.18	6.71	64.9	63.5	71.2
	B, 8	65.5	63.6	31.6	99.9	101.6	142.0
Qwen3-4B	A, 1	4.37	4.62	5.10	194	185	171

Table 5: Online serving ( vllm bench serve , closed loop, 1,024 output tokens). Set A: 256 LongBench prompts of 4k–64k tokens; set B: 64 prompts of 64k–126k tokens.

c

: concurrency. FI: dense FlashInfer. Bold: best per row.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Shape	Rows	CV	Moira	varlen	pad	pow2	kernel	plan	Dense
				FA3			FlashInfer
8k $\times$ 32	256	0	70.1	68.3	68.5	75.7	65.5	668	248
		1	73.5	84.2	168.4	107.5	105.0	672
		real	120.5	139.2	228.5	161.7	149.0	683
32k $\times$ 16	128	0	120.6	118.6	118.7	132.1	120.8	455	489
		1	123.2	151.3	337.3	207.3	119.4	467

Appendix

Table 6: Attention time per layer (µs) on the same page lists (Llama-3.1-8B geometry; data of Figure 4 ). CV 0: all rows of equal length; CV 1: nominal spread (measured 0.84–0.88); real: page counts from decode steps of Moira at about 8k tokens, applied as fractions of each context. FlashInfer needs its host plan at every step (plan: page counts copied to the host, index arrays built and plan() called, per layer); “–”: not measured. Bold: fastest method per row, FlashInfer counted with its plan.

	Llama-3.1-8B				Qwen3-4B
Shape	dense	Moira	$\times$	reads	dense	Moira	$\times$	reads
8k $\times$ 32	12.84	10.29	1.25	0.55	12.37	10.02	1.23	0.60
16k $\times$ 32	19.57	13.10	1.49	0.49	20.76	14.00	1.48	0.54
32k $\times$ 16	16.65	9.91	1.68	0.43	17.15	10.73	1.60	0.50
64k $\times$ 8	19.70	9.12	2.16	0.24	20.66	10.22	2.02	0.34
130k $\times$ 4	19.50	8.59	2.27	0.22	20.49	9.03	2.27	0.28

Appendix

Table 7: TPOT (ms) with

\gamma{=}0.99

and dense FA3, same shapes and protocol as Table 3 (mean of two interleaved pairs), and page reads of Moira . Bold: lower TPOT.

Figure 5: Online serving with vllm bench serve (closed loop, prefix caching off, 1,024 output tokens, one run per point). Set A: LongBench prompts of 4k–64k tokens; set B: 64k–126k tokens.

Figure 6: TPOT in vLLM ( Moira at

\gamma{=}0.98

vs. dense FA3; data of Table 3 ). Median of 5 repetitions, two interleaved pairs per shape; the error bars span the two pairs.

Figure 7: Sparse decode methods vs. dense FA3, attention only (Section 5.5.1 ). (a) One decode step (32 layers) of Llama-3.1-8B at 4

\times

130k. Twilight: the authors’ selection chain and FlashInfer attention on its own token sets (GPU), plus host metadata and plan (hatched). (b) Speedup against the kept fraction at the same geometry ( Moira : 64-token pages; FFD: 128/256-token blocks at a preset kept fraction; shaded band: FFD’s kept fraction on RULER). Moira uses page counts from a real decode step, scaled to the reads of

\gamma{=}0.98

and

0.99

on RULER and, labelled “scaled”, to an intermediate fraction. (c) Across shapes: FFD at about 15% kept; Moira on real

\gamma{=}0.98

steps (kept fraction in parentheses).

Language Model Decoding

Language model decoding studies how models choose and generate outputs, balancing quality, reliability, diversity, and inference cost.

See all 168 papers

Jun 3, 2026

SparDA: Sparse Decoupled Attention for Efficient Long-Context LLM Inference

Sparse attention reduces compute and memory bandwidth for long-context LLM inference. However, two key challenges remain: (1) KV cache capacity still grows with sequence length, and offloading to CPU memory introduces a PCIe transfer bottleneck; (2) the sparse selection step itself retains

O(T^2)

complexity and can dominate attention cost at long contexts. We propose SparDA, a decoupled sparse attention architecture that introduces a fourth per-layer projection, the Forecast, alongside Query, Key, and Value. The Forecast predicts the KV blocks needed by the next layer, enabling lookahead selection that overlaps CPU-to-GPU prefetch with current-layer execution. Because Forecast is decoupled from the attention query, our GQA implementation uses one Forecast head per GQA group, reducing selection overhead versus the original multi-head selector. SparDA adds

<

0.5% parameters and trains only the Forecast projections by matching the original selector's attention distribution. On two sparse-pretrained 8B models, SparDA matches or slightly improves accuracy and delivers up to 1.25

\times

prefill speedup and 1.7

\times

decode speedup over the sparse-attention offload baseline. By enabling larger feasible batch sizes on a single GPU, SparDA further reaches up to 5.3

\times

higher decode throughput than the non-offload sparse baseline. Our source code is available at https://github.com/NVlabs/SparDA.

Yaosheng Fu, Guangxuan Xiao, Xin Dong +2

1NVIDIA · 2Thinking Machines Lab · 3ByteDance Seed +1

KV-Cache Offloading Long-Context Language Model Inference LLM Inference Acceleration Sparse Attention

Ask PDF

Oct 6, 2026

SPIN: Shadow Predictive Indexer for Sparse Attention

Indexer-based sparse attention reduces the cost of core attention by passing only a fixed, small number of important tokens to it. However, the indexer must still score the entire KV cache at every decoding step. This scoring overhead becomes a major bottleneck as the context length grows. We propose SPIN (Shadow Predictive Indexer) to reduce this indexer overhead. SPIN uses lightweight, history-based prediction to identify important KV blocks, avoiding the need to score the full KV cache at every decoding step. SPIN treats KV blocks and speculative decoding as first-class design and implementation considerations. Across extensive evaluations on long-context and agentic benchmarks, SPIN achieves 30-40% sparsity while preserving task quality. In end-to-end vLLM serving, SPIN improves output throughput by up to 14.9% and reduces median inter-token latency by up to 13.2%.

Yao Fu, Jiahan Chang, Ritchie Zhao +11

LLM Inference Acceleration Sparse Attention Long-Context Language Model Inference Block-Sparse Attention

Ask PDF

Sep 16, 2026

Elastic Threshold Attention: Learned Contextual Sparsity for Long-Context Decoding

Massive KV caches can cause severe memory-bandwidth bottlenecks during long-context decoding. Sparse attention methods mitigate this problem, but often drop necessary context, leading to quality degradation. We introduce Elastic Threshold Attention (ETA), an end-to-end trainable architecture that rivals dense model quality under hardware-aligned block-sparse decoding. ETA predicts dynamic, contextual thresholds directly from query representations, adjusting context retention depending on the task at hand. To learn this policy from scratch while enabling near lossless KV cache pruning at inference time, we filter attention logits through a shifted SiLU gate during training. We show theoretically and empirically that this creates a smooth, near-uniform attention floor that neutralizes sub-threshold value contributions while simultaneously causing localized attention sinks on initial tokens to disappear. To materialize these advantages, we implement a fused inference-time kernel in Triton that screens KV blocks in

O(1)

time using cached geometric-probabilistic bounds. Across language modeling, reasoning, and RULER benchmarks, our 1.45B ETA model matches or exceeds dense quality, outperforming alternative fast decoding methods (Quest, H

_2

O, NSA) while achieving higher sparsity levels. Our kernel also achieves up to

2.15\times

end-to-end speedups over FlashAttention-2 at context lengths of up to

512

K tokens. Finally, we introduce an offline calibration algorithm for domain-specific deployments that freezes per-head constant thresholds, cutting attention compute by an additional 27% at no quality cost.

Themistoklis Haris, Henry Li, Maryam Karimzadehgan

Google · Boston University

Long-Context Language Model Inference Sparse Attention Block-Sparse Attention KV-Cache Eviction

Ask PDF

Abstract

Language Model Decoding

Explore similar work

SparDA: Sparse Decoupled Attention for Efficient Long-Context LLM Inference

SPIN: Shadow Predictive Indexer for Sparse Attention

Elastic Threshold Attention: Learned Contextual Sparsity for Long-Context Decoding