Accelerating Retrieval Augmented Language Model via PIM and PNM Integration

---

Review Form

Reviewer: The Guardian (Adversarial Skeptic)

Summary

The paper proposes MNM, a heterogeneous computing architecture integrating Processing-In-Memory (PIM) on the HBM core die and Processing-Near-Memory (PNM) on the logic die to accelerate Retrieval-Augmented Language Models (RALMs). The authors claim that by offloading memory-bound GEMV operations of the language model to PIM and retrieval-specific tasks (vector search, sorting) to PNM, their architecture can overcome memory bandwidth bottlenecks inherent to conventional GPU-based systems. They supplement this hardware with a novel scheduling strategy, "selective batching and early generation," which purportedly overlaps retrieval and generation to reduce idle cycles. The authors report significant performance speedups (up to 29.2x) and energy savings (up to 71.5%) over an NVIDIA H100 GPU baseline.

Strengths

1. Problem Characterization: The initial workload analysis in Section 3 (Figure 4, page 5) is thorough and correctly identifies the memory-bound nature of key RALM components. The roofline model analysis effectively motivates the need for a non-conventional architectural solution by demonstrating that both the language model's MHA kernels and the retriever's PQ code scan/top-k selection are far from the compute-bound regime on a state-of-the-art GPU.

2. Task Partitioning Rationale: The architectural decision to partition tasks—assigning simple, regular GEMV operations to PIM units and more complex, specialized retrieval logic to PNM—is well-reasoned (Section 4.1, page 6). This approach correctly acknowledges the fabrication and area constraints of logic on a DRAM die versus the flexibility of a logic die.

Weaknesses

My primary concerns with this work lie in the realism of the evaluation methodology, the downplaying of critical trade-offs, and the lack of rigor in overhead analysis.

1. Unjustified Quality-Performance Trade-off: The proposed "Early Generation" scheduling scheme is presented as a key innovation, yet it comes at the cost of model accuracy. Figure 8 (page 9) explicitly shows an increase in perplexity as batch size and nprobe scale. While the authors frame this as a smaller degradation compared to a GPU-based equivalent, it is a degradation nonetheless. The paper fails to adequately justify this trade-off. An architecture that achieves speedup by compromising the correctness of the model's output is fundamentally flawed. The core premise of RALM is to improve generation quality with retrieved data; this scheduler actively works against that goal by using stale data.

2. Questionable Area and Power Overheads: The overhead analysis in Section 6.6 (page 13) relies on highly suspect assumptions. The area of PIM logic is synthesized for a 14nm process and then scaled by a factor of "10x" to approximate a DRAM process node. This scaling factor is a crude approximation based on a single reference and lacks the necessary justification for a rigorous hardware paper. Consequently, the claim of a "15.0% overhead" on the core die is built on a weak foundation. This figure is not minor; a 15% reduction in memory cell area is a commercially prohibitive cost that the authors dismiss too readily. The power figures in Table 3 (page 10) also appear optimistic, and it is unclear if they account for all sources of leakage and activity under realistic concurrent operation.

3. Oversimplified Thermal Modeling: The thermal analysis presented in Section 4.1 (page 6) to justify PNM logic placement is superficial. The use of a simple 1D compact layered-conduction model is inadequate for a complex 3D-stacked device like HBM. This model ignores lateral heat spreading, the cumulative effect of hotspots from underlying PIM logic, and the thermal impact of the high-density TSV arrays. Presenting a temperature rise of less than 0.5°C based on such a model is unconvincing and potentially misleading. Real-world thermal throttling could easily negate the claimed performance benefits.

4. Unquantified System-Level Complexity: The proposed architecture introduces significant complexity that is not accounted for in the evaluation.

   • The MNM Controller (Figure 6, page 7), which translates GPU instructions into MNM commands, is a critical component whose design, latency, and area/power overheads are completely ignored.
   • The scheme to reorder PQ codeword IDs to solve memory alignment issues requires a host-side mapping table. The memory footprint and lookup latency of this table are never quantified. For a large database, this could be a non-trivial overhead.

5. Speculative Scalability Claims: The model and database scaling analysis in Section 6.5 (page 13) is based entirely on projection, not simulation or measurement. While such analyses can be illustrative, the claims made here are presented with undue confidence. The argument that MNM suffers from lower communication overhead than a baseline multi-GPU system is unsubstantiated, as the authors fail to quantify the command and data traffic their own scaled-up system would require between the host, GPUs, and multiple MNM-enabled HBM stacks.

Questions to Address In Rebuttal

1. Provide a rigorous justification for the 10x area scaling factor used for PIM logic. How does this factor account for differences in cell libraries, routing density, and process design rules between a mature logic process and a modern DRAM process?

2. The "Early Generation" scheduler trades model accuracy (perplexity) for latency. At what point does this degradation become unacceptable for a production-level RALM? Please provide a sensitivity analysis showing how perplexity scales with more extreme retrieval latencies and characterize the point at which the model's output quality is critically compromised.

3. Elaborate on the cost of the 15.0% PIM area overhead on the HBM core die. Quantify this in terms of lost memory capacity per die and per stack. How does this impact the commercial viability of such a memory product compared to standard HBM?

4. Address the limitations of the 1D thermal model. Acknowledge the potential impact of 3D thermal effects and explain why the risk of thermal throttling from the proposed PNM logic is negligible, especially when placed adjacent to high-activity PHY and TSV regions.

5. Provide a detailed architectural design for the "MNM Controller" on the GPU side, including its projected area, power, and the latency it adds to the command path. Furthermore, quantify the memory footprint and access overhead of the host-side mapping table required for PQ codeword ID reordering for the largest dataset used.

Review Form

Reviewer: The Synthesizer (Contextual Analyst)

---

Summary

This paper presents MNM, a heterogeneous computing architecture designed to accelerate Retrieval-Augmented Language Models (RALMs). The core contribution is the insightful co-design of a hardware system and a scheduling policy that addresses the two distinct, memory-bound bottlenecks inherent in RALMs. The authors propose integrating both Processing-In-Memory (PIM) on the HBM core die and Processing-Near-Memory (PNM) on the HBM logic die. The PIM units are leveraged to accelerate the GEMV-heavy attention operations in the language model, while the more flexible PNM logic is dedicated to the lookup- and sort-intensive tasks of the vector search-based retriever. This hardware architecture is complemented by a novel "selective batching and early generation" scheduling strategy that exploits the hardware's capabilities to maximize the overlap between token generation and retrieval, mitigating the idle cycles common in batched RALM inference.

---

Strengths

1. Excellent Problem Characterization and Motivation: The authors perform a thorough workload analysis (Section 3, page 5) that correctly identifies the fundamental challenge in RALM acceleration: it is not a monolithic problem. They astutely recognize that both the language model and the retriever are memory-bound, but for different reasons. The language model's attention is limited by bandwidth for GEMV operations, while the IVF-PQ retriever is constrained by irregular memory access patterns (LUTs) and sorting. This clear diagnosis provides a strong foundation for their proposed solution.

2. Holistic, System-Level Approach: This work is a superb example of system-level co-design. Rather than focusing on a narrow hardware optimization, the authors consider the entire RALM pipeline. They propose a hardware solution (MNM) and a software scheduling policy that are synergistic. The hardware enables more efficient, concurrent operations, and the scheduling policy is explicitly designed to exploit this new capability. This holistic view is a significant strength and reflects a mature understanding of the problem domain.

3. Elegant and Well-Justified Architectural Mapping: The core architectural idea—mapping the two distinct RALM workloads to different parts of an HBM stack—is elegant and compelling. Placing simple, massively parallel MAC units (PIM) in the core die for the regular GEMV operations is a natural fit. Simultaneously, using the logic die for a more specialized PNM accelerator for retrieval tasks is equally wise, as it allows for more complex logic (e.g., sorters, arbiters) without the process technology constraints of the DRAM die. This "right tool for the right job" approach is the paper's central technical insight.

4. Strong Contextual Fit within Current Research Trends: This work is exceptionally well-positioned within the broader landscape of computer architecture and AI systems. It directly engages with several critical, contemporary research threads:

   • The Memory Wall: It is a direct attack on the memory wall for a critical emerging workload.
   • Heterogeneous Computing: It moves beyond the traditional CPU-GPU dichotomy and embraces a more specialized, heterogeneous model integrating PIM and PNM.
   • System-Level AI Acceleration: It recognizes that accelerating AI is about more than just speeding up matrix multiplies; it involves optimizing the entire data-to-answer pipeline, including data retrieval. This paper serves as a valuable case study for the future of AI system design.

---

Weaknesses

While the core idea is strong, the paper could be improved by broadening its contextual discussion to better situate its specific design choices.

1. Limited Discussion of Alternative Heterogeneous Designs: The proposed PIM/PNM integration within a single HBM stack is a compelling design point. However, the academic landscape includes other heterogeneous proposals. For instance, the recent work on HeterRAG [60] proposes using HBM-PIM for generation (low latency) and DIMM-based PIM for retrieval (high capacity). A discussion contrasting the MNM approach (all-in-HBM, maximizing bandwidth) with an approach like HeterRAG's would strengthen the paper by exploring the fundamental trade-off between retrieval database capacity and access latency. This would help readers understand the specific part of the design space that MNM is optimizing for.

2. Clarity on the Interdependence of Hardware and Software Contributions: The "Early Generation" scheduling scheme appears highly effective in conjunction with the MNM hardware. However, its performance on conventional systems is less clear. The perplexity analysis in Figure 8 (page 9) is a good start, but the paper could more explicitly articulate whether the scheduling strategy is a general contribution applicable to any RALM system, or if its benefits are primarily unlocked by the dramatically reduced retrieval latency provided by the MNM architecture. Clarifying this relationship would help delineate the impact of the hardware and software components of the work.

---

Questions to Address In Rebuttal

1. The choice to integrate both PIM and PNM accelerators within the HBM stack prioritizes bandwidth and low-latency communication between the two components. Could the authors elaborate on the trade-offs of this approach compared to a disaggregated design, such as the one proposed in HeterRAG [60], which might offer much larger retriever database capacity by using DIMM-based memory for retrieval? In what scenarios or scales of RALM deployment would the MNM design be most advantageous?

2. The proposed "Early Generation" scheduling is shown to be highly effective with MNM. Could you further elaborate on its applicability to conventional GPU-based systems? Given that GPUs would have a much longer retrieval latency, would the window for "early generation" shrink to the point of being ineffective, or does the selective batching component still provide significant benefits on its own?

3. The current PNM design is tailored for IVF-PQ retrieval. Looking forward, retrieval methods are evolving to include more complex logic, such as multi-hop reasoning or graph traversal. How extensible is the proposed PNM architecture to support these more compute-intensive retrieval tasks? Would it require a more general-purpose programmable core on the logic die, and what would be the implications for area and power?

---

Review Form

Reviewer: The Innovator (Novelty Specialist)

Summary

This paper proposes MNM, a heterogeneous PIM/PNM architecture integrated within a High Bandwidth Memory (HBM) stack to accelerate Retrieval-Augmented Language Models (RALMs). The core idea is to partition the RALM workload according to its distinct computational characteristics and map each part to a specialized memory-centric compute unit. Specifically, the authors propose using Processing-In-Memory (PIM) on the HBM core die to accelerate the GEMV-heavy attention operations of the language model, while using Processing-Near-Memory (PNM) on the HBM logic die to accelerate the lookup- and sort-intensive tasks of the vector search retriever. This architecture is coupled with a novel scheduling policy, "Early Generation," which builds upon selective batching to overlap retrieval and generation phases, aiming to maximize throughput.

Strengths

The primary strength of this work lies in its elegant synthesis and co-design of existing concepts into a cohesive and well-motivated architecture.

1. Novelty in Integration: The central novel claim is the tight, heterogeneous integration of both PIM and PNM accelerators within a single HBM stack, with each component tailored to a different part of the RALM workload. While PIM and PNM have been proposed for these tasks separately, their co-location and co-design within one memory device to serve a single application is a distinct architectural proposition.
2. Workload-Specific Partitioning: The motivation for the heterogeneous design is well-argued. The paper correctly identifies the distinct bottlenecks of RALMs—GEMV-bound attention in the language model and irregular memory access patterns in the retriever (Section 3, page 5). The proposed mapping of these workloads to PIM (for its high internal bandwidth) and PNM (for its more complex logic capability) is logical and architecturally sound.
3. Hardware-Software Co-Design: The proposed "Early Generation" scheduling scheme appears to be a direct consequence of the underlying hardware's capabilities. The ability to perform concurrent PIM and PNM operations, enabled by features like the dual row buffer (Section 4.1, page 6), allows for a more aggressive overlapping than would be possible in systems with physically separate accelerators. This synergy between the proposed hardware and software is a notable strength.

Weaknesses

My primary concern is that while the integration is novel, the fundamental building blocks and conceptual approaches are largely derivative of prior art. The paper could do a better job of positioning its contribution against these closely related works.

1. Constituent Ideas are Not New: The core ideas, taken in isolation, are well-established.

   • PIM for Attention: Using HBM-PIM to accelerate GEMV operations in transformer attention is not a new idea. Prior works like AttAcc [80] and NeuPIM [24] have already established this approach. The PIM component of MNM appears functionally similar to these proposals.
   • PNM for Vector Search: Accelerating vector search using near-memory processing is also a known technique. Works like Chameleon [34] and FAANS [32] have proposed dedicated accelerators (on FPGAs or DIMMs) for IVF-PQ and other approximate nearest neighbor search algorithms. The PNM component of MNM addresses the same problem with a similar approach.

2. Conceptual Overlap with Prior Heterogeneous Systems: The concept of a heterogeneous PIM architecture for RALMs has been recently proposed. HeterRAG [60] specifically proposes combining HBM-based PIM for generation with DIMM-based PIM for retrieval. MNM's core conceptual contribution—partitioning the RALM workload across different PIM/PNM technologies—is therefore not entirely novel. The delta between MNM and HeterRAG is primarily in the implementation: MNM integrates both units within a single HBM stack, whereas HeterRAG uses separate memory systems. The paper fails to discuss this crucial piece of prior art and articulate the specific benefits of its tightly-coupled approach over HeterRAG's disaggregated one.

3. Incremental Novelty in Scheduling: The proposed scheduling scheme is an intelligent adaptation of existing techniques. The idea of overlapping retrieval and generation was explored in PipeRAG [35]. The concept of dynamically managing requests to improve utilization is the basis of continuous/selective batching, as seen in systems like Orca [99]. The authors' "Early Generation" scheme is a clever combination of these ideas, but its fundamental novelty is tied exclusively to its co-design with the MNM hardware rather than being a new scheduling paradigm in itself.

Questions to Address In Rebuttal

The authors should address the following points to clarify the novelty and significance of their contribution.

1. Comparison with HeterRAG [60]: Please provide a detailed, quantitative comparison against HeterRAG. What are the specific architectural and performance advantages of integrating both PIM and PNM within a single HBM stack versus HeterRAG's approach of using separate HBM-PIM and DIMM-PIM systems? Does the tight coupling provide benefits beyond reduced physical distance, such as a more unified programming model or lower-latency coordination?

2. Clarifying Scheduling Novelty: The "Early Generation" scheduling is enabled by the MNM architecture. Can the authors pinpoint the exact architectural feature(s) that make this scheduling scheme uniquely effective or even possible? For instance, is it the dual row buffer that allows for contention-free concurrent access, which would not be possible with other PIM designs? How would a PipeRAG-style scheduling perform on the MNM architecture, and vice-versa?

3. Generality of the PNM Accelerator: The PNM design is highly optimized for IVF-PQ-based retrieval. How adaptable is this design to other, more modern vector search algorithms like HNSW? A truly novel contribution would ideally show a path toward broader applicability. Is the PNM component a fixed-function unit for IVF-PQ, or is it a more general-purpose programmable engine for near-memory retrieval acceleration?