IANUS: Integrated Accelerator based on NPU-PIM Unified Memory System

LLM inference is two workloads wearing one coat — compute-hungry prefill and bandwidth-starved generation — so IANUS bolts a commercial NPU to commercial processing-in-memory chips and, crucially, makes the PIM the NPU's main memory, since 90% of the weights are shared; a scheduler referees the one catch, that in-memory compute and ordinary reads cannot happen at once.

Why this problem is worth solving

Transformer inference at a datacenter has two phases with opposite personalities. The summarization stage (often called prefill) processes all input tokens at once — big matrix-matrix multiplications, compute-bound, exactly what GPUs and

Two further characters complicate the play. Vector operations (layer normalization, residual additions) are trivial FLOP-wise — under 0.06% — yet consume 13.2% of latency without dedicated vector hardware. And non-computing operations (transposes, data rearrangement in attention) occupy 66.1% of self-attention latency on the GPU — moving data, not crunching it.

What people did before, and where it broke

Stage specialists. DFX (the prior state of the art here) is a multi-FPGA appliance tuned for the generation stage — effective there, but with too little compute for prefill. Attention accelerators (A³, ELSA-style) accelerate one operator, not the end-to-end path. NPUs alone (the paper’s NPU-MEM baseline: the same NPU with ordinary GDDR6) run prefill well and crawl through generation — ending up 24% slower than DFX overall.

PIM as an accelerator with partitioned memory. Prior xPU+PIM systems treat PIM as a device with its own memory, separate from the host’s. For LLMs this is uniquely wasteful: the fully-connected weights — about 90–91% of all parameters in GPT-2 — are needed by both sides (matrix-matrix in prefill on the NPU, matrix-vector in generation on the PIM). Partitioned memory means either duplicating nearly the whole model or shuttling it back and forth every stage transition. (NeuPIMs, covered elsewhere on this site, takes the partitioned path and buys concurrency with dual row buffers; IANUS bets the other way.)

The unified-memory catch. Making PIM chips serve as the NPU’s main memory eliminates duplication — but a PIM chip executing an in-memory computation cannot simultaneously answer ordinary load requests. Naive scheduling either serializes the NPU behind PIM work or corrupts the overlap that makes heterogeneity worthwhile. That conflict is the paper’s central engineering problem.

The idea

Match each operation to the engine its arithmetic intensity demands. Prefill’s matrix-matrix FCs → NPU matrix unit. Generation’s matrix-vector FCs and feed-forward layers → PIM, where per-bank units consume weights at internal bandwidth. Normalizations and residuals → NPU vector processors. Transposes and reshapes → DMA engines, overlapped with compute. The mapping is not always static: for FC layers, an adaptive algorithm picks PIM or matrix unit by input token count — matrix-vector work belongs in memory, but as tokens accumulate the operation becomes matrix-matrix and the crossover flips to the NPU.

One subtle anti-mapping: attention’s score and context multiplies stay on the NPU even in generation, though they are matrix-vector shaped. Two reasons the paper quantifies. First, geometry: a PIM DRAM row holds 1,024 BF16 elements, but an attention head’s dimension is 64 — computing QKᵀ in PIM would use 6.25% of each activated row, a bandwidth disaster. Second, parallelism: if PIM did everything, the NPU would idle; keeping self-attention on the NPU lets it run concurrently with PIM’s FC work. Heterogeneity pays only when both faces are busy.

Unify the memory, then schedule the conflict away. The unified system’s constraint — PIM compute and normal memory access cannot overlap on the same device — is handled by PIM Access Scheduling as a joint mapping-and-ordering problem. Weights for Q/K/V are partitioned head-wise across PIM chips, aligning each NPU core with its own chips so parallel reads don’t collide; other FC weights are split column-wise so the large weight matrices never move between cores. Scheduling then interleaves: while PIM chips compute one head’s FC, the NPU prefetches another head’s keys and values (ordinary reads to different banks/chips), runs softmax, or processes the previous head — inter-head pipelining that hides each engine’s work behind the other’s.

Prefetch what generation actually re-reads. In generation, offloading Q/K/V FCs to PIM frees the NPU’s DMA to prefetch the growing key/value history during softmax — turning what NPU-MEM serializes into overlap, and yielding a 4.3× self-attention speedup without offloading a single attention operation to PIM.

LayerNorm → vector unit
input: one token vector (batch 1)

Why it works

• Generation’s cost is weight traffic, not arithmetic → computing FCs at the banks converts internal bandwidth (4× external) into token latency → the memory-bound stage stops being bound by the bus.
• ~90% of parameters serve both stages → one copy in unified memory halves footprint and deletes stage-transition transfers → the heterogeneous pair costs no more memory than either engine alone.
• The unified conflict is schedulable because the workload is known → head-wise/column-wise weight placement aligns engines with disjoint devices, and PAS orders the rest → NPU and PIM overlap instead of alternating.
• Mapping follows geometry, not category: attention’s matvecs stay on the NPU (64-wide data vs 1,024-wide rows), FCs flip engines with token count → each operation lands where its efficiency is highest, not where its shape superficially suggests.
• Commercial parts on both sides plus an FPGA prototype with real AiM chips → the claims rest on buildable silicon, not idealized PIM.

What it costs

• Serialization risk is managed, not eliminated. PIM compute blocks normal reads per device; PAS hides this for transformer structure specifically — irregular workloads would re-expose it.
• PIM compute energy. In-DRAM operations cost ~3× a DRAM read in the model; the net energy win comes from traffic eliminated, not cheaper arithmetic.
• Batch-1 framing. Evaluation targets non-batched datacenter serving (following DFX); large-batch serving shifts generation toward matrix-matrix and would re-balance the mapping toward the NPU.
• Simulator-centered results. Performance and energy come from detailed simulation (with measured A100/DFX baselines); the FPGA prototype validates functionality and accuracy, not the headline speedups — and the NPU in the prototype is itself simulated.
• Model scale. GPT-2 class (up to 2.5B) on one device; larger models need multiple IANUS devices (shown to 30B with 8), and modern many-billion-parameter LLMs would lean on that scaling.

The evaluation

The simulator models the full system in cycle-level detail (including the FPGA-side PIM control unit, pipelined with PIM commands) with energy accounting. Baselines: an NVIDIA A100-SXM (PyTorch 2.0, HuggingFace and Megatron-LM codepaths), DFX with four FPGAs (the generation-stage state of the art), and NPU-MEM — the identical NPU with standard GDDR6 instead of PIM, isolating the memory system’s contribution. Workloads: BERT-Large and GPT-2 M/L/XL/2.5B in BF16, inputs of 128–512 tokens, outputs of 1–512, batch size 1.