Klotski: Efficient Mixture-of-Expert Inference via Expert-Aware Multi-Batch Pipeline

Zhiyuan Fang, Yuegui Huang, Zicong Hong, Yufeng Lyu, Wuhui Chen, Yue Yu and 2 others

A Mixture-of-Experts language model has enormous parameters but activates only a couple of experts per token, so to run one on a modest machine you offload the idle experts to CPU and disk and stream them back on demand — but loading a single expert takes far longer than computing it, so the GPU stalls in idle bubbles; this work from Sun Yat-sen University, HKUST, Peng Cheng Laboratory, and Huawei builds Klotski, which runs several batches together so the long computation of the popular hot experts hides the slow loading of the rare cold experts, shrinking the bubbles to almost nothing and lifting throughput up to 85 times over prior offloading engines.

The argument in one glance

Problem

Mixture-of-Experts models are huge, but only sparsely used

A Mixture-of-Experts (MoE) model swaps each dense feed-forward layer for many expert sub-networks and lets a small gate route each token to just one or two of them. This makes the model cheap to compute yet vast in knowledge — but the sheer parameter count no longer fits in GPU memory. DeepSeek-V2, at 236 billion parameters, needs seven top-end GPUs to hold. To run such a model on a small server or a single consumer card, systems offload the mostly-idle experts to CPU memory and disk, and stream each one back to the GPU only when the gate actually picks it.
What people did

Prefetch the right experts more accurately, or compute them on the CPU

Offloading engines built for MoE try to hide the loading cost. The dominant line of work sharpens expert prefetching: MoE-Infinity watches activation traces to guess which experts a layer will need, and SiDA trains an offline predictor that hits over 90 percent accuracy, so the right expert is on its way before the gate asks for it. A different tack, Fiddler, keeps some experts on the CPU and runs them there to avoid moving them at all.
Why it fell short

Even perfect prefetching cannot beat the loading-versus-compute gap

The bottleneck is physical, not predictive. Loading a single expert takes longer than computing it — on an RTX 3090 an attention layer computes in about 2.6 milliseconds while one expert transfer takes about 21 — so even with 100 percent accurate prefetching the GPU still stalls, waiting on I/O, in what are called bubbles. Worse, the dense-model trick of batching more work to fill the wait backfires here: more batches make the gate activate more distinct experts, so the loading time grows and the bubbles multiply.
This paper's idea

Overlap the compute of popular experts with the loading of rare ones

Klotski notices that a few hot experts handle most of the tokens while many cold experts handle few. Run several batches together as one group, and the hot experts now carry a large, time-consuming pile of computation, while the cold experts each need only a little compute but a lot of loading. So schedule the hot experts first, right after prefetching just them, and their long computation buys exactly the time needed to stream in the cold experts. The two demands are complementary, and pairing them squeezes the bubbles to almost nothing.

Why this problem is worth solving

The starting point is what makes MoE both attractive and awkward to run, so it is worth being precise about the sparsity.

A Mixture-of-Experts model replaces the feed-forward network inside each transformer block with an MoE layer: a small gating network plus many experts, where each expert is itself a feed-forward network. For every token the gate computes routing weights with a softmax and activates only the top-k experts, usually one or two, rather than all of them. The final output is a weighted sum of just those experts. This is what lets MoE scale to hundreds of billions or even trillions of parameters without a matching rise in compute — most of the weights sit unused for any given token. But that same design is a memory nightmare for inference: in Switch Transformers the experts can be up to 99 percent of all parameters, and a model like DeepSeek-V2 with 236 billion parameters needs seven state-of-the-art 80-gigabyte GPUs merely to hold its weights. On a personal computer or a small server, the gap between the model size and the available GPU memory is enormous. So the sparsity that makes MoE cheap to compute is exactly what makes it expensive to store, and running one in a resource-constrained setting is an open problem.

The natural fix, offloading, then runs into an arithmetic wall that no amount of clever prediction can climb.

Offloading keeps only the tensors needed right now in GPU memory and parks the rest in the much larger CPU memory and disk, streaming each piece back just before it is used. For MoE this is a great fit, because sparsely activated experts are idle most of the time and can be offloaded. The catch is that the link from CPU to GPU is far slower than the GPU computes, and an expert is large: transferring one expert takes much longer than the tiny amount of computation it performs on a token. The result is bubbles — stretches where the GPU sits idle, waiting for data. These come in two flavors. Inter-layer bubbles appear between the attention layer and the MoE layer, because the attention layer computes in a couple of milliseconds but the expert layer needs tens of milliseconds to load. Intra-layer bubbles appear within the expert layer, because each activated expert processes only a slice of the tokens yet you still pay the full cost of transferring its whole feed-forward network. Sharpening prefetch accuracy does not remove either bubble; the loading is simply slower than the compute, so the wait remains. Breaking that wall needs a way to give the GPU enough work to do while the data is in flight — which is what makes rethinking the pipeline worthwhile.

What people did before, and where it broke

Make expert prefetching as accurate as possible. MoE-Infinity prefetches and caches experts from activation traces; SiDA trains an offline data-aware predictor reaching over 90 percent accuracy. Where it broke. Even flawless prediction only decides which expert to load, not how long loading takes — and since one expert’s I/O outlasts its compute, the GPU still stalls waiting for it.
Run experts on the CPU to avoid moving them. Fiddler orchestrates CPU and GPU together, using the CPU’s compute so some experts never have to be transferred to the GPU at all. Where it broke. It cuts data movement but leans on limited CPU compute, and it can only offload experts, so a large batch’s key-value cache overflows GPU memory and it caps out at small batch sizes.
Cache and quantize a subset of experts. Mixtral-offloading keeps a working set of experts on the GPU with an LRU cache and quantizes them, so fewer transfers are needed. Where it broke. A cache helps only when the same experts recur; misses still trigger slow loads, and the fundamental I/O-versus-compute imbalance inside a layer is untouched.
Borrow the dense-model multi-batch overlap trick. FlexGen, built for dense models, runs several batches together so the shared weights of one layer are reused across all of them, extending compute to overlap the next layer’s loading. Where it broke. In an MoE, more batches feed more varied tokens to the gate, which activates more distinct experts — so the loading time grows with the batching, and applying the dense trick directly to a MoE gives far less speedup than it does to a dense model of the same size.
The gap. Every approach either predicted better without changing the loading-versus-compute imbalance, or borrowed a dense-model overlap that MoE’s data-dependent routing actively fights. What was missing was a pipeline that is aware of which experts are hot and which are cold, and that arranges the multi-batch computation around that structure — which is exactly what efficient MoE offloading needs.

How Klotski works

Arranging the multi-batch pipeline around hot and cold experts

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

The idea

Attack the loading imbalance, not the prediction accuracy. The field spent its effort guessing which experts to load, but Klotski’s measurements show that even a perfect guess leaves the GPU stalled, because a single expert takes longer to load than to compute. The real enemy is the imbalance between I/O and compute, so Klotski gives the GPU enough computation to chew on while data is in flight, rather than trying to make the flight shorter through better prediction. This is the same lesson the deep-learning-on-GPUs retrospective covered elsewhere on this site draws for training hardware — that once the arithmetic units are fast enough, the ceiling is set by moving the data, so the useful work is hiding or shrinking that movement rather than accelerating compute that is already idle.

Extend the computation so it hides the loading, then keep it extended. The lever is running many batches together: because a loaded layer’s weights are shared across every batch in the group, one transfer now serves n computations, and the stacked compute of those n batches stretches long enough to overlap the loading of the next layer. Klotski cuts the number of transfers to roughly one over n of the original and tucks each one behind computation. Overlapping a slow transfer behind useful compute so the wait never shows on the critical path is exactly the discipline the communication-overlapping scheduler covered elsewhere on this site applies to distributed training, where the trick is to keep the expensive movement running quietly underneath the math that was going to happen anyway.

Predict what is hot, prefetch only that, and let the rest arrive just in time. Klotski does not need a highly accurate predictor, because the multi-batch overlap already hides most of the loading; it needs only to know which handful of experts are hot. It builds an expert correlation table from a short pre-run, recording how a token’s expert choices in one layer predict its choices in the next, and uses it to prefetch just the gate and the top hot experts during attention, updating the table as inference proceeds. Prefetching only the popular items and reusing warm state so the rare items load lazily behind ongoing work is the same instinct the multi-turn LLM-serving cache covered elsewhere on this site follows, keeping the frequently-needed context ready while letting cold data stream in under cover of computation.

Following a Mixtral batch through Klotski

Running Mixtral on a single consumer GPU

Tracing one batch group shows how hot and cold experts trade places in the pipeline.

Step 1 of 5

A group of batches enters the MoE block

Klotski gathers several request batches into one batch group and feeds them through a transformer block. The attention layer runs first, over all the batches at once, while the gate and a small set of hot experts are prefetched from CPU memory in the background.

Loading one layer's weights once and sharing them across the whole group is what makes the compute long enough to hide the loading.

Why it works

Sharing one layer’s weights across many batches turns a single load into many computations. A plain pipeline loads a layer and computes one batch, so the load dwarfs the compute → but reusing the loaded weights across a group of n batches stacks n computations behind that one load → so the compute finally stretches long enough to overlap the next layer’s loading and the inter-layer bubble closes.
Hot and cold experts have opposite demands, so pairing them cancels the wait. Hot experts carry heavy compute but load rarely, while cold experts load fully but compute little → but running the hot experts’ long computation while the cold experts stream in overlaps the two demands → so the cold transfers finish under cover of the hot compute and the intra-layer bubble shrinks toward zero.
Prefetching only the gate and hot experts makes the transfer small enough to hide. Loading the entire MoE layer during attention is too much data for the brief attention compute to cover → but prefetching just the gate and the one or two hot experts is a small, achievable transfer → so the multi-batch attention computation actually covers it instead of stalling behind a full-layer load.
Reordering to compute hot experts first keeps the GPU from stalling on data in transit. Computing experts in gate order can hit a cold expert whose transfer has not arrived, freezing the GPU → but scheduling the already-loaded hot experts first defers the cold ones until their transfers complete → so the GPU always has ready work and the cold loads finish just in time.
Solving for the smallest sufficient batch count avoids paying for overlap you do not need. Too few batches leaves compute too short to hide the loading, while too many balloons the key-value cache and risks running out of memory → but measuring the hardware and solving the overlap inequalities pins n to the smallest value that erases the bubbles → so Klotski gets a full pipeline without a wasteful memory footprint.

What it costs (trade-offs and scope)

Multi-batching inflates the key-value cache, and too large a batch count runs out of memory. The overlap works by processing many batches together, but each batch carries its own growing key-value cache; on a single 3090 the computed batch count for Mixtral-8x22B was large enough to cause out-of-memory, so the authors manually dialed it down to n equal to 10.
The whole scheme leans on the hot-expert phenomenon. If a workload’s routing were uniform — every token spread across cold experts — the hot experts would carry no extra compute, the prefetching would be useless, and the bubbles would return. The authors argue this worst case is very unlikely in practice, but it is the assumption the design rests on.
The correlation-aware prefetcher needs a pre-run and stays deliberately shallow. Building the expert correlation table requires a warm-up pass over sample text, and the authors set the lookback length to just one layer to keep the table small and fast, accepting lower prefetch accuracy because the multi-batch overlap already hides most loading.
The strongest numbers use extra options and hand-tuned settings. The headline throughput uses the maximum batch count of 15 to show the best result rather than the default the planner computes, and the memory-saving figures rely on optional 4-bit quantization (via Half-Quadratic Quantization) and StreamingLLM-style sparse attention, each of which trades some accuracy for smaller transfers.
It is a prototype on a narrow hardware and model set. Klotski is about 3,000 lines of Python built on PyTorch, Hugging Face Transformers, and FlexGen, evaluated on two machines (an RTX 3090 and an H800) with the Mixtral-8x7B and Mixtral-8x22B models; the multiplicative speedups are up to numbers against specific baselines and depend heavily on batch size, model, and hardware.

The evaluation

Four ways to run an offloaded MoE, and where the bubbles go

Property	Single-batch prefetch	Dense multi-batch (FlexGen)	Accurate prefetch (MoE-Infinity)
How it hides loading	Prefetch next layer while computing current	Share weights across many batches	Predict and cache likely experts
Inter-layer bubble	Large, attention too short to cover experts	Reduced by longer compute	Reduced when prediction is right
Intra-layer bubble	Large, expert I/O outlasts its compute		Unaddressed, still I/O-bound
Mixtral-8x7B throughput on a 3090	Baseline offloading	Multi-batch, about 18 tokens per second	Prefetch-limited
Best reported gain	Reference point	Klotski beats it by 2.23 times	Klotski beats it by 19.06 times

Ways to run an offloaded Mixture-of-Experts model, evaluated on Mixtral-8x7B and Mixtral-8x22B across an RTX 3090 and an H800. Klotski raises throughput over Accelerate, DeepSpeed-FastGen, FlexGen, MoE-Infinity, and Fiddler by up to 85.12, 15.45, 2.23, 19.06, and 9.53 times respectively, and an ablation on Mixtral-8x7B shows the pipeline climbing from about 5.7 tokens per second for a simple pipeline to 18.2 with multi-batching, 19.1 with hot-expert prefetching, and 22.4 after reordering the expert computations — because it attacks the loading-versus-compute imbalance instead of the prediction.Underlined values explain themselves when clicked.

Where this stops working

What to keep in mind

Klotski is a clean systems insight backed by a careful pipeline. It observes that offloaded MoE inference is limited not by how well you predict experts but by the plain fact that loading one expert outlasts computing it, and it responds by running many batches together and pairing the long computation of the hot experts with the slow loading of the cold ones, so both the inter-layer and intra-layer bubbles nearly vanish. The reported gains are large — up to 85 times over the weakest baseline and 2.23 times over the strongest dense-multi-batch engine, with over 94 percent less GPU memory in the extreme offloading configuration. Read it with its scope. The design rests on the hot-expert phenomenon holding, so a uniform routing pattern would erase the benefit; the multi-batch overlap inflates the key-value cache, forcing a manually reduced batch count on the smallest GPU to avoid running out of memory; the best throughput uses a hand-picked maximum batch count and optional quantization and sparse attention that trade some accuracy; and it is a roughly 3,000-line prototype tested on two machines and the Mixtral family. The durable idea is the reframing: when loading dominates compute, the win is not a better guess about what to load but a schedule that always has computation ready to hide the load behind.

Background

Why offloaded inference keeps fighting the same imbalance

It helps to see Klotski as one more round in a long fight between fast compute and slow data movement. Modern accelerators can do arithmetic far faster than they can pull the operands in from anywhere but their own tiny on-chip memory, so as models outgrow that memory the story of running them becomes the story of hiding or avoiding transfers. MoE sharpens the tension: its sparsity is precisely what makes the weights too big to keep resident, so most experts live in slower memory and have to be fetched, and each fetch dwarfs the little computation it feeds. Klotski’s answer is to manufacture enough computation — by grouping batches and front-loading the busy hot experts — that the fetches always have something to hide behind. The in-storage serverless-acceleration work covered elsewhere on this site fights the same imbalance from the other end, moving the computation down to where the data already sits rather than hauling the data up to the compute; both are refusals to let the slow link set the pace. The broad lesson is that for memory-bound workloads, throughput is decided less by raw compute than by how completely you can keep the movement off the critical path.

If you remember only three things

1. A Mixture-of-Experts model activates only a couple of its many experts per token, so it is cheap to compute but its full parameter set (up to hundreds of billions of weights) will not fit in GPU memory; the standard fix is to offload idle experts to CPU and disk and stream them back, but loading one expert takes far longer than computing it (about 21 milliseconds to transfer versus 2.6 to compute an attention layer on a 3090), so the GPU stalls in inter-layer and intra-layer bubbles, and sharpening prefetch accuracy cannot fix an imbalance that is about speed, not prediction.
2. Klotski runs several batches together as a group so one weight load serves many computations, then exploits that a few hot experts handle most tokens while the rest are cold: it prefetches only the gate and hot experts during attention (closing the inter-layer bubble), computes the already-loaded hot experts first so their long run hides the streaming of the cold experts (compressing the intra-layer bubble), and a constraint-sensitive planner measures the hardware and solves for the smallest batch count that keeps the pipeline full without an oversized key-value cache, aided by adaptive tensor placement, a correlation-table prefetcher, and optional 4-bit quantization and sparse attention.
3. On Mixtral-8x7B and Mixtral-8x22B across an RTX 3090 and an H800, Klotski raises throughput over Accelerate, DeepSpeed-FastGen, FlexGen, MoE-Infinity, and Fiddler by up to 85.12, 15.45, 2.23, 19.06, and 9.53 times, with an ablation showing the climb from 5.7 to 22.4 tokens per second as multi-batching, hot-expert prefetching, and reordering are added, and over 94 percent less GPU memory in the extreme case — but the scheme relies on the hot-expert phenomenon holding, the multi-batch key-value cache can force a hand-reduced batch count to avoid running out of memory, the best numbers use a hand-picked batch count and optional lossy compression, and it is a roughly 3,000-line prototype on two machines and the Mixtral family.

Check yourself

Check your understanding

1. Why does sharpening expert-prefetch accuracy fail to remove the bubbles in offloaded MoE inference?
2. What is the core insight behind Klotski's expert-aware multi-batch pipeline?
3. Why does naively applying the dense-model multi-batching trick to a MoE make things worse?
4. Which statement best reflects the limits of Klotski?

Sources

The paper is “Klotski: Efficient Mixture-of-Expert Inference via Expert-Aware Multi-Batch Pipeline” by Zhiyuan Fang, Yuegui Huang, Zicong Hong, Yufeng Lyu, Wuhui Chen, Yue Yu, Fan Yu, and Zibin Zheng, published at ASPLOS 2025, DOI 10.1145/3676641.3716261. The authors are at Sun Yat-sen University, the Hong Kong University of Science and Technology, Peng Cheng Laboratory, and Huawei Technologies.

This explainer is written from the full 15-page text. Everything above comes from the paper’s own sections: the introduction framing MoE’s memory pressure and the inter-layer and intra-layer bubbles that offloading suffers (with the 2.6-millisecond attention versus 21-millisecond expert-transfer measurement on an RTX 3090); the background and related-work survey of MoE architecture, offloading systems (ZeRO, FlexGen, HeteGen, STI), and MoE-specific engines (Mixtral-offloading, Fiddler, MoE-Infinity); the motivation showing that even accurate prefetching leaves bubbles and that the dense multi-batch trick underperforms on MoE (Table 1); the strawman construction and the expert-aware multi-batch pipeline paradigm with its hot and cold experts, prefetch-only-hot rule, and reordering to compute hot experts first; the tensor management with adaptive placement across GPU, CPU, and disk and the correlation-aware expert prefetcher’s expert correlation table; the constraint-sensitive I/O-compute planner solving for the smallest batch count n, plus the optional Half-Quadratic Quantization and StreamingLLM sparse-attention compression; and the evaluation on Mixtral-8x7B and Mixtral-8x22B across an RTX 3090 and an H800 reporting throughput gains of up to 85.12, 15.45, 2.23, 19.06, and 9.53 times over Accelerate, DeepSpeed-FastGen, FlexGen, MoE-Infinity, and Fiddler, over 94.1 percent memory reduction, and the ablation climbing from 5.721 to 22.604 tokens per second.