A mixture-of-experts model can hold hundreds of gigabytes of parameters while activating only a few per token, which makes it cheap to run and impossible to fit — so systems stream the weights from host memory a layer at a time and use enormous batches to amortise the transfer, at which point host memory rather than the accelerator becomes the limit; this paper (Berkeley and Stanford) extends the roofline model to reason about both processors at once, discovers that attention belongs on the host, and chops the weight transfer into pieces so it stops blocking everything else — reaching the same ceiling with two to three times less host memory.
The argument in one glance
-
Problem
A model that computes small and stores large
A mixture-of-experts model holds many expert networks and routes each token to only one or two, so it does roughly the arithmetic of a small model while holding the parameters of a very large one. That trade is excellent when memory is plentiful and exactly wrong when it is not — one popular model needs over 260 gigabytes of parameters, which no affordable accelerator has, even though the work per token would fit comfortably.
-
What people did
Keep the weights in host memory and stream them in
The standard answer is offloading. Store the parameters in the host's memory, and as each layer is needed, copy its weights across the interconnect to the accelerator. Since that copy is expensive, process an enormous batch of requests while the layer is resident, so one transfer serves thousands of tokens rather than one.
-
Why it fell short
The amortisation trick moves the bottleneck rather than removing it
A large batch needs somewhere to hold the attention state of every request in it, and that lives in host memory alongside the weights — so the batch size you can afford, and therefore the throughput you can reach, is set by how much host memory you have. Meanwhile that state has to travel to the accelerator and back, competing for the same interconnect that should be fetching weights. Both processors end up waiting.
-
This paper's idea
Model both processors properly, then schedule for what the model says
This paper (Cao, Liu, Griggs, Schafhalter, Liu, Sheng, Gonzalez, Zaharia and Stoica) extends the roofline model so it describes a machine with two processors and two memories rather than one of each, uses it to decide which computation belongs where, and finds that attention should run on the host. It then rebuilds the schedule around the fact that weight transfers and activation transfers have opposite urgency.
Why this problem is worth solving
The analysis tool the paper extends is worth understanding, because the conclusions come from it.
THE
ROOFLINE MODEL asks one question about a computation — how many arithmetic operations does it
perform per byte it must fetch, a ratio called OPERATIONAL INTENSITY. Plot achievable
performance against that ratio and you get two limits. Below a threshold the computation is
memory-bound, and performance rises with intensity because the processor is waiting for data.
Above it the computation is compute-bound and performance flattens at the hardware’s peak.
Where a computation sits tells you which resource to spend effort on. The classical model
assumes one processor and one memory, which is not the machine here — there is an accelerator
with its own fast memory, a host with a much larger and slower one, and an interconnect
between them. The extension gives each level of the hierarchy its own processor, its own peak
performance, and its own bandwidth, which produces two turning points the original does not
have: one where a computation is better performed on the host than the accelerator, and one
where the whole system is limited by host memory rather than accelerator memory.
The second thing to see is why throughput becomes a question about host memory at all.
WITH
the weights in host memory, each layer’s parameters must cross the interconnect before that
layer can run. That transfer is expensive, so the way to make it worthwhile is to have a great
many tokens waiting to use the layer while it is resident — the cost is amortised across the
batch. But a language model generating text must remember, for every request, the keys and
values computed for every token so far, and for a batch of thousands of requests that state is
large. It cannot live on the accelerator, whose capacity is the original problem, so it lives
in host memory beside the weights. The chain is therefore: more host memory allows a bigger
batch, a bigger batch amortises the weight transfer better, and better amortisation means
higher throughput. A system that reaches its ceiling only at very large batches is a system
that demands a great deal of host memory, which is the axis this paper competes on.
What people did before, and where it broke
-
Offload weights and attention state, run everything on the accelerator. The state-of-the-art batch inference system for this setting keeps both in host memory, fetches the attention state for the next micro-batch while computing on the current one, and runs all computation on the accelerator.
-
The gap. Moving attention state across the interconnect costs bandwidth that could have been fetching weights, and it is more traffic than doing the attention on the host would have required at all.
-
The same system, with attention on the host. That system does offer a mode where attention runs on the host instead, which removes the state transfers.
-
The gap. Its schedule for that mode is the least optimised of the alternatives, and the paper measures it as sometimes worse than keeping attention on the accelerator — the idea is right and the schedule around it is not.
-
Overlap host attention with accelerator computation. A separate line of work does schedule host-side attention concurrently with accelerator work, which is the correct structure.
-
The gap. It does not target memory-constrained machines, so it never considers weight transfer at all. The hardest scheduling problem here is the one it does not have.
-
Fit a cost model to measured runs. The prior system searches for a good policy using a model fitted to profiling data.
-
The gap. The fitting is expensive, taking hours or days, and must be redone per model and machine. An analytical model built from operation counts and measured peak bandwidths gives usable answers immediately.
How the schedule gets rebuilt
From an offloaded model to a saturated pipeline
-
-
new
-
new
-
new
-
new
-
new
The idea
Extend the analysis before optimising the system. The paper’s structure is unusual in that a third of it is analysis with no system in sight. The roofline model is generalised to a machine with several levels of memory, each paired with a processor, and this immediately produces questions the original could not express — for a given computation, which processor should perform it, and which memory is the one actually limiting the system. Applied to a mixture-of-experts layer during generation, the answer for attention is the host. Attention at this stage has very low operational intensity, meaning it moves a great deal of state per arithmetic operation, so an accelerator’s superior compute is worth little while the cost of shipping that state across the interconnect is substantial. Running it where the state already sits removes the transfer entirely. Deciding placement from the arithmetic character of each phase rather than from where the fast hardware happens to be is the same reasoning as the prefill-decode overlap work covered elsewhere on this site, which pairs a compute-hungry phase with a bandwidth-hungry one precisely because their appetites differ.
Two transfers with opposite urgency should not share a queue. With attention state no longer moving, what remains crossing the interconnect is weights for the next layer and small intermediate results for the current one, and these behave nothing alike. An activation is needed immediately — the computation that consumes it is next in line, so any delay is a stall. A layer’s weights are not needed until every micro-batch of the current layer has been processed, so they have slack. But the weights are much larger and take much longer to move, so issuing them as one block leaves the small urgent transfer stuck behind them, and the pipeline develops a bubble at every layer boundary. The fix is to page the weights: cut the transfer into as many chunks as there are micro-batches and interleave them with the urgent traffic. Nothing is transferred that was not transferred before, and nothing arrives late. Splitting a large low-priority transfer so it stops obstructing a small urgent one is the scheduling instinct behind the offloaded fine-tuning work covered elsewhere on this site, which hides its swaps behind computation on a similarly starved interconnect.
A model built from arithmetic beats one fitted from measurements. The policy — batch size, micro-batch size, where each computation runs, how much weight to keep resident — is chosen by searching against a performance model, and the paper’s is analytical. Operation counts and byte counts come from the model architecture, peak bandwidths and peak throughput come from a one-time microbenchmark, and the estimate follows. This is less accurate in absolute terms than fitting to measured runs, and the authors concede it, but they argue correctly that a policy search only needs to rank two candidates rather than predict either one’s runtime. The payoff is that there is no fitting phase costing hours or days per model and machine. The ablation shows the policy alone is worth a great deal: feeding it to the competing system, without any of the scheduling work, raises that system’s throughput by 1.77 times, and with a larger batch by 2.17. Trading absolute accuracy for a model cheap enough to use everywhere is the same judgement as the mixture-of-experts training work covered elsewhere on this site, which fits two-parameter cost models per operation and schedules from them.
Why a big transfer must be cut up
How weight paging removes the bubble at every layer boundary
Following the interconnect through one layer of generation, with several micro-batches in flight.
Step 1 of 5
Two kinds of traffic want the same link
During a layer, small intermediate results move between host and accelerator for each micro-batch, and the weights for the next layer must arrive before that layer starts.
They differ in every respect that matters to a scheduler. The activations are small and urgent; the weights are large and have slack. A queue that treats them alike will serve them in the wrong order.
Why it works
- Attention in generation is bandwidth-hungry and compute-light. Its operational intensity is very low → an accelerator’s peak arithmetic buys almost nothing while shipping the state costs a great deal → running it on the host removes traffic entirely rather than merely overlapping it.
- Freed bandwidth goes where it is scarce. Attention state no longer crossing the interconnect → the link is available for weight transfer, which is the transfer that cannot be avoided → the amortisation that offloading depends on gets cheaper.
- Urgency, not size, should order the queue. Paging the weights so each chunk moves during one micro-batch → a small urgent activation waits behind one page instead of a whole layer → the bubble at each layer boundary disappears without moving any additional bytes.
- Ranking is easier than predicting. An analytical model from operation counts and measured peak bandwidths → good enough to compare two policies even where absolute latency is off → no fitting phase, where the alternative costs hours or days per model and machine.
- Less host memory for the same ceiling. Better interconnect utilisation → the weight transfer is amortised at a smaller batch → the throughput upper bound is reached with two to three times less host memory, which is the resource that was binding.
The evaluation
| Property | What was measured | The result | What it tells you |
|---|---|---|---|
| Throughput on a single small accelerator | |||
| Host memory required | |||
| Where the gains come from | |||
| Scaling across several cards | |||
| Breadth |
Check yourself
Check your understanding
1. Why does offloading make host memory the thing that limits throughput?
2. What does the extended roofline analysis conclude about where attention should run, and why?
3. Why does issuing a layer's weights as one transfer create a bubble, and how does paging fix it?
4. The performance model uses theoretical operation counts rather than being fitted to measured runs. Predict why that is defensible.
Sources
MoE-Lightning: High-Throughput MoE Inference on Memory-constrained GPUs. Shiyi Cao, Shu Liu, Tyler Griggs, Peter Schafhalter, Xiaoxuan Liu, Ying Sheng, Joseph E. Gonzalez, Matei Zaharia, Ion Stoica. ASPLOS 2025, Volume 1, pages 715 to 730. DOI 10.1145/3669940.3707267. Full text read from the authors’ preprint, arXiv 2411.11217.