Ouroboros: Wafer-Scale SRAM CIM with Token-Grained Pipelining for Large Language Model Inference

Yiqi Liu, Yudong Pan, Mengdi Wang, Shixin Zhao, Haonan Zhu, Yinhe Han and 2 others

Running a large language model spends most of its energy not on math but on shuttling weights, activations, and the growing key-value cache up and down a deep memory hierarchy — this paper (Chinese Academy of Sciences — Yiqi Liu, Ying Wang, and colleagues) proposes Ouroboros, a wafer-scale chip that does every LLM operation right inside its first-level SRAM (computing in memory, so data never leaves), and makes the scarce on-chip capacity pay off with token-by-token pipelining, KV caching tucked into leftover memory, and a communication- and defect-aware mapping — reaching up to 9.1 times the throughput and 17 times the energy efficiency of prior systems.

The argument in one glance

Problem

LLM inference burns most of its energy moving data, not computing

A large language model runs by pushing enormous amounts of data through hardware: the model's weights, the activations flowing between layers, and the Key-Value (KV) cache that stores information about every token seen so far. Conventional accelerators keep these in a deep memory hierarchy — a small, fast first-level SRAM on the chip, backed by large, slower off-chip DRAM or High Bandwidth Memory (HBM). The trouble is that, as models scale into the tens and hundreds of billions of parameters, inference stops being limited by arithmetic and becomes limited by moving data across those memory tiers. The paper calls this the "hardware scaling tax": the energy spent shuttling weights, activations, and KV data between memory levels, and across the interconnects between them, comes to dominate total energy consumption, far outweighing the cost of the actual computation. Because LLM inference is so memory-intensive, this constant data movement — not the multiply-accumulates — is what fundamentally caps energy efficiency and throughput. So the real problem is not building faster compute; it is stopping the data from having to travel at all.
What people did

Add more fast memory and more accelerators, and stream data to compute

The standard response is to scale the conventional design: use more accelerators (several GPUs each with their own HBM), pack in more high-bandwidth memory, and rely on the deep hierarchy — stream weights and KV data from off-chip HBM into on-chip SRAM and registers, compute, and write results back. Each layer's work is done by fetching its parameters and activations to a compute unit, then moving the outputs onward. The hope is that fast enough memory and enough parallel accelerators keep the compute units fed.
Why it fell short

Streaming across tiers cannot escape the movement it depends on

Scaling the conventional architecture does not remove the tax; it pays it more times. Every layer still fetches weights and KV data from off-chip memory to the compute units and writes activations back, so the more the model grows, the more data crosses the tiers and the interconnects — and that traversal is exactly what dominates energy and limits throughput. Faster HBM and more GPUs raise the ceiling but do not change the fundamental structure: compute and storage are in separate places, so data must move between them constantly. The only way to truly eliminate the movement is to stop separating compute from storage — to perform the computation in the very memory that holds the data, so weights and activations never have to travel to a distant compute unit. But doing all of an LLM's operations in the first-level, on-chip memory means abandoning the large off-chip tiers entirely, and on-chip SRAM has far lower density (much less capacity per area) than DRAM or HBM. What was missing was an architecture that computes entirely in fast on-chip memory to kill data movement, yet manages that scarce, dual-purpose memory — which must serve as both storage and compute substrate — carefully enough to actually fit and efficiently run a large model.
This paper's idea

Compute the whole model inside a wafer of SRAM, and allocate it wisely

Ouroboros is a wafer-scale, SRAM-based Computing-in-Memory (CIM) architecture for LLM inference — the name evokes the snake eating its own tail, because computation happens inside the memory itself rather than fetching data out to a separate engine. Instead of cutting a silicon wafer into many small chips, a wafer-scale chip keeps the whole wafer as one enormous chip holding thousands of CIM cores connected by an on-wafer network; each core's SRAM does double duty as both storage and compute. Ouroboros performs all LLM operations in situ within this first-level SRAM, eliminating off-chip data migration and, with it, the hardware scaling tax — yielding, the paper reports, unprecedented energy efficiency. The hard part is that SRAM capacity is scarce, so the model, its activations, and its KV cache must be packed and scheduled with great care. Ouroboros contributes three innovations to do this. Token-Grained Pipelining schedules the model as a pipeline at the granularity of individual tokens rather than whole sequences, so cores stay busy despite varying input lengths. Distributed Dynamic KV Cache Management tucks the KV cache into fragmented, underused SRAM inside cores already allocated to the model, instead of reserving dedicated memory for it. And Communication-Aware, Fault-Tolerant Core Mapping places transformer blocks onto the wafer's cores to minimize communication and to route around manufacturing defects. Together these reach up to 9.1 times the throughput and 17 times the energy efficiency of state-of-the-art systems (4.1 times and 4.2 times on average).

Why this problem is worth solving

Data movement, not computation, is now the dominant cost of running large models, so an architecture that removes the movement attacks the real bottleneck.

For decades the intuition was that computation is the expensive thing, but for LLM inference that has inverted. Moving a number from off-chip memory to a compute unit costs far more energy than the arithmetic done on it, and an LLM does a staggering amount of moving: streaming billions of weights, passing activations between layers, and reading and writing an ever-growing KV cache, all across a deep hierarchy of memory tiers connected by power-hungry interconnects. As the paper shows, this data movement dominates total energy and caps throughput, and it gets worse as models scale — the “hardware scaling tax.” That means the highest-leverage thing you can do is not make the multipliers faster but stop the data from traveling. Computing in memory does exactly that: if the multiply-accumulate happens inside the memory array that already holds the weights, the weights never move. For a workload where movement is the majority of the energy, eliminating it is not an incremental tuning but a change in the dominant term, which is why the paper reports order-of-magnitude energy-efficiency gains. Solving this well matters because inference is where deployed models spend almost all their lifetime energy and cost.

The catch is that computing in on-chip SRAM trades a movement problem for a capacity problem, and the whole design is about winning that trade.

Eliminating off-chip memory sounds unambiguously good until you notice what you gave up: capacity. Off-chip DRAM and HBM are dense — lots of storage per unit area — precisely because they are optimized for storage, not computation. SRAM, especially SRAM that also performs computation, is far less dense, so a wafer of SRAM CIM holds much less than a stack of HBM. Now every scarce byte of that SRAM is doing two jobs at once — storing data and acting as the compute substrate — so wasting it on idle cores, on a poorly-placed KV cache, or on data that must be shuffled between distant cores directly costs you either capacity or performance. This is why the architecture cannot succeed on the CIM idea alone; it lives or dies on how well it allocates the memory. Every one of Ouroboros’s three innovations is an answer to a specific way that naive use would waste this SRAM: pipeline bubbles waste compute, dedicated KV storage wastes capacity, and bad core placement wastes bandwidth. The reason the problem is deep, and worth a paper, is that removing the data-movement tax only pays off if you can then fit and schedule a huge model inside a fundamentally capacity-limited substrate.

What people did before, and where it broke

Scale the conventional GPU-plus-HBM design. Use more accelerators and more high-bandwidth memory, streaming weights and KV data from off-chip into on-chip compute. Where it broke. Compute and storage stay separate, so data must keep crossing the memory tiers — the movement that dominates energy grows with the model rather than shrinking.
Deepen the memory hierarchy. Add more levels (SRAM over DRAM over HBM) to keep hot data close. Where it broke. It only reduces how far data travels on average; the traversals between tiers and their interconnects still dominate energy at LLM scale.
Pipeline whole sequences across stages. Split the model into pipeline stages, each processing different input sequences. Where it broke. LLM computational intensity varies with sequence length and batching, so sequence-grained stages become load-imbalanced, leaving some cores idle (pipeline bubbles) while others are overloaded.
Manage the KV cache with centralized, dedicated memory. Reserve memory and manage the KV cache from a central controller. Where it broke. Centralized management does not suit a large, distributed array of thousands of CIM cores and fails to balance their load, while dedicated KV memory wastes scarce SRAM.
The gap. No design both eliminated off-chip data movement by computing entirely in on-chip memory and solved the resulting capacity-allocation problems — keeping cores busy, fitting the KV cache, and placing work to minimize communication and tolerate defects.

How Ouroboros works

From streaming across memory tiers to computing inside a wafer of SRAM

↓ →
new
↓ →
new
↓ →
new
↓ →
new
↓ →
new

The idea

Keep the pipeline full by scheduling tokens, not whole sequences. Ouroboros splits the model into pipeline stages across CIM cores, but schedules work at the granularity of individual tokens so that the wildly varying amount of computation per sequence does not leave some stages idle and others swamped. Reworking pipeline parallelism so that stages stay balanced and cores are not left waiting — adapting the schedule to the workload’s uneven shape rather than assuming uniform stages — is the same problem the work on optimizing pipeline parallelism with adaptive recomputation (AdaPipe) covered elsewhere on this site tackles, keeping a pipelined accelerator busy. Ouroboros’s token-grained pipelining is its answer for the fluctuating computational intensity of LLM inference, and it also shrinks the activation memory each stage must hold.

Reclaim leftover on-chip memory to hold the KV cache. The KV cache is a major memory consumer, and on a capacity-starved SRAM wafer there is no room to reserve dedicated storage for it, so Ouroboros decouples memory from compute and slots the KV cache into the fragmented, underused SRAM inside cores that are already allocated to the model. Treating memory as a first-class resource to be carefully managed and packed for the growing KV cache, rather than statically over-reserving it, is the same discipline the work on dynamic memory management for serving LLMs (vAttention) covered elsewhere on this site brings to KV-cache storage, squeezing more serving out of fixed memory. Ouroboros does this in a distributed way across thousands of cores, with a mapping that keeps their loads balanced.

Place computation on the wafer to minimize the communication it forces. Because a CIM core can only compute where its data physically lives, Ouroboros treats mapping transformer blocks to cores as an optimization — using Mixed Integer Quadratic Programming to balance inter-stage against intra-stage communication, and accounting for defective cores. Co-exploring how work is mapped onto hardware to jointly satisfy memory-capacity and communication constraints, rather than optimizing compute placement in isolation, is the same co-design the work on hardware mapping co-exploration toward memory capacity and communication (CoCCO) covered elsewhere on this site pursues, treating placement and communication as one problem. Ouroboros adds wafer-scale fault tolerance, mapping around the manufacturing defects that are unavoidable at wafer scale.

Following a request through Ouroboros

Serving an LLM request on the wafer

A trace, drawn from the paper, of how Ouroboros runs inference without moving data off-chip.

Step 1 of 5

The model is mapped onto the wafer once

Before serving, Ouroboros uses Mixed Integer Quadratic Programming to place the transformer blocks onto the thousands of CIM cores so that neighboring pipeline stages are physically close and large layers are spread for compute locality — while avoiding cores with manufacturing defects.

Weights now live in the SRAM where they will also be computed on.

Why it works

Computing in memory removes the dominant energy term. Data movement across memory tiers dominates LLM inference energy → but performing every operation inside the SRAM that holds the data → so weights and activations never travel, eliminating the tax rather than shrinking it.
Token-grained pipelining keeps scarce cores utilized. Sequence-grained stages idle some cores while overloading others → but scheduling at token granularity to smooth per-stage work → so CIM cores stay busy and less activation memory is needed, using the capacity efficiently.
Reusing fragmented SRAM fits the KV cache for free. Reserving dedicated KV memory would waste scarce SRAM → but tucking the KV cache into underused memory in already-allocated cores → so KV storage costs almost no additional capacity while staying load-balanced.
Communication-aware mapping respects that compute is pinned to data. In CIM, work must run where its data lives, so placement sets communication cost → but optimizing the block-to-core mapping to balance inter- and intra-stage traffic → so the on-wafer network is not the bottleneck.
Defect-aware mapping makes wafer-scale actually manufacturable. A whole-wafer chip inevitably has some faulty cores → but folding core-level defects into the mapping decision → so the model runs across the working cores, making the design fault-tolerant in practice.

What it costs (trade-offs and scope)

SRAM capacity is the hard limit. Abandoning off-chip DRAM and HBM removes data movement but caps how much model, activation, and KV data can reside on-chip; the whole design exists to cope with this scarcity, and models must fit within the wafer’s SRAM budget.
It is a specialized wafer-scale chip, not commodity hardware. The gains rest on a custom SRAM-CIM wafer-scale architecture with an on-wafer network, not on running the idea on existing GPUs — a substantial hardware commitment.
The scheduling and mapping add complexity. Token-grained pipelining, distributed KV management, and Mixed Integer Quadratic Programming plus Dynamic Programming for mapping are non-trivial machinery that must be gotten right for the utilization gains to materialize.
Benefits depend on model and workload shape. The peak 9.1-times and 17-times figures are for the 13B model; average gains are 4.1 times and 4.2 times, and the advantage depends on how well a given model and request mix fit and utilize the wafer.
Reported figures are architecture-evaluation results. The throughput and energy numbers come from the paper’s evaluation against specific state-of-the-art systems and configurations; real-world gains depend on the models, sequence lengths, batching, and the manufactured wafer’s defect rate.

The comparison

Running LLM inference: conventional hierarchy versus wafer-scale CIM

Property	GPU + HBM hierarchy	Naive SRAM CIM	Ouroboros
Where computation happens
Data movement energy
Core utilization
KV cache storage
Throughput and energy

A comparison of ways to run LLM inference. A GPU-plus-HBM hierarchy is bottlenecked by data movement across memory tiers; naive SRAM computing-in-memory removes movement but wastes its scarce capacity on pipeline bubbles and dedicated KV storage. Ouroboros computes every operation inside first-level SRAM on a wafer-scale chip and allocates that capacity carefully — token-grained pipelining, KV caching in fragmented memory, and communication- and defect-aware mapping — reaching up to 9.1 times throughput and 17 times energy efficiency over state-of-the-art systems on the 13B model. Exact figures should be read from the paper.Underlined values explain themselves when clicked.

Where this stops working

What to keep in mind

This explainer is based on the paper’s full text (the authors’ arXiv copy, 2603.02737, of the ASPLOS 2026 paper). The described system is the paper’s: Ouroboros is a wafer-scale SRAM-based Computing-in-Memory (CIM) architecture that performs all LLM operations in situ within first-level SRAM to eliminate off-chip data migration. Because SRAM capacity is limited and each core’s SRAM serves as both storage and compute, the paper introduces three innovations: (1) Token-Grained Pipelining, which partitions sequences into token-level units scheduled across pipeline stages to avoid the load imbalance and pipeline bubbles of sequence-grained pipelining and to reduce activation storage; (2) Distributed Dynamic KV Cache Management, which decouples CIM memory from compute assignment and repurposes fragmented, underused SRAM in already-allocated cores for KV storage, with a load-balancing mapping; and (3) Communication-Aware and Fault-Tolerant Core Mapping, which uses Mixed Integer Quadratic Programming (and Dynamic Programming for intra-core pressure) to place transformer blocks so as to balance inter-stage and intra-stage communication while routing around core-level defects. The reported results are up to 9.1 times throughput and 17 times energy efficiency for the 13B model, with 4.1 times and 4.2 times average improvements over state-of-the-art systems. Background on the memory hierarchy, computing-in-memory, wafer-scale chips, transformer blocks (attention and FFN), prefill versus decoding, and the KV cache is standard context provided to make the approach intuitive. The step-by-step walkthrough is an illustration of the described mechanism. Exact architectural details, the scheduling and mapping formulations, and the full evaluation should be read from the paper.

Background

Why computing in memory keeps returning as the answer to the data-movement wall

It helps to see Ouroboros within a long-running response to a fundamental problem: in modern computing, and especially in machine learning, moving data costs far more than computing on it, and the gap keeps widening. The classic architecture separates memory from compute and pays, over and over, to ferry data between them — a cost that now dominates the energy and latency of running large models. The recurring structural answer is to collapse that separation: bring the computation to the data, whether by processing near memory, inside memory arrays, or, as here, by making a vast expanse of memory that also computes. Each incarnation faces the same follow-on question, because memory optimized for computing is scarcer than memory optimized only for storage: once you have eliminated the movement, how do you fit and schedule a huge workload inside the capacity you have left? Ouroboros is a clean, extreme instance — a whole wafer of SRAM that computes in place — and its real content is the allocation machinery (token-grained pipelining, fragmented KV storage, communication- and defect-aware mapping) that makes the scarce capacity pay off. This drive to unleash in- or near-memory computation for transformer inference, so that data movement stops throttling large models, is the same terrain the work on unleashing processing-in-memory for batched transformer generation (AttAcc) covered elsewhere on this site occupies, moving the computation to where the data already sits. Ouroboros’s particular contribution is to push this to wafer scale and show that, with careful management of the scarce on-chip memory, computing an entire LLM in situ delivers order-of-magnitude energy and throughput gains.

If you remember only three things

1. Large language model inference has become memory-bound: as models scale, most of the energy and latency go not to arithmetic but to moving weights, activations, and the growing Key-Value (KV) cache across a deep memory hierarchy (on-chip SRAM backed by off-chip DRAM or HBM) — a "hardware scaling tax" where data movement across tiers and interconnects dominates cost. Scaling the conventional GPU-plus-HBM design does not remove this tax; because compute and storage stay separate, more model size means more data crossing the tiers. The only way to truly eliminate the movement is to stop separating compute from storage and compute inside the memory that holds the data.
2. Ouroboros is a wafer-scale SRAM-based Computing-in-Memory (CIM) architecture — thousands of CIM cores on one wafer, each core's SRAM doubling as storage and compute — that performs all LLM operations in situ in first-level SRAM, so data never goes off-chip. The challenge it must then solve is that SRAM is low-density and scarce, so it contributes three allocation innovations: Token-Grained Pipelining schedules the model as a pipeline at the token level (not whole sequences) so varying input lengths do not leave cores idle in pipeline bubbles; Distributed Dynamic KV Cache Management tucks the KV cache into fragmented, underused SRAM in cores already allocated to the model instead of reserving dedicated memory; and Communication-Aware, Fault-Tolerant Core Mapping uses Mixed Integer Quadratic Programming to place transformer blocks to minimize inter- and intra-stage communication while routing around manufacturing defects.
3. By eliminating off-chip data movement and using the scarce on-chip SRAM efficiently, Ouroboros reports up to 9.1 times throughput and 17 times energy efficiency over state-of-the-art systems on the 13B model, and 4.1 times and 4.2 times on average. The trade-offs: SRAM capacity is a hard limit on what fits, it is a specialized wafer-scale chip rather than commodity hardware, the token-grained scheduling and optimization-based mapping add real complexity, and the gains depend on how well a given model and workload fit and utilize the wafer. This explainer is based on the paper's full text; the architectural and scheduling details should be confirmed in the paper.

Check yourself

Check your understanding

1. Why is large language model inference described as memory-bound, and what is the "hardware scaling tax"?
2. What is the core architectural idea of Ouroboros, and what new problem does it create?
3. What are Ouroboros's three innovations for using the scarce SRAM efficiently?
4. What does Ouroboros achieve, and what are the limits of the approach?

Sources

The paper is “Ouroboros: Wafer-Scale SRAM CIM with Token-Grained Pipelining for Large Language Model Inference” by Yiqi Liu, Yudong Pan, Mengdi Wang, Shixin Zhao, Haonan Zhu, Yinhe Han, Lei Zhang, and Ying Wang (SKLP, Institute of Computing Technology, Chinese Academy of Sciences, and University of Chinese Academy of Sciences), published at ASPLOS 2026, DOI 10.1145/3779212.3790197. It is available on arXiv as 2603.02737.

This explainer is written from the paper’s full text (the authors’ arXiv copy). Ouroboros is a wafer-scale SRAM Computing-in-Memory architecture that performs all LLM operations in situ within first-level SRAM to eliminate off-chip data movement, and manages the resulting scarce capacity with three innovations: token-grained pipelining to keep CIM cores utilized despite varying sequence lengths, distributed dynamic KV cache management that reuses fragmented on-chip memory, and communication-aware, fault-tolerant core mapping (via Mixed Integer Quadratic Programming) that minimizes inter- and intra-stage communication while routing around defects. The paper reports up to 9.1 times throughput and 17 times energy efficiency over state-of-the-art systems on the 13B model (4.1 times and 4.2 times on average). Background on the memory hierarchy, computing-in-memory, wafer-scale chips, transformer inference, and the KV cache is standard context. The numbers are the paper’s own; exact mechanisms and full results should be confirmed in the paper.