跳到正文
Cerebras:Blog·· 21 小时前AI 评分42

Cerebras 如何用 59 片 WSE-3T 晶圆部署 3T 参数 Kimi K3

How we deploy Trillion-parameter frontier models on Cerebras

AI 导读

Cerebras 详解如何把约 3T 参数的 Kimi K3 拆分到多片 WSE-3T 晶圆上运行,实现 1000-2000+ tokens/s 的多万亿参数模型推理。

正文 · 原文

One of the most frequently asked questions we get is how we deploy large models on Cerebras wafer-scale chips. A single Cerebras WSE-3T wafer has 44 GB of on-chip memory, and trillion-parameter frontier models need terabytes for their weights alone.

In this post, we’ll explain how we split a large frontier model across multiple WSE-3T wafers, to run multi-trillion parameter models at 1000-2000+ tokens per second.

We’ll use the open source model Kimi K3, with ~3 trillion parameters, as an illustrative example.

By the end, you'll understand:

  • How we fit the full model into on-chip memory
  • How we map the model across multiple chips
  • How we handle a high # of concurrent requests

Breaking it down with Kimi K3

Kimi K3 is an open source, 3T parameter model with 69 linear-attention layers and 24 dense layers.

For this illustrative example, we use:

#1: Size of Each layer

To calculate the total number of Cerebras WSE-3T wafers we need to support Kimi K3, let’s first calculate the size of each model layer.

Kimi K3 has 3 trillion parameters. If all 3T parameters are stored at 4 bits each, the raw parameter payload is:

Dividing that payload by 93 layers then gives:

#2: Number of Wafers

Now that we know each layer needs about 16 GB, we can figure out how many layers fit on each 44 GB wafer. The answer depends on the layer type:

  • Dense layers need additional room for a KV cache, because each dense (full-attention) layer stores the keys and values for every token in every active request. So, we place one dense layer per wafer and leave the remaining SRAM for its KV cache.
  • Linear-attention layers do not require additional room for a KV cache. Instead of storing keys and values for every token, a linear-attention layer keeps a small, fixed-size state per request, regardless of context length. That leaves room for two linear layers per wafer (2 × 16 GB = 32 GB).

Therefore, with Kimi K3’s 69 linear layers and 24 dense layers, we get:

This gives us a total wafer count of 35 + 24 = 59 wafers.

#3: Connecting the wafers together

Cerebras executes models in a pipeline. After the Cerebras compiler maps the model layers across regions of wafers, the runtime maps the path for tokens to flow between wafers. For ultrafast latency, only activations need to be passed between the wafers, and this can be done via RDMA over Ethernet. This is roughly 2 microseconds per hop.

It’s important to emphasize that because each Cerebras wafer has 44 GB of on-chip SRAM and each layer is only 16 GB, we are typically able to fit entire layers of the model directly on the chip.

This means that in most cases, the model weights never have to move, and layers don’t need to be split across wafers. Only activations cross from one wafer to the next.

#4: How we serve 1 request

Now that we have our wafers connected together and the model laid out, we can think of each wafer as one stage in an assembly line: the model’s weights for each stage stay in place, and the request’s intermediate results move through the layer stages.

For example, Wafer 1 runs model layers 1–2 and sends its activations to wafer 2. Wafer 2 runs layers 3–4 and passes its results to Wafer 3, and so on until the last wafer produces the next token.

#5: How we serve many requests

This raises a common misconception. If a token has to travel through each wafer in order, doesn’t each wafer sit idle while it waits for its turn? In effect, wouldn't the system be running at a batch size of one?

No! If we only sent one request down the line at a time and waited for it to complete before starting the next request, this would be an inefficient use of the pipeline. Instead, we keep the entire pipeline full with different requests. Because the model is split into pipeline stages already, each stage can work on a different request at the same time.

While stage 3 processes a token for request A, stage 2 processes one for request B, and stage 1 one for request C. In our example, each stage is a layer. In production, Cerebras uses much finer-grained stages, and many more of them.

Video thumbnail

Loading player...

Each stage can also process a micro-batch, a small group of requests, rather than a single request. So the number of requests in flight is roughly the number of stages times the micro-batch size.

Video thumbnail

Loading player...

To understand the tradeoff, smaller micro-batches make each request even faster, while larger micro-batches increase total throughput.

#6: How many requests can we serve?

The amount of on-chip memory allocated to KV cache determines the maximum number of concurrent requests. Every request in flight needs its own KV cache on each dense-layer wafer, so the size of on-chip memory sets the ceiling on concurrent requests. As explained above, we only need to calculate KV cache requirements for dense layers.

For each 44 GB wafer that already has one dense layer, we can calculate the remaining space for the KV cache with:

Total wafer storage - size of one dense layer. In this case:

Calculating the # of Concurrent Requests

To find the number of concurrent requests supported for any particular model mapping, we first calculate how much cache a single 256k request takes up:

To then calculate the total number of concurrent requests this configuration would support, we divide:

The total number of concurrent requests represents the number of simultaneous requests moving through the pipeline at any point in time. Unlike a GPU batch, those requests don’t wait for one another to finish, they stream through the pipeline continuously.

Of course, neither concurrency nor batch size tells you exactly how efficiently an inference system performs. When evaluating the throughput of an inference system, the more meaningful metric is total token throughput per unit of power (e.g. tokens per second per megawatt), when measured at the same output generation speed.

How this scales beyond 3T parameters

The same approach extends to larger models: calculate the memory needed for weights and request state, distribute the model across wafers, and connect them into a pipeline. Each additional wafer contributes compute and on-chip memory, allowing us to accommodate models that exceed the capacity of any single chip.

Thanks for learning with us! Interested in deploying a model on Cerebras or building with our inference cloud?

Designs by Halley Chang. Thanks to Atharva Talpade, Zhenwei Gao, Angela Yueng, Jessica Liu, Joyce er, Emma Call, and Rohan for their review and feedback.

来源:Cerebras:Blog · cerebras.ai