NeuPIMs: NPU-PIM Heterogeneous Acceleration for Batched LLM Inferencing

Pairing a matrix-math accelerator with compute-in-memory only pays off if both run at the same time — NeuPIMs makes that possible with a second row buffer inside each memory bank and a schedule that splits every batch in two counter-phased halves.

Why this problem is worth solving

Serving a large language model means running its stack of decoder blocks over and over, once per generated token, for every request. Each block has three parts, and they do not resemble each other computationally.

QKV generation and the feed-forward networks multiply activations by large trained weight matrices. When requests are batched, the weights are shared, so these become GEMM operations — general matrix-matrix multiplications — with high arithmetic intensity: many operations per byte fetched. This is exactly what machine-learning accelerators are built for. The paper uses NPU (neural processing unit) as the umbrella term for such devices — TPUs and their kin, built around systolic arrays, grids of multipliers that stream matrices through.

Multi-head attention is different. It multiplies the current token’s vector against the request’s stored key and value matrices — activations against activations. Every request has its own matrices, so batching shares nothing, and each element fetched is used essentially once. These are GEMV operations — matrix-vector multiplications — and they are bandwidth-bound: the paper’s roofline analysis of GPT-3 places the whole generation phase deep in memory-bound territory while the summarization phase (reading the prompt) is compute-bound.

The cost of this mismatch is measured, not hypothetical. Running four production models (GPT-NeoX, LLaMA-2, OPT, MPT) on real GPUs — an RTX 3090 and an A100 with 1,555 gigabytes per second of memory bandwidth — memory capacity is nearly 100% used, but compute utilization never reaches 40%. Operators pay for arithmetic that spends most of its time waiting for bytes.

Processing-in-memory (PIM) attacks exactly this: put small compute units inside the memory banks, where data can be read at internal bandwidth far beyond what the external bus offers. PIM is excellent at GEMV and poor at GEMM. So the obvious system pairs an NPU (for GEMM) with PIM (for GEMV) — and the obvious system disappoints. Understanding why is the heart of this paper.

What people did before, and where it broke

The serving-software layer is not the problem, and NeuPIMs keeps it: Orca’s iteration-level scheduling (requests join and leave the batch at every token boundary) and selective batching (batch the batchable layers, serialize attention), plus vLLM’s paged memory management for the key-value cache. These techniques fill the batch; they do not change what hardware the math runs on.

On the hardware side, PIM-based GEMV accelerators exist — the paper builds on Newton, an industry design. Whole-transformer PIM was also tried: TransPIM runs everything in memory. But GEMM in memory wastes the NPU’s strength; TransPIM targets single-request, encoder-style models, and in the authors’ re-implementation NeuPIMs outruns it by 79× to 431× on batched decoder inference. A companion ASPLOS 2024 paper, AttAcc, accelerates the attention layer with PIM as well; NeuPIMs’ distinct contribution is making the NPU and the PIM run at the same time, plus the scheduling that feeds them both.

Why doesn’t the naïve NPU+PIM combination already win? Two independent walls:

The microarchitectural wall: blocked mode. A DRAM bank reads data through its row buffer — a small latch holding the currently open row. Existing PIM designs use that single row buffer for both jobs: ordinary reads and writes, and operand staging for in-memory computation. The two modes cannot coexist, so while PIM computes, the whole memory is deaf to the NPU — it cannot even fetch weights. The devices are physically forced to take turns.

The algorithmic wall: dependency. Inside a decoder block, the order is fixed: QKV generation (GEMM) produces what attention (GEMV) consumes, and attention produces what the projection and feed-forward layers (GEMM) consume. Even with perfect hardware, within one batch there is nothing concurrent to run — when the batch is in its GEMV phase, the NPU legitimately has no work.

Measured on the naïve integration: NPU utilization 28%, PIM utilization 17%. Two accelerators, both mostly idle, taking turns being the bottleneck.

The idea

A second row buffer in every bank

The blocked mode has a mundane cause — one latch doing two jobs — and NeuPIMs removes it with a mundane fix: dual row buffers. Each bank gets a MEM row buffer, used exclusively by ordinary reads and writes, and a PIM row buffer, used exclusively by in-memory GEMV, each with its own data path. A DRAM array can activate different rows independently, so while the PIM buffer holds a slice of someone’s key matrix mid-computation, the NPU can stream weights out through the MEM buffer of the very same bank.

The cost is deliberately small — the design principle was to push complexity out of the silicon and into the command protocol. Measured with CACTI at a 22-nanometer node, the doubled buffers add 3.11% area. The command protocol absorbs three extensions: PIM_HEADER announces an upcoming GEMV’s dimensions — attention’s size varies with sequence length, and the controller must know the latency in advance to schedule around DRAM refresh; PIM_GEMV bundles many dot-products and the result readback into one command, keeping the command bus from becoming the new bottleneck; and PIM_PRECHARGE closes the PIM buffer’s row without touching the MEM buffer. The memory controller interleaves the two command streams, giving PIM commands priority because they are few and long-running.

Two sub-batches, half a phase apart

Dual buffers make concurrency possible; they do not create concurrent work. Within one batch, the GEMM-GEMV-GEMM chain is a hard dependency. The schedule manufactures independence instead: split the batch into two sub-batches. Requests in different sub-batches share nothing, so sub-batch 1’s attention (GEMV, on PIM) can legally run while sub-batch 2’s QKV generation or feed-forward layers (GEMM, on NPU) run — and vice versa, alternating through all the decoder blocks like two pistons.

For the pipeline to stay full, the two halves must take similar time, which raises two balancing problems. Across memory channels: a channel’s attention latency is set by the longest sequences it holds, so an estimator (built from the PIM tile layout) predicts each request’s attention latency, and a greedy bin-packing pass assigns the longest-sequence requests first, always to the least-loaded channel. Across sub-batches: each channel’s requests are split evenly between the halves, so neither sub-batch drags.

One more overlap rides on the same hardware: attention’s softmax runs on the NPU’s vector units between the two PIM multiplications (scores, then weighting). With dual buffers, partial results flow out while the PIM is still computing other attention heads, so the vector units and the memory overlap at head granularity too.

NPU-S: [SB1: QKV gen, block 1]
PIM:   [idle]

Why it works

Because the model’s layers alternate between weight-shared matrix math and per-request vector math, each layer type has a device that suits it — but the chain of dependencies means one batch can only occupy one device at a time. → Because two halves of a batch share no data, splitting the batch manufactures exactly one unit of independent work — enough to keep two devices in counter-phase. → Because each bank has separate row buffers for memory traffic and for in-memory computation, the counter-phase is physically executable: the NPU streams weights from the same memory whose banks are computing attention. → Because channel loads and sub-batch sizes are balanced using a latency estimator, neither piston out-runs the other, and the overlap stays tight through all the decoder blocks. → Therefore both devices’ utilization rises together — the NPU from 28% to 65%, the PIM from 17% to 26%, memory bandwidth from 27% to 85% — and throughput follows: 1.6× over the naïve integration, 2.4× over an NPU alone, 3× over a GPU, at the same memory capacity.

What it costs

• More power, less energy. Running both devices concurrently and holding a second row buffer raises average memory power 1.8× (635 milliwatts versus 364). The 2.4× speedup more than compensates: energy per unit of work falls about 25%.
• Small batches lose. Splitting a small batch starves the systolic arrays: the ablation shows sub-batch interleaving only wins at batch sizes of 256 and up. Below that, the technique can cost more than it saves — this is a large-batch, datacenter design.
• Custom memory silicon. The 3.11% area overhead is modest, but it is an overhead on a modified HBM die plus a new command protocol — nothing you can deploy on stock hardware. All results are cycle-level simulation (ONNXim plus DRAMsim3; the simulator is open-source), with a real A100 as the GPU reference.
• Training gains nothing. Training uses fixed-length sequences and is GEMM end to end; PIM has no GEMV to accelerate.
• Parallelism choices are constrained. Scaling to many devices prefers tensor parallelism; pipeline parallelism shrinks the per-device batch, which the previous points make expensive.

The evaluation

Cycle-accurate simulation of GPT-3 variants at 7B, 13B, 30B and 175B parameters, driven by two real request-length distributions: ShareGPT (conversational; average 80 input, 296 output tokens) and Alpaca (short instructions; 12 and 56). Batch sizes 64 to 512. The prototype device: 8 systolic arrays of 128×128, 8 vector units, 32 HBM channels of 1 gigabyte, Newton-style PIM.

The ablation assigns credit cleanly. Dual row buffers alone contribute the most — 69.7% average throughput gain over the naïve baseline — because they unlock all concurrency. Channel load balancing always helps. Sub-batch interleaving is the conditional piece: a loss below batch 256, the top configuration at and above it. And against TransPIM — the run-everything-in- memory approach — NeuPIMs averages 228× higher throughput, a reminder that PIM earns its keep only on the operations that are actually bandwidth-bound.