Title: Hardware-aligned Hierarchical Sparse Attention for Efficient Long-term Memory Access
Venue: NeurIPS 2025
Accepted: yes
Keywords: long-context modeling

language modeling

attention mechanism

sparse attention

RNN
Abstract: A key advantage of Recurrent Neural Networks (RNNs) over Transformers is their linear computational and space complexity enables faster training and inference for long sequences. However, RNNs are fundamentally unable to randomly access historical context, and simply integrating attention mechanisms may undermine their efficiency advantages.
To overcome this limitation, we propose \textbf{H}ierarchical \textbf{S}parse \textbf{A}ttention (HSA), a novel attention mechanism that enhances RNNs with long-range random access flexibility while preserving their merits in efficiency and length generalization. HSA divides inputs into chunks, selecting the top-$k$ chunks and hierarchically aggregates information.
The core innovation lies in learning token-to-chunk relevance based on fine-grained token-level information inside each chunk. This approach enhances the precision of chunk selection across both in-domain and out-of-domain context lengths.
To make HSA efficient, we further introduce a hardware-aligned kernel design.
By combining HSA with Mamba, we introduce RAMba, which achieves perfect accuracy in passkey retrieval across 64 million contexts despite pre-training on only 4K-length contexts, and significant improvements on various downstream tasks, with nearly constant memory footprint. These results show RAMba's huge potential in long-context modeling.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper claims that current RNNs fail to randomly access historical context and questions whether the efficiency advantages still hold when attention mechanisms are integrated into RNNs. To this end, the authors propose Hierarchical Sparse Attention and combine it with Mamba2, leading to a hybrid RAMba architecture. HSA consists of two stages, including a bidirectional encoder within each chunk and a weighted summation for the chunks selected by top-k mechanisms. Moreover, they design a hardware-aligned Triton kernel to efficiently implement HSA. RAMba achieves remarkable performance in length generalization, performance, and efficiency, pushing ahead with the boundaries of long-context modeling.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\n1. A sparse attention design with a two-stage hierarchical structure. Chunk-level information is captured in the first stage and then fused in the second stage, which helps the length generalization. The working memory and long-term memory from a psychological perspective are interesting.\n\n2. A hardware-aligned algorithm to achieve efficient parallel computation.\n\n3. Extensive experiments on long-range language modeling, Passkey retrieval, LongBench, RULER, and fine-tunable datasets. Results of 2.7B models on LongBench and RULER validate the scalability. The result that the running time and memory consumption outperform NSA is encouraging, thanks to the shared chunk memory design. Additionally, the ablation studies regarding memory reset and sticking weights are thorough.\n\nweakness\n\n1. The motivation and the solution of the paper are not tightly aligned. My main concern is that the authors use \"Random Long-context Access for Mamba\" as the central motivation, while the main focus of the paper is actually on proposing the HSA method. In fact, the authors do not design any sparse mechanism specifically for the Mamba block itself, nor do they provide an in-depth discussion on the relationship between HSA and Mamba. Instead, HSA is simply applied within a hybrid architecture. As a result, this narrative may dilute the core contribution of the work. Moreover, I believe that using \"random access\" to describe chunk selection may be misleading, although the phrase contributes to the formation of the acronym \"RAMba\".\n\n2. Line 29-30: Is there any literature supporting the claim that augmenting RNNs with attention mechanisms leads to poor length extrapolation? The authors appear to assume that hybrid models extrapolate worse than pure RNNs, but this assumption is questionable. To the best of my knowledge, hybrid architectures do not necessarily hinder length extrapolation [1].\n\n3. Lines 139 and 157: Does the introduction of the transformer-based bi-directional encoder and the mean-pooled representation compromise the model's causality? MoBA [2] points out: \"The routing to the current block could also violate causality, since mean pooling across the entire block can inadvertently include information from future tokens.\" To address this, they ensure that each token is routed to its respective current block and apply a causal mask during the current block attention. I would be glad to see the authors' perspective on this issue.\n\n4. The results in Figure 3 are excellent. However, I have one question: What mechanism enables RAMba to achieve better extrapolation performance than both Mamba-2 and NSA? Even when the encoder is removed, the results remain strong. If the authors could provide a reasonable explanation for this significant performance gain, it might offer useful insights for future architecture design.\n\n5. Can HSA operate independently of the Mamba blocks? If a model were designed using only HSA, would it still be faster than NSA?\n\n[1] MiniMax-01: Scaling Foundation Models with Lightning Attention.\n[2] MoBA: Mixture of Block Attention\n\nQUESTIONS:\n1. Why does HSA adopt the stick-breaking attention mechanism for constructing chunk weights, rather than other forms of attention with positional encoding? I would like to see a theoretical justification or supporting ablation studies. (Table 1 presents RAMba w/o m.r. w/o s.b., which uses softmax without positional encoding, but it does not include ablation related to positional encoding)\n\n2. Line 174: Does the equation for the HSA layer in this line omit the Mamba layers? I ask because the formula shows H_t^(l+1)=H_t^l+O_t^l, but in the architecture, there are components other than the HSA layer between H_t^(l+1) and H_t^l, including Mamba layers.\n\n3. Line 273: Why does NSA's implementation currently not support generation?\n\nLIMITATIONS:\nI agree that experiments with RNNs besides Mamba can be explored further.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces RAMba, which combines Mamba (a Linear RNN) with a novel attention method, called Hierarchical Sparse Attention (HSA), whose goal is to optimize the tradeoffs between efficiency, length generalization, and random access in long-context modeling. \nThe key innovation is HSA, a two-stage hierarchical attention mechanism that learns chunk-aware token-to-chunk relevance. \nThe authors implement hardware-aligned kernels for efficiency and introduce a memory reset mechanism during training to improve length generalization.\nRamba achieves perfect accuracy on context passkey retrieval up to 64M context length despite pre-training on only 4K contexts. \nOn language modeling tasks, it performs slightly outperforms with Transformers and Mamba on IID data (4K length), and performs on par with Mamba on OOD samples (16K and 64K context length).\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n- I believe the hierarchical sparse attention mechanism represents a meaningful advance over existing chunk-based methods. The distinction between \"chunk-aware\" and \"chunk-unaware\" learning seems intuitive and well-motivated to me, with the two-stage attention design enabling end-to-end optimization of chunk selection.\n\n- The 64M context passkey retrieval with perfect accuracy is quite impressive. That's 16000x beyond pre-training length, which I think demonstrates exceptional length extrapolation. The consistent improvements across language modeling benchmarks within their evaluated scope also support the method's effectiveness.\n\n- I appreciate that the authors provided hardware-aligned kernel design with detailed algorithmic descriptions (Algorithms 1-3), which is very important for reproducibility and practical deployment. \n\n- The efficiency-generalization-access trilemma is well-articulated and represents a fundamental challenge in long-context modeling.\n\n**Weakness**\n\nWhile I largely enjoying reading this paper, I believe it still has room for improvement. Here are some of them:\n\n- The paper states that existing methods \"often suffer from inaccurate chunk selection\" (Line 37) without citing any supporting evidence. I think this unverified claim serves as the primary motivation for their approach, which I consider a significant methodological flaw.\n\n- I believe the multiple interacting components (Mamba + HSA + chunk encoder + memory reset) make it difficult to isolate individual contributions and understand failure modes. A more objective presentation and evaluation of HSA sounds more compelling.\n\n- The paper reports NSA results for language modeling (Table 1), passkey retrieval (Figure 3), and other tasks using Mamba with NSA, but later states \"NSA's implementation currently does not support generation.\" Since language modeling IS generation (next-token prediction), I think this reveals either implementation inconsistencies or conceptual confusion (possibly on my side). \n\n- I find it hard to accept the claim that Mamba \"excels\" (Line 292) at RULER tasks based on 20-40% accuracy rates. I'd say, instead, that all models perform poorly after a context length of 64K.\n\n- I believe that a critical issue lies in the efficiency comparison using \"Physics Processing Units (PPUs)\", which, to the best of my knowledge, lack tensor cores. In turn, this makes me think the comparisons with FlashAttention-2 (specifically optimized for tensor core utilization) are misleading, making the efficiency claims invalid or overstretched.\n\n- Finally, the paper lacks detailed computational complexity analysis. Based on my calculations (disregarding sorting), for the settings from section 4.1 (K=8 chunks, S=64 chunk size), HSA requires approximately $O(n \\times \\lceil n/S \\rceil + n \\times K \\times S^2)$ FLOPs per sequence. For a 16K sequence with $d=1024$, we get ~540M FLOPs per token, and an arithmetic intensity of ~33 FLOPs/byte. This indicates that HSA is likely compute-bound rather than memory-bound on modern hardware (e.g., GPUs), which could better elucidate how memory offloading can provide empirical benefits.\n\nQUESTIONS:\n- Can you provide concrete evidence (citations, experiments) supporting the claim that existing chunk-based methods \"suffer from inaccurate chunk selection\"? I really think this is crucial because without this evidence, the fundamental motivation for HSA remains somewhat obscure.\n\n- Do the PPUs used in efficiency evaluation support tensor cores? If not, I don't see how you can validly compare against FlashAttention-2, which is optimized specifically for tensor core utilization. \n\n- Are the perplexity results in Table 1 on validation or test sets? \n\n- How and where is the sorting operation applied? If it's inside the Triton kernel, how is it implemented?\n\n- Can you provide a detailed breakdown of the FLOPs required for HSA and its arithmetic intensity? Specifically, I'd like to understand the computational cost per token including chunk selection, token-level attention, and chunk-level fusion, and whether your method is memory-bound or compute-bound.\n\nLIMITATIONS:\nI think the authors provide a limitations section in the Appendix acknowledging the focus on Mamba and computational resource constraints. However, I believe some limitations could be addressed in order to improve the paper:\n\n- Discussion of how findings might change at larger model/context scales\n\n- I don't see any acknowledgment of the experimental inconsistencies or hardware evaluation issues that concern me (see weakness above)\n\n- I believe there's limited theoretical analysis of when/why the approach might fail.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a novel Hierarchical Sparse Attention (HSA) mechanism to give RNN-based models like Mamba random-access capability while maintaining linear complexity. By integrating HSA with a hardware-aligned kernel and memory offloading, the resulting RAMba model shows exceptional length generalization and efficiency. Notably, it achieves perfect accuracy on a 64M-token passkey retrieval task despite being trained on only a 4K context.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1.\tThe paper’s core contribution is a novel, well-motivated Hierarchical Sparse Attention (HSA) mechanism that directly addresses the critical trade-off between efficiency and random-access capability in long-context models.\n2.\tThe claims are backed by extensive and compelling experiments across a wide range of tasks. The standout performance on a 64M-token passkey retrieval task strongly validates the method’s length extrapolation capabilities.\n3.\tThe work extends beyond theory to a practical, system-level solution, including a hardware-aligned kernel design and an efficient inference strategy, which significantly enhances its real-world applicability.\nWeaknesses:\n1.\tWhile powerful, the RAMba architecture introduces significant complexity and multiple new hyperparameters (e.g., chunk size S, selected chunks K, HSA layer frequency G). The paper lacks ablation studies on these crucial design choices, making it difficult to assess their individual impact or the model’s sensitivity.\n2.\tThe choice of a randomized segment for the memory reset mechanism is not sufficiently justified. Although shown to be effective, the lack of comparison against simpler strategies (like resetting to zero) makes it unclear whether the randomization itself is a key ingredient or an unnecessary complexity.\n3.\tThe comparison against the NSA baseline is confounded. As shown in Fig. 3, RAMba has a significant advantage over RAMba-w/o-enc, which proves the chunk encoder is critical. Given this, the comparative experiments should have equipped the NSA baseline with the same component to fairly test the merits of the hierarchical attention design itself. Without this control, the source of the performance gain remains ambiguous.\n\nQUESTIONS:\n1.\tRegarding the model’s architectural complexity and its new hyperparameters (S, K, G): could the authors comment on their sensitivity? How were the reported values chosen, and what is the expected impact of varying them?\n2.\tRegarding the memory reset mechanism: could the authors provide an ablation study against simpler strategies (e.g., resetting to zero or to the last state) to better justify the choice of a randomized approach?\n3.\tTo validate the true advantage of the HSA hierarchical design, we suggest the authors provide results for an NSA baseline augmented with the same chunk encoder. This experiment is crucial for a fair comparison.\n\nLIMITATIONS:\nyes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces Hierarchical Sparse Attention (HSA), a novel attention mechanism designed to provide RNNs with efficient, long-range random access capabilities, addressing the 'trilemma' of balancing efficiency, random-access flexibility, and length generalization. \n\nHSA operates by dividing the input into chunks, selecting the most relevant chunks, and then applying a two-stage attention process: first within tokens of a chunk, and then across the selected chunks. The authors integrate HSA with a Mamba-2 backbone to create RAMba. To manage memory, RAMba employs techniques like hardware-aligned kernels and offloading the KV cache to CPU memory. \n\nThe paper's experiments demonstrate that RAMba achieves near state-of-the-art performance on long-context tasks, such as perfect accuracy on a 64-million-token passkey retrieval task , while maintaining high throughput and a near-constant memory footprint during inference.\n\nSTRENGTHS AND WEAKNESSES:\nRAMba is a novel architecture that integrates HSA with a Mamba backbone. The goal is to equip Mamba with efficient, long-range random access capabilities, a direction of significant interest for sequence modeling. While the approach is promising and the results are strong, I have a few concerns regarding the architectural design and the justification for its components.\n\n* On the Source of Performance Gains in HSA: A central claim of HSA is its \"chunk-aware\" learning mechanism, which is presented as its core innovation over prior work. However, the mechanics of how this is achieved beyond standard backpropagation through the hierarchical structure could be elaborated upon. This point is crucial because the ablation results in Table 1 for RAMba w/o s.b. show a significant performance degradation when stick-breaking attention is removed. This raises the question of whether the primary performance gains stem from the specific properties of the stick-breaking mechanism itself, rather than the proposed two-stage hierarchical attention. Further clarification or a more targeted ablation could help disentangle these effects and solidify the paper's core claim.\n* Justification for the Transformer-based Encoder: The choice of a Transformer-based bi-directional encoder for chunk memory is another key design decision. What was the specific motivation for this choice over other alternatives, such as using another Mamba layer? Given that RAMba is a Mamba-centric architecture, integrating a Transformer component introduces additional complexity. An ablation study comparing different encoder designs would significantly strengthen the justification for this specific architectural choice and its contribution to the model's impressive length generalization capabilities.\n\nGenerally, I think the the overall design of RAMba is quite sophisticated, with several novel components (HSA, chunk encoder, memory reset) interleaved with the Mamba backbone. \nWhile the results on very long contexts are impressive, a discussion on the model's practical efficiency (latency and throughput) on shorter sequences—which are common in many real-world scenarios—would be beneficial. It is unclear how the overhead from the added components impacts performance at smaller scales.\n\nTo make the contribution more convincing and the design principles clearer, the authors could add a more detailed architectural justification. For instance, explaining the rationale behind alternating one HSA layer with G Mamba layers in the upper decoder could help readers understand the intended synergy between the components. A more minimalist design, supported by careful ablations that justify each component's inclusion, would make the final architecture even more compelling.\n\nQUESTIONS:\n1. Can the authors also reports the results for pure NSA w/o SWA as used in MiniCPM4? I think it is an important baseline besides transformer.\n2. Could you elaborate more on BPTT? How did you process the last hidden states after param optimized, If I understand correctly, the saved states does not align with the ones produced by the updated model weights, will this hurt the performance?\n3. I did not figure out the difference between the alg in  Fig. 2 (b) and Algorithm 1 with NSA.\n4. How about the performance of MK-NIAH on RULER?\n5. I think some more results beyond RULER and long-bench could be present to verify the claims of random long-context access abilities. Wondering if it's possible to report the results of Phonebook lookup used in \"Repeat After Me\", which is training-free\n\nMiniCPM4: https://arxiv.org/abs/2506.07900\n\nRepeat After Me: https://arxiv.org/pdf/2402.01032\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-13"}