Pimba: A Processing-in-Memory Acceleration for Post-Transformer Large Language Model Serving

---

Review Form

Reviewer: The Guardian (Adversarial Skeptic)

Summary

The paper proposes PIMBA, a Processing-in-Memory (PIM) accelerator designed to serve both transformer and "post-transformer" Large Language Models (LLMs), such as those based on State Space Models (SSMs). The authors' central thesis is that a common "state update" operation in post-transformer models is memory-bandwidth-bound, similar to the attention mechanism in transformers. PIMBA's architecture is based on two primary ideas: (1) a State-update Processing Unit (SPU) shared between two DRAM banks using an interleaved access pattern to improve utilization, and (2) the use of MX8 quantized arithmetic to achieve a better accuracy-area tradeoff compared to other low-precision formats. The authors claim significant throughput improvements over GPU and GPU+PIM baselines with minimal accuracy loss.

Strengths

1. The paper correctly identifies a relevant and timely research direction: accelerating the emerging class of post-transformer LLMs.
2. The workload characterization in Section 3, which identifies the state update operation as a potential bottleneck under batched inference, provides a reasonable starting point for the investigation.
3. The analysis considers the critical trade-off between quantization-induced accuracy degradation and hardware area overhead (Section 4.2), which is a necessary component of any proposal involving low-precision arithmetic.

Weaknesses

1. Unsupported Foundational Motivation: The paper's motivation hinges on the claimed superiority of post-transformer models. Figure 1(a) presents a comparison where Mamba-2 achieves 4.5% higher accuracy than a baseline transformer. This result is cited from an external source [15] without providing the necessary context to validate its fairness (e.g., equivalent training compute, data, and model tuning). Without this validation, the entire premise that the community should invest in specialized hardware for these models is built on a potentially specious claim.

2. Oversimplification of the "State Update" Primitive: The paper generalizes the core operations of diverse architectures (SSMs, linear attention, RNNs) into a single "state update" primitive, formalized in Equation 2 (Page 4). This abstraction is a significant simplification. It glosses over key differences, such as the scalar decay in Mamba-2 versus the vector-based gating in GLA. The paper provides no evidence or sensitivity analysis to demonstrate that this single, generalized hardware implementation can efficiently serve these varied computational patterns without significant performance penalties or architectural compromises for one model type versus another.

3. Insufficient Justification for the Choice of MX8: The argument for using MX8 over int8 rests on the claim that int8 incurs "substantial area overhead" due to the need for dequantization and requantization for addition operations (Section 4.2, Page 6). This claim is asserted but not rigorously substantiated. The paper fails to provide a quantitative, apples-to-apples area comparison between the components of their proposed MX Adder (Figure 9b), which includes shifters and comparison logic, and a standard int8 datapath for the same function. Without this direct comparison, the claim that MX8 is Pareto-optimal remains an unproven assertion.

4. Weak and Potentially Biased Baselines: The experimental comparison relies on baselines that appear to be deliberately weakened. The GPU+PIM baseline is described as a "time-multiplexed design" that explicitly lacks the "access interleaving technique of PIMBA" (Section 6.1, Page 10). This constructs a strawman; a fair comparison would be against a more aggressively pipelined PIM baseline, allowing the reader to properly assess the incremental benefit of PIMBA's specific design choices. By comparing against a seemingly self-crippled baseline, the reported speedups of up to 2.1x are likely inflated.

5. Questionable Architectural Novelty: The core architectural proposal of sharing a processing unit between two banks using "access interleaving" (Section 5.2, Page 7) is presented as a key innovation. However, both resource sharing to amortize area cost and interleaving to mitigate memory bank contention are foundational techniques in computer architecture and parallel systems. The paper fails to adequately position this contribution with respect to prior art, thereby overstating its novelty.

Questions to Address In Rebuttal

1. Regarding the core motivation in Figure 1: Can you provide concrete evidence that the 2.7B parameter transformer and Mamba-2 models are fairly compared in terms of training data, total training FLOPs, and hyperparameter tuning? If not, how does this uncertainty affect the premise of your work?

2. Regarding the generalized state update primitive (Equation 2): Please provide a quantitative analysis of the performance and area efficiency of your proposed SPU when executing the specific, non-generalized operations of RetNet, GLA, and HGRN2. How much efficiency is lost by forcing these distinct operations onto your unified hardware?

3. Regarding the choice of MX8 over int8: Please provide a detailed, post-synthesis area breakdown of the MX Adder and MX Multiplier versus an equivalent int8 datapath that includes the necessary dequantization/requantization logic. This data is required to substantiate the claim of Pareto optimality made in Figure 6.

4. Regarding the GPU+PIM baseline: Please justify the decision to use a non-interleaved, "time-multiplexed" design as your primary PIM baseline. Why was a more competitive, pipelined PIM design not used for comparison? How much of your reported speedup is attributable to simply having a better pipeline structure versus the baseline you designed?

5. Regarding the accuracy results in Table 2: Was any form of quantization-aware training or post-quantization fine-tuning used to achieve the reported near-zero accuracy degradation? If only post-training quantization was used, the results are unusually strong and require further explanation and validation.

---

Review Form

Reviewer: The Synthesizer (Contextual Analyst)

Summary

This paper presents PIMBA, a Processing-in-Memory (PIM) accelerator designed for serving both existing Transformer-based Large Language Models (LLMs) and the emerging class of "post-transformer" models (e.g., State Space Models like Mamba-2, linear attention).

The authors' core contribution is founded on a crucial unifying insight: despite their algorithmic differences, the performance of both model classes during batched inference is fundamentally bottlenecked by memory bandwidth. For Transformers, this is the well-known attention mechanism; for post-transformers, it is a newly identified bottleneck in the "state update" operation.

Based on this, the authors propose a unified PIM architecture that can efficiently execute both types of operations. The design is guided by two key principles derived from their analysis: (1) The varied primitives in state update operations make per-bank PIM logic area-inefficient, motivating a shared "State-update Processing Unit" (SPU) that interleaves access between two banks. (2) Post-transformer models are sensitive to quantization, and the authors identify Microsoft's MX format (specifically MX8) as a Pareto-optimal choice for balancing accuracy and hardware area in a PIM context. The resulting system, PIMBA, demonstrates significant throughput improvements over GPU and GPU+PIM baselines while maintaining model accuracy and adhering to practical area constraints.

Strengths

1. Timeliness and Important Vision: This work is exceptionally timely. As the research community actively seeks alternatives to the quadratically scaling Transformer, the question of how to build hardware that supports this transition is critical. The paper's vision of a unified serving system that bridges the gap between today's Transformers and tomorrow's post-transformers is both ambitious and highly valuable. It provides a practical roadmap for evolving our hardware infrastructure.

2. Excellent Unifying Abstraction: The paper's primary strength lies in its abstraction of the core performance problem. By identifying the "state update" as the conceptual parallel to "attention" and demonstrating through roofline analysis (Figure 1b, page 2) and latency breakdowns (Figure 3, page 4) that both are memory-bound, the authors distill a complex landscape into a single, addressable hardware challenge. This insight forms a powerful foundation for their entire work.

3. Principled, Data-Driven Hardware Design: The design of PIMBA is not arbitrary; it follows directly from well-articulated principles and thorough analysis.

   • The decision to share an SPU between two banks is a clever solution to the area-throughput tradeoff identified in their analysis of pipelined vs. time-multiplexed PIM designs (Section 4.1, page 6).
   • The selection of the MX8 format is convincingly justified through a detailed accuracy-area tradeoff analysis (Figure 6, page 6), which clearly shows its superiority over other formats like int8 (too much area) and low-precision floating point (poor accuracy) for this specific workload. This is an excellent piece of co-design.

4. Comprehensive Scope and Evaluation: The authors evaluate their proposal against a broad set of modern architectures—four distinct post-transformer models, a hybrid model (Zamba2), and a traditional Transformer (OPT). The evaluation across multiple scales (up to 70B parameters) and on a wide range of metrics (throughput, latency, energy, accuracy, and area) lends significant credibility to their claims. The minimal accuracy degradation shown in Table 2 (page 11) is particularly compelling evidence for the viability of their chosen quantization strategy.

Weaknesses

My concerns are less about flaws in the existing work and more about the boundaries of its claims and its integration into the broader systems landscape.

1. Robustness of the "State Update" Generalization: The authors unify several post-transformer operations into a single generalized state update form (Equation 2, page 4). This is elegant and effective for the models studied. However, the post-transformer field is nascent and evolving rapidly. It is conceivable that a future dominant architecture might introduce primitives (e.g., complex non-linearities, different data dependencies) that do not map cleanly to this structure. This could limit the future-proofing of the PIMBA design.

2. System-Level Integration Challenges: The paper effectively addresses the accelerator microarchitecture but is lighter on its integration into a full-fledged, dynamic serving system. Modern LLM serving schedulers (like those in vLLM or Orca) use sophisticated techniques like continuous batching and preemption to maximize GPU utilization. The paper acknowledges that PIMBA operates in a "blocked manner" with the GPU (Section 8, page 13), which inherently leads to utilization gaps. While the authors suggest leveraging techniques from NeuPIMs [28], a more detailed discussion of how PIMBA's deterministic, command-based execution model would coexist with the highly dynamic and asynchronous nature of a production-level scheduler would strengthen the paper's system-level contribution.

3. The Software Hurdle: As with all novel hardware, the software and programmability aspect is a major barrier to adoption. The paper mentions extending CUDA and defining custom DRAM commands (Section 5.1, page 7), which is a non-trivial engineering effort. A deeper contextualization of the required software changes would provide a more complete picture of the path to real-world deployment.

Questions to Address In Rebuttal

1. Regarding the generalized state update operation (Equation 2): Could the authors comment on the potential limitations of this formulation? Have they considered any emerging post-transformer algorithms that might challenge this abstraction, and how might PIMBA be adapted to accommodate them?

2. Could the authors elaborate on the system-level scheduling vision? How would a system using PIMBA-enabled memory efficiently manage the pipeline bubbles created during the hand-offs between GPU and PIM execution, especially in a dynamic, multi-user environment managed by a sophisticated scheduler?

3. The Pareto-optimality of the MX8 format was compellingly demonstrated for the Mamba-2 model (Figure 6). Was this detailed tradeoff analysis also performed for the other SU-LLMs (e.g., RetNet, GLA)? Confirming that MX8 remains the optimal choice across this diverse set of models would further strengthen this key design decision.

---

Review Form

Reviewer: The Innovator (Novelty Specialist)

Summary

This paper presents PIMBA, a Processing-in-Memory (PIM) accelerator designed to serve both transformer and, more critically, "post-transformer" Large Language Models (LLMs) like State Space Models (SSMs) and their variants. The authors first identify a common, memory-bound "state update" operation that unifies various post-transformer architectures. They then propose a novel PIM architecture to accelerate this operation alongside standard attention. The core architectural claims of novelty are (1) a State-update Processing Unit (SPU) that is shared between two memory banks, using an "access interleaving" technique to mitigate read/write hazards and maximize throughput within a constrained area budget, and (2) a State-update Processing Engine (SPE) within the SPU that uses custom microarchitecture for element-wise multiplication and addition on the MX low-precision data format, extending its use beyond its original dot-product design.

Strengths

The primary strength of this paper lies in its well-defined and well-defended novel contributions at multiple levels of abstraction.

1. Novel Workload Abstraction: The identification and generalization of the "state update" operation (Equation 2, Section 3.1, page 4) across a diverse set of emerging post-transformer models (Mamba-2, GLA, RetNet, HGRN2) is a significant conceptual contribution. While prior work has analyzed individual models, this paper is the first I have seen to propose a unified hardware target based on this common algorithmic pattern. This insight alone provides a strong foundation for the work.

2. Novel Architectural Technique: The "access interleaving" mechanism where a single SPU is shared between two banks (Section 5.2, page 7) is an elegant and novel solution to a practical PIM design problem. The authors correctly identify the area-throughput tradeoff between naive time-multiplexed and fully-pipelined per-bank designs (Section 4.1, page 6). Their proposed solution, which pipelines operations by alternating reads and writes between two banks into a single SPU, effectively achieves the throughput of a per-bank design with roughly half the area overhead. This is a clever application of pipeline hazard mitigation in a new domain.

3. Novel Microarchitectural Design: The paper's novelty extends to the microarchitecture of the SPE (Section 5.3, page 8). While the MX format itself is not new [16], its application has been predominantly for GEMM/dot-product operations. The authors’ contribution is the design of bespoke MX Multiplier and MX Adder units for element-wise operations, which are critical for the "state update" primitive. This extension of the MX format's utility is a non-trivial and novel piece of engineering that is directly justified by their compelling area-accuracy tradeoff analysis (Figure 6, page 6).

4. Novel Empirical Insights: The quantization analysis (Section 3.2, page 5) provides a genuinely new finding. The observation that post-transformer state updates are highly susceptible to the "swamping effect" with standard low-precision floating-point formats (e.g., e4m3, e5m2) is a crucial insight that distinguishes this workload from transformer KV cache quantization. The subsequent identification of MX8 with stochastic rounding as a Pareto-optimal solution in the context of PIM area constraints is a strong, data-backed novel claim.

Weaknesses

My criticisms are primarily focused on contextualizing the novelty and identifying elements that are more evolutionary than revolutionary.

1. Incremental System Integration: The overall system design, particularly the software stack and host interface, heavily leverages prior art. The authors explicitly state their system architecture is "similar to the existing PIM-based LLM serving systems [28, 54, 55, 67]" and that the software stack is "based on a prior work, HBM-PIM [40]" (Section 5.1, page 6). The proposed custom DRAM commands (ACT4, REG_WRITE, etc., in Section 5.5, page 8) are logical extensions of command-based PIM interfaces seen in prior works and do not represent a fundamental shift in the PIM execution model. The novelty is clearly in the PIM logic itself, not its system-level integration.

2. Application of a Known Principle: The concept of interleaving accesses between memory banks to improve functional unit utilization is a classic computer architecture principle. While its application to solve the state update read/write hazard in PIM is novel and well-executed, the underlying idea of hiding latency by finding parallelism between independent resources is not fundamentally new. The paper would be strengthened by acknowledging this principle and more clearly articulating how its constraints and implementation differ in the PIM context.

Questions to Address In Rebuttal

1. On the Novelty of Access Interleaving: The proposed SPU shared between two banks is a cornerstone of your architectural contribution. Can you clarify if any prior work in the broader near-data processing or PIM literature has employed a similar shared-unit, interleaved-bank access scheme to resolve read-after-write or write-after-read hazards, even if for a different application (e.g., database operations, graph processing)?

2. On the Generality of the "State Update" Primitive: Your design is predicated on the "state update" abstraction (Equation 2). This holds for the current crop of post-transformer models. However, this field is evolving rapidly. How robust is your architecture to potential future post-transformer models that may deviate from this specific pattern of decay ⨀ state + outer(k, u)? Is the SPE's functionality general enough, or is PIMBA's novelty tightly coupled to this specific formulation?

3. On the Microarchitectural Delta of the MX SPE: The paper proposes custom MX multipliers and adders (Figure 9, page 8). Beyond the logic for handling the shared exponents and microexponents, how much of the core datapath (i.e., the mantissa computation) differs from standard integer multiplication/addition? Please quantify the "delta" in hardware complexity between your proposed units and a hypothetical design that dequantizes MX to a shared internal format (e.g., FP16), performs standard FP operations, and then re-quantizes. This would help solidify the claimed area benefits.