采用分层树状索引器的次二次复杂度注意力机制

# Subquadratic Attention with a Hierarchical Tree-Structured Indexer
Source: [https://arxiv.org/html/2610.42112](https://arxiv.org/html/2610.42112)
Cutter Dawes\*Alexander SpeicherKamesh Krishnamurthy†Tomás Figliolia†Beren Millidge†Affiliation:1Zyphra Research\* Equal contribution† Senior AuthorAffiliation:Correspondence:\{leon,tom\}@zyphra\.com

###### Abstract

The quadratic scaling of attention is a fundamental bottleneck limiting the long\-context capabilities of large language models \(LLMs\)\. A variety of approaches have attempted to circumvent this, but they either fail to match attention’s expressiveness \(e\.g\., sliding\-window or linear attention variants with fixed\-size state\) or only achieve constant\-factor speedups \(e\.g\., sparse attention methods whose indexers themselves scale quadratically\)\. In this paper, we derive an asymptotically subquadratic sparse attention method based on a hierarchical, tree\-structured indexer, which we call Zyphra Hierarchical Attention \(ZHA\)\. ZHA constructs a tree of recursively compressed key summaries over the KV cache, which future queries traverse from the top down to retrieve a sparse subset of keys\. Crucially, by iteratively growing the tree’s depth, the total cost including routing and final attention grows sublinearly in sequence length, resulting in overall subquadratic \(but superlinear\) scaling for attention\. We demonstrate that ZHA can be cheaply applied to existing pretrained models with only a short KL distillation training stage over a few billion tokens, and maintains performance on long\-context retrieval and challenging reasoning, mathematical, and agentic coding evaluation tasks\. Moreover, we show that our method is highly robust to tree geometry including very deep trees of up to 6 levels\. Finally, we develop a detailed performance model and show that our approach scales subquadratically, dramatically outperforms competing methods beyond 1 million context, and ultimately has the potential to extend effective context windows to the hundreds of millions or billions of tokens\.

## IIntroduction

Attention is the fundamental sequence\-mixing primitive in modern transformer architectures\[[41](https://arxiv.org/html/2610.42112#bib.bib41)\], and it is the basis for the highly flexible and expressive sequence modeling that underpins the recent rapid and dramatic advances in AI\. Attention expresses an all\-to\-all mapping between sequence elements, where each token in the sequence is compared to every other token to identify those which are relevant\. It is highly expressive and unstructured, yet parallelizable and efficiently implementable on GPUs\. These properties have made it the predominant sequence\-mixing architecture and enabled massive scaling of model size and sequence lengths\.

However, the all\-to\-all nature of attention leads to fundamentally unfavorable scaling in computational and memory costs\. In prefill, its quadratic computational cost dominates other operations, while during decode, re\-reading the whole key\-value \(KV\) cache per new token imposes memory bandwidth and storage costs which linearly increase in the sequence length\. This has prohibited extremely long LLM contexts, despite continuing gains from longer reasoning and input lengths across a wide range of increasingly challenging tasks\[[9](https://arxiv.org/html/2610.42112#bib.bib9),[2](https://arxiv.org/html/2610.42112#bib.bib2)\]\.

This has motivated a variety of approaches that reduce the computational cost of attention\[[38](https://arxiv.org/html/2610.42112#bib.bib38)\]\. State\-space models and sliding\-window attention\[[6](https://arxiv.org/html/2610.42112#bib.bib6),[15](https://arxiv.org/html/2610.42112#bib.bib15),[24](https://arxiv.org/html/2610.42112#bib.bib24),[12](https://arxiv.org/html/2610.42112#bib.bib12)\]replace the full KV cache with a fixed\-size state that is updated along the sequence\. While these eliminate the asymptotic cost of attention with linear prefill and constant decode cost, their fixed\-size state saturates at long contexts, leading to worse performance\[[21](https://arxiv.org/html/2610.42112#bib.bib21)\]\. Practical “hybrid” methods\[[14](https://arxiv.org/html/2610.42112#bib.bib14),[26](https://arxiv.org/html/2610.42112#bib.bib26),[13](https://arxiv.org/html/2610.42112#bib.bib13),[39](https://arxiv.org/html/2610.42112#bib.bib39)\]thus often interleave a fixed\-state mixer with full attention, reaping the computational benefits in most layers while restricting the attention costs to only a few, necessary layers\. However, at sufficiently long contexts, the few layers using full attention still dominate\.

Other approaches aim to compress the sequence directly by online\-compressing multiple tokens into a single token during prefill or decode \(e\.g\., CSA\[[8](https://arxiv.org/html/2610.42112#bib.bib8)\], NSA\[[46](https://arxiv.org/html/2610.42112#bib.bib46)\], HNet\[[20](https://arxiv.org/html/2610.42112#bib.bib20)\]\), or by compressing the KV cache by reducing the number of independent heads\[[1](https://arxiv.org/html/2610.42112#bib.bib1),[28](https://arxiv.org/html/2610.42112#bib.bib28),[11](https://arxiv.org/html/2610.42112#bib.bib11)\]\. A related set of methods focuses on offline KV cache eviction, where tokens deemed to be of little value are evicted periodically from the KV cache, reducing the cost of all future queries\[[23](https://arxiv.org/html/2610.42112#bib.bib23)\]\. However, although head\-compression methods such as GQA and MLA have become ubiquitous, they are fundamentally limited in possible compression rate by the number of heads\. Meanwhile, sequence\-compression and KV\-eviction methods theoretically enable large compression rates but have fared less well empirically\. Many of these methods both only obtain a constant factor speedup, and also are usually brittle and scale poorly in practice\. A core reason for this brittleness is that it is practically impossible to know ahead of time, or on average, whether a token is useful or not\.

This leads directly to a family of*fine\-grained sparse attention*methods that enable each query to retrieve its own sparse subset of keys and values from the context \(e\.g\., DSA\[[29](https://arxiv.org/html/2610.42112#bib.bib29)\], MSA\[[25](https://arxiv.org/html/2610.42112#bib.bib25)\]\)\. In these methods, a lightweight indexer attention scores every KV pair’s relevance to each query, and full attention is performed only on the top\-KKmost relevant pairs\. SinceKKis fixed, the full attention has a constant cost\. Crucially, this approach still permits each query to select an independent subset of keys most important to it, which enables higher precision for retrieval than compressing the same set of keys for all queries\. However, this flexibility comes with two costs\. First, the full KV cache must be stored, even if not computed over every token\. Second, and most importantly, it relocates the cost of attention into the indexer\. In practice, one can obtain significant speedups by making the indexer cheap with fewer heads or lower precision, but at contexts near 1M tokens, 60\-90% of time in these layers is still spent on indexing\[[3](https://arxiv.org/html/2610.42112#bib.bib3),[47](https://arxiv.org/html/2610.42112#bib.bib47)\]\. A subfamily of*chunk\-sparse*attention methods\[[46](https://arxiv.org/html/2610.42112#bib.bib46),[18](https://arxiv.org/html/2610.42112#bib.bib18),[19](https://arxiv.org/html/2610.42112#bib.bib19)\]attempts to address this by compressing the indexer itself, and sparsely selecting chunks from the KV cache over which to perform attention\. Of interest is[Huang et al\. \[19\]](https://arxiv.org/html/2610.42112#bib.bib19)which uses an adaptive sparsity functionα\\alpha\-entmax, similar to our top\-pprouting in spirit, and[Hu et al\. \[18\]](https://arxiv.org/html/2610.42112#bib.bib18), which constructs chunk summaries as an attention\-weighted sum, similar to our own key construction method\. However, these methods tend to be empirically brittle and only perform chunking over a single level, thus also only obtaining constant factor speedup\.

This leads to the conclusion that to construct an effective attention replacement yielding better than constant\-factor speedups, two properties must be satisfied\. The first property, derived from sparse attention, is*query\-dependent sparsity*, since different queries can require substantially different information from the KV cache\. The second property is a*growing but sublinear active state size*\. New sequence elements add predictive information indefinitely, so a fixed\-size state will inevitably reach capacity and degrade in performance\. By contrast, full attention can be viewed as a non\-parametric state that grows linearly with sequence length\. This linear growth is maximally expressive but highly redundant\[[48](https://arxiv.org/html/2610.42112#bib.bib48),[30](https://arxiv.org/html/2610.42112#bib.bib30),[4](https://arxiv.org/html/2610.42112#bib.bib4),[10](https://arxiv.org/html/2610.42112#bib.bib10)\], since the predictive information of each new token diminishes: the millionth token contains information, but likely less than the first\. The challenge is to design a state that can grow dynamically, but sublinearly, at the rate that predictive information accumulates in the sequence\.

Motivated by these considerations, we introduce a novel hierarchical sparse attention method, which we call Zyphra Hierarchical Attention \(ZHA\), that can achieve better than constant\-factor speedups over dense attention\. By analyzing the structure of attention’s representations, we find that the size of its effective state is sublinear and that it is amenable to a hierarchical indexer\. Moreover, we show that we can adapt pretrained models to use our hierarchical indexer with distillation training of only a few billion tokens and very few additional parameters, without changing any of the base model’s weights\. We convert models up to 35B parameters and observe virtually no degradation on challenging evaluations covering mathematical reasoning, agentic coding, and long\-context retrieval\. We also perform a detailed study of various indexer geometries and demonstrate that ZHA is highly robust to such variations and, crucially, can maintain performance at large depths\.

Finally, we develop a detailed theoretical performance model of our method as well as comprehensive empirical sparsity results for a variety of indexer geometries\. Our results demonstrate that ZHA can remain asymptotically subquadratic, subject to careful top\-ppscaling, when extrapolated to extremely long sequences of tens of millions to billions of tokens, and can thus theoretically strongly outperform sparse attention methods at these context lengths in our performance model\. We model ZHA’s speedups against dense attention and MSA with a roofline model\[[45](https://arxiv.org/html/2610.42112#bib.bib45)\], which predicts a crossover point against the latter around 0\.5M tokens for a performant geometry withp=0\.75p=0\.75\.

## IIMotivation

### II\-AThe Natural Sparsity of Attention

Our work starts from the observation that attention’s representations on language data already admit a sublinear state\. Importantly, sublinear compressibility is not fundamental to the attention architecture but rather determined by the data and task distribution\. Adversarial sequences consisting of random tokens from which arbitrary tokens must be retrieved uniformly throughout the sequence would require full quadratic attention to process\. However, we observe that natural language sequences appear, on average, to follow a sublinear growth of predictive information with sequence length\.

We operationalize this notion by studying the size of the top\-ppKV subset of global attention layers in a number of language models on a dataset of long\-context fiction books\. The top\-ppset is the smallest subset of keys whose attention probabilities sum to at leastpp\. Top\-ppcan be thought of as a denoiser for identifying attention’s “true” receptive field\. We can write the output of attention as

Attention⁡\(q,K,V\)\\displaystyle\\operatorname\{Attention\}\(q,K,V\)=\\displaystyle=\{\}p\\displaystyle pAttention⁡\(q,Kp,Vp\)\\displaystyle\\operatorname\{Attention\}\(q,K\_\{p\},V\_\{p\}\)\(1\)\+\(1−p\)\\displaystyle\{\}\+\(1\-p\)Attention⁡\(q,Kp¯,Vp¯\)\.\\displaystyle\\operatorname\{Attention\}\(q,K\_\{\\bar\{p\}\},V\_\{\\bar\{p\}\}\)\.where\(Kp,Vp\)\(K\_\{p\},V\_\{p\}\)denotes the top\-ppsubset for queryqq, and\(Kp¯,Vp¯\)\(K\_\{\\bar\{p\}\},V\_\{\\bar\{p\}\}\)contains the remaining KV pairs\. When performing top\-ppattention, we keep only the first term in[equation1](https://arxiv.org/html/2610.42112#S2.E1), without the prefactorpp\. If it is the case that only a subset of tokens are relevant for a query, it is still difficult for attention to set the other tokens’ probabilities precisely to zero, but they should still have far smaller scores than the relevant subset\. So, we can interpret the first term in[equation1](https://arxiv.org/html/2610.42112#S2.E1)as the true “signal” and the second term as noise\. By choosing the appropriatep<1p<1, we can extract the desired signal\.

We find the size of the top\-ppsubset in each attention layer grows as a sublinear power law in the context length, with exponent depending onpp\(see[Figure1](https://arxiv.org/html/2610.42112#S2.F1)\)\. Appendix[A\-A](https://arxiv.org/html/2610.42112#A1.SS1)shows top\-ppcurves broken down by layer and head\.111Assessing whether an entire layer or model admits a power law is more conservative than studying a single attention head\. Here, we wish to see that all the layers with significant attention exhibit sublinear growth in their top\-ppsubset sizes, establishing that compressibility extends across the model rather than being confined to isolated heads\.[TableI](https://arxiv.org/html/2610.42112#A1.T1)provides a complementary comparison on ZHA rollout selection counts, where power laws are preferred in many configurations but not uniformly across checkpoint geometries and query\-tile widths\. These fits describe router\-selected leaf unions, rather than dense top\-ppsets\.

That this sublinear growth is a power law already suggests hierarchical structure, although does not prove this structure is efficiently searchable in practice\. Power laws are typically associated with scale invariance, that is, an underlying representation that remains the same at every level of analysis\. Here, this would likely correspond to the statistics of some summary of tokens remaining the same as we summarize groups of tokens of geometrically increasing size\. If we could learn such statistics and use them for routing, then perhaps we could route recursively through these geometric groups at increasing levels of precision according to a simple decision rule\.

We note that this power\-law structure does not always hold, and it seems to depend on positional encoding\. Qwen3\.6\-35B\-A3B uses partial RoPE, whereas Nemotron\-3\.5\-Lightning\-30B\-A3B, which we study in[Figure12](https://arxiv.org/html/2610.42112#A1.F12), uses NoPE and does not show as robust power\-law top\-ppscaling, though it still performs comparably with ZHA \(see Appendix[B\-D](https://arxiv.org/html/2610.42112#A2.SS4)\)\.

Figure 1:Dense top\-ppset size versus causal context length in Qwen3\.6\-35B\-A3B on six fiction books \(averaged across books, layers, and heads\)\. We observe that at allp<1p<1\(shades of blue\), the top\-ppset size increases sublinearly as a power law, indicating that longer sequences become increasingly compressible\.
### II\-BSublinear Retrieval through Hierarchical Search

That attention for an individual query is sparse does not imply that most KV entries can be permanently discarded\. The relevant subset can change with the query such that keys receiving negligible attention for a given query may nevertheless become important for a later query\. Sparse individual attention distributions can therefore coexist with broad coverage of the cache across the union of future queries\.

To assess whether sparse attention supports permanent KV eviction, we measure the union of the top\-ppsubsets over increasing numbers of future queries, repeating this analysis for several values ofpp\. More precisely, among sequences from long fiction books of lengthT=261,548T=261\{,\}548tokens, we consider the window of theWWmost\-recent queries \(1≤W<T1\\leq W<T\) and study

Up​\(W\)=1T−W​\|⋃q=T−WT−1Sp​\(q\)\|,U\_\{p\}\(W\)=\\frac\{1\}\{T\-W\}\\left\|\\bigcup\_\{q=T\-W\}^\{T\-1\}S\_\{p\}\(q\)\\right\|,whereSp​\(q\)S\_\{p\}\(q\)is the top\-ppsubset of indices \(using full causal attention\) for the query at positionqq, but this subset of indices is restricted to the prefix\{0,…,T−W−1\}\\\{0,\\ldots,T\-W\-1\\\}after selection, without renormalization\. Thus,UUis the fraction of the prefix covered by the union of all the top\-ppsets across allWWmost\-recent queries\.

IfUp​\(W\)U\_\{p\}\(W\)is near 1, this suggests that nearly all keys prior to timeT−WT\-Ware in the top\-ppset for some query at or after timeT−WT\-W\. We see in[Figure2](https://arxiv.org/html/2610.42112#S2.F2)that this is inevitably the case asWWgrows large\. That is, the union expands to cover the entire cache, indicating that entries receiving negligible attention from some queries may still be relevant to others\. These results motivate retaining the full KV cache while retrieving a query\-dependent subset for each attention computation\.

Figure 2:Fraction of KV entries of a 256K sequence covered by the union of top\-ppattention subsets as the number of future queries increases, for several values ofpp\. We observe that for the higherppvalues we use in practice, a substantial fraction of the total KV cache is included in the union of the queries\.If we are to identify the top\-ppsubset asymptotically faster than quadratic attention, then the work needed to identify relevant entries must grow sublinearly with context length, i\.e\., we cannot compare every query directly to every prior key\. A hierarchical, geometric tree\-structured indexer offers a natural mechanism for addressing this search problem\. Coarse summaries guide the query toward promising groups of keys\. If a group of keys is deemed relevant via its summary, then those keys are scored again at a higher resolution, while if they are not deemed relevant, then they can all be discarded together without needing to examine them individually\. This can be an effective way to reduce the number of query\-key comparisons, provided that the summaries can sufficiently capture relevance\.

It is not obvious a\-priori that a tree is expressive enough to identify the top\-ppsubset with sufficient accuracy, since exact computation of this subset requires scoring the entire KV cache\. However, we saw in[SectionII\-A](https://arxiv.org/html/2610.42112#S2.SS1)that the top\-ppsubset specifically follows a power\-law, which suggests an underlying tree\-like structure that we aim to learn\.

## IIIArchitecture

ZHA is designed to approximate quadratic attention despite subquadratic computational complexity\. Our approach involves first constructing a tree of recursively compressed keys, each of which summarizes a chunk of keys in the level below it \(which may themselves be compressed\)\. Compression is performed using chunk\-wise attention scores from trained reference queries\. Then, once the tree is constructed, we route through the tree with a given query by selecting the top\-ppnodes at a given level to expand\.

We train just the hierarchical indexer, whose only learnable components are a small set of reference queries\.*ZHA may therefore be applied to existing pretrained models with only a small distillation phase*; in fact, we find that a short KL warmup of only a few billion tokens is sufficient to match the base model’s performance\. ZHA does not require finetuning the full model, and can support a wide range of attention configurations such as different positional encodings\[[37](https://arxiv.org/html/2610.42112#bib.bib37)\], gating\[[34](https://arxiv.org/html/2610.42112#bib.bib34)\], and head\-structures such as GQA, MLA, or CCA\[[1](https://arxiv.org/html/2610.42112#bib.bib1),[28](https://arxiv.org/html/2610.42112#bib.bib28),[11](https://arxiv.org/html/2610.42112#bib.bib11)\]\. An overview of the ZHA architecture is shown in[Figure3](https://arxiv.org/html/2610.42112#S3.F3)\.

Figure 3:Overview of the ZHA architecture: \(a\) tree construction and \(b\) sparse routing\. During tree construction, for each chunk of contiguous key nodes, we construct a compressed key node at the higher level by an attention\-weighted sum of the keys, where the attention is computed with respect to the learned reference queries\. In practice, we use multiple reference queries and store the key nodes for each in parallel\. This process continues recursively with higher level key nodes being constructed from the summaries of lower levels\. During routing, we start at the top of the tree and compute dot\-product similarity scores between our query and the key nodes of that level\. We select the set of nodes to expand using top\-pp, and then recursively descend the tree until we reach the leaf nodes corresponding to raw KV cache entries\. The union of all the leaves we reach defines the sparse KV cache subset that we dispatch to the main attention\.### III\-ATree Construction

To construct the tree, ZHA recursively creates key summaries that compress contiguous chunks of keys in the level below it \(which themselves are either key summaries or, at the bottom of the tree, the original keys\)\. During decode, we construct the key summaries immediately upon generation of a complete chunk of keys, maintaining a buffer until sufficient keys have accumulated to construct that chunk’s summary\. Supposing the chunk size isCC, this implies that we add a compressed key to the first level of the tree everyCCdecode steps, to the second level everyC2C^\{2\}steps, and so on\. Our approach is intentionally lightweight, so that tree construction adds negligible latency amortized across the sequence \(around 1% or less in our roofline model\)\.

To perform well, it is vital that the compressed keys are a faithful summarization of the underlying chunk of keys\. Thus, we aim to directly approximate a chunk’s total attention mass over a wide distribution of queries\. To perform this approximation, we use the observation that the log unnormalized attention mass of a chunk is a convex function of the query that is naturally approximated by a piecewise\-affine tiling of tangent planes\. In particular, we take the maximum over a set of tangents centered at some point in the query distribution \(these become our reference queries\), each of which is cheaply computable by appropriately attention\-weighting the keys and measuring the entropy of the attention weights at that point\.

To make this concrete, denote a contiguous chunk of keys byKC=\[k0,…,kC−1\]K\_\{C\}=\[k\_\{0\},\\dots,k\_\{C\-1\}\]with chunk sizeCC, and letRRbe the number of tangents andddthe head dimension\. Its log unnormalized attention mass for a given query is the log\-sum\-exp \(LSE\) of its scaled attention scores:

FC\(q\):=LSEj∈\[C\]\(q⊤kj/d\)=log∑j∈\[C\]exp\(q⊤kj/d\)\.F\_\{C\}\(q\):=\\operatorname\{LSE\}\_\{j\\in\[C\]\}\\\!\\left\(q^\{\\top\}k\_\{j\}/\\sqrt\{d\}\\right\)=\\log\\sum\_\{j\\in\[C\]\}\\exp\(q^\{\\top\}k\_\{j\}/\\sqrt\{d\}\)\.The first\-order Taylor approximation ofFCF\_\{C\}at some reference queryq^r\\hat\{q\}\_\{r\}gives the tangenttr​\(q\)=q⊤​k^r/d\+b^rt\_\{r\}\(q\)=q^\{\\top\}\\hat\{k\}\_\{r\}/\\sqrt\{d\}\+\\hat\{b\}\_\{r\}, where the compressed key and bias are

k^r\\displaystyle\\hat\{k\}\_\{r\}=Attn⁡\(q^r,KC,KC\)=∑j∈\[C\]exp⁡\(q^r⊤​kj/d\)​kj∑j′∈\[C\]exp⁡\(q^r⊤​kj′/d\)\\displaystyle=\\operatorname\{Attn\}\(\\hat\{q\}\_\{r\},K\_\{C\},K\_\{C\}\)=\\sum\_\{j\\in\[C\]\}\\frac\{\\exp\(\\hat\{q\}\_\{r\}^\{\\top\}k\_\{j\}/\\sqrt\{d\}\)k\_\{j\}\}\{\\sum\_\{j^\{\\prime\}\\in\[C\]\}\\exp\(\\hat\{q\}\_\{r\}^\{\\top\}k\_\{j^\{\\prime\}\}/\\sqrt\{d\}\)\}b^r\\displaystyle\\hat\{b\}\_\{r\}=FC​\(q^r\)−q^r⊤​k^r/d\\displaystyle=F\_\{C\}\(\\hat\{q\}\_\{r\}\)\-\\hat\{q\}\_\{r\}^\{\\top\}\\hat\{k\}\_\{r\}/\\sqrt\{d\}Provided a set of tangents representing a given chunk, its LSE for a future queryqqis approximated bymaxr∈\[R\]⁡tr​\(q\)\\max\_\{r\\in\[R\]\}t\_\{r\}\(q\)\. To summarize a complete node at levelℓ\>1\\ell\>1in the tree, which representsCℓC^\{\\ell\}original keys, we apply the same construction to itsCCchild summaries for each reference query, including the child biases in the attention scores and adding their attention\-weighted biases to the entropy of the merge weights\.222When the same reference query is used throughout a subtree, this merge recovers the summary computed directly from its original keys\. With chunk\-dependent RoPE rotations, it instead gives an affine lower bound that need not be tangent to the original chunk’s log attention mass\.For models with GQA or other KV\-head groupings, we construct one tree per KV head, shared by the query heads in its group\. To provide some intuition for the variety of geometries,[Figure4](https://arxiv.org/html/2610.42112#S3.F4)illustrates trees varying chunk size, depth, and number of reference queries\.

Figure 4:Illustrative ZHA tree geometries of the three core tree parameters\. We can trade off depth and width for a given level of compression in the tree, since compression scales asCLC^\{L\}\. Secondly, at the cost of additional compute, we can increase the number of reference queries\.
### III\-BTraining the Module

Our chunk summarization technique relies upon a suitable choice of reference queries\. Our approach is to learn them during a short KL\-warmup phase; these are in fact the only learnable component in our ZHA mechanism\. We find it sufficient to learn a single set ofRRreference queries per tree, shared by all nodes across the sequence, at all levels in the tree, and, in the case of GQA, by the query heads sharing the tree\. With a reasonable number of reference queriesRR, the ZHA parameters add a negligible increase to the parameter count of the models we test\.333For example, withR=16R=16and one tree per KV head in Qwen3\.6\-35B\-A3B, the parameter count increases by approximately0\.00024%0\.00024\\%\(81,920 / 34\.66B\)\.For models that incorporate RoPE, we rotate the shared reference queries to the last token position of each node’s span, at every tree level\. Letq^r,c\\hat\{q\}\_\{r,c\}denote the reference queryrrrotated for nodecc\. Before merging the children of a parent nodePP, we evaluate all child summaries at the parent’s rotated reference queryq^r,P\\hat\{q\}\_\{r,P\}\. Each childccwas originally summarized at its own rotated reference queryq^r,c\\hat\{q\}\_\{r,c\}\. To evaluate its affine approximation at the parent’s query, we add1d​⟨k^r,c,q^r,P−q^r,c⟩\\tfrac\{1\}\{\\sqrt\{d\}\}\\langle\\hat\{k\}\_\{r,c\},\\hat\{q\}\_\{r,P\}\-\\hat\{q\}\_\{r,c\}\\rangleto its estimated log attention mass\. This correction is simply the compressed key’s score at the parent’s reference query minus its score at the child’s reference query\.

From a trained dense\-attention model, the ZHA conversion procedure is remarkably simple and compute\-efficient\. In particular, we freeze the base model \(including all native attention components\) and learn the reference query parameters by matching the predicted chunk probabilities to the true attention distribution pooled into chunks at each level\. For a queryqq, letaj​\(q\)a\_\{j\}\(q\)denote its dense attention weight on keyjj, and let𝒫ℓ​\(q\)\\mathcal\{P\}\_\{\\ell\}\(q\)partition the original key indices visible toqqinto chunks𝒞\\mathcal\{C\}at tree levelℓ\\ell\. The target and predicted chunk probabilities are

pℓ,𝒞​\(q\)=∑j∈𝒞aj​\(q\),p^ℓ,𝒞​\(q\)=exp⁡sℓ,𝒞​\(q\)∑𝒞′∈𝒫ℓ​\(q\)exp⁡sℓ,𝒞′​\(q\),p\_\{\\ell,\\mathcal\{C\}\}\(q\)=\\sum\_\{j\\in\\mathcal\{C\}\}a\_\{j\}\(q\),\\quad\\hat\{p\}\_\{\\ell,\\mathcal\{C\}\}\(q\)=\\frac\{\\exp s\_\{\\ell,\\mathcal\{C\}\}\(q\)\}\{\\sum\_\{\\mathcal\{C\}^\{\\prime\}\\in\\mathcal\{P\}\_\{\\ell\}\(q\)\}\\exp s\_\{\\ell,\\mathcal\{C\}^\{\\prime\}\}\(q\)\},wheresℓ,𝒞​\(q\)=maxr∈\[R\]⁡tℓ,𝒞,r​\(q\)s\_\{\\ell,\\mathcal\{C\}\}\(q\)=\\max\_\{r\\in\[R\]\}t\_\{\\ell,\\mathcal\{C\},r\}\(q\)is the chunk’s approximate log attention mass\.444For the current incomplete chunk, we use its exact log attention mass in place ofsℓ,𝒞​\(q\)s\_\{\\ell,\\mathcal\{C\}\}\(q\)\.We supervise the shared reference queries at every tree level by minimizing

ℒKL=∑ℓ=1L𝔼q​\[DKL​\(pℓ​\(q\)∥p^ℓ​\(q\)\)\]\.\\mathcal\{L\}\_\{\\mathrm\{KL\}\}=\\sum\_\{\\ell=1\}^\{L\}\\mathbb\{E\}\_\{q\}\\\!\\left\[D\_\{\\mathrm\{KL\}\}\\\!\\left\(p\_\{\\ell\}\(q\)\\,\\middle\\\|\\,\\hat\{p\}\_\{\\ell\}\(q\)\\right\)\\right\]\.HereLLis the number of tree levels, and the expectation averages over eligible queries, attention heads, and model layers at each tree level\. In practice, learning the reference query parameters requires only a few billion tokens, and can be done at sequence lengths significantly shorter than pretrained backbone’s full context length\.555In the open\-source models we tested, which range from 4B to 35B parameters and were trained at context lengths from 256K to 512K, we trained for 4B tokens at 32K context length, with the reported Qwen and Nemotron runs using a peak learning rate of3\.2×10−23\.2\\times 10^\{\-2\}and the WSD schedule described in Appendix[B\-A](https://arxiv.org/html/2610.42112#A2.SS1)\.

### III\-CRouting Through the Tree

Having constructed the tree, we must now sparsely traverse it during inference\. A given query retrieves from the tree by scoring nodes and selecting a subset of nodes at each level, recursively descending by repeating this process for the selected nodes’ children, until terminating at a sparse subset of the original KV, from which we compute the attention output\.

In addition to the low\-latency yet expressive attention mass approximation and the multiple levels of hierarchy, a core innovation of this work is that we utilize top\-ppas our selection mechanism \(described in[SectionII](https://arxiv.org/html/2610.42112#S2)\)\. Top\-ppallows the selection size to adapt to the attention distribution and sequence length, in contrast to fixed\-selection methods such as top\-KK\[[27](https://arxiv.org/html/2610.42112#bib.bib27)\]\. If only a few nodes contain most of the mass, top\-ppis not obliged to select marginal nodes to satisfy a fixedKKbudget\. Conversely, if attention mass is near\-uniform, a large fraction of the nodes at a given level can be selected for expansion\. We posit that this property enables length extrapolation beyond that seen in other methods, which we confirm for a variety of open\-source models in[SectionIV](https://arxiv.org/html/2610.42112#S4)\. The hierarchy also keeps the ordering cost of top\-ppselection manageable\. Pruning at each level limits the number of candidate nodes considered at the next, reducing the size of subsequent ordering operations\. In practice, we find that these candidate sets remain small enough to make top\-ppselection feasible at every level of the tree\.

The concrete routing procedure is as follows\. Eachqqdenotes one query head’s query; selection and child expansion remain independent across query heads, including those sharing a KV tree\. Let𝒜ℓ​\(q\)=\{𝒞1,…,𝒞nℓ\}\\mathcal\{A\}\_\{\\ell\}\(q\)=\\\{\\mathcal\{C\}\_\{1\},\\ldots,\\mathcal\{C\}\_\{n\_\{\\ell\}\}\\\}denote the complete candidate chunks considered at levelℓ\\ell, indexed in decreasing order of scoresℓ,𝒞​\(q\)s\_\{\\ell,\\mathcal\{C\}\}\(q\)\. Note that the ordering andnℓ=\|𝒜ℓ​\(q\)\|n\_\{\\ell\}=\|\\mathcal\{A\}\_\{\\ell\}\(q\)\|vary depending onqq, and there may be incomplete chunks at each level in the tree depending on the exact sequence length\. At each level, we pick the top\-ppnodes; i\.e\., for chosenp∈\(0,1\)p\\in\(0,1\), we select chunks

𝒮ℓ​\(q\)=\{𝒞1,…,𝒞mℓ​\(q\)\},\\mathcal\{S\}\_\{\\ell\}\(q\)=\\\{\\mathcal\{C\}\_\{1\},\\ldots,\\mathcal\{C\}\_\{m\_\{\\ell\}\(q\)\}\\\},where

mℓ​\(q\)=min⁡\{m∈\[nℓ\]:∑i=1mexp⁡sℓ,𝒞i​\(q\)∑i=1nℓexp⁡sℓ,𝒞i​\(q\)≥p\}\.m\_\{\\ell\}\(q\)=\\min\\left\\\{m\\in\[n\_\{\\ell\}\]:\\frac\{\\sum\_\{i=1\}^\{m\}\\exp s\_\{\\ell,\\mathcal\{C\}\_\{i\}\}\(q\)\}\{\\sum\_\{i=1\}^\{n\_\{\\ell\}\}\\exp s\_\{\\ell,\\mathcal\{C\}\_\{i\}\}\(q\)\}\\geq p\\right\\\}\.Thus, the selected chunks form the smallest prefix containing at least a fractionppof the estimated attention mass among the current candidates\. Let𝒰ℓ​\(q\)\\mathcal\{U\}\_\{\\ell\}\(q\)denote the complete chunks at levelℓ\\ellwith no complete parent\. Starting with all complete top\-level chunks,𝒜L​\(q\)=𝒰L​\(q\)\\mathcal\{A\}\_\{L\}\(q\)=\\mathcal\{U\}\_\{L\}\(q\), we pass the selected chunks’ children and any parentless chunks to the next level:

𝒜ℓ−1\(q\)=\(⋃𝒞∈𝒮ℓ​\(q\)ch\(𝒞\)\)∪𝒰ℓ−1\(q\),ℓ=L,…,2,\\mathcal\{A\}\_\{\\ell\-1\}\(q\)=\\left\(\\bigcup\_\{\\mathcal\{C\}\\in\\mathcal\{S\}\_\{\\ell\}\(q\)\}\\operatorname\{ch\}\(\\mathcal\{C\}\)\\right\)\\cup\\mathcal\{U\}\_\{\\ell\-1\}\(q\),\\qquad\\ell=L,\\ldots,2,wherech⁡\(𝒞\)\\operatorname\{ch\}\(\\mathcal\{C\}\)denotes the child chunks of𝒞\\mathcal\{C\}\. We exclude the first key \(index00\) from the tree construction and instead always retain it666In many models, the first token acts as an attention sink, receiving a large share of attention mass regardless of content; we found it beneficial to always attend to it\.together with the visible keys in the current chunk,𝒞cur​\(q\)\\mathcal\{C\}\_\{\\mathrm\{cur\}\}\(q\), giving the token indices

ℐ⁡\(q\)=\{0\}∪\(⋃𝒞∈𝒮1​\(q\)𝒞\)∪𝒞cur​\(q\)\.\\mathcal\{I\}\(q\)=\\\{0\\\}\\cup\\left\(\\bigcup\_\{\\mathcal\{C\}\\in\\mathcal\{S\}\_\{1\}\(q\)\}\\mathcal\{C\}\\right\)\\cup\\mathcal\{C\}\_\{\\mathrm\{cur\}\}\(q\)\.During decode, query heads sharing a KV head jointly load the union of their final selected leaf chunks; this sharing does not merge their upper\-level routing selections or attention masks\.777For block prefill, loads are shared across query positions within each head’s tile, with per\-query masks retained\.Each head computes standard scaled dot\-product attention over its own selected KV:

o⁡\(q\)=Attn⁡\(q,Kℐ⁡\(q\),Vℐ⁡\(q\)\)\.o\(q\)=\\operatorname\{Attn\}\\\!\\left\(q,K\_\{\\mathcal\{I\}\(q\)\},V\_\{\\mathcal\{I\}\(q\)\}\\right\)\.

## IVEvaluations

We evaluate ZHA on a suite of reasoning and retrieval tasks with context lengths extending to hundreds of thousands of tokens, demonstrating that the module can match the dense model while attending to a small fraction of the total KV cache\. Across these settings, ZHA is impressively robust to choice of geometry, whose effect on retrieval accuracy is nonetheless highly predictable by a simple model\. These results suggest that, not only does ZHA enable efficient retrieval from past context, this retrieval remains effective out to very long contexts\.

### IV\-AReasoning and Retrieval

For the primary suite of evaluations, we convert Qwen3\.6\-35B\-A3B\[[35](https://arxiv.org/html/2610.42112#bib.bib35)\]with ZHA routers of varying tree geometry\. To start, we train both deep \(chunk size 4, six tree levels; C4 L6\) and wide \(chunk size 256, one tree level; C256 L1\) routers, each withR=16R=16reference queries and trained on 4B tokens with the backbone frozen \(for more training details, please see Appendix[B\-A](https://arxiv.org/html/2610.42112#A2.SS1)\)\. For our evaluations, we varyp∈\{0\.8,0\.9,0\.95\}p\\in\\\{0\.8,0\.9,0\.95\\\}, constant at every tree level\. Both routers perform within noise of the dense baseline on reasoning, agentic coding, and long\-context benchmarks\. The exception is RULER\-13 for the deep router, which scores 93\.3% versus 96\.0% for dense; however, this is just a 3% drop in performance while attending to 1%\-6% of the original KV cache\. See[Figure5](https://arxiv.org/html/2610.42112#S4.F5)for the full suite of evaluation results for ZHA \(deep and wide\) and dense\.

Figure 5:\(a\) Reasoning and \(b\) retrieval evaluation scores with dense attention, MSA, and deep \(C4 L6\) or wide \(C256 L1\) ZHA \(p=0\.95p=0\.95,R=16R=16, converted from Qwen3\.6\-35B\-A3B\)\. Appendix[B](https://arxiv.org/html/2610.42112#A2)gives evaluation details and results across routing thresholds\. We observe that ZHA with both wide and deep geometries maintains performance even at high sparsity levels, while MSA suffers performance degradation especially on retrieval tasks such as RULER and LongBench v2\.We also compare ZHA against Minimax Sparse Attention \(MSA\)\[[25](https://arxiv.org/html/2610.42112#bib.bib25)\], which selects the top 16 blocks of 128 tokens using a learned indexer\. We use the same Qwen backbone and train MSA’s indexer on 40B tokens using only the KL warmup stage, with the backbone frozen\. MSA stays close to dense attention on AIME 2026, MMLU\-Pro, and GPQA Diamond, but falls well behind both ZHA routers on long\-context retrieval\. On LongBench v2, where MSA’s fixed budget amounts to roughly 1% of the context, the deep router attends to a similar 1\.1% of decode keys yet scores 14 points higher, even though MSA was trained on 10x the token budget of ZHA\. For comparison, MSA is also included in[Figure5](https://arxiv.org/html/2610.42112#S4.F5)\.

Finally, we also train ZHA onto Nemotron\-3\.5\-Lightning\-30B\-A3B\[[33](https://arxiv.org/html/2610.42112#bib.bib33)\], to assess applicability to an architecture without explicit positional encodings\. Generally Nemotron performs comparably to its baseline on the agentic coding benchmarks, but suffers moderately on the RULER benchmark\. Full results for Nemotron are reported in Appendix[B\-D](https://arxiv.org/html/2610.42112#A2.SS4)\.

### IV\-BEffect of Geometry and Sparsity

Thus far we have shown that ZHA is competitive with dense attention, whether the tree is wide \(with chunk size 256\) or deep \(with 6 levels\), and that these variants remain competitive even while attending to a fraction of the native KV cache\. However, the space of tree geometries, spanned by choice of reference countRR, chunk sizeCC, tree depthLL, and routing thresholdpp, is combinatorially complex\. Despite preliminary evidence of robustness, it is easy to see why variation along any of the axes above may have a direct effect on performance: each changes the effective compression and sparsity in different ways, and each may introduce error in the router’s approximation to the true top\-ppset\. Moreover, depth and width introduce potentially qualitatively different kinds of error into the approximation\. Here, we explore the effect of geometry and sparsity along a few particularly revealing slices of the design space; then, to demonstrate the predictability of performance as a function of these variables, we develop a simple yet accurate model of RULER accuracy in terms of\(R,C,L,p\)\(R,C,L,p\)that we fit with just four parameters\.

Figure 6:Geometry and sparsity have predictable effects on RULER accuracy \(mean across tasks\)\. Greater sparsity without degrading performance is achieved by: \(a\) decreasing compressionCLC^\{L\}, \(b\) increasing depth at fixed compression \(CL=4096C^\{L\}=4096\), or \(c\) making key summaries more expressive \(with increasing reference\-query numberRR\) at fixed chunk size and depth\. \(The horizontal dotted lines in panels \(a\)–\(c\) show dense attention\.\) \(d\) For the simple geometry\-to\-performance model, predictions versus measurements across 64K, 128K, and 256K, holding out each length in turn \(RMSE: 1\.44 percentage points\)\.##### Robustness across geometries\.

To more thoroughly probe the effect on performance, we sweep a wide range of geometries by varyingRR,CC, andLL\(this yields 28 different configs; see[Figure18](https://arxiv.org/html/2610.42112#A3.F18)in Appendix[C](https://arxiv.org/html/2610.42112#A3)for the 28 used for fitting\)\. For each, we use the same 4B\-token KL warmup and then evaluate on RULER \(at lengths of 64K, 128K, and 256K\) while varyingp∈\{0\.7,0\.8,0\.9,0\.95,0\.99\}p\\in\\\{0\.7,0\.8,0\.9,0\.95,0\.99\\\}\. Across 13 tasks at 256K, the aggregate RULER accuracy as a function of percentage of decode keys attended is shown in[Figure6](https://arxiv.org/html/2610.42112#S4.F6)\(a\), color\-coded by the top\-level node spanCLC^\{L\}, our nominal compression factor\. The primary observation is that, when attending to between 10–30% of keys, retrieval performance is robust to choice of geometry\. Beyond this, performance generally degrades with increasing key sparsity, with the rate depending on both top\-level compression and its decomposition intoCCandLL\.[Figure6](https://arxiv.org/html/2610.42112#S4.F6)\(c\) isolates the effect of reference countRRat fixedC=8C=8andL=2L=2\.[Figure20](https://arxiv.org/html/2610.42112#A3.F20)in Appendix[C](https://arxiv.org/html/2610.42112#A3)shows the 256K RULER Pareto frontiers across geometries for decode keys attended and modeled speedup over dense\.

##### Depth at fixed compression\.

Because ZHA’s subquadratic scaling properties rely on increasing tree depth, a critical question for ZHA’s promise is whether depth specifically leads to sharp degradation in retrieval\. To isolate the effect of tree depthLLfrom top\-level compressionCLC^\{L\},[Figure6](https://arxiv.org/html/2610.42112#S4.F6)\(b\) shows RULER accuracy versus decode keys attended for a slice of ZHA geometries with fixedCL=4096C^\{L\}=4096\. In fact, it appears that increasing depth enables sustaining effective retrieval out to greater sparsity and that the deeper geometries are more performant at a fixed sparsity level\. However, this alone does not suggest that ZHA can indefinitely increase depth while scaling to longer contexts \(as this increases top\-level compressionCLC^\{L\}at fixedCC\), but it does provide strong evidence that depth itself is not the limiting factor\. Notably, to our knowledge, our approach is the first hierarchical attention variant to have demonstrated convincing scaling to depths greater than 1 while maintaining close parity to dense attention\.

##### Predicting retrieval accuracy\.

To quantify and generalize the accuracy\-sparsity tradeoff, we develop a simple model predicting RULER accuracy from input\(R,C,L,p\)\(R,C,L,p\)\. The model is

dℓ​\(R,C\)\\displaystyle d\_\{\\ell\}\(R,C\)=log⁡\(1\+CℓRr\),\\displaystyle=\\log\\\!\\left\(1\+\\frac\{C^\{\\ell\}\}\{R^\{r\}\}\\right\),\(2\)A^\\displaystyle\\widehat\{A\}=Adense​\(N\)​∏ℓ=1Lexp⁡\[−a​\(1−p\)b​dℓh\]\.\\displaystyle=A\_\{\\mathrm\{dense\}\}\(N\)\\,\\prod\_\{\\ell=1\}^\{L\}\\exp\\\!\\left\[\-a\(1\-p\)^\{b\}d\_\{\\ell\}^\{h\}\\right\]\.Intuitively, the termlog⁡\(1\+Cℓ/Rr\)\\log\(1\+C^\{\\ell\}/R^\{r\}\)represents compression difficulty at levelℓ\\ell, increasing with node span and decreasing with reference capacity\. The product compounds per\-level accuracy retention over the tree\. The factor\(1−p\)b\(1\-p\)^\{b\}accounts for the effect of sparsity, whose penalty can grow nonlinearly\.888We constrain1≤b≤21\\leq b\\leq 2to retain flexibility while avoiding unstable low\-ppextrapolation for near\-ceiling tasks\.We fit four parameters:aa, which specifies the overall penalty strength;bb, which controls sensitivity to sparsity;hh, which controls sensitivity to each level’s difficulty; andrr, which controls reference capacity\. Each exponential maps a level’s penalty to an accuracy\-retention factor; their product scales the dense\-attention accuracyAdense​\(N\)A\_\{\\mathrm\{dense\}\}\(N\)at that sequence length \(NN\)\. Aside from length\-dependence inherited from this baseline, sequence length does not enter this simple model\.

We fit the model using 28 geometries from the sweep acrossR∈\{1,2,4,16\}R\\in\\\{1,2,4,16\\\}, evaluated for 100 RULER examples per task atp∈\{0\.7,0\.8,0\.9,0\.95\}p\\in\\\{0\.7,0\.8,0\.9,0\.95\\\}and lengths \(64K, 128K, 256K\), together withp=0\.99p=0\.99evaluations for 13 of these geometries\. Fitting uses settings complete at all three lengths; the empirical panels include complete 256K settings\. We fit the four parameters \(constraininga\>0a\>0,1≤b≤21\\leq b\\leq 2,0\.2≤h≤60\.2\\leq h\\leq 6, and0≤r≤60\\leq r\\leq 6\) by nonlinear least squares on accuracy \(in percentage points\)\. Here we fit to the 13\-task mean, and Appendix[C](https://arxiv.org/html/2610.42112#A3)fits each task independently with the same form\. To evaluate held\-out prediction accuracy \(for the aggregate and per\-task fits separately\) and particularly length transfer, we measure predicted versus actual RULER accuracy on a held\-out sequence length while fitting on the other two lengths; we do this separately for the three sequence lengths, and for each we fit on 250 cells \(125 geometry\-threshold settings at 2 lengths\) and evaluate on 125 cells\. Each fold reuses the same geometries and thresholds at the other lengths and uses the measured dense baseline at the held\-out length, so this tests transfer of the relative accuracy penalty across lengths conditional on that baseline\.[Figure6](https://arxiv.org/html/2610.42112#S4.F6)\(d\) shows all 150R=16R=16leave\-one\-length\-out predictions pooled across the three lengths, with an RMSE of 1\.44 percentage points \(the 256K subset has an RMSE of 1\.65 percentage points\)\. The relatively low prediction\-error of such a simple model suggests that the accuracy varies surprisingly smoothly with sparsity and tree geometry over the fitted range\.

## VSpeedup

ZHA is designed to achieve fundamentally subquadratic latency scaling, in contrast to the variety of chunk\-sparse methods that yield significant but constant\-factor speedups\. Yet, however principled this design, its practicality depends on whether it can run efficiently on GPUs\. While we do not release optimized kernels in this work, we do investigate the question of efficient GPU implementation via a detailed roofline model that enables extrapolation of ZHA’s latency scaling in comparison to dense attention and other baselines\.

Figure 7:Per\-level selection scales sublinearly in the ZHA router with predictable power\-law exponent\. \(a\)–\(e\): For ZHA geometries with depthsL=2L=2–66\(C=4C=4,R=16R=16,p=0\.8p=0\.8\), per\-level selection counts \(normalized by first\-bin count per\-level\) on AIME problems fitted by power\-law scaling curves\. \(f\): Per\-level exponents fitted by cumulative top\-down model, demonstrating that the across\-depth variation in their exponents is well\-described by a simple model that accumulates top\-down contributions\.### V\-ARoofline Model

Although we leave the full details to Appendix[D](https://arxiv.org/html/2610.42112#A4)\(which provides a treatment of the top\-ppselection algorithm among other specifics\), we briefly describe the roofline model we use to estimate how ZHA latency scales with sequence length\. SupposeTTis the context length,L⁡\(T\)L\(T\)the ultimate tree depth, andℓ\\ellthe routing level\. The modeled ZHA attention\-subsystem latency𝒕ZHA\\bm\{t\}\_\{\\mathrm\{ZHA\}\}is given by a sum reflecting its serial execution through the tree:

𝒕ZHA​\(T\)\\displaystyle\\bm\{t\}\_\{\\mathrm\{ZHA\}\}\(T\)=ttree​\(T\)\+λA\+𝚽attn​\(T\)\\displaystyle=t\_\{\\mathrm\{tree\}\}\(T\)\+\\lambda\_\{A\}\+\\bm\{\\Phi\}\_\{\\mathrm\{attn\}\}\(T\)\(3\)\+∑ℓ=1L⁡\(T\)\[λℓ\+𝚽routing,ℓ\(T\)\+𝚽scan,ℓ\(T\)\]\\displaystyle\+\\sum\_\{\\ell=1\}^\{L\(T\)\}\\left\[\\lambda\_\{\\ell\}\+\\bm\{\\Phi\}\_\{\\mathrm\{routing\},\\ell\}\(T\)\+\\bm\{\\Phi\}\_\{\\mathrm\{scan\},\\ell\}\(T\)\\right\]Above,ttreet\_\{\\mathrm\{tree\}\}accounts for prefill construction or amortized decode maintenance, including launches\.λℓ\\lambda\_\{\\ell\},λA\\lambda\_\{A\}measure the per\-level and final\-attention launch overheads, respectively\. The three kinds of execution stages are reference scoring𝚽routing,ℓ\\bm\{\\Phi\}\_\{\\mathrm\{routing\},\\ell\}, node selection𝚽scan,ℓ\\bm\{\\Phi\}\_\{\\mathrm\{scan\},\\ell\}, and retained\-token exact attention𝚽attn\\bm\{\\Phi\}\_\{\\mathrm\{attn\}\}\.

Each constituent kernel or pass𝚽x\\bm\{\\Phi\}\_\{x\}\(indexed byxx\) follows a three\-term roofline model accounting for compute and memory bandwidth along with a worker\-wave term𝒕x,worker​waves\\bm\{t\}\_\{x,\\mathrm\{worker\\ waves\}\}that accounts for finite block concurrency:

𝚽x​\(T\)=max⁡\{𝐅𝐋𝐎𝐏𝐬x​\(T\)/Px,𝐛𝐲𝐭𝐞𝐬x​\(T\)/BWx,𝒕x,worker​waves​\(T\)\}\.\\bm\{\\Phi\}\_\{x\}\(T\)=\\max\\left\\\{\\begin\{aligned\} &\\mathbf\{FLOPs\}\_\{x\}\(T\)/P\_\{x\},\\\\ &\\mathbf\{bytes\}\_\{x\}\(T\)/\\mathrm\{BW\}\_\{x\},\\\\ &\\bm\{t\}\_\{x,\\mathrm\{worker\\ waves\}\}\(T\)\\end\{aligned\}\\right\\\}\.\(4\)Above,𝐅𝐋𝐎𝐏𝐬x\\mathbf\{FLOPs\}\_\{x\}denotes the floating\-point operation count andPxP\_\{x\}the effective FLOP/s \(scalar for scan\); similarly,𝐛𝐲𝐭𝐞𝐬x\\mathbf\{bytes\}\_\{x\}denotes the HBM traffic andBWx\\mathrm\{BW\}\_\{x\}the effective bytes/s\. For stages with sequential kernels or passes, the execution times are summed\.

Though the compute and memory terms are familiar to any roofline model, it’s worth noting why the worker\-wave term is of particular importance for ZHA\. This term is

𝒕x,worker​waves​\(T\)=𝝉x​\(T\)​⌈𝑵tiles,x​\(T\)NSM​cx⌉,\\bm\{t\}\_\{x,\\mathrm\{worker\\ waves\}\}\(T\)=\\bm\{\\tau\}\_\{x\}\(T\)\\left\\lceil\\frac\{\\bm\{N\}\_\{\\mathrm\{tiles\},x\}\(T\)\}\{N\_\{\\mathrm\{SM\}\}c\_\{x\}\}\\right\\rceil,\(5\)where𝑵tiles,x\\bm\{N\}\_\{\\mathrm\{tiles\},x\}is the number of thread blocks,NSM​cxN\_\{\\mathrm\{SM\}\}c\_\{x\}the number that can execute concurrently, and𝝉x\\bm\{\\tau\}\_\{x\}the per\-block execution time at that concurrency\. The ceiling therefore counts the sequential waves needed to execute the stage\. This matters for ZHA because pruning at one level reduces work at the next, potentially leaving individual stages with limited parallelism and thus idle SMs in the higher levels of the tree\.

### V\-BSparsity Scaling

Predicting latency therefore requires understanding how the number of nodes selected at each level scales with sequence length\. Extrapolating this behavior to longer sequences requires observing how the selected top\-ppset scales at each level, and then modeling how this scaling varies across levels, particularly as we increase the tree’s ultimate depth\.

##### Empirical per\-level exponents\.

Inspired by the robustly sublinear scaling observed in the top\-ppselection set for dense models \(see[SectionII](https://arxiv.org/html/2610.42112#S2)\), we fit a power law to the number of nodes selected per level\. We fit these curves with per\-level KV\-group union counts collected across heads and layers of ZHA generations to AIME 2026 questions \(response lengths 16K–64K\), with geometries ofC=4C=4,R=16R=16, and depthsL=2L=2–66\(a subset of the Qwen3\.6\-35B\-A3B conversions from[SectionIV\-B](https://arxiv.org/html/2610.42112#S4.SS2)\)\.

For a fixed geometry and routing thresholdpp, letSℓ​\(T\)S\_\{\\ell\}\(T\)denote the mean selected\-node count per query at context lengthTTand levelℓ\\ell, counting the union within each KV group at every level and averaging over groups \(withℓ=L\\ell=Lat the top andℓ=1\\ell=1at the leaf chunks\)\. Our power\-law summaryS^ℓ​\(T\)\\widehat\{S\}\_\{\\ell\}\(T\)has the form

S^ℓ​\(T\)=cℓ​Tαℓ,\\widehat\{S\}\_\{\\ell\}\(T\)=c\_\{\\ell\}T^\{\\alpha\_\{\\ell\}\},\(6\)wherecℓc\_\{\\ell\}is a constant factor andαℓ\\alpha\_\{\\ell\}is the scaling exponent that we fit to the empirical selection counts\.999In particular, we estimate observed exponentsαℓ\\alpha\_\{\\ell\}by fitting each problem’s log counts against log context length with1/T1/Tweights, then taking the median slope across problems at each level\. See Appendix[E](https://arxiv.org/html/2610.42112#A5)for more details\.[Figure7](https://arxiv.org/html/2610.42112#S5.F7)\(a\)–\(e\) shows the empirical number of selected nodes for the \(color\-coded\) levels across fixed\-CC, different\-depth geometries\. We also plot lines denoting each level’s fitted depth\-model power law, demonstrating good quality\-of\-fit to the selection counts\.

##### Modeled per\-level exponents\.

In order to extrapolate ZHA’s latency to longer sequences with increasing tree depth, we develop a model for how the exponent varies per\-level\. Using the same variable\-depth geometries as a case study, we begin by modeling per\-level exponents given some fixed ultimate depthLL\. Notice that, in[Figure7](https://arxiv.org/html/2610.42112#S5.F7), the scaling exponent is highest at the top level and decreases proceeding down each tree, consistently for each geometry\.101010[Figure7](https://arxiv.org/html/2610.42112#S5.F7)showsp=0\.8p=0\.8; the corresponding data forp∈\{0\.7,0\.9,0\.95\}p\\in\\\{0\.7,0\.9,0\.95\\\}may be found in[Figure22](https://arxiv.org/html/2610.42112#A5.F22)in Appendix[E](https://arxiv.org/html/2610.42112#A5)and broadly corroborates these same trends\. The exception is the deepest geometry \(C=4C=4,L=6L=6,R=16R=16\) atp=0\.7p=0\.7, whose exponents rise again toward the leaf levels\. This setting achieves only 44\.0% accuracy on AIME, compared with 93\.1% for dense attention, so its anomalous scaling should be interpreted in the context of severe accuracy degradation\.This pattern inspires our approach to modeling the per\-level exponentα^ℓ\\widehat\{\\alpha\}\_\{\\ell\}as a product of top\-down contributions by top\-level exponentβ\\betaand per\-level factorsβℓ\\beta\_\{\\ell\}:

εℓ\\displaystyle\\varepsilon\_\{\\ell\}=2​ε​\(1−2−\(L−ℓ\)\),\\displaystyle=2\\varepsilon\\left\(1\-2^\{\-\(L\-\\ell\)\}\\right\),\(7\)βℓ\\displaystyle\\beta\_\{\\ell\}=β\+εℓ,εL=0,\\displaystyle=\\beta\+\\varepsilon\_\{\\ell\},\\qquad\\varepsilon\_\{L\}=0,α^ℓ\\displaystyle\\widehat\{\\alpha\}\_\{\\ell\}=∏k=ℓLβk,α^L=β\.\\displaystyle=\\prod\_\{k=\\ell\}^\{L\}\\beta\_\{k\},\\qquad\\widehat\{\\alpha\}\_\{L\}=\\beta\.This product includes the current level and every level above it:α^L−1=βL−1​β\\widehat\{\\alpha\}\_\{L\-1\}=\\beta\_\{L\-1\}\\betaandα^L−2=βL−2​βL−1​β\\widehat\{\\alpha\}\_\{L\-2\}=\\beta\_\{L\-2\}\\beta\_\{L\-1\}\\beta\. To account for sparsity introduced by higher\-level pruning and KV\-group unioning,εℓ\\varepsilon\_\{\\ell\}is a saturating additive correction to the factorβℓ\\beta\_\{\\ell\}\. We fit the overall shape parameters\(β,ε\)\(\\beta,\\varepsilon\)separately for each geometry andpp\.111111Specifically, we fit\(β,ε\)\(\\beta,\\varepsilon\)for each geometry by equally weighted least squares against the empirical exponentsαℓ\\alpha\_\{\\ell\}, allowing signedε\\varepsilonwhile keeping everyβℓ\\beta\_\{\\ell\}positive, with an independent amplitudecℓc\_\{\\ell\}calibrated to pooled mean counts by1/T1/T\-weighted least squares in log space at each level\. See Appendix[E](https://arxiv.org/html/2610.42112#A5)for more details\.In Appendix[E](https://arxiv.org/html/2610.42112#A5), we compare to more general additive corrections, from a constant additive factor to increasing linear drift\.

[Figure7](https://arxiv.org/html/2610.42112#S5.F7)\(f\) shows the empirically\-estimated per\-level exponentsαℓ\\alpha\_\{\\ell\}along with their fitted cumulative formα^ℓ\\widehat\{\\alpha\}\_\{\\ell\}for all five depths; note that this form accurately captures the across\-level exponent trends atp=0\.8p=0\.8\. Appendix[E](https://arxiv.org/html/2610.42112#A5)completes the comparison across all four routing thresholds and gives the problem counts, full\-resolution fitting, display binning, and bootstrap details\.

![Refer to caption](https://arxiv.org/html/2610.42112v1/hatt-latency-constant-linear.png)Figure 8:For a fixed top\-p, and hence sparsity level, ZHA \(here,R=1R=1,C=8C=8,p=0\.75p=0\.75\) scales subquadratically according to our decode roofline model, eventually dramatically outperforming both dense attention and the MSA sparse attention variant\. Importantly, we have not empirically validated the scaling at these ultra\-long context lengths due to the lack of existing LLMs which can operate at these contexts\. Because of this, while we can obtain subquadratic latency, it is unclear to what extent reasoning and retrieval performance degrades at the level of depths required for these ultra\-long contexts\. Latencies are modeled per\-token decode times for the attention subsystem of one layer at batch size 1\.Crucially, to predict sparsity scaling for longer sequences, we must also extrapolate the within\-tree profiles above to increasingly deeper trees\. Given the sparsity of available data \(20 measurements, spanning five depths perpp\), we present upper and lower bounds rather than attempt a tight fit\. In particular, we parameterize the top exponent as

β^​\(L,p\)=b0\+bp​p\+\[d0\+dp​p\]​L\.\\widehat\{\\beta\}\(L,p\)=b\_\{0\}\+b\_\{p\}p\+\\left\[d\_\{0\}\+d\_\{p\}p\\right\]L\.\(8\)The constant\-top scenario setsd0=dp=0d\_\{0\}=d\_\{p\}=0; the linear\-top scenario fits all four coefficients directly to the observed top\-level exponentsαL\\alpha\_\{L\}, independently of the within\-tree profile fit\. To isolate the effect of ultimate depth, we fit this model on the top\-level exponents across the variable\-depth geometries \(the same 20 measurements referred to above\); Appendix[E](https://arxiv.org/html/2610.42112#A5)reports the exact coefficients and validation procedure\. From this limited data, linear\-top is better than constant\-top, though neither provide a particularly good fit\. However, they jointly provide a loose bound on the top\-level exponent with increasing depth, which we use to extrapolate ZHA’s latency out to very long contexts\.

### V\-CLatency Comparison

Given this roofline model and selection scaling framework, we compare the theoretical latency scaling of ZHA to dense attention and MSA for sequence lengths up to10910^\{9\}tokens\. For this analysis, we restrict to a light\-weight ZHA geometry \(R=1R=1,C=8C=8,p=0\.75p=0\.75; Qwen 35B with BF16\) that retains high performance across our long\-context reasoning and retrieval evaluation suite \(see[Figure15](https://arxiv.org/html/2610.42112#A2.F15)in Appendix[B](https://arxiv.org/html/2610.42112#A2)\), and compare its modeled latency to dense attention and MSA\. To model the per\-level exponent extrapolation with depth, we transfer theC=4C=4,R=16R=16AIME exponent fits and calibrate count amplitudes toR=1R=1,C=8C=8,L=2L=2AIME measurements\. At each sequence length, we choose the ultimate tree depth that minimizes latency\. As shown in[Figure8](https://arxiv.org/html/2610.42112#S5.F8), ZHA becomes faster than dense attention at approximately 48K tokens and MSA at 0\.48M tokens, reaching about373373–417×417\\timesand5252–58×58\\timesspeedup, respectively, at10910^\{9\}tokens\. The selected depth isL=5L=5at10910^\{9\}tokens in both scenarios\.[Figure24](https://arxiv.org/html/2610.42112#A5.F24)in Appendix[E](https://arxiv.org/html/2610.42112#A5)comparesp=0\.7,0\.75,0\.8p=0\.7,0\.75,0\.8\.

## VIAnalysis

Several design choices in the development of ZHA rely upon assumptions about the underlying attention distribution, and therefore how best to approximate it with subquadratic compute\. Here, we empirically investigate two such phenomena: the dynamic sparsity introduced by top\-pp, and the learned distribution and usage dynamics of the router’s reference queries\.

### VI\-AVariable Sparsity

In contrast to fixed top\-KKthresholds such as MSA, ZHA’s top\-ppthreshold dynamically varies the number of attended keys\. Not only does this facilitate proper scaling of retrieval with sequence length, this also enables adaptive retrieval depending on the query and preceding context\. The utility of the latter depends on the extent to which the underlying attention dynamics have differing sparsity depending on the setting and task, and whether the model can adaptively route through the tree according to these dynamics\.

First, we compare attention sparsity across the evaluation domains, which require retrieval patterns ranging from narrow lookup to diffuse aggregation\. At context lengths of approximately 64K tokens, 90% of a query head’s dense attention mass is concentrated in a median of 2\.3% of 16\-key chunks on RULER, 11\.7% on SWE\-bench Verified, and 15\.0% on fiction books\.121212We exclude key 0 and the current partial chunk, which ZHA always attends, and record dense attention along the router’s own decode\. Medians are taken over decode steps, the 10 full\-attention layers, and query heads within each sequence, then across sequences\.A representative ZHA geometry \(R=16R=16,C=16C=16,L=3L=3\), also atp=0\.9p=0\.9, retains 0\.44%, 5\.0%, and 7\.4% of chunks, respectively\. These chunks contain a median of 80%, 78%, and 78% of dense attention mass, with individual RULER tasks ranging from 74–82%\. Though the router is consistently sparser than the dense top\-ppset \(reflecting iterative pruning across the three levels\), the pattern across tasks is consistent\. Furthermore, the variation for both dense and ZHA matches our intuition: in RULER, its synthetic filler carries little that the answer depends on; in agentic contexts like SWE\-bench, the relevant information is spread over code, tool outputs, and earlier reasoning; and in fiction books, the information is diffuse throughout the entire context\.

RULER provides a particularly controlled setting in which to isolate how specific tasks induce different attention distributions such as either diffuse or concentrated\. As such, we compare the size of ZHA’s retrieved key set to the geometric mean over all 13 RULER subtasks\. The tasks fall into three sparsity tiers: tasks with a single target in repetitive filler \(variable\_tracking,niah\_single\_1\) read about half as many keys as the typical task; needles in essay text and document QA sit near the middle \(0\.82–1\.26×\\times\); and tasks that must scan many candidates, either aggregating items \(cwe,fwe\) or distinguishing many look\-alike needles \(niah\_multikey\_2/3\), read about 1\.6×\\timesas many\. To concretely illustrate how ZHA achieves variable sparsity,[Figure25](https://arxiv.org/html/2610.42112#A6.F25)depicts a representative attention head’s traversal through the tree forniah\_single\_1andniah\_multikey\_2; both traversals ultimately attend to the correct answer\-containing chunk, but the former tree is far more aggressively pruned than the latter\. Note that this adaptive sparsity is impossible for a fixed\-budget top\-KKmechanism, which instead forces a tradeoff between a large budget wasting keys on concentrated inputs, and a small one missing context needed for diffuse retrieval\. ZHA’s top\-ppthreshold avoids this choice entirely\.

### VI\-BReference\-Query Structure

Another question concerns the nature of the learned structure in the reference queries, which comprise the only learned components in the router and therefore are critical for its efficacy\. Recall that each reference query produces a compressed key and bias defining an affine approximation to the chunk’s log attention mass\. Taking the maximum over these approximations yields a piecewise\-affine approximation to the true function\. In this setup, each reference query serves as a tangent plane’s center point, and therefore it might reflect some central tendency in the query distribution\. However, this maximum over the reference\-query\-mediated scores in the KL training objective does not explicitly enforce balanced usage or distinct reference query representations\. In the worst case, they might simply collapse while still adding a constant\-factor computational cost\.

Thus, a reference query is only useful insofar as it is selected by the max\-pooling operation\. We therefore measure the selection rates of reference queries across levels of the tree, using a sample of length\-32K natural language documents from our training mix\.131313133 documents of at least length\-32K in one training\-mix shard; these are chat\-formatted, mostly long mathematical reasoning, plus general chat, agentic coding, and long\-document QA\. At every 8th query position, we count which reference query wins each kept node\. As the routers may have seen this data,[AppendixG](https://arxiv.org/html/2610.42112#A7)repeats the analysis on RULER\.The reference queries develop a coherent pattern of utilization both within and between tree levels; this is summarized in[Figure9](https://arxiv.org/html/2610.42112#S6.F9)\. In particular, the selection rate among reference queries within a level follows a gradual decay, which becomes increasingly concentrated among a few reference queries at higher tree levels and more balanced toward the leaves\. Furthermore, comparing the rank correlation in the selection rate orderings between levels, the higher levels consistently use the same reference queries most and least commonly, whereas the lowest level’s usage pattern diverges, suggesting different selection dynamics depending on the level’s compression factor\. The across\-level usage patterns are further broken down in Appendix[G](https://arxiv.org/html/2610.42112#A7), demonstrating rich selection\-rate structure in which reference query usage varies smoothly across levels\.

Moreover, we investigate the learned structure, using a deep geometry \(R=16R=16,C=4C=4,L=6L=6\) with moderate sparsity \(p=0\.9p=0\.9\) as a case study\. First,[Figure10](https://arxiv.org/html/2610.42112#S6.F10)summarizes their geometric variation\. Comparing the reference queries \(here, 16\) in a KV group; their significant pairwise cosine similarity, with a mean pairwise cosine similarity of 0\.44 \(averaged over all 20 KV groups\), show they occupy a cone of directions rather than collapsing to identical vectors\. This structure emerges across geometries and independently trained routers with varying chunk sizes and depths learn closely matching reference queries \([AppendixH](https://arxiv.org/html/2610.42112#A8)\)\. This similarity across geometries is clearest for the most frequently selected references\. The reference query norms have some spread while remaining fairly concentrated\. Interestingly, selection also favors smaller\-norm reference queries\. Since smaller norms lead to smoother weighting, it may indicate that more uniform averaging of keys becomes advantageous at higher levels\. Together, the distribution in cosine similarity and norms suggests that the reference queries have organized into a coherent and non\-collapsed geometric structure\.

![Refer to caption](https://arxiv.org/html/2610.42112v1/qref_selection_rank.png)Figure 9:Reference\-query selection rates per tree level for the C4 L6 ZHA router on natural text\. \(a\) Within each KV group and level, reference queries are independently ranked by selection frequency\. Each point shows the level’s share of selections at that rank, averaged over layers and KV groups\. Higher levels concentrate selections among fewer reference queries, while lower levels use them more uniformly\. \(b\) Spearman correlation between levels’ selection rates\. Higher levels have strongly aligned reference\-query rankings\. Agreement generally weakens as levels become further separated, with the chunk level showing the greatest divergence from the upper levels\.Figure 10:\(a\) Cosine similarity distribution of every pair of reference queries within a KV group; \(b\) norm distribution of every reference query\. Reference queries within a KV group become slightly more similar and their norms more concentrated as the tree gets shallower and its chunks larger\.

## VIIDiscussion

In this paper, we have introduced ZHA, a tree\-structured hierarchical sparse attention method that can achieve asymptotically subquadratic attention scaling, subject to specific constraints on the tree geometry and the top\-ppvalue\. Moreover, we have shown that ZHA can be applied to existing base models in a short distillation phase with only a few billion tokens, and it maintains performance on challenging long\-context reasoning, agentic coding, and retrieval tasks\.

While empirically ZHA appears to perform very well for contexts up to 256k, performance does degrade at increasingly higher levels of sparsity, and thus ultimately at longer context lengths\. Here, we present perhaps the simplest possible version of the tree\-structured indexer idea, which only learns reference queries in a separate independent KL\-distillation phase\. There are potentially many improvements that could be made to the core algorithm or training recipe to improve its performance\. For instance, chunk boundaries could be dynamically learned instead of set arbitrarily; we could introduce a recency bias into the tree so that it becomes more detailed closer to the current token, rather than uniformly distributing mass over the sequence; and we could finetune the base model parameters on long\-context training sets with a highly aggressive ZHA geometry to let the rest of the model learn to compensate for the degradation caused by the sparsity\. We could also augment the tree with compressed value summaries, allowing nodes pruned during routing to contribute coarse attention outputs\. This would preserve approximate visibility over the full KV cache span, combining token\-level attention over selected regions with summary\-level contributions from pruned subtrees\.

##### Parameter scaling\.

Our method is surprisingly simple, requiring only a single KL\-warmup phase in which we train a set of reference queries \(these are the only learnable parameters we add\) to match the attention mass induced by regular dense attention\. Although we initially tried using ideas from[Hu et al\. \[18\]](https://arxiv.org/html/2610.42112#bib.bib18)such as a reference\-query LoRA and an explicit factorization of the attention probabilities across different levels of the tree, we found that these methods added complexity and did not improve performance in our preliminary ablations\.

One interesting trend we noted during our initial development of ZHA, is that the performance of ZHA improves with model scale\. We found that small models under 1B parameters often struggled significantly with RULER and other retrieval tasks during conversion to ZHA because the original attention scores we were distilling from were very noisy themselves\. Our success with distillation improved dramatically when moving to larger models in the 4B scale and up\. This makes us optimistic that ZHA will continue to improve with further increases to model scale\.

##### Positional encoding\.

Another interesting question concerns the impact of different positional encoding schemes on the effectiveness of ZHA\. The primary experiments we ran were on the Qwen model series which uses p\-RoPE, as were our experiments on the structure of naturalistic data in[Figure1](https://arxiv.org/html/2610.42112#S2.F1)\. However, we find that RoPE induces a very particular structure into the attention scores over long contexts, where most heads represent a strong recency bias while a few heads appear to primarily attend to the lowest RoPE frequencies and are sensitive much more uniformly across the context\. NoPE attention, by contrast, can induce a more diffuse attention structure, and thus may extrapolate or give different asymptotic scalings for ZHA than RoPE models\. In Appendix[B\-D](https://arxiv.org/html/2610.42112#A2.SS4), we also tested the hybrid\-NoPE Nemotron3 models and found that ZHA was able to be trained and maintain evaluation performance for these models too\. However, the long\-context extrapolation capabilities of ZHA with NoPE remain uncertain\.

##### Asymptotic length scaling\.

The most important limitation in our results is that, although we have extrapolated performance out to millions of tokens of context and beyond, we have not empirically validated the evaluation performance of ZHA at these context lengths\. The primary limitation is finding models with sufficiently long contexts that are small enough to adapt without substantial computational expense\. In future work, we plan to scale up the method to frontier open\-source models and determine whether our method continues to scale to 1 million context and beyond\.

Of course, this problem of evaluating performance at contexts much beyond 1 million tokens is not unique to ZHA, and is in fact shared by all attention\-based models\. Though this in part derives from data limitations, it also points to the inspiration for sparse attention methods: attention’s quadratic scaling makes long\-context training and inference prohibitively expensive\. Beyond this computational bottleneck, however, it is still unclear whether dense attention itself is even performant at longer sequence lengths\. Prior work has in fact found that softmax can be “dispersive”, in the sense that it upweights noise within the preceding context, leading to increasingly uniform attention scores that eventually prevent effective retrieval\[[42](https://arxiv.org/html/2610.42112#bib.bib42)\]\.

Therefore, an argument may be made that sparsity is not only necessary for computational efficiency, but also to continue to scale language modelling to ultra\-long contexts\. It may be the case that sparsity’s pruning effectively “denoises” the past context, sidestepping the dispersiveness of dense attention and uncovering the truly relevant history to be retrieved\. Indeed, our finding that top\-ppgrows like a sublinear power law is highly suggestive that the subset of information that is truly relevant for predicting the next token has fundamentally different scaling properties than dense retrieval of the entire context\.

A core assumption of our work is that this top\-ppset contains sufficient information to enable effective language modeling\. ZHA’s empirical success in part validates this assumption; however, it is admittedly unclear whether this will remain the case indefinitely\. Similarly, it is unclear whether top\-pp, and the underlying predictive information that we suppose it picks up, grows sublinearly to indefinite context lengths\.

Supposing that retrieving the true top\-ppset is “all you need”, the remaining source of error is properly routing to that set\. Due to asymptotic divergence of different power\-law curves, precisely matching the true scaling of the top\-ppset is a surprisingly subtle point\. In short, it turns out that our method of applying top\-ppiteratively down a growing tree requires carefully tuning the per\-levelpp, and failing to do so results in two possibilities\. First, fixingp<1p<1as we do currently, the “cumulative”peff=pL→0p\_\{\\text\{eff\}\}=p^\{L\}\\to 0as depthL→∞L\\to\\infty\. In this case, the retrieval state asymptotes towards some constant size and performance necessarily degrades\. Alternatively, if we rescaled per\-levelpptop1/Lp^\{1/L\}to counteract this effect, we could then observe the opposite pathology wherepptends to 1 as depth increases, leading to all of the nodes at the higher levels being selected and the upper levels of the tree effectively becoming redundant\. If this occurs then we return to quadratic scaling and hence the diffusive pathology of fully dense attention\.

To remain subquadratic in the limit while also maintaining performance, then, we need some method to scaleppdynamically with depth such that the effectiveppat the leaf nodes decreases, but at a slow enough rate that it tracks the decrease in predictive information in the sequence with increasing context length, similar to that we found in[Figure1](https://arxiv.org/html/2610.42112#S2.F1)\. This was roughly a power\-law scaling\. However, this depends not on the architecture, but on the intrinsic distribution of the data and thus is deeply dependent upon the statistical structure of ultra\-long context data\. Our analysis suggests that naturalistic data does exhibit an amenable scaling, but we only investigated data up to 256K context, and thus extrapolating this to hundreds of millions or billions of tokens could fail\. Whether or not naturalistic text data continues to exhibit power\-law self\-similar structure at these ultra\-long scales thus remains an important avenue for future work\.

## References

- \[1\]Joshua Ainslie, James Lee\-Thorp, Michiel De Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai\.GQA: Training generalized multi\-query transformer models from multi\-head checkpoints\.In*Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing*, pp\. 4895–4901, 2023\.
- \[2\]Yushi Bai, Shangqing Tu, Jiajie Zhang, Hao Peng, Xiaozhi Wang, Xin Lv, Shulin Cao, Jiazheng Xu, Lei Hou, Yuxiao Dong, et al\.LongBench v2: Towards deeper understanding and reasoning on realistic long\-context multitasks\.*arXiv preprint arXiv:2412\.15204*, 2024\.
- \[3\]Yushi Bai, Qian Dong, Ting Jiang, Xin Lv, Zhengxiao Du, Aohan Zeng, Jie Tang, and Juanzi Li\.IndexCache: Accelerating sparse attention via cross\-layer index reuse\.*arXiv preprint arXiv:2603\.12201*, 2026\.
- \[4\]Zefan Cai, Yichi Zhang, Bofei Gao, Yuliang Liu, Yucheng Li, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Junjie Hu, et al\.PyramidKV: Dynamic KV cache compression based on pyramidal information funneling\.*arXiv preprint arXiv:2406\.02069*, 2024\.
- \[5\]Neil Chowdhury, James Aung, Chan Jun Shern, Oliver Jaffe, Dane Sherburn, Giulio Starace, Evan Mays, Rachel Dias, Marwan Aljubeh, Mia Glaese, et al\.Introducing SWE\-bench Verified, 2024\.URL[https://openai\.com/index/introducing\-swe\-bench\-verified/](https://openai.com/index/introducing-swe-bench-verified/)\.
- \[6\]Tri Dao and Albert Gu\.Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality, 2024\.URL[https://arxiv\.org/abs/2405\.21060](https://arxiv.org/abs/2405.21060)\.
- \[7\]Tri Dao, Daniel Y\. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré\.FlashAttention: Fast and memory\-efficient exact attention with IO\-awareness\.*Advances in Neural Information Processing Systems*, 35:16344–16359, 2022\.
- \[8\]DeepSeek\-AI\.DeepSeek\-V4: Towards highly efficient million\-token context intelligence, 2026\.URL[https://arxiv\.org/abs/2606\.19348](https://arxiv.org/abs/2606.19348)\.
- \[9\]DeepSeek\-AI et al\.DeepSeek\-R1: Incentivizing reasoning capability in LLMs via reinforcement learning\.*arXiv preprint arXiv:2501\.12948*, 2025\.
- \[10\]Yuan Feng, Junlin Lv, Yukun Cao, Xike Xie, and S Kevin Zhou\.Ada\-KV: Optimizing KV cache eviction by adaptive budget allocation for efficient LLM inference\.*Advances in Neural Information Processing Systems*, 38:113152–113188, 2026\.
- \[11\]Tomas Figliolia, Nicholas Alonso, Rishi Iyer, Quentin Anthony, and Beren Millidge\.Compressed convolutional attention: Efficient attention in a compressed latent space\.*arXiv preprint arXiv:2510\.04476*, 2025\.
- \[12\]Zichuan Fu, Wentao Song, Yejing Wang, Xian Wu, Yefeng Zheng, Yingying Zhang, Derong Xu, Xuetao Wei, Tong Xu, and Xiangyu Zhao\.Sliding window attention training for efficient large language models, 2025\.URL[https://arxiv\.org/abs/2502\.18845](https://arxiv.org/abs/2502.18845)\.
- \[13\]Paolo Glorioso, Quentin Anthony, Yury Tokpanov, Anna Golubeva, Vasudev Shyam, James Whittington, Jonathan Pilault, and Beren Millidge\.The Zamba2 suite: Technical report\.*arXiv preprint arXiv:2411\.15242*, 2024a\.
- \[14\]Paolo Glorioso, Quentin Anthony, Yury Tokpanov, James Whittington, Jonathan Pilault, Adam Ibrahim, and Beren Millidge\.Zamba: A compact 7B SSM hybrid model\.*arXiv preprint arXiv:2405\.16712*, 2024b\.
- \[15\]Albert Gu and Tri Dao\.Mamba: Linear\-time sequence modeling with selective state spaces, 2024\.URL[https://arxiv\.org/abs/2312\.00752](https://arxiv.org/abs/2312.00752)\.
- \[16\]Cheng\-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg\.RULER: What’s the real context size of your long\-context language models?*arXiv preprint arXiv:2404\.06654*, 2024\.
- \[17\]Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, et al\.MiniCPM: Unveiling the potential of small language models with scalable training strategies\.*arXiv preprint arXiv:2404\.06395*, 2024\.URL[https://arxiv\.org/abs/2404\.06395](https://arxiv.org/abs/2404.06395)\.
- \[18\]Xiang Hu, Xinyu Wei, Hao Gu, Minshen Zhang, Tian Liang, Huayang Li, Lei Zhu, Yan Wang, Sirui Han, Yushi Bai, et al\.Hierarchical sparse attention done right: Toward infinite context modeling\.*arXiv preprint arXiv:2607\.02980*, 2026\.
- \[19\]Yuxiang Huang, Nuno MT Gonçalves, Federico Alvetreti, Lei Li, Xu Han, Edoardo M Ponti, André FT Martins, and Marcos V Treviso\.DashAttention: Differentiable and adaptive sparse hierarchical attention\.*arXiv preprint arXiv:2605\.18753*, 2026\.
- \[20\]Sukjun Hwang, Brandon Wang, and Albert Gu\.Dynamic chunking for end\-to\-end hierarchical sequence modeling\.*arXiv preprint arXiv:2507\.07955*, 2025\.
- \[21\]Samy Jelassi, David Brandfonbrener, Sham M\. Kakade, and Eran Malach\.Repeat after me: Transformers are better than state space models at copying\.In*Proceedings of the 41st International Conference on Machine Learning*, volume 235, pp\. 21502–21521, 2024\.URL[https://proceedings\.mlr\.press/v235/jelassi24a\.html](https://proceedings.mlr.press/v235/jelassi24a.html)\.
- \[22\]Carlos E\. Jimenez, John Yang, Alexander Wettig, Shunyu Yao, Kexin Pei, Ofir Press, and Karthik R\. Narasimhan\.SWE\-bench: Can language models resolve real\-world GitHub issues?In*The Twelfth International Conference on Learning Representations*, 2024\.URL[https://openreview\.net/forum?id=VTF8yNQM66](https://openreview.net/forum?id=VTF8yNQM66)\.
- \[23\]Jang\-Hyun Kim, Dongyoon Han, and Sangdoo Yun\.Fast KVzip: Efficient and accurate LLM inference with gated KV eviction, 2026\.URL[https://arxiv\.org/abs/2601\.17668](https://arxiv.org/abs/2601.17668)\.
- \[24\]Aakash Lahoti, Kevin Y\. Li, Berlin Chen, Caitlin Wang, Aviv Bick, J\. Zico Kolter, Tri Dao, and Albert Gu\.Mamba\-3: Improved sequence modeling using state space principles, 2026\.URL[https://arxiv\.org/abs/2603\.15569](https://arxiv.org/abs/2603.15569)\.
- \[25\]Xunhao Lai, Weiqi Xu, Yufeng Yang, Qiaorui Chen, Yang Xu, Lunbin Zeng, Xiaolong Li, Haohai Sun, Haichao Zhu, Vito Zhang, et al\.MiniMax sparse attention\.*arXiv preprint arXiv:2606\.13392*, 2026\.
- \[26\]Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev\-Shwartz, et al\.Jamba: A hybrid transformer\-Mamba language model\.*arXiv preprint arXiv:2403\.19887*, 2024\.
- \[27\]Chaofan Lin, Jiaming Tang, Shuo Yang, Hanshuo Wang, Tian Tang, Boyu Tian, Ion Stoica, Song Han, and Mingyu Gao\.Twilight: Adaptive attention sparsity with hierarchical top\-pppruning\.In*Advances in Neural Information Processing Systems*, volume 38, 2025\.URL[https://arxiv\.org/abs/2502\.02770](https://arxiv.org/abs/2502.02770)\.
- \[28\]Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi Dengr, Chong Ruan, Damai Dai, Daya Guo, et al\.DeepSeek\-V2: A strong, economical, and efficient mixture\-of\-experts language model\.*arXiv preprint arXiv:2405\.04434*, 2024\.
- \[29\]Aixin Liu, Aoxue Mei, Bangcai Lin, Bing Xue, Bingxuan Wang, Bingzheng Xu, Bochao Wu, Bowei Zhang, Chaofan Lin, Chen Dong, et al\.DeepSeek\-V3\.2: Pushing the frontier of open large language models\.*arXiv preprint arXiv:2512\.02556*, 2025\.
- \[30\]Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava\.Scissorhands: Exploiting the persistence of importance hypothesis for LLM KV cache compression at test time\.*Advances in Neural Information Processing Systems*, 36:52342–52364, 2023\.
- \[31\]Ilya Loshchilov and Frank Hutter\.Decoupled weight decay regularization\.In*International Conference on Learning Representations*, 2019\.URL[https://arxiv\.org/abs/1711\.05101](https://arxiv.org/abs/1711.05101)\.
- \[32\]MathArena\.AIME 2026 dataset, 2026\.URL[https://huggingface\.co/datasets/MathArena/aime\_2026/tree/d2de22f3c656b4f56cf8981212186377d1e23bc3](https://huggingface.co/datasets/MathArena/aime_2026/tree/d2de22f3c656b4f56cf8981212186377d1e23bc3)\.
- \[33\]NVIDIA\.NVIDIA Nemotron 3\.5 Super and Lightning, 2026\.URL[https://research\.nvidia\.com/labs/nemotron/Nemotron\-3\.5\-Super\-and\-Lightning/](https://research.nvidia.com/labs/nemotron/Nemotron-3.5-Super-and-Lightning/)\.
- \[34\]Zihan Qiu, Zekun Wang, Bo Zheng, Zeyu Huang, Kaiyue Wen, Songlin Yang, Rui Men, Le Yu, Fei Huang, Suozhi Huang, et al\.Gated attention for large language models: Non\-linearity, sparsity, and attention\-sink\-free\.*Advances in Neural Information Processing Systems*, 38:100092–100118, 2026\.
- \[35\]Qwen Team\.Qwen3\.6\-35B\-A3B model card, 2026\.URL[https://huggingface\.co/Qwen/Qwen3\.6\-35B\-A3B](https://huggingface.co/Qwen/Qwen3.6-35B-A3B)\.
- \[36\]David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Dirani, Julian Michael, and Samuel R\. Bowman\.GPQA: A graduate\-level Google\-proof Q&A benchmark\.*arXiv preprint arXiv:2311\.12022*, 2023\.
- \[37\]Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu\.RoFormer: Enhanced transformer with rotary position embedding\.*Neurocomputing*, 568:127063, 2024\.
- \[38\]Bohao Tang, Zhen Qin, Yuqi Pan, Zheng Li, and Pengfei Liu\.Block sparse attention with log\-linear complexity\.*arXiv preprint arXiv:2609\.31093*, 2026\.
- \[39\]Kimi Team, Yu Zhang, Zongyu Lin, Xingcheng Yao, Jiaxi Hu, Fanqing Meng, Chengyin Liu, Xin Men, Songlin Yang, Zhiyuan Li, et al\.Kimi Linear: An expressive, efficient attention architecture\.*arXiv preprint arXiv:2510\.26692*, 2025\.
- \[40\]Terminal\-Bench Team\.Terminal\-Bench 2\.1, 2026\.URL[https://www\.tbench\.ai/news/terminal\-bench\-2\-1](https://www.tbench.ai/news/terminal-bench-2-1)\.
- \[41\]Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N\. Gomez, Łukasz Kaiser, and Illia Polosukhin\.Attention is all you need\.In*Advances in Neural Information Processing Systems*, volume 30, 2017\.URL[https://arxiv\.org/abs/1706\.03762](https://arxiv.org/abs/1706.03762)\.
- \[42\]Petar Veličković, Christos Perivolaropoulos, Federico Barbero, and Razvan Pascanu\.Softmax is not enough \(for sharp size generalisation\)\.*arXiv preprint arXiv:2410\.01104*, 2024\.
- \[43\]Yubo Wang, Xueguang Ma, Ge Zhang, Yuansheng Ni, Abhranil Chandra, Shiguang Guo, Weiming Ren, Aaran Arulraj, Xuan He, Ziyan Jiang, et al\.MMLU\-Pro: A more robust and challenging multi\-task language understanding benchmark\.*arXiv preprint arXiv:2406\.01574*, 2024\.
- \[44\]Robert Washbourne, Rishi Iyer, Tomas Figliolia, Henry Zheng, Ryan Lorig\-Roach, Sungyeon Yang, Pritish Yuvraj, Quentin Anthony, Yury Tokpanov, Xiao Yang, et al\.ZAYA1\-8B technical report\.*arXiv preprint arXiv:2605\.05365*, 2026\.
- \[45\]Samuel Williams, Andrew Waterman, and David Patterson\.Roofline: An insightful visual performance model for multicore architectures\.*Communications of the ACM*, 52\(4\):65–76, 2009\.doi:10\.1145/1498765\.1498785\.
- \[46\]Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, Yuxing Wei, Lean Wang, Zhiping Xiao, et al\.Native sparse attention: Hardware\-aligned and natively trainable sparse attention\.In*Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\)*, pp\. 23078–23097, 2025\.
- \[47\]Wen Zan, Jiaqi Zhang, Jianchao Tan, Hong Liu, Cunguang Wang, Xiang Li, Duyue Ma, Guanyu Wu, Yifan Lu, Fengcun Li, et al\.LongCat sparse attention: Taming the lightning via streaming\-aware hierarchical cross\-layer indexing\.*arXiv preprint arXiv:2608\.01662*, 2026\.
- \[48\]Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al\.H2O: Heavy\-hitter oracle for efficient generative inference of large language models\.*Advances in Neural Information Processing Systems*, 36:34661–34710, 2023\.

Appendix

## Appendix ATop\-ppScaling by Layer and Head

### A\-ADense Top\-ppSets by Layer and Head

In[SectionII](https://arxiv.org/html/2610.42112#S2), we average dense top\-ppset sizes over layers and query heads\.[Figure11](https://arxiv.org/html/2610.42112#A1.F11)shows all 10 full\-attention layers and 16 query heads of Qwen3\.6\-35B\-A3B separately, averaging only across the six books\. Common logarithmic axes allow comparison with the linearp=1p=1reference, which retains all keys\. Sublinear trends appear throughout the full\-attention stack, with substantial variation in set size and growth rate\. The aggregate trend is thus not confined to a few sparse heads, but does not imply a common exponent\.

![Refer to caption](https://arxiv.org/html/2610.42112v1/dense-top-p-heads.png)Figure 11:Dense top\-ppset size for individual query heads in Qwen3\.6\-35B\-A3B, using the same six\-book cohort and thresholds as[Figure1](https://arxiv.org/html/2610.42112#S2.F1)\. Rows identify full\-attention layers and columns identify query heads, with zero\-based indices\. Each panel averages only across books\. All panels share logarithmic axis limits; colors denoteppand the dotted black line gives thep=1p=1reference\.
### A\-BAffine Versus Power\-Law Fits on Rollouts

[TableI](https://arxiv.org/html/2610.42112#A1.T1)provides a complementary comparison using ZHA leaf\-level selection counts, not dense top\-ppset sizes\. Each count unions selected leaf nodes across the eight query heads sharing a KV head and excludes mandatory sink and local tokens\. Here,CCis the checkpoint’s base KV chunk size,QQis the query\-tile width, andppis the routing mass threshold\.

We fit each rollout–layer–KV\-head trace separately, averaging tile counts in 256\-query bins\. Both models minimize squared log\-count residuals on the same observations\. Raw and cumulative\-maximum targets are fitted separately, with running maxima computed within each rollout before bin averaging\. The Raw and Cum\-max columns each give the affine win rate among 9,600 paired fits\.

The preferred model depends on the configuration\. Atp=0\.9p=0\.9and\(C,Q\)=\(8,16\)\(C,Q\)=\(8,16\), the affine win rate is3\.76%3\.76\\%for raw counts and5\.36%5\.36\\%for cumulative\-maximum counts\. For\(C,Q\)=\(256,256\)\(C,Q\)=\(256,256\)at the same threshold, the corresponding rates are99\.97%99\.97\\%and100%100\\%\. Power laws are therefore preferred in many configurations, but not uniformly across checkpoint geometries and query\-tile widths\.

Table I:Percentage of individual rollout–layer–head fits for which the affine model beats the power law,ρ=\(1−Rlin2\)/\(1−Rpow2\)<1\\rho=\(1\-R^\{2\}\_\{\\mathrm\{lin\}\}\)/\(1\-R^\{2\}\_\{\\mathrm\{pow\}\}\)<1, evaluated in log\-yyspace at the leaf level\. Raw and Cum\-max denote raw and cumulative\-maximum targets\. Each percentage uses 9,600 paired fits \(30 problems×\\times16 rollouts×\\times10 layers×\\times2 KV heads\), without averaging across rollouts before fitting\.\(a\)Q=1Q=1
\(b\)Q=16Q=16
\(c\)Q=32Q=32
\(d\)Q=64Q=64
\(e\)Q=128Q=128
\(f\)Q=256Q=256

### A\-CNemotron Top\-ppScaling

[Figure12](https://arxiv.org/html/2610.42112#A1.F12)repeats the dense analysis for Nemotron\-3\.5\-Lightning\-30B\-A3B on the same six books\. Here, we find that while top\-ppset sizes again grow sublinearly, the stronger log\-log curvature and late\-context fluctuations make a single power law less descriptive\. Sublinear growth thus need not follow a clean power law\. Although the models differ in positional encoding, this comparison does not isolate it as the cause\.

![Refer to caption](https://arxiv.org/html/2610.42112v1/nemotron35-dense-top-p-mean.png)Figure 12:Dense top\-ppset size for Nemotron\-3\.5\-Lightning\-30B\-A3B using the same six\-book cohort as[Figure1](https://arxiv.org/html/2610.42112#S2.F1)\.

## Appendix BTraining and Evaluation Details

### B\-ATraining Setup and Compute

[TableII](https://arxiv.org/html/2610.42112#A2.T2)summarizes the training regime for the four ZHA routers and the MSA indexer used in our evaluations\. Qwen denotes Qwen3\.6\-35B\-A3B and Nemotron denotes Nemotron\-3\.5\-Lightning\-30B\-A3B\.

Table II:Training setup and compute for the evaluated ZHA routers and MSA index\. Backbones are frozen; ZHA trains only reference queries\. Deep and wide ZHA useR=16R=16\.##### Training recipe\.

ZHA is trained by minimizing KL to dense attention over tree nodes at each level\. This is a local objective applied to each attention block independently\. By default, unless otherwise specified, ZHA shares 16 reference queries per KV group\. The same reference queries are used for all levels of the tree\. We train at 32,768 context using document\-masked attention\. Unless otherwise specified, the backbone is in BF16 and the routing parameters are in FP32\. We train for 1,900 steps, totaling approximately 4 billion tokens, and we use the final checkpoint of each training run for evaluation\.

For MSA, we use 128\-token blocks and select the top 16 blocks per query\. We perform warmup by minimizing KL over the key blocks\.

##### Optimization\.

We train with the AdamW optimizer\[[31](https://arxiv.org/html/2610.42112#bib.bib31)\]usingβ=\(0\.9,0\.95\)\\beta=\(0\.9,0\.95\)and gradient\-norm clipping at 1\.0\. ZHA’s reference\-query learning rate is3\.2×10−23\.2\\times 10^\{\-2\}, with no weight decay\. We use the warmup–stable–decay \(WSD\) learning rate schedule\[[17](https://arxiv.org/html/2610.42112#bib.bib17)\]with 20 warmup steps and linear decay over the final 20% of steps to 1% of the peak learning rate\. MSA uses a peak learning rate of3×10−43\\times 10^\{\-4\}, 200 warmup steps, and cosine decay to 10% of peak; weight decay is 0\.1 on projections and zero on norm gains\. The MSA paper does not report its optimizer settings\. We chose ours from the best of three 40B\-token pilot runs that varied the learning rate, batch size and schedule\. Nemotron’s stored FP32 MoE correction biases are frozen and excluded from the trainable parameter count\.

##### Data and compute accounting\.

For our Qwen and Nemotron training runs, we sample approximately 4B tokens \(13\.8M documents\) from a 100B\-token mid\-training mixture we used for the ZAYA1\-8B model\[[44](https://arxiv.org/html/2610.42112#bib.bib44)\]for each model\. We retokenized the data with the appropriate tokenizer\. We trained for a single epoch over the corpus\. We chose mid\-training data since we believed it was likely similar to the data distribution used for final mid\-training and SFT stages of the Qwen and Nemotron models prior to RL\.

##### Training KL\.

When we report the final KL, we are averaging the last 100 steps\. We report ZHA’s total sums over tree levels;[TableIII](https://arxiv.org/html/2610.42112#A2.T3)gives the individual values, with level 0 denoting the raw KV cache chunks\. The MSA value is its block\-level objective\. These differently aggregated objectives should not be compared directly across depths or methods\.

Table III:Final per\-level ZHA training KL, averaged over logged points in the last 100 steps\.

### B\-BEvaluation Protocols

Table IV:Evaluation datasets, scoring metrics, and inference protocols\.1Uses the pinned MathArena dataset revision containing the problem\-10 statement error\.

2Uses the Qwen3\.6 model\-card protocol: flat 3 h per trial and 32 CPUs / 48 GB per container\.

3Prompts exceeding the cap retain their first and last halves\.

##### Evaluation protocols\.

AIME 2026 reports avg@16 over 30 problems, grading the final boxed answer after reasoning\. SWE\-bench Verified reports resolved tasks out of 500\. Terminal\-Bench 2\.1 averages five attempts per task over 89 tasks with a three\-hour agent limit\. MMLU\-Pro uses 12,032 questions, and GPQA Diamond uses 198 questions with eight samples each\. RULER\-13 averages 13 official task scores at 256K tokens, with 100 examples per task\. LongBench v2 uses 205 questions with contexts≥\\geq128K tokens, a 32K\-token reasoning budget, and a 4K\-token answer budget\.

### B\-CQwen3\.6\-35B\-A3B

[Figure13](https://arxiv.org/html/2610.42112#A2.F13)extends the main\-text dense vs ZHA comparisons atp=0\.95p=0\.95to routing thresholdsp∈\{0\.95,0\.9,0\.8\}p\\in\\\{0\.95,0\.9,0\.8\\\}\. Each threshold is applied uniformly at each tree level\. The figure reports both benchmark scores and, where available, the share of keys attended during decoding, showing the tradeoff between accuracy and keys attended\. Furthermore it demonstrates how sparsity compounds with depth: at the same routing thresholdpp, the deep router achieves greater sparsity than the wide router\. We also compare the sparsity of both ZHA geometries with MSA in[Figure14](https://arxiv.org/html/2610.42112#A2.F14)\. For both figures; blue bars denote the deep router and green bars the wide router; darker shades indicate higher thresholds\.

Figure 13:Qwen3\.6\-35B\-A3B scores and decode key shares across routing thresholds for deep \(C4 L6, blue\) and wide \(C256 L1, green\) ZHA, withR=16R=16\. Darker shades indicate higherpp\. \(a\) Reasoning; \(b\) retrieval\. Labels above bars show the benchmark scores and labels inside bars give key shares where available\.Figure 14:Qwen3\.6\-35B\-A3B decode key shares for deep \(C4 L6, blue\) and wide \(C256 L1, green\) ZHA \(R=16R=16\) atp∈\{0\.95,0\.9,0\.8\}p\\in\\\{0\.95,0\.9,0\.8\\\}, compared with MSA \(top\-16 blocks of 128 tokens\)\. \(a\) Reasoning; \(b\) retrieval\. Darker shades indicate higherpp; the dashed line marks dense attention \(100%\)\.##### Small\-reference router\.

Figure 15:Dense attention vs\. a ZHA router with one reference query per KV group, 8\-token chunks, fan 8 and two levels \(R1 C8 L2\) on Qwen3\.6\-35B\-A3B, with uniform routing top\-pp0\.8, 0\.75 and 0\.7\. \(a\) Reasoning \(b\) Retrieval\.
##### Router precision\.

Figure 16:Dense attention vs\. the wide ZHA router \(16 reference queries per KV group, 256\-token chunks, one level; top\-pp0\.95\) on Qwen3\.6\-35B\-A3B, with the router in BF16 and in NVFP4 \(router operands, reference\-key cache and scoring in FP4; backbone, KV cache and the sparse attention in BF 16\)\. \(a\) Reasoning \(b\) Retrieval

### B\-DNemotron\-3\.5\-Lightning\-30B\-A3B

On Nemotron\-3\.5\-Lightning\-30B\-A3B, which has no explicit positional encodings, ZHA retains performance close to the dense baseline on the reasoning and coding benchmarks\. On RULER\-13 however atpp=0\.95, the deep and wide geometry scores 25\.3 and 10\.2 percentage points below the dense baseline, respectively\.[Figure17](https://arxiv.org/html/2610.42112#A2.F17)shows results across routing thresholds\.

Figure 17:Nemotron\-3\.5\-Lightning\-30B\-A3B scores and decode key shares across routing thresholds for deep \(C4 L6, blue\) and wide \(C256 L1, green\) ZHA, withR=16R=16\. Darker shades indicate higherpp\. \(a\) Reasoning; \(b\) retrieval\. Labels above bars show the benchmark scores and labels inside bars give key shares where available\.

## Appendix CGeometry, Retrieval Accuracy, and Efficiency

Here we provide supplementary material for the performance analysis of a variety of ZHA geometries and sparsities, including further distinguishing their performance per retrieval task and studying which geometries are on the sparsity–performance Pareto frontier\.

### C\-AGeometry Sweep

We are careful to cover a wide range of geometries to obtain enough coverage on the effect of each variable in\(R,C,L,p\)\(R,C,L,p\)\.[Figure18](https://arxiv.org/html/2610.42112#A3.F18)summarizes the 28 geometries used in our fits, spanning reference countsR∈\{1,2,4,16\}R\\in\\\{1,2,4,16\\\}, chunk sizesCCfrom 4 to 256, and depthsLLfrom 1 to 6\. We focus on coverage of top\-level compressionCLC^\{L\}, ranging from 16 to 4096; matched\-compression comparisons varying the decomposition intoCCandLL; and repeated\(C,L\)\(C,L\)pairs acrossRRto isolate the effect of reference\-query expressiveness\. Each geometry is evaluated atp∈\{0\.7,0\.8,0\.9,0\.95,0\.99\}p\\in\\\{0\.7,0\.8,0\.9,0\.95,0\.99\\\}and sequence lengths of 64K, 128K, and 256K\.

Figure 18:The 28 geometries used in the RULER fits\. Each panel fixes the reference countRR; dotted curves mark constant top\-level comparessionCLC^\{L\}\. These geometries are evaluated atp∈\{0\.7,0\.8,0\.9,0\.95,0\.99\}p\\in\\\{0\.7,0\.8,0\.9,0\.95,0\.99\\\}and sequence lengths of 64K, 128K, and 256K\.
### C\-BEffect of Geometry on Per\-Task RULER Accuracy

[SectionIV\-B](https://arxiv.org/html/2610.42112#S4.SS2)fitted retrieval accuracy in terms of the ZHA geometry and sparsity, but performance was measured as the aggregate accuracy across the 13 RULER tasks, but there is enormous variation in the qualitative retrieval dynamics and difficulty of the RULER subtasks\. Some tasks require not only sharp but narrow retrieval in order to ignore a large set of distractor needles, while other tasks require multi\-needle or multi\-hop retrieval\. As such, it is meaningful to study the accuracy of ZHA geometries per RULER task, as this may reveal particular weaknesses of different geometries and inform which provide the best tradeoff of speed and performance\.

[Figure19](https://arxiv.org/html/2610.42112#A3.F19)\(a\)–\(c\) show the same views of geometry’s effect on retrieval accuracy, that separate the effect of varying compression \(CLC^\{L\}\), depth \(LL\), and per\-node expressiveness \(RR\); panel \(d\) shows the displays the simple geometry\-to\-performance model’s quality\-of\-fit to the empirical RULER accuracy data \(in this case, per task\)\. The first observation is the variety in the “loss\-less sparsity” \(i\.e\., without accuracy loss\) permitted by each task; some admit good performance with a fraction of a percent of dense keys attended \(e\.\.g, single needle 1\) whereas others see decay by 10% of keys attended \(e\.g\., common\-word extraction\)\. This reflects that ZHA’s efficacy is bounded by the inherent compressibility of the underlying data\. Nevertheless, note also that the same trends with geometry remain, and are consistent across geometries: greater compression can hurt performance, and increasing depth or expressiveness enables greater sparsity at fixed performance\. Furthermore, though there is some variability in the fitting quality per task \(often when there is little spread in accuracy, whether high or low\), the geometry\-to\-performance model is broadly a good fit to the empirical RULER accuracy across tasks and sequence lengths\.

Figure 19:Per\-task geometry, sparsity, and prediction accuracy for Qwen3\.6\-35B\-A3B with reference\-query\-only routing; the final group shows the direct aggregate fit\. Panels \(a\), \(b\), and \(d\) fixR=16R=16\. \(a\) Measured accuracy at 256K versus decode keys attended across 11 geometries, colored byCLC^\{L\}\. \(b\) The matchedCL=4096C^\{L\}=4096geometries, colored by depthLL\. \(c\) TheC=8C=8,L=2L=2slice at 256K, colored by reference countRR\. Curves sweepp∈\{0\.7,0\.8,0\.9,0\.95\}p\\in\\\{0\.7,0\.8,0\.9,0\.95\\\}, withp=0\.99p=0\.99included where available; markers denote depth in \(a,b\) andRRin \(c\)\. Key shares use each task’s native decode counters \(pooled counters for the aggregate\), and the key\-share axes are logarithmic\. Dotted horizontal lines show dense attention\. \(d\) All 150 leave\-one\-length\-out predictions per response over 64K, 128K, and 256K, colored and marked by held\-out length; inset values give RMSE in percentage points and dashed diagonals show equality\. Each task has four independent parameters in[equation2](https://arxiv.org/html/2610.42112#S4.E2), fitted at the other two lengths acrossR∈\{1,2,4,16\}R\\in\\\{1,2,4,16\\\}using the measured dense baseline at the held\-out length\. Accuracy uses 100 examples per task; vertical\-axis ranges vary by task\.
### C\-CGeometry–Performance Pareto Frontier

To obtain a holistic view of the aggregate effect of different geometries on retrieval performance, we study the 28 geometries above together in[Figure20](https://arxiv.org/html/2610.42112#A3.F20), comparing accuracy with decode keys attended and modeled decode\-attention speedup\. At 256K, the high\-accuracy end of the sparsity frontier favorsR=16,C=4,L=2R=16,C=4,L=2, which achieves95\.8%95\.8\\%accuracy atp=0\.7p=0\.7while attending to0\.38%0\.38\\%of keys\. However, the deeperL=5L=5–66variants reach still lower key fractions but with substantially lower accuracy\. The modeled speedup frontier instead favors shallower or lower\-RRtrees: atp=0\.7p=0\.7,R=4,C=16,L=1R=4,C=16,L=1gives96\.4%96\.4\\%accuracy at1\.93×1\.93\\times,R=4,C=64,L=1R=4,C=64,L=1gives95\.5%95\.5\\%at2\.15×2\.15\\times, andR=1,C=8,L=2R=1,C=8,L=2gives92\.7%92\.7\\%at2\.28×2\.28\\times\. That these shallower or more lightweight trees are favored over the deeper variants reflects that minimizing the attended key count alone does not maximize speedup, as routing and selection also contribute to attention cost\.

Figure 20:Measured RULER\-13 accuracy at 256K versus \(a\) measured decode key share and \(b\) modeled decode\-attention speedup over dense for Qwen3\.6\-35B\-A3B \(115 settings, 28 geometries\)\. Speedups are H100 BF16 attention\-layer projections at batch size one, not measured end\-to\-end timings; missing higher\-level union counts use the slower modeled scenario\. Colors denoteRR, markersCC, and curves sweepppwithin each geometry\. Black lines mark Pareto frontiers; dotted horizontal lines show dense accuracy, and the vertical guide marks1×1\\times\.

## Appendix DRoofline Model

We use roofline models to compute the theoretical latencies of various attention mechanisms to estimate their optimized real\-world speeds\. Specifically, we build models for dense attention, ZHA, and MSA\. We study prefill and \(amortized\) decode\.

### D\-ADense

We begin with dense attention, using the same execution variables as ZHA below\. At fixed workload dimensions, these quantities are deterministic\.

Table V:Variables for the dense\-attention roofline model\.##### Latency\.

For each kernel, the roofline execution time is

Φx​\(T\)\\displaystyle\\Phi\_\{x\}\(T\)=max⁡\{FLOPsx​\(T\)Px,bytesx​\(T\)BWx\}\.\\displaystyle=\\max\\left\\\{\\frac\{\\mathrm\{FLOPs\}\_\{x\}\(T\)\}\{P\_\{x\}\},\\frac\{\\mathrm\{bytes\}\_\{x\}\(T\)\}\{\\mathrm\{BW\}\_\{x\}\}\\right\\\}\.\(9\)Dense attention has no tree construction, routing, or selection\. Its unsplit latency is simply

tdense​\(T\)=λA\+Φattn​\(T\)\.t\_\{\\mathrm\{dense\}\}\(T\)=\\lambda\_\{A\}\+\\Phi\_\{\\mathrm\{attn\}\}\(T\)\.\(10\)For split\-KV decode or noncausal attention, the partial and merge kernels give

tdense,split​\(T\)=λpartial\+Φpartial​\(T\)\+λmerge\+Φmerge​\(T\)\.t\_\{\\mathrm\{dense\},\\mathrm\{split\}\}\(T\)=\\lambda\_\{\\mathrm\{partial\}\}\+\\Phi\_\{\\mathrm\{partial\}\}\(T\)\+\\lambda\_\{\\mathrm\{merge\}\}\+\\Phi\_\{\\mathrm\{merge\}\}\(T\)\.\(11\)Each stage uses its own effective compute throughput and memory bandwidth\.

##### Work and traffic\.

With rectangular query/head tiling, the query\-side count is

Nqtiles,x​\(T\)\\displaystyle N\_\{\\mathrm\{qtiles\},x\}\(T\)=B​HKV​⌈Q⁡\(T\)κQ,x⌉​⌈GκH,x⌉=\{B​HKVgrouped\-head decode\(Q⁡\(T\)=1,κH,x=G\)B​Hq​⌈TκQ,x⌉prefill\(Q⁡\(T\)=T,κH,x=1\)\\displaystyle=BH\_\{\\mathrm\{KV\}\}\\left\\lceil\\frac\{Q\(T\)\}\{\\kappa\_\{Q,x\}\}\\right\\rceil\\left\\lceil\\frac\{G\}\{\\kappa\_\{H,x\}\}\\right\\rceil=\\begin\{cases\}BH\_\{\\mathrm\{KV\}\}&\\text\{grouped\-head decode $\(Q\(T\)=1,\\kappa\_\{H,x\}=G\)$\}\\\\ BH\_\{q\}\\left\\lceil\\frac\{T\}\{\\kappa\_\{Q,x\}\}\\right\\rceil&\\text\{prefill $\(Q\(T\)=T,\\kappa\_\{H,x\}=1\)$\}\\end\{cases\}\(12\)Theκ∙,x\\kappa\_\{\\bullet,x\}are kernel tile sizes\.

For unsplit attention over allTTkeys, the stage counts are

Ntiles,attn​\(T\)\\displaystyle N\_\{\\mathrm\{tiles\},\\mathrm\{attn\}\}\(T\)=Nqtiles,attn​\(T\),\\displaystyle=N\_\{\\mathrm\{qtiles\},\\mathrm\{attn\}\}\(T\),\(13\)FLOPsattn​\(T\)\\displaystyle\\mathrm\{FLOPs\}\_\{\\mathrm\{attn\}\}\(T\)=2​B​Q​\(T\)​Hq​\(dk\+dv\)​T,\\displaystyle=2BQ\(T\)H\_\{q\}\(d\_\{k\}\+d\_\{v\}\)T,\(14\)bytesattn​\(T\)\\displaystyle\\mathrm\{bytes\}\_\{\\mathrm\{attn\}\}\(T\)=sQKV​\(dk\+dv\)​\[Nqtiles,attn​\(T\)​T\+B​Q​\(T\)​Hq\]\.\\displaystyle=s\_\{\\mathrm\{QKV\}\}\(d\_\{k\}\+d\_\{v\}\)\\left\[N\_\{\\mathrm\{qtiles\},\\mathrm\{attn\}\}\(T\)T\+BQ\(T\)H\_\{q\}\\right\]\.\(15\)
For causal prefill with one\-head query tiles, letn⁡\(T\)=⌈T/κQ,attn⌉n\(T\)=\\lceil T/\\kappa\_\{Q,\\mathrm\{attn\}\}\\rceil\. SetFLOPsattn​\(T\)=B​Hq​\(dk\+dv\)​T​\(T\+1\)\\mathrm\{FLOPs\}\_\{\\mathrm\{attn\}\}\(T\)=BH\_\{q\}\(d\_\{k\}\+d\_\{v\}\)T\(T\+1\)and replace the KV\-read countNqtiles,attn​\(T\)​TN\_\{\\mathrm\{qtiles\},\\mathrm\{attn\}\}\(T\)Tinbytesattn​\(T\)\\mathrm\{bytes\}\_\{\\mathrm\{attn\}\}\(T\)by

B​Hq​\[κQ,attn​n​\(T\)​\(n⁡\(T\)−1\)2\+T\]\.BH\_\{q\}\\left\[\\frac\{\\kappa\_\{Q,\\mathrm\{attn\}\}n\(T\)\(n\(T\)\-1\)\}\{2\}\+T\\right\]\.\(16\)The query and output traffic is unchanged\.

For split\-KV decode or noncausal attention, the counts are

Ntiles,partial​\(T\)\\displaystyle N\_\{\\mathrm\{tiles\},\\mathrm\{partial\}\}\(T\)=Nqtiles,partial​\(T\)​J​\(T\),Ntiles,merge​\(T\)=Nqtiles,merge​\(T\),\\displaystyle=N\_\{\\mathrm\{qtiles\},\\mathrm\{partial\}\}\(T\)J\(T\),\\qquad N\_\{\\mathrm\{tiles\},\\mathrm\{merge\}\}\(T\)=N\_\{\\mathrm\{qtiles\},\\mathrm\{merge\}\}\(T\),\(17\)FLOPspartial​\(T\)\\displaystyle\\mathrm\{FLOPs\}\_\{\\mathrm\{partial\}\}\(T\)=2​B​Q​\(T\)​Hq​\(dk\+dv\)​T,FLOPsmerge​\(T\)=B​Q​\(T\)​Hq​\[2​J​\(T\)​dv\+3​J​\(T\)−1\],\\displaystyle=2BQ\(T\)H\_\{q\}\(d\_\{k\}\+d\_\{v\}\)T,\\qquad\\mathrm\{FLOPs\}\_\{\\mathrm\{merge\}\}\(T\)=BQ\(T\)H\_\{q\}\\left\[2J\(T\)d\_\{v\}\+3J\(T\)\-1\\right\],\(18\)bytespartial​\(T\)\\displaystyle\\mathrm\{bytes\}\_\{\\mathrm\{partial\}\}\(T\)=sQKV​\(dk\+dv\)​Nqtiles,partial​\(T\)​T\+B​Q​\(T\)​Hq​\[sQKV​dk\+saccum​J​\(T\)​\(dv\+2\)\],\\displaystyle=s\_\{\\mathrm\{QKV\}\}\(d\_\{k\}\+d\_\{v\}\)N\_\{\\mathrm\{qtiles\},\\mathrm\{partial\}\}\(T\)T\+BQ\(T\)H\_\{q\}\\left\[s\_\{\\mathrm\{QKV\}\}d\_\{k\}\+s\_\{\\mathrm\{accum\}\}J\(T\)\(d\_\{v\}\+2\)\\right\],\(19\)bytesmerge​\(T\)\\displaystyle\\mathrm\{bytes\}\_\{\\mathrm\{merge\}\}\(T\)=B​Q​\(T\)​Hq​\[saccum​J​\(T\)​\(dv\+2\)\+sQKV​dv\]\.\\displaystyle=BQ\(T\)H\_\{q\}\\left\[s\_\{\\mathrm\{accum\}\}J\(T\)\(d\_\{v\}\+2\)\+s\_\{\\mathrm\{QKV\}\}d\_\{v\}\\right\]\.\(20\)For the merge,PmergeP\_\{\\mathrm\{merge\}\}is an effective scalar arithmetic rate\. Query loads are assumed cached across splits\. Each partial writes advd\_\{v\}\-element vector and two softmax statistics atsaccums\_\{\\mathrm\{accum\}\}bytes per element; the merge reads them and writes the final output atsQKVs\_\{\\mathrm\{QKV\}\}bytes per element\. Both launches are included explicitly\.

The existing worker\-aware calibration was checked against native FlashAttention\[[7](https://arxiv.org/html/2610.42112#bib.bib7)\]on576576held\-out H100 cases \(144144/row\)\. Each usesB=1B=1,Hq=32H\_\{q\}=32;dk=dv∈\{64,128,256\}d\_\{k\}=d\_\{v\}\\in\\\{64,128,256\\\}andHKV∈\{1,4,8,32\}H\_\{\\mathrm\{KV\}\}\\in\\\{1,4,8,32\\\}, withT=29,…,220T=2^\{9\},\\ldots,2^\{20\}\. The error ise=100​\(tmodel/tFlash−1\)e=100\(t\_\{\\mathrm\{model\}\}/t\_\{\\mathrm\{Flash\}\}\-1\); calibration uses separate cases\.

Decode includes the profiled split/merge geometry\. These are attention\-kernel timings, not end\-to\-end model timings\.

\(a\)BF16 latency\.\(b\)FP16 latency\.\(c\)BF16 timing/prediction\.\(d\)FP16 timing/prediction\.
Figure 21:Held\-out dense\-attention latency and timing divided by prediction\. Within each subfigure, rows are MHA, GQA4, GQA8, and MQA; columns are causal prefill and decode\. Colours distinguishd∈\{64,128,256\}d\\in\\\{64,128,256\\\}\. Dots show all recorded timing samples and short ticks mark medians; solid lines in the latency views show worker\-aware predictions\. Each CUDA\-event sample averages a replay batch, not an individual invocation\. The lower views showtFlash/tmodelt\_\{\\mathrm\{Flash\}\}/t\_\{\\mathrm\{model\}\}; the dashed1×1\\timesline denotes exact prediction, and values below it indicate overprediction\. Matching views retain the same layout, colours, marks, and axis limits across precisions; the ratio views share the same ratio scale\.

### D\-BZHA

We now extend the dense model to ZHA, including tree construction, routing, and selection\. We model it as

𝒕ZHA​\(T\)\\displaystyle\\bm\{t\}\_\{\\mathrm\{ZHA\}\}\(T\)=ttree​\(T\)\+∑ℓ=1L⁡\(T\)\[λℓ\+𝚽routing,ℓ​\(T\)\+𝚽scan,ℓ​\(T\)\]⏟Levelℓof tree\+λA\+𝚽attn​\(T\)⏟Final attention\.\\displaystyle=t\_\{\\mathrm\{tree\}\}\(T\)\+\\sum\_\{\\ell=1\}^\{L\(T\)\}\\underbrace\{\\left\[\\lambda\_\{\\ell\}\+\\bm\{\\Phi\}\_\{\\mathrm\{routing\},\\ell\}\(T\)\+\\bm\{\\Phi\}\_\{\\mathrm\{scan\},\\ell\}\(T\)\\right\]\}\_\{\\text\{Level $\\ell$ of tree\}\}\+\\underbrace\{\\lambda\_\{A\}\+\\bm\{\\Phi\}\_\{\\mathrm\{attn\}\}\(T\)\}\_\{\\text\{Final attention\}\}\.\(21\)We useλ\\lambdato represent kernel launch overhead and𝚽\\bm\{\\Phi\}to capture execution time, accounting for compute, memory bandwidth, and available parallel workers\. Boldface indicates data\-dependent random quantities, including deterministic special cases\. Workload dimensions, kernel choices, and calibration coefficients are fixed, not random, in this model\. Context\-length dependence is explicit; in particular,Q⁡\(T\)Q\(T\)andL⁡\(T\)L\(T\)are deterministic functions set by the workload and tree configuration\. Dependence on the input and query is implicit, with routing parameters fixed\. For fusion, replace routing and selection by one execution term with its own worker and launch parameters\. Variable definitions are provided in[TablesVI](https://arxiv.org/html/2610.42112#A4.T6),[VII](https://arxiv.org/html/2610.42112#A4.T7)and[VIII](https://arxiv.org/html/2610.42112#A4.T8)\. We choose this model for a few reasons:

1. 1\.The sum over levels captures the inherent*serial*processing through the depth of tree\.
2. 2\.It separates kernel launch overhead from execution time, i\.e\., it does not assume perfect pipelining, which may be challenging with our architecture\.
3. 3\.It captures data\-dependent sparsity and latency\.
4. 4\.It accounts for potential under\-utilization of a GPU at higher levels of the tree\.

Table VI:Variables for the ZHA roofline model\.WithCCchildren per retained node, expansion gives

𝑲ℓ−1​\(T\)\\displaystyle\\bm\{K\}\_\{\\ell\-1\}\(T\)=C​𝑺ℓ​\(T\),1≤ℓ≤L⁡\(T\)\.\\displaystyle=C\\,\\bm\{S\}\_\{\\ell\}\(T\),\\qquad 1\\leq\\ell\\leq L\(T\)\.\(22\)Atℓ=1\\ell=1, the children are raw tokens\. Partially filled nodes contribute only their actual children\.

#### D\-B1Worker\-aware Execution

For a sparse stagexx, with the routing levelℓ\\ellimplicit when applicable, we use the same worker\-aware maximum:

𝚽x​\(T\)\\displaystyle\\bm\{\\Phi\}\_\{x\}\(T\)=max⁡\(𝒕x,worker​waves​\(T\),𝐅𝐋𝐎𝐏𝐬x​\(T\)Px,𝐛𝐲𝐭𝐞𝐬x​\(T\)BWx\)\.\\displaystyle=\\max\\left\(\\bm\{t\}\_\{x,\\mathrm\{worker\\ waves\}\}\(T\),\\frac\{\\mathbf\{FLOPs\}\_\{x\}\(T\)\}\{P\_\{x\}\},\\frac\{\\mathbf\{bytes\}\_\{x\}\(T\)\}\{\\mathrm\{BW\}\_\{x\}\}\\right\)\.\(23\)This is a standard roofline model except for the𝒕x,worker​waves​\(T\)\\bm\{t\}\_\{x,\\mathrm\{worker\\ waves\}\}\(T\)term\. In𝐛𝐲𝐭𝐞𝐬x​\(T\)\\mathbf\{bytes\}\_\{x\}\(T\), each stored Q/K/V, compressed\-key, reference\-bias, or routing\-score element contributessQKVs\_\{\\mathrm\{QKV\}\}bytes per transfer\. Accumulation precision is separate from storage width; intermediates kept on\-chip incur no HBM traffic\. The throughput, occupancy, and block\-latency parameters correspond to each stage’s operand and accumulation types\. For worker waves, we use

𝒕x,worker​waves​\(T\)\\displaystyle\\bm\{t\}\_\{x,\\mathrm\{worker\\ waves\}\}\(T\)=𝝉x​\(T\)​⌈𝑵tiles,x​\(T\)NSM​cx⌉\.\\displaystyle=\\bm\{\\tau\}\_\{x\}\(T\)\\left\\lceil\\frac\{\\bm\{N\}\_\{\\mathrm\{tiles\},x\}\(T\)\}\{N\_\{\\mathrm\{SM\}\}c\_\{x\}\}\\right\\rceil\.\(24\)For fixed work per block,𝝉x​\(T\)=τx\\bm\{\\tau\}\_\{x\}\(T\)=\\tau\_\{x\}is constant in this model\.

Table VII:Kernel and worker variables for the roofline model\.For separate kernels, we use the provisional occupancy assumptions

cx\\displaystyle c\_\{x\}=\{2,x∈\{routing,scan\},1,x=attn\.\\displaystyle=\\begin\{cases\}2,&x\\in\\\{\\mathrm\{routing\},\\mathrm\{scan\}\\\},\\\\ 1,&x=\\mathrm\{attn\}\.\\end\{cases\}\(25\)These values are not GPU\-independent constants: a different GPU or kernel can changecxc\_\{x\}, and𝝉x​\(T\)\\bm\{\\tau\}\_\{x\}\(T\)must be calibrated at that concurrency\.

Using the query\-side tile counts defined above, useκrouting\\kappa\_\{\\mathrm\{routing\}\}compressed keys per routing block andκscan\\kappa\_\{\\mathrm\{scan\}\}node scores per scan block:

𝑵tiles,routing,ℓ​\(T\)\\displaystyle\\bm\{N\}\_\{\\mathrm\{tiles\},\\mathrm\{routing\},\\ell\}\(T\)=Nqtiles,routing​\(T\)​⌈R​𝑲ℓ​\(T\)κrouting⌉,\\displaystyle=N\_\{\\mathrm\{qtiles\},\\mathrm\{routing\}\}\(T\)\\left\\lceil\\frac\{R\\bm\{K\}\_\{\\ell\}\(T\)\}\{\\kappa\_\{\\mathrm\{routing\}\}\}\\right\\rceil,\(26\)𝑵tiles,scan,ℓ​\(T\)\\displaystyle\\bm\{N\}\_\{\\mathrm\{tiles\},\\mathrm\{scan\},\\ell\}\(T\)=Nqtiles,scan​\(T\)​⌈𝑲ℓ​\(T\)κscan⌉,\\displaystyle=N\_\{\\mathrm\{qtiles\},\\mathrm\{scan\}\}\(T\)\\left\\lceil\\frac\{\\bm\{K\}\_\{\\ell\}\(T\)\}\{\\kappa\_\{\\mathrm\{scan\}\}\}\\right\\rceil,\(27\)𝑵tiles,attn​\(T\)\\displaystyle\\bm\{N\}\_\{\\mathrm\{tiles\},\\mathrm\{attn\}\}\(T\)=Nqtiles,attn​\(T\)​\{1,unsplit attention,⌈𝑲0​\(T\)/κattn⌉,split\-KV partials\.\\displaystyle=N\_\{\\mathrm\{qtiles\},\\mathrm\{attn\}\}\(T\)\\begin\{cases\}1,&\\text\{unsplit attention\},\\\\ \\left\\lceil\\bm\{K\}\_\{0\}\(T\)/\\kappa\_\{\\mathrm\{attn\}\}\\right\\rceil,&\\text\{split\-KV partials\}\.\\end\{cases\}\(28\)Hereκattn\\kappa\_\{\\mathrm\{attn\}\}is retained tokens per split\-KV block; split\-KV also needs a separately timed merge pass withNqtiles,attn​\(T\)N\_\{\\mathrm\{qtiles\},\\mathrm\{attn\}\}\(T\)blocks\. Unsplit attention instead loops over retained keys within each block, so𝑲0​\(T\)\\bm\{K\}\_\{0\}\(T\)affects𝝉attn​\(T\)\\bm\{\\tau\}\_\{\\mathrm\{attn\}\}\(T\), not its block count\. FLOPs and block latency include all packed query positions and heads\. These are per\-pass counts; additional selection or reduction passes add their own execution and launch costs\.

Gather costs are reflected in the block latencies161616Is this the best place to account for them?\.

##### FLOPs\.

The dominant matrix work is

𝐅𝐋𝐎𝐏𝐬routing,ℓ​\(T\)\\displaystyle\\mathbf\{FLOPs\}\_\{\\mathrm\{routing\},\\ell\}\(T\)≃2​B​Q​\(T\)​Hq​R​dk​𝑲ℓ​\(T\),\\displaystyle\\simeq 2BQ\(T\)H\_\{q\}Rd\_\{k\}\\,\\bm\{K\}\_\{\\ell\}\(T\),\(29\)𝐅𝐋𝐎𝐏𝐬attn​\(T\)\\displaystyle\\mathbf\{FLOPs\}\_\{\\mathrm\{attn\}\}\(T\)≃2​B​Q​\(T\)​Hq​\(dk\+dv\)​𝑲0​\(T\)\.\\displaystyle\\simeq 2BQ\(T\)H\_\{q\}\(d\_\{k\}\+d\_\{v\}\)\\,\\bm\{K\}\_\{0\}\(T\)\.\(30\)The routing contribution is𝑸⁡\(T\)​𝑲c​\(T\)⊤\\bm\{Q\}\(T\)\\bm\{K\}\_\{c\}\(T\)^\{\\top\}; selection is accounted for separately\. The two attention contributions are𝑸⁡\(T\)​𝑲​\(T\)⊤\\bm\{Q\}\(T\)\\bm\{K\}\(T\)^\{\\top\}and multiplication of the attention weights by𝑽⁡\(T\)\\bm\{V\}\(T\)\.

For radix selection,𝐅𝐋𝐎𝐏𝐬scan,ℓ​\(T\)\\mathbf\{FLOPs\}\_\{\\mathrm\{scan\},\\ell\}\(T\)counts the ordinary floating\-point arithmetic of the chosen implementation\.PscanP\_\{\\mathrm\{scan\}\}is a scalar arithmetic rate, not Tensor Core throughput\. Exponentials, bit operations, comparisons, and index handling are reflected in𝝉scan,ℓ​\(T\)\\bm\{\\tau\}\_\{\\mathrm\{scan\},\\ell\}\(T\)\.

##### Bytes\.

For traffic, assume attention outputs and reference biases use the same storage precisionsQKVs\_\{\\mathrm\{QKV\}\}, measured in bytes\. For streaming radix selection, we have to doκradix\\kappa\_\{\\mathrm\{radix\}\}full\-frontier passes171717Is this really true? Is it possible to keep most of this on\-chip or in L2?, each reading one stored score per candidate\. The byte counts are

𝐛𝐲𝐭𝐞𝐬routing,ℓ​\(T\)\\displaystyle\\mathbf\{bytes\}\_\{\\mathrm\{routing\},\\ell\}\(T\)=R​sQKV​\(dk\+1\)​Nqtiles,routing​\(T\)​𝑲ℓ​\(T\)⏟compressed keys \+ biases\+sQKV​dk​B​Q​\(T\)​Hq⏟reference queries\\displaystyle=\\underbrace\{Rs\_\{\\mathrm\{QKV\}\}\(d\_\{k\}\+1\)N\_\{\\mathrm\{qtiles\},\\mathrm\{routing\}\}\(T\)\\bm\{K\}\_\{\\ell\}\(T\)\}\_\{\\text\{compressed keys \+ biases\}\}\+\\underbrace\{s\_\{\\mathrm\{QKV\}\}d\_\{k\}BQ\(T\)H\_\{q\}\}\_\{\\text\{reference queries\}\}\(31\)𝐛𝐲𝐭𝐞𝐬scan,ℓ​\(T\)\\displaystyle\\mathbf\{bytes\}\_\{\\mathrm\{scan\},\\ell\}\(T\)=sQKV​κradix​B​Q​\(T\)​Hq​𝑲ℓ​\(T\),\\displaystyle=s\_\{\\mathrm\{QKV\}\}\\kappa\_\{\\mathrm\{radix\}\}BQ\(T\)H\_\{q\}\\bm\{K\}\_\{\\ell\}\(T\),\(32\)𝐛𝐲𝐭𝐞𝐬attn,unsplit​\(T\)\\displaystyle\\mathbf\{bytes\}\_\{\\mathrm\{attn\},\\mathrm\{unsplit\}\}\(T\)=sQKV​\(dk\+dv\)​\[Nqtiles,attn​\(T\)​𝑲0​\(T\)\+B​Q​\(T\)​Hq\],\\displaystyle=s\_\{\\mathrm\{QKV\}\}\(d\_\{k\}\+d\_\{v\}\)\\left\[N\_\{\\mathrm\{qtiles\},\\mathrm\{attn\}\}\(T\)\\bm\{K\}\_\{0\}\(T\)\+BQ\(T\)H\_\{q\}\\right\],\(33\)𝐛𝐲𝐭𝐞𝐬attn,split​\(T\)\\displaystyle\\mathbf\{bytes\}\_\{\\mathrm\{attn\},\\mathrm\{split\}\}\(T\)=𝐛𝐲𝐭𝐞𝐬attn,unsplit​\(T\)\+2​saccum​B​Q​\(T\)​Hq​⌈𝑲0​\(T\)κattn⌉​\(dv\+2\)\.\\displaystyle=\\mathbf\{bytes\}\_\{\\mathrm\{attn\},\\mathrm\{unsplit\}\}\(T\)\+2s\_\{\\mathrm\{accum\}\}BQ\(T\)H\_\{q\}\\left\\lceil\\frac\{\\bm\{K\}\_\{0\}\(T\)\}\{\\kappa\_\{\\mathrm\{attn\}\}\}\\right\\rceil\(d\_\{v\}\+2\)\.\(34\)Here⌈𝑲0​\(T\)/κattn⌉\\lceil\\bm\{K\}\_\{0\}\(T\)/\\kappa\_\{\\mathrm\{attn\}\}\\rceilis the number of KV splits, andsaccums\_\{\\mathrm\{accum\}\}denotes the bytes per stored accumulator/statistic element\.

Split\-KV writes and rereads advd\_\{v\}\-element partial vector and two softmax statistics per query/head and split\. Query loads are assumed cached; partial and merge kernels are timed separately\.

These leading payload counts include reference biases but omit index/mask reads, retained\-ID or mask writes, quantization scales, and algorithm\-specific scratch\.

#### D\-B2Tree construction

In addition to routing through the tree, we must consider the cost of constructing it\. Use the baseline stable\-state merge, without Taylor/RoPE re\-anchoring, with reference queries already prepared\. A leaf computes reference/key scores and weighted key sums:4​R​C​dk4RCd\_\{k\}leading FLOPs per query head\. A parent rescales and sumsCCchild weighted sums independently for each reference:2​R​C​dk2RCd\_\{k\}leading scalar FLOPs\.

Per reference, retain a maximum, weight sum, entropy, anddkd\_\{k\}\-vector, usingsaccums\_\{\\mathrm\{accum\}\}bytes per element for subsequent merges\. Emit normalized reference keys atsQKVs\_\{\\mathrm\{QKV\}\}bytes per element in the same kernel, rather than an additional mean key or a later normalization pass\. The counts below include the leaf kernel’s two key reads and the parent kernel’s repeated maximum read\.

Table VIII:Variables for tree construction and maintenance\.Tree kernels use the free parameters in[TableVII](https://arxiv.org/html/2610.42112#A4.T7), withx=summx=\\mathrm\{summ\}orx=\(tree,merge\)x=\(\\mathrm\{tree\},\\mathrm\{merge\}\)\.

Forn\>0n\>0full nodes per sequence, withG=Hq/HKVG=H\_\{q\}/H\_\{\\mathrm\{KV\}\}, the kernel counts are

Ntiles,summ​\(n\)\\displaystyle N\_\{\\mathrm\{tiles\},\\mathrm\{summ\}\}\(n\)=B​HKV​n,Ntiles,tree,merge​\(n\)=B​Hq​n,\\displaystyle=BH\_\{\\mathrm\{KV\}\}n,\\qquad N\_\{\\mathrm\{tiles\},\\mathrm\{tree\},\\mathrm\{merge\}\}\(n\)=BH\_\{q\}n,\(35\)FLOPssumm​\(n\)\\displaystyle\\mathrm\{FLOPs\}\_\{\\mathrm\{summ\}\}\(n\)=4​B​Hq​n​R​C​dk,FLOPstree,merge​\(n\)=2​B​Hq​n​R​C​dk,\\displaystyle=4BH\_\{q\}nRCd\_\{k\},\\qquad\\mathrm\{FLOPs\}\_\{\\mathrm\{tree\},\\mathrm\{merge\}\}\(n\)=2BH\_\{q\}nRCd\_\{k\},\(36\)bytessumm​\(n\)\\displaystyle\\mathrm\{bytes\}\_\{\\mathrm\{summ\}\}\(n\)=B​Hq​n​\[2​sQKV​dk​\(C/G\+R\)\+saccum​R​\(dk\+3\)\],\\displaystyle=BH\_\{q\}n\\left\[2s\_\{\\mathrm\{QKV\}\}d\_\{k\}\(C/G\+R\)\+s\_\{\\mathrm\{accum\}\}R\(d\_\{k\}\+3\)\\right\],\(37\)bytestree,merge​\(n\)\\displaystyle\\mathrm\{bytes\}\_\{\\mathrm\{tree\},\\mathrm\{merge\}\}\(n\)=B​Hq​n​R​\[saccum​C​\(dk\+4\)\+saccum​\(dk\+3\)\+sQKV​dk\]\.\\displaystyle=BH\_\{q\}nR\\left\[s\_\{\\mathrm\{accum\}\}C\(d\_\{k\}\+4\)\+s\_\{\\mathrm\{accum\}\}\(d\_\{k\}\+3\)\+s\_\{\\mathrm\{QKV\}\}d\_\{k\}\\right\]\.\(38\)The kernel costs are

tsumm​\(n\)\\displaystyle t\_\{\\mathrm\{summ\}\}\(n\)=λsumm\+max⁡\{τsumm​⌈Ntiles,summ​\(n\)NSM​csumm⌉,FLOPssumm​\(n\)Psumm,bytessumm​\(n\)BWsumm\},\\displaystyle=\\lambda\_\{\\mathrm\{summ\}\}\+\\max\\left\\\{\\tau\_\{\\mathrm\{summ\}\}\\left\\lceil\\tfrac\{N\_\{\\mathrm\{tiles\},\\mathrm\{summ\}\}\(n\)\}\{N\_\{\\mathrm\{SM\}\}c\_\{\\mathrm\{summ\}\}\}\\right\\rceil,\\frac\{\\mathrm\{FLOPs\}\_\{\\mathrm\{summ\}\}\(n\)\}\{P\_\{\\mathrm\{summ\}\}\},\\frac\{\\mathrm\{bytes\}\_\{\\mathrm\{summ\}\}\(n\)\}\{\\mathrm\{BW\}\_\{\\mathrm\{summ\}\}\}\\right\\\},\(39\)tmerge​\(n\)\\displaystyle t\_\{\\mathrm\{merge\}\}\(n\)=λtree,merge\+max⁡\{τtree,merge​⌈Ntiles,tree,merge​\(n\)NSM​ctree,merge⌉,FLOPstree,merge​\(n\)Ptree,merge,bytestree,merge​\(n\)BWtree,merge\}\.\\displaystyle=\\lambda\_\{\\mathrm\{tree\},\\mathrm\{merge\}\}\+\\max\\left\\\{\\tau\_\{\\mathrm\{tree\},\\mathrm\{merge\}\}\\left\\lceil\\tfrac\{N\_\{\\mathrm\{tiles\},\\mathrm\{tree\},\\mathrm\{merge\}\}\(n\)\}\{N\_\{\\mathrm\{SM\}\}c\_\{\\mathrm\{tree\},\\mathrm\{merge\}\}\}\\right\\rceil,\\frac\{\\mathrm\{FLOPs\}\_\{\\mathrm\{tree\},\\mathrm\{merge\}\}\(n\)\}\{P\_\{\\mathrm\{tree\},\\mathrm\{merge\}\}\},\\frac\{\\mathrm\{bytes\}\_\{\\mathrm\{tree\},\\mathrm\{merge\}\}\(n\)\}\{\\mathrm\{BW\}\_\{\\mathrm\{tree\},\\mathrm\{merge\}\}\}\\right\\\}\.\(40\)Leaf blocks share keys acrossGGquery heads; parent blocks are per query head\. Softmax, entropy, and normalization enter the calibrated block latencies;Ptree,mergeP\_\{\\mathrm\{tree\},\\mathrm\{merge\}\}is a scalar rate, not a Tensor Core rate\. Both costs are zero forn=0n=0\.

For the balanced full\-node model, letnℓ​\(T\)=⌊T/Cℓ⌋n\_\{\\ell\}\(T\)=\\lfloor T/C^\{\\ell\}\\rfloor\. Prefill batches all nodes at each level; decode averages one completed node per sequence everyCℓC^\{\\ell\}tokens:

ttreeprefill​\(T\)\\displaystyle t\_\{\\mathrm\{tree\}\}^\{\\mathrm\{prefill\}\}\(T\)=tsumm​\(n1​\(T\)\)\+∑ℓ=2L⁡\(T\)tmerge​\(nℓ​\(T\)\),\\displaystyle=t\_\{\\mathrm\{summ\}\}\(n\_\{1\}\(T\)\)\+\\sum\_\{\\ell=2\}^\{L\(T\)\}t\_\{\\mathrm\{merge\}\}\(n\_\{\\ell\}\(T\)\),\(41\)t¯treedecode​\(T\)\\displaystyle\\overline\{t\}\_\{\\mathrm\{tree\}\}^\{\\mathrm\{decode\}\}\(T\)=tsumm​\(1\)C\+∑ℓ=2L⁡\(T\)tmerge​\(1\)Cℓ\\displaystyle=\\frac\{t\_\{\\mathrm\{summ\}\}\(1\)\}\{C\}\+\\sum\_\{\\ell=2\}^\{L\(T\)\}\\frac\{t\_\{\\mathrm\{merge\}\}\(1\)\}\{C^\{\\ell\}\}\(42\)=tsumm​\(1\)C\+tmerge​\(1\)C⁡\(C−1\)​\(1−C−\(L⁡\(T\)−1\)\)\.\\displaystyle=\\frac\{t\_\{\\mathrm\{summ\}\}\(1\)\}\{C\}\+\\frac\{t\_\{\\mathrm\{merge\}\}\(1\)\}\{C\(C\-1\)\}\\left\(1\-C^\{\-\(L\(T\)\-1\)\}\\right\)\.\(43\)These counts exclude partial boundary nodes\. The decode expression is an average, not a tail latency; an unchanged resident tree hasttree​\(T\)=0t\_\{\\mathrm\{tree\}\}\(T\)=0\.

### D\-CMSA

MSA uses one index\-query head per KV group and*one shared index\-key head*, both of widthdId\_\{I\}\. It scores every visible token, max\-pools withinCC\-token blocks, then selectskkblocks per group\. The local block occupies one of these slots\.

The shared workload and hardware variables retain the definitions in[TablesV](https://arxiv.org/html/2610.42112#A4.T5)and[VII](https://arxiv.org/html/2610.42112#A4.T7);[TableIX](https://arxiv.org/html/2610.42112#A4.T9)defines the additional notation used here\. The block sizeCCand selection counts below are local to MSA\.

Table IX:Additional variables for the MSA roofline model\.For decode \(Q=1Q=1\), define

M⁡\(T\)\\displaystyle M\(T\)=⌈T/C⌉,\\displaystyle=\\left\\lceil T/C\\right\\rceil,kT\\displaystyle k\_\{T\}=min⁡\(k,M⁡\(T\)\),\\displaystyle=\\min\(k,M\(T\)\),\(44\)r⁡\(T\)\\displaystyle r\(T\)=T−C⁡\(M⁡\(T\)−1\),\\displaystyle=T\-C\(M\(T\)\-1\),K⁡\(T\)\\displaystyle K\(T\)=C⁡\(kT−1\)\+r⁡\(T\)\.\\displaystyle=C\(k\_\{T\}\-1\)\+r\(T\)\.\(45\)ThusK⁡\(T\)=TK\(T\)=Twhen all blocks fit, andK⁡\(T\)≤k​CK\(T\)\\leq kCotherwise\. The selected addresses vary, but these counts are deterministic\. The decode stage counts are:

The two index projections are fused into one launch; cached historical index keys are not reprojected\. Projection weights and inputs, and index queries across key tiles, are assumed cached within their stage\. Only pooled block scores reach HBM, not aTT\-element token\-score vector per index query\. Pooling comparisons and exp\-free top\-kkwork are charged through their block latencies; they are not Tensor Core FLOPs\. The top\-kkrow assumes each group’s candidate list fits on\-chip\.

Using the same worker\-aware maximum, with stage\-specific parameters,

Φx​\(T\)=max⁡\{τx​\(T\)​⌈Nx​\(T\)NSM​cx⌉,Fx​\(T\)Px,Vx​\(T\)BWx\},\\displaystyle\\Phi\_\{x\}\(T\)=\\max\\left\\\{\\tau\_\{x\}\(T\)\\left\\lceil\\tfrac\{N\_\{x\}\(T\)\}\{N\_\{\\mathrm\{SM\}\}c\_\{x\}\}\\right\\rceil,\\frac\{F\_\{x\}\(T\)\}\{P\_\{x\}\},\\frac\{V\_\{x\}\(T\)\}\{\\mathrm\{BW\}\_\{x\}\}\\right\\\},\(46\)tMSA​\(T\)=∑x∈\{proj,I,top​k\}\[λx\+Φx​\(T\)\]\+tsparseMSA​\(K⁡\(T\)\)\.\\displaystyle\\boxed\{t\_\{\\mathrm\{MSA\}\}\(T\)=\\sum\_\{x\\in\\\{\\mathrm\{proj\},I,\\mathrm\{top\}k\\\}\}\[\\lambda\_\{x\}\+\\Phi\_\{x\}\(T\)\]\+t\_\{\\mathrm\{sparse\}\}^\{\\mathrm\{MSA\}\}\(K\(T\)\)\.\}\(47\)OmitFx/PxF\_\{x\}/P\_\{x\}for top\-kk, whose scalar work is included inτtop​k​\(M,kT\)\\tau\_\{\\mathrm\{top\}k\}\(M,k\_\{T\}\)\. There is one top\-kklaunch, notkTk\_\{T\}launches\. For query\-outer main attention, the unsplit counts are

NA\\displaystyle N\_\{A\}=B​HKV​⌈G/κH,A⌉,\\displaystyle=BH\_\{\\mathrm\{KV\}\}\\left\\lceil G/\\kappa\_\{H,A\}\\right\\rceil,\(48\)FA\\displaystyle F\_\{A\}=2​B​Hq​\(dk\+dv\)​K​\(T\),\\displaystyle=2BH\_\{q\}\(d\_\{k\}\+d\_\{v\}\)K\(T\),\(49\)VA\\displaystyle V\_\{A\}=s⁡\(dk\+dv\)​\[NA​K​\(T\)\+B​Hq\]\.\\displaystyle=s\(d\_\{k\}\+d\_\{v\}\)\\,\[N\_\{A\}K\(T\)\+BH\_\{q\}\]\.\(50\)For split\-KV, use the partial/merge counts in[equations17](https://arxiv.org/html/2610.42112#A4.E17)to[20](https://arxiv.org/html/2610.42112#A4.E20), replacing the*key length*byK⁡\(T\)K\(T\), not the query length\. The same two\-statistic partial\-state convention applies\. As elsewhere, final\-attention index reads, gather penalties, and padding are not separately counted; they require kernel\-specific calibration\.

For causal prefill, sum the visible per\-position counts:

Fprojprefill\\displaystyle F\_\{\\mathrm\{proj\}\}^\{\\mathrm\{prefill\}\}=2​B​T​dmodel​WI,\\displaystyle=2BTd\_\{\\mathrm\{model\}\}W\_\{I\},\(51\)FIprefill\\displaystyle F\_\{I\}^\{\\mathrm\{prefill\}\}=B​HKV​dI​T​\(T\+1\),\\displaystyle=BH\_\{\\mathrm\{KV\}\}d\_\{I\}T\(T\+1\),\(52\)FAprefill\\displaystyle F\_\{A\}^\{\\mathrm\{prefill\}\}=2​B​Hq​\(dk\+dv\)​∑i=1TK⁡\(i\)\.\\displaystyle=2BH\_\{q\}\(d\_\{k\}\+d\_\{v\}\)\\sum\_\{i=1\}^\{T\}K\(i\)\.\(53\)Pooled\-score and selected\-ID counts becomeB​HKV​∑iM⁡\(i\)BH\_\{\\mathrm\{KV\}\}\\sum\_\{i\}M\(i\)andB​HKV​∑ikiBH\_\{\\mathrm\{KV\}\}\\sum\_\{i\}k\_\{i\}, respectively\. If index query tilejjends attjt\_\{j\}, its key\-read count usestjt\_\{j\}and its key\-grid count uses⌈tj/κI⌉\\lceil t\_\{j\}/\\kappa\_\{I\}\\rceil; do not multiply the final\-position decode cost byTT\.

Native MSA prefill instead uses KV\-outer scheduling\. Let𝒏b\\bm\{n\}\_\{b\}be the number of queries selecting blockbbwithin its KV group, and𝒂b=⌈G​𝒏b/κpacked⌉\\bm\{a\}\_\{b\}=\\lceil G\\bm\{n\}\_\{b\}/\\kappa\_\{\\mathrm\{packed\}\}\\rceilits packed query tiles\. Sums below also include batch and KV group;CbC\_\{b\}is the stored block length andE=∑b𝒏b=B​HKV​∑ikiE=\\sum\_\{b\}\\bm\{n\}\_\{b\}=BH\_\{\\mathrm\{KV\}\}\\sum\_\{i\}k\_\{i\}\. For one CTA per packed query tile, the partial\-attention grid and traffic are

𝑵partial\\displaystyle\\bm\{N\}\_\{\\mathrm\{partial\}\}=∑b𝒂b,\\displaystyle=\\sum\_\{b\}\\bm\{a\}\_\{b\},\(54\)𝐛𝐲𝐭𝐞𝐬partial\\displaystyle\\mathbf\{bytes\}\_\{\\mathrm\{partial\}\}=s⁡\(dk\+dv\)​∑bCb​𝒂b\+s​dk​G​E\+saccum​\(dv\+2\)​G​E\.\\displaystyle=s\(d\_\{k\}\+d\_\{v\}\)\\sum\_\{b\}C\_\{b\}\\bm\{a\}\_\{b\}\+sd\_\{k\}GE\+s\_\{\\mathrm\{accum\}\}\(d\_\{v\}\+2\)GE\.\(55\)Add reverse\-index/scheduling launches and scratch traffic, then combinekik\_\{i\}partials per query using the merge model\. Native chunking or persistent schedules must use their actual launch grids, partial counts, and traffic\. Block popularity therefore affects worker waves and KV reuse, even at fixedkk; the decode comparison below does not claim to reproduce this native prefill schedule\.

## Appendix ESparsity Scaling Across Routing Thresholds

The latency extrapolations in[SectionV](https://arxiv.org/html/2610.42112#S5)assume a sublinear power\-law scaling at each tree level that changes with depth, and therefore we fit those exponents and model their dependence on depth in\. Here we extend the main\-textp=0\.8p=0\.8results to other thresholds, explain how these trends are estimated and validated, and connect them to the latency model\.

### E\-AEstimating Per\-Level Growth

Estimating the per\-level sublinear exponent from empirical selection rates requires carefully controlling for outliers in sequence length and example count\. We seek typical problem\-level growth, limiting the influence of unusually long generations or extreme slopes\. The Qwen3\.6\-35B\-A3B sweep uses the same 30 AIME 2026 problems, with 16 generations per setting:C=4C=4,R=16R=16,L=2L=2–66, andp∈\{0\.7,0\.8,0\.9,0\.95\}p\\in\\\{0\.7,0\.8,0\.9,0\.95\\\}\(9,600 completions\)\.

We replay saved generations with single\-query tiles to measure shared KV footprints: each selected complete node is counted once across the eight query heads in its group, separately at each position, layer, and tree level\. Counts are averaged over two groups, 10 full\-attention layers, and complete 256\-token bins, excluding the always\-retained first key and current partial chunk\.

For each problem, we average available generations within each bin and fit a free\-intercept log–log slope over 16K–64K \(192 native bins\)\. Objective weights1/T1/Tapproximate equal weighting in log context length\. Requiring three complete bins leaves 29–30 problems per setting\. Their median slope definesαℓ\\alpha\_\{\\ell\}, weighting problems equally without imputing missing tails\.

### E\-BModeling Levels Within a Tree

Provided empirically estimated exponents per level, we may fit our depthwise exponent model\. For a fixed tree,β\\betasets the top exponent andε\\varepsilonits correction toward the leaves in[equation7](https://arxiv.org/html/2610.42112#S5.E7)\. Equally weighted least squares fits both to the observedαℓ\\alpha\_\{\\ell\}separately for each geometry andpp, with signedε\\varepsilon,0\.01≤β≤10\.01\\leq\\beta\\leq 1, andβℓ≥0\.01\\beta\_\{\\ell\}\\geq 0\.01\. Note that joint fitting can shiftβ\\betafrom the observed top slope, and that the factors describe combined pruning and KV\-group union effects\.

The parameterrrdetermines how much of the accumulated correction carries into the next lower level:

εL=0,εℓ−1=ε\+r​εℓ,εℓ=\{ε⁡\(1−rL−ℓ\)/\(1−r\),r≠1,ε⁡\(L−ℓ\),r=1\.\\varepsilon\_\{L\}=0,\\qquad\\varepsilon\_\{\\ell\-1\}=\\varepsilon\+r\\varepsilon\_\{\\ell\},\\qquad\\varepsilon\_\{\\ell\}=\\begin\{cases\}\\varepsilon\(1\-r^\{L\-\\ell\}\)/\(1\-r\),&r\\neq 1,\\\\ \\varepsilon\(L\-\\ell\),&r=1\.\\end\{cases\}\(56\)
Herer=0r=0gives a constant correction below the top,r=1/2r=1/2approaches2​ε2\\varepsilon, andr=1r=1gives linear drift\. Across 20 settings \(80 exponent summaries\), their RMSEs, weighting settings equally, are0\.02180\.0218,0\.01300\.0130, and0\.01560\.0156\. We favorr=1/2r=1/2as a compact descriptive choice that balances accuracy across settings: paired bootstraps favorr=1/2r=1/2overr=0r=0, though not clearly overr=1r=1; its advantage comes mainly fromp=0\.7p=0\.7,L=6L=6\.

[Figure22](https://arxiv.org/html/2610.42112#A5.F22)shows decreasing exponent profiles atp=0\.9,0\.95p=0\.9,0\.95; the deepestp=0\.7p=0\.7tree instead rises toward the leaves\. We normalize each level and geometry by the first\-bin count to make growth rates comparable\. For visualizations, we further average 16 native\-bin means into each of 12 display bins; fits retain native resolution\. With profile slopes fixed, amplitudescℓc\_\{\\ell\}fit pooled arithmetic\-mean counts by1/T1/T\-weighted least squares in log space, weighting available generations equally within each bin\.

The results for these otherppvalues broadly corroborate the main\-textp=0\.8p=0\.8findings; scaling exponents decrease smoothly and predictably proceeding top\-down the tree\. The exception is the deepest geometry \(C=4C=4,L=6L=6,R=16R=16\) atp=0\.7p=0\.7, whose exponents rise again toward the leaf levels\. This setting achieves only 44\.0% accuracy on AIME, compared with 93\.1% for dense attention, so its anomalous scaling should be interpreted in the context of severe accuracy degradation\.

Figure 22:Per\-level selection scaling forC=4C=4,R=16R=16, andp=0\.7,0\.9,0\.95p=0\.7,0\.9,0\.95, extending[Figure7](https://arxiv.org/html/2610.42112#S5.F7)\. The first five columns show depths 2–6; the sixth compares observed median exponents with ther=1/2r=1/2profile\. Bars show pointwise 95% whole\-problem bootstrap intervals\.
### E\-CAdding Tree Depth

In order to estimate the selection exponents not only within fitted trees, but to extrapolate to larger depths, we must determine what happens when another level is added\. To isolate this effect,[Figure23](https://arxiv.org/html/2610.42112#A5.F23)fits only the five observed top slopesαL\\alpha\_\{L\}at eachpp, independently of lower levels:

β^const​\(L,p\)=ap,β^lin​\(L,p\)=ap\+bp​\(L−4\)\.\\widehat\{\\beta\}\_\{\\mathrm\{const\}\}\(L,p\)=a\_\{p\},\\qquad\\widehat\{\\beta\}\_\{\\mathrm\{lin\}\}\(L,p\)=a\_\{p\}\+b\_\{p\}\(L\-4\)\.\(57\)
The latency scenarios share the dependence onpp\. Equally weighted least squares on all 20 top slopes gives this form of[equation8](https://arxiv.org/html/2610.42112#S5.E8):

β^const​\(p\)\\displaystyle\\widehat\{\\beta\}\_\{\\mathrm\{const\}\}\(p\)=0\.3736\+0\.5536​p,\\displaystyle=0\.3736\+0\.5536p,\(58\)β^lin​\(L,p\)\\displaystyle\\widehat\{\\beta\}\_\{\\mathrm\{lin\}\}\(L,p\)=0\.3429\+0\.4129​p\+\[0\.007697\+0\.03518​p\]​L\.\\displaystyle=0\.3429\+0\.4129p\+\[0\.007697\+0\.03518p\]L\.
Comparing a linear\-top model to a constant\-top, the linear\-top follows the observed trend more closely, reducing top\-exponent RMSE from0\.05740\.0574to0\.02260\.0226\. On heldout depths or thresholds beyond the fitted settings, linear\-top again improves on constant\-top: RMSE0\.03170\.0317versus0\.07060\.0706for held\-out depths, and0\.04420\.0442versus0\.05970\.0597for held\-out thresholds\.

Figure 23:Observed top\-level exponents versus total tree depth forC=4C=4,R=16R=16\. Dotted lines fit constant \(gray\) and linear \(blue\) depth dependence separately at eachpp, using only the five top\-level observations\. Bars show pointwise 95% whole\-problem bootstrap intervals\.
### E\-DRobustness of the Depth Trend

Though the trends identified appear robust, it is important to qualify that we fit in a low\-data regime\. The deepest trees and longest contexts have the least support:L=6L=6has only 4–15 complete top\-level candidates over 16K–64K, and fewer generations reach 64K\. Many positional bins therefore do not imply many independent examples\.

We quantify uncertainty with 2,000 whole\-problem bootstrap draws, keeping each sampled problem’s generations, bins, and levels together and recomputing the median slopes and model fits\. Model comparisons pair draws across depths and thresholds and repeat the holdouts\. We take the median which reduces sensitivity to extreme slopes, while still not removing differences in the length ranges covered by each problem\.

### E\-EFrom Selection Counts to Latency

Having fitted a model for how the sparsity exponents grow with depth, we can plug in the per\-level work that implies to the latency model\. We transfer theC=4C=4,R=16R=16exponent scenarios by depth and calibrate count amplitudes to nativeR=1R=1,C=8C=8,L=2L=2AIME measurements \(480 generations perpp\), using pooled counts and the same weighted log fit\. ExistingC=8C=8checks find that compression\-based transfer can improve the top exponent while worsening the full profile\. They also use different texts and, forR=1R=1, a different reference count, so we do introduce error by transferring between adjacent geometries\.

The plugged\-in selection rates model the per\-level expansion in the latency model\. That is, the model expands parent unions to obtain candidate counts, adds nodes without complete parents, and includes the first key and partial chunk in final token reads\.181818At the internal 128K reference, extrapolated from 16K–64K, added\-parent occupancy followsqparent=1−\(1−qchild\)8q\_\{\\mathrm\{parent\}\}=1\-\(1\-q\_\{\\mathrm\{child\}\}\)^\{8\}, whereqqis the selected fraction of nodes\.[Figure24](https://arxiv.org/html/2610.42112#A5.F24)extends the main\-textp=0\.75p=0\.75comparison top=0\.7,0\.8p=0\.7,0\.8; note that these are tightly clustered relative to the baselines, reflecting their different scaling exponent\.

![Refer to caption](https://arxiv.org/html/2610.42112v1/hatt-latency-aime-three-p.png)Figure 24:Modeled BF16 decode latency for dense attention, MSA, and ZHA \(R=1R=1,C=8C=8\) atp=0\.7,0\.75,0\.8p=0\.7,0\.75,0\.8, with the H100 assumptions and nondecreasing latency\-minimizing depth of[Figure8](https://arxiv.org/html/2610.42112#S5.F8)\. Latencies are modeled per\-token decode times for the attention subsystem of one layer at batch size 1\.

## Appendix FRetrieval Traversal Examples

Aggregate key shares show how much of the context ZHA attends to, but not where that budget goes within the tree\.[Figure25](https://arxiv.org/html/2610.42112#A6.F25)follows one retrieval head through each level of a three level router on two 64K RULER prompts\. The router’s search widens depending on the task difficulty: on the multi\-key prompt, whose distractor needles resemble the target, it expands more nodes at multiple regions of the context, while for the single\-needle prompt the router is more targeted\. These examples come from a single layer, head and decode step; they are meant to illustrate the mechanism, not to represent typical behaviour\. Nevertheless, these worked examples demonstrate the power of hierarchical routing in action, where the model is able to accurately find the needle while only activating a small fraction of the total underlying KV cache\. Our RULER scores in the main text imply that this retrieval is highly robust in practice, even to challenging multi\-key \+ distractor RULER subtasks\.

![Refer to caption](https://arxiv.org/html/2610.42112v1/traversal_strip_64k_combined.png)Figure 25:Retrieval traversals at 64K tokens for \(a\) single\-needle and \(b\) multi\-key 2 RULER tasks \(Qwen3\.6\-35B\-A3B, C16 L3,R=16R=16,p=0\.9p=0\.9\)\. Both show layer 35, query head 8, while emitting the first answer digit\. Dark blue blocks are scored and expanded nodes; light blue blocks scored but pruned nodes; grey blocks are nodes which were never scored by the router\. Dense\-attention rows show chunk\-level mass on a logarithmic scale\. The green marker indicates the ground\-truth needle location\. Boxed regions are enlarged below to reveal finer traversal details\. These examples illustrate task\-dependent routing with top\-pp: for single\-needle, the router explores few branches before reaching the answer\-containing chunk; for multi\-key 2, whose context contains many distractor needles, it explores substantially more upper\-level nodes\. Both traversals ultimately attend to a single answer\-containing chunk\.
## Appendix GReference\-Query Agreement Across Tree Levels

We extend the reference\-query usage analysis in[Figure9](https://arxiv.org/html/2610.42112#S6.F9)to all four geometries and compare aggregate level agreement with individual layer/KV\-group examples\. Levelℓ\\ellcontains nodes spanningCℓC^\{\\ell\}keys for the geometries shown\. Our analysis demonstrates that a similar pattern of rank\-correlation between reference queries across levels holds across a range of tree geometries\.

![Refer to caption](https://arxiv.org/html/2610.42112v1/qref_level_spearman_geoms.png)Figure 26:Agreement between tree levels in reference\-query selection for Qwen3\.6\-35B\-A3B ZHA \(R=16R=16\)\. Columns show C4 L6, C8 L4, C16 L3 and C64 L2; rows show RULER at 64K tokens and natural text at 32K tokens\. For each layer/KV group, Spearman correlation is computed between the 16 reference\-query selection shares at each pair of levels, then averaged over 20 groups \(10 layers×\\times2 KV groups\)\. Parenthesized axis labels give node spans in keys\. Adjacent levels are similar in their reference selection rate while further apart levels have a different selection priority\. This behavior holds across all geometries\.### G\-AIndividual layer and KV\-group examples

We show the actual reference selection behavior per level for a typical KV group in[Figure27](https://arxiv.org/html/2610.42112#A7.F27), and two outlier behaviors in[Figure28](https://arxiv.org/html/2610.42112#A7.F28)and[Figure29](https://arxiv.org/html/2610.42112#A7.F29)\. The reference queries are shared across levels\. In general, selection becomes increasingly concentrated at higher levels, while remaining more uniform at the lowest level\.

![Refer to caption](https://arxiv.org/html/2610.42112v1/qref_level_rates_L19kv0.png)Figure 27:Typical selection behavior across levels: layer 19, KV group 0 of the C4 L6 ZHA router on RULER at 64K tokens\. Adjacent levels have similar reference rankings, but agreement weakens over larger level separations; the chunk and top\-level rankings are nearly unrelated \(L1 versus L6:ρ=0\.12\\rho=0\.12\)\.![Refer to caption](https://arxiv.org/html/2610.42112v1/qref_level_rates_L39kv1.png)Figure 28:Outlier selection across levels: layer 39, KV group 1 of the C4 L6 ZHA router on RULER at 64K tokens\. Upper levels agree closely, while L1 reference rankings are negatively correlated with the top levels \(L1 versus L6:ρ=−0\.35\\rho=\-0\.35\)\.![Refer to caption](https://arxiv.org/html/2610.42112v1/qref_level_rates_L3kv1.png)Figure 29:Outlier selection across levels: Similar reference rankings across all levels: layer 3, KV group 1 of the C4 L6 ZHA router on RULER at 64K tokens\. Chunk and top\-level rankings remain strongly aligned \(ρ=0\.86\\rho=0\.86\)\.

## Appendix HReference\-Query and Reference\-Key Similarity

ZHA learnsRRreference queries per KV group without explicitly regularizing their structure\. Different geometries impose different chunk sizes, depths, and compression levels, raising the question of whether a common reference structure nonetheless emerges\.[Figures30](https://arxiv.org/html/2610.42112#A8.F30)to[32](https://arxiv.org/html/2610.42112#A8.F32)compare the 16 references across four geometries, examining both the query vectors and the reference keys they produce\. Across most layers, matched references are similar across geometries, while reference keys remain differentiated within each geometry\. This suggests that training finds a largely geometry\-independent, specialized basis\. All plots use the four independently trainedR=16R=16geometries C4 L6, C8 L4, C16 L3, and C64 L2; a top level node covers exactly 4096 keys in each geometry\.

![Refer to caption](https://arxiv.org/html/2610.42112v1/qref_cos_r16.png)Figure 30:Reference\-query cosine similarities across the four Qwen3\.6\-35B\-A3B ZHA geometries\. The C16 L3 references are ordered by average\-linkage clustering with cosine distance; references in each other geometry are Hungarian\-matched to them by cosine similarity\. Matched references therefore occupy corresponding positions\. Diagonal structure within the off\-diagonal blocks indicates similar reference directions across geometries\.![Refer to caption](https://arxiv.org/html/2610.42112v1/qref_cos_r16_byusage_ruler.png)Figure 31:Reference\-query cosine similarities with references independently ordered by top\-level selection frequency within each geometry and KV group, most to least frequent\. Frequencies were measured during routed prefill on 16 RULER prompts at 64K tokens andp=0\.9p=0\.9\. No cross\-geometry matching is applied; diagonal structure therefore indicates correspondence in both direction and selection rank\.![Refer to caption](https://arxiv.org/html/2610.42112v1/refkey_cos_r16_n4096_ruler.png)Figure 32:Mean cosine similarities of centered reference keys across geometries, evaluated on the same 60 top\-level nodes from four RULER prompts at 64K tokens\. Each node covers 4096 keys\. Within each node and KV group, the mean over the 16 reference keys is subtracted before computing cosine similarities, which are then averaged over nodes\. C16 L3 references are clustered using reference\-key cosine similarity; other geometries are Hungarian\-matched using the same measure\. The cross\-geometry correspondence extends to the pooled key representations\.