Cerebras runs Kimi K3 on 59 wafers at 1,000 tokens per second

A chip the size of a dinner plate still has only 44 GB of memory on it. A 3 trillion parameter model needs more than that for weights alone. Cerebras has published its own worked example of the split.

Sarah Chieng wrote the post on the Cerebras blog. It uses Kimi K3, an open 3 trillion parameter model from Moonshot, as the example. The model has 69 linear-attention layers and 24 dense layers. Kimi K3 scores 88.3 percent on Terminal-Bench 2.1 and 91.2 percent on BrowseComp, both Moonshot's own figures. The Hacker News thread has 147 points.

How the model fits on the chips

Cerebras assumes 4-bit weight quantisation and an 8-bit KV cache. At 4 bits each, 3 trillion weights come out to a layer of about 16 GB. One wafer has 44 GB of on-chip memory.

A dense layer needs room for the KV cache. The KV cache stores keys and values for every token in every active request. So each dense-layer wafer holds one layer and 28 GB of KV cache. A linear-attention layer keeps a fixed-size state per request. Two linear layers fit on one wafer, with 12 GB left over.

Kimi K3 has 69 linear layers over 35 wafers and 24 dense layers over 24 wafers. The total is 59 wafers. Each layer sits on one chip, so the weights never move during inference.

How the pipeline keeps the chips busy

Each wafer passes only the activations to the next wafer, over Ethernet with RDMA. The latency between hops is about 2 microseconds. One request runs through 59 stages in order.

The pipeline does not wait for a request to finish before the next one starts. Each stage works on a different request at the same time. Cerebras runs finer-grained stages in production, each on a small group of requests at once. The system reaches 1,000 to 2,000 plus tokens a second this way.

What limits the number of live requests

The 28 GB of KV cache on each dense-layer wafer sets the ceiling. Cerebras cites 512 live requests at an assumed average size of 256 thousand tokens. That number is per-wafer and per-model. A smaller average request size would allow more.

The post is written as an engineering explanation of a question Cerebras says it gets often. The company says the better measure is tokens per second per megawatt. Both systems have to run at the same output speed. Cerebras does not publish that figure in the post.

The method shown here is pipeline parallelism, where each stage holds its own weights. The step-by-step arithmetic names every number used: 4-bit weights, 8-bit KV cache, 44 GB per wafer, 16 GB per layer, and the 256 thousand token request size. A team planning to serve a frontier open model can repeat the same steps on its own hardware.

Source

How we deploy trillion parameter frontier models on Cerebras, by Sarah Chieng, Cerebras. Discussion: Hacker News.