AttAcc! Unleashing the Power of PIM for Batched Transformer-based Generative Model Inference

Batching makes the weight-heavy layers of a large language model cheap but leaves attention untouched, so AttAcc moves the attention layer into the memory chips themselves and lets the GPU keep the part it is good at.

Why this problem is worth solving

A Transformer-based generative model — the paper’s term for systems like GPT — answers a request in two phases. The summarization stage reads the whole input prompt at once and builds an understanding of it. Then the generation stages produce the answer one token (word piece) at a time, each new token feeding back as input for the next step.

During summarization, the model computes, for every input token, a key vector and a value vector — a compressed record of what that token contributes to the context. These are stored as the KV matrices. Every later generation step reads the entire KV matrices to decide what to attend to, and appends one new row for the token it just produced. Write once, read once per generated token, grow forever.

The authors’ earlier letter quantifies what this costs for GPT-3 with 175 billion parameters. One inference over a 2,048-token input and a 2,048-token output takes 1,425 tera-floating-point operations — about 365,000 times the work of classifying one image with ResNet-50. The generation stages dominate: even with a long input of 2,048 tokens and a short output of 128, they take over 88% of the time; for more balanced shapes, over 96%. The reason is their sequential nature — every token requires re-reading the model’s weights, and weights for this model fill 326 gigabytes in 16-bit precision.

And the KV matrices are not small either: at 2,048 input plus 2,048 output tokens, they occupy 18 gigabytes per request — so a batch of 64 requests needs about 1,152 gigabytes for KV matrices alone, roughly 3.5 times the size of the entire model.

What people did before, and where it broke

The standard remedy is batching: group many requests and run them through the model together, with continuous scheduling that slips a new request into the batch the moment another finishes.

Batching works because of a quantity worth defining carefully. Operations per byte measures how many arithmetic operations a computation performs for every byte it fetches from memory. Hardware has a matching ratio: an NVIDIA DGX A100 server delivers 2,500 tera-operations per second against 16 terabytes per second of memory bandwidth — about 156 operations per byte. A computation below that ratio leaves compute idle; the memory bus is the bottleneck.

The fully-connected layers — the weight-heavy matrix multiplications that make up most of the model — batch perfectly. The weights are shared by all requests, so one fetch feeds the whole batch: multiply the batch size, and operations per byte rise in step. The letter’s measurements show throughput scaling almost linearly while the energy spent on memory traffic falls.

The attention layer does not batch, for a structural reason: its “weights” are the KV matrices, and every request has its own. There is nothing to share. Each generation step multiplies one small query vector against the request’s large KV matrices — about one operation per byte, two orders of magnitude below what the hardware wants, and batching does not move it. The measured consequences on the DGX, from the letter:

• At batch size 64, attention consumes more than 60% of generation time.
• Even at batch size 256, GPU compute utilization stays below 20%.
• With a 50-millisecond-per-token latency target, batch 64 misses the deadline (108 milliseconds), forcing the batch down to 16 or fewer — surrendering the very throughput batching was meant to buy.
• And the KV capacity wall stands regardless: 1.5 terabytes wanted, 640 gigabytes present.

Buying more GPUs scales the machine out, but every added GPU brings expensive idle compute along with its memory — paying for arithmetic to obtain bandwidth. The mismatch is architectural, and the fix has to be too.

The idea

Attention is exactly the computation processing-in-memory wants

Processing-in-memory (PIM) places small processing elements inside a memory device, next to the DRAM banks — the independent sub-arrays a memory chip is divided into. Reading within a bank is cheap and parallel; it is the journey across the external bus to the processor that is slow and energy-hungry. If many banks compute simultaneously on their own data, the internal bandwidth exposed to computation is many times the external bandwidth the chip could ever ship out — in the letter’s configuration, 64 terabytes per second internal against a 4× smaller external figure for the whole DGX.

PIM is no free lunch: data must be laid out in the right banks in advance, and results must be small enough to ship out cheaply. Two properties make the attention layer an unusually good fit.

Write once, read many. A request’s KV matrices are written during summarization, then read again at every generation step — once per output token. Placing them in PIM pays the write once and harvests internal-bandwidth reads hundreds or thousands of times.

Tiny inputs, tiny outputs, huge middles. Each attention step brings in one query vector and returns one context vector; in between, it reads the entire KV matrices. For GPT-3’s dimensions (embedding width 12,288, 96 attention heads), the traffic across the external bus is about 1/128th of the data read internally. The expensive reads never need to leave the chip at all.

The AttAcc architecture

From the letter’s baseline design: AttAcc is built from high-bandwidth memory (HBM) stacks with processing elements coupled to the DRAM banks. The multiplication work — the score step (query against all keys) and the context step (attention weights against all values) — runs at the banks. The softmax between them needs exponentiation, which is too costly to replicate per bank, so it runs on the HBM stack’s buffer die, the logic layer at the base of the stack. Each attention head maps to one HBM stack; each bank holds a slice of the head’s key matrix (split by columns) and value matrix (split by rows). Because requests, heads and tokens all offer independent work, every bank stays busy.

The ASPLOS paper generalizes this into a design space: its simulator models processing units placed per bank, per bank group, or per pseudo-channel on the buffer die, over HBM3 memory — three rungs of a ladder trading parallelism against area and integration cost. Per its abstract, the paper also adds system-level optimizations for using the xPU and PIM together; their details are in the full text, which this page could not access.

The division of labor

The end system uses each side for what it is structurally good at. The xPU keeps summarization and the fully-connected layers, where batching drives operations per byte up into compute-bound territory. AttAcc takes the attention layer, whose operations per byte no batch size can raise, and supplies bandwidth instead. The interconnect between them carries only query and output vectors. And the capacity problem dissolves as a side effect: the KV matrices now live in AttAcc’s memory, so growing the KV budget means adding memory devices that are the accelerator, not idle GPU compute.

xPU: batched GEMM (weights shared by batch)
per request: Q, K_new, V_new — three vectors of width 12,288

Why it works

Because fully-connected layers share weights across a batch while attention shares nothing, batching splits the model into a part that scales and a part that cannot — so the two parts deserve different hardware. → Because a request’s KV matrices are written once but read at every generated token, placing them inside memory pays a one-time cost for a per-token saving. → Because attention’s inputs and outputs are vectors while its reads are matrices, only about 1/128th of the traffic must ever cross a chip boundary — the expensive external bus is almost removed from the loop. → Because banks compute in parallel on local data, attention runs at internal bandwidth several times what any external bus provides, and each bit read costs about 3.5× less energy than a bus crossing. → Therefore attention stops throttling the batch: latency per token falls, bigger batches fit under the latency target, the KV capacity wall moves from the GPU’s memory to the accelerator’s own, and both throughput and energy improve at once — up to 2.81× and 2.67× in the ASPLOS paper’s headline, against a conventional system with the same memory capacity.

What it costs

• It is new silicon in the memory itself. PIM-enabled HBM stacks are not commodity parts; the results come from simulation (the released simulator builds on Ramulator 2). The letter’s authors point to a commercial HBM-PIM precedent, but deployment at serving scale is an assumption, not a fact.
• Only attention accelerates. The fully-connected layers still run on the xPU and now bound the end-to-end gain — the familiar Amdahl ceiling. The headline 2.81× is far below the attention layer’s isolated potential.
• Capacity is committed, not pooled. Memory inside AttAcc holds KV matrices; it is no longer general-purpose GPU memory. The split between weight memory and KV memory must be provisioned for the workload mix.
• The two-stage structure must hold. The design leans on generation dominating and on KV reuse. Workloads dominated by the summarization stage — including ordinary training — get little; the letter says this directly.
• A dependency on interconnect behavior. Every generated token makes a round trip between xPU and AttAcc. The abstract’s “strategic” co-use and optimizations exist precisely because that boundary is not free; their mechanics are in the full paper.

The evaluation, as far as the sources report it

From the ASPLOS abstract: the GPU-PIM system outperforms a conventional system with the same memory capacity, improving performance and energy efficiency of serving a 175-billion-parameter model by up to 2.81× and 2.67× respectively. Same-capacity framing matters — it removes “you just added more memory” as an explanation.

From the authors’ earlier letter (in-house simulator, GPT-3 175B, 16-bit precision, DGX A100 baseline, simulator validated against a real A100 on a 6-billion-parameter model):