Running 125B MoE on Two 32 GB Cards - vLLM Optimization Edition

· 22 min · llm, vllm, qwen, gpu, cuda

In the previous article, I brought Qwen3.8-Flash-Next (a 125B-A6B MoE plus a 51B n-gram table) to 95-99 tok/s single-stream decode in vLLM on three RTX 3090s. After that, a user with two RTX 5090s asked about the published quantized model, in short: “Does it run on two 32 GB cards? If it does, how many experts go to RAM, and how much slower is decode?”

This article describes the 2x32GB release configuration I made as the answer. The weights do not fit in two 32 GB cards, so the least used routed experts stay in host RAM and are read over PCIe only when the router picks one.

I do not have a 32 GB card. Every number in this article was measured on RTX 3090s. The measurement conditions:

  • Most measurements were on three 3090s (PP=3), with the experts chosen for the 2-card configuration (3 GiB per card, 6 GiB total) in host RAM. In single-stream decode the PP stages run one after another, so the time of one token is the sum over all stages. Which experts are in host RAM does not depend on how the layers are split, so the number and time of host reads per token are the same as in the 2-card configuration.
  • All three 3090s are on PCIe 4.0 x16. That is the same bandwidth as a 5090 in an x8 slot (PCIe 5.0 x8).
  • I also ran the model on two 3090s (PP=2). To fit in 24 GB, that needs 10.5 GiB per card in host RAM (section 6).
  • I did not measure the compute and memory bandwidth of the 5090 itself, the effect of PP going from 3 stages to 2, or whether the model actually fits in 32 GB (section 8).

The results first.

3x3090 (no offload)2x32GB first version (offload-02)2x32GB now (offload-03)
Decode, 1 stream, average of 16 new prompts97.7 tok/s81.8 tok/s91.5 tok/s
Decode, 1 stream, code / English story / Japanese83-84 / 85 / 8193-97 / 94-95 / 93-95
Prefill 122k (2 runs)5,154 / 5,196 tok/s5,083 / 5,292 tok/s
Startup time750 s335 s
Routed experts in host RAMnone3 GiB per card3 GiB per card

All on three 3090s (PP=3), --max-num-batched-tokens 2048, --kv-cache-memory-bytes 550000000, --max-num-seqs 2, prefix caching on. The first-version column is the current image with VLLM_EXPERT_ONE_TABLE=0, which runs the offload-02 path (with the fix from section 5). All decode runs use temperature 0, ignore_eos, thinking off and 512 output tokens. “16 new prompts” is the average over prompts that alternate between subjects (code, Japanese, English story, math, other), each sent once to a freshly started server. “Code / English story / Japanese” sends the same prompt twice in a row.

On top of the 12 patches of the previous release, there are 3 new ones.

PatchWhat it doesEffect
offload-01-expert-host-layersKeeps all routed experts of chosen layers in host RAM1.2 GiB VRAM per layer, about 1 ms/token of decode
offload-02-cold-expertsKeeps only the experts listed in a file in host RAMOnly the least used experts leave the GPU
offload-03-expert-one-tableOne expert table that spans VRAM and host RAM, and moving experts at run timeDecode 81.8 -> 91.5 tok/s

The sections below follow the order of the work.

1. What does not fit

After mem-01 from last time, the weights are 62.27 GiB. With PP=2 and 24 layers per card, that is 31.36 / 31.21 GiB per card.

Besides the weights, a card needs about 1.6 GiB (including 0.34 GiB of KV), measured on 3090s. A 2048-token prefill chunk adds another 0.4-0.7 GiB (section 3.5). A 5090 has 32,607 MiB. If about 31.4 GiB per card is usable, at least about 2.2 GiB per card has to leave VRAM.

I chose the routed experts for this.

  • There are 48 layers x 512 routed experts, 58.01 GiB, 93% of the weights.
  • One expert is 2.417 MiB in INT4 g128 (w13 [1280, 2560] + w2 [2560, 640] + BF16 scales).
  • One token selects 48 layers x top-10 = 480 experts, 2% of all experts.

An expert that is not selected is not used for that token. If the least used experts are in host RAM, PCIe reads happen only when one of them is selected.

As before, the QSA K/V of the KV cache is in host RAM. With PP=2 each rank has 6 QSA layers, and the attention block size is 8,096 tokens (12,144 with 3 cards). The release setting --kv-cache-memory-bytes 360000000 gives 112 blocks, 262,144 x 1.04 requests, and 20.8 GiB of pinned host memory.

2. Whole layers in host RAM (offload-01)

The first implementation keeps all routed experts of chosen layers in host RAM. The layer numbers go in VLLM_EXPERT_HOST_LAYERS.

  • The expert weights are allocated in pinned host memory when the block is built, and the GPU reads them through a UVA view. The MoE kernels are unchanged, and only the selected experts cross PCIe.
  • The Marlin repack after loading returns new GPU tensors, so each result is moved back to the host right after the repack. Only one layer is on the GPU at a time, so a model whose weights do not fit in VRAM still loads.

Each layer frees 1.2 GiB of VRAM. But decode reads about 24 MB per token per layer (the top-10 experts) over PCIe, about 1 ms per token per layer. Prefill reads nearly the whole layer (1.2 GiB) per chunk.

3 x 3090, chunk 512Code decodePrefill 72k
No offload99-100 tok/s5,414 / 5,429 tok/s
6 layers in host RAM61-62 tok/s2,165 tok/s

Six layers cost about 40% of decode speed. With 18 layers in host RAM on two 3090s, decode was 35 tok/s. The release keeps this mode, but it is not the default.

3. Only the least used experts in host RAM (offload-02)

Choosing per expert instead of per layer lets only the rarely selected experts go to host RAM. Below, the set of experts in host RAM is called the cold set (VLLM_EXPERT_COLD_FILE / cold-experts/ in the release).

3.1 How uneven the expert selection is

I added a counting hook to vLLM that counts the gate’s top-10 per layer, and ran the following data through it. The experts selected at a position depend only on the context up to that position. So running an already generated token sequence through prefill gives the same selections as decode did.

SetContentsSize
codeCode from real repositories (several languages), 8k x 32260k tokens
LCB-A / LCB-BLiveCodeBench generations (reasoning and answer), split into two halves by alternating problems1.13M / 1.26M tokens
G1 / G2Generations (with reasoning) for 32 requests to implement, debug, review, test and refactor code in 8 languages24.5k decode tokens each

With 3.5 GiB (1,482 experts, 6.0% of all) in host RAM, I computed h, the share of selections on the evaluation set that fall on a host expert.

How the cold set is chosenEvaluated on G2 (decode)Evaluated on LCB-B
Uniformly at random6.0%5.5%
From LCB-A only3.19%0.16% (same kind)
From code only5.51%4.15%
From LCB-A and code with equal weight0.55%
+ G10.16%

A cold set chosen from one kind of data does almost nothing for a different kind of task. Competitive programming Python with reasoning text and code from real repositories use different experts. Mixing the two keeps only experts that are rarely used by both, and on G2, which was not used for the choice, h is about 1/11 of the uniform case.

The unevenness itself is large. On LCB-B, the bottom 10% of experts in a layer receive 0.29% of the selections (median over layers).

3.2 The cost of one host read

On three 3090s with the production configuration, I made an experimental build that calls the fused decode Marlin MoE twice, once for GPU experts and once for host experts, and measured the cost per read. All weights stayed on the GPU, and the host experts were read from a UVA view of pinned memory. The inputs were 5 LiveCodeBench problems (chat) and 3 continuations of 30k of code, with temperature 0, ignore_eos and 768 output tokens.

Configuration (3.5 GiB in host RAM)Short prompt30kHost reads/token
Production (no split)96.9 tok/s94.6 tok/s
Two calls, host experts also on GPU (control)89.988.1
Cold set chosen from LCB-A + code + G183.883.05.5 (h 1.16%)
Uniformly random cold set69.668.629.1 (h 6.06%)

Dividing the ITL difference from the control by the number of reads gives 111-136 µs per host read (19-23 GB/s effective). With the uniform cold set, reads spread evenly over the tokens, and p90 stays around p50 + 0.8 ms. With the chosen cold set, reads concentrate on a few tokens, and p90 goes up to 15.8 ms.

Splitting into two calls alone costs +0.8 ms. This experiment ran the reordering and the Marlin call twice in each of the 48 layers. Section 4 removes this cost.

3.3 Mixing subjects in the choice

The cold set in 3.1 was chosen from code only. With it, decode of an English story slowed down a lot.

3 x 3090, chunk 512CodeEnglish storyPrefill 72k
No offload99-100 tok/s95-96 tok/s5,414 / 5,429 tok/s
Cold set chosen from code only (3.5 GiB total)91-9357-613,934

So I counted data other than code and chose the cold set again. The added data: Japanese (160 samples of the quantization calibration data, 620k tokens), arXiv papers and science documents (32 each, 260k tokens each), and C1, generations for miscellaneous chats in English and Japanese (16 chats, 18k decode tokens). The seven sets LCB-A, code, G1, Japanese, arXiv, science and C1 are each normalized to a sum of 1 and added, and experts are chosen from the lowest sum up. The choice fills the budget of each card separately: card 1 (layers 0-23) and card 2 (layers 24-47) for PP=2.

For evaluation I used G2 (coding decode) and C2 (decode of miscellaneous chats in English and Japanese), neither used for the choice. Host reads per token:

Per cardCode only, G2 / C2General, G2 / C2Uniformly random
1.5 GiB0.6 / 24.92.4 / 4.124.8
2.5 GiB1.5 / 43.85.4 / 9.741.4
3 GiB2.2 / 51.37.1 / 12.949.7

The code-only cold set reads more than a uniformly random one on C2. The experts that code does not use include experts that chat uses often. The general cold set reads more than the code-only one on code, but on both evaluations it is about 1/4 to 1/10 of the uniform case. The release has this general cold set in five sizes: 1.5, 2, 2.5, 3 and 4 GiB per card.

3.4 Implementation

  • The GPU-side RoutedExperts has fewer rows, and its expert map sets the host experts to -1 (the same form as EP).
  • A second module of the same class holds the host experts. It is not a child of the model, so parameter names and the loaded-weight bookkeeping do not change. The loader wraps weight_loader and sends host expert ids to the second module.
  • In the fused decode path, the router kernel writes two block tables, one for the GPU experts and one for the host experts, and Marlin is called twice.

A point to note: Marlin MoE does not skip blocks with expert id -1, and reads out of range. A unit test reproduced the illegal address. The first implementation called Marlin twice on one block table containing -1, and crashed here. A block table must not contain experts that its module does not hold.

On output, the logprob difference when scoring the same generations (generated part) was 0.005-0.024, about the same as between two servers without the offload (0.004-0.021).

3.5 Prefill and chunk size

In prefill, each chunk reads the host experts it uses over PCIe. A larger chunk reads less per prompt token.

3 x 3090, general cold setChunkPrefill 72k
No offload5125,414 / 5,429 tok/s
2.5 GiB per card5123,395 / 3,427
2.5 GiB per card20485,561 / 5,661
3 GiB per card20485,340 / 5,378

With chunk 2048, prefill is about the same as without the offload (chunk 512). In exchange, VRAM goes up by 0.4-0.7 GiB per card. The estimate in section 1 includes this. The release uses chunk 2048 and 3 GiB per card.

3.6 Results of offload-02

On three 3090s with the general cold set at 3 GiB per card and chunk 2048, decode of code / English story / Japanese was 80-83 / 83-84 / 78-80 tok/s. This was the first release. On the 16 new prompts under the same conditions, it was 81.9 tok/s against 97.7 tok/s without the offload, about 16% lower. Section 4 reduces this.

4. Decode: one table and moving experts at run time (offload-03)

4.1 What was slow in offload-02

On the 16 new prompts, no offload was 97.7 tok/s and offload-02 was 81.9 tok/s. The difference is about 2 ms per token, and it has two parts.

  • Extra calls. For each layer, offload-02 runs another set of Marlin w13, activation and w2 for the host experts. Even in a step with no host read, that adds 144 kernels per token.
  • Host reads. The general cold set was chosen across subjects, so every conversation reads from the host at some rate.

4.2 One table that spans VRAM and host RAM

If the first rows of one expert tensor are in VRAM and the last rows are in pinned host RAM, one Marlin call is enough. I implemented this with the CUDA virtual memory management (VMM) API, calling libcuda directly through ctypes.

  1. cuMemAddressReserve reserves one contiguous virtual address range.
  2. cuMemCreate creates physical memory on the device, and physical memory on the host with CU_MEM_LOCATION_TYPE_HOST_NUMA.
  3. cuMemMap maps the two next to each other.

The GPU reads the host part at 25.1 GB/s, the same as the earlier pinned-memory UVA.

  • The GPU experts go in the first rows and the host experts in the last rows, and the rows of the router’s gate are reordered the same way. The id the router outputs is the row of the table, so vLLM’s MoE path needs no change, no expert map is needed, and Marlin runs once.
  • Loading uses the offload-02 code (two modules) as is, and the two are joined into one table after the repack. moe_quant_config holds references to the scales, so _setup_kernel is called again.
  • The VRAM part is rounded down to 2 MiB, so VRAM use does not grow compared with offload-02 (within +40 MiB in nvidia-smi, same KV capacity).
  • With VLLM_EXPERT_ONE_TABLE=0, or when the driver refuses the host mapping (tried with 4 MiB at startup), the server runs the offload-02 path.
  • Because the gate rows are reordered, the expert selected when two BF16 logits are equal can differ from the original (on a tie, the lower id wins). The effect is smaller than vLLM’s run-to-run nondeterminism.

Two points to note. The Model loading took line in vLLM’s startup log does not count memory allocated through VMM, so it shows less than the real amount. Also, the host memory allocated through VMM is not mapped into the CPU’s virtual address space, so it does not appear in the process’s RssShmem (it does appear in used of free). The automatic KV sizing uses measured free VRAM, so the first point does not affect it.

At this stage decode went from 81.9 to 86.6 tok/s. About 0.66 ms/token was the extra calls.

4.3 Moving experts at run time: simulation

What remains is the host reads. The general cold set is chosen so that every subject reads from the host at some rate. But within one conversation, the experts in use are uneven. Moving the experts a conversation uses to the GPU at run time, and moving unused ones to the host, reduces the reads.

I checked the effect with a simulation before implementing it. I ran 26 sessions (code, explanations of existing code, Japanese, English stories, math, other; about 2k tokens each) plus 8 real LiveCodeBench generations, 137k tokens in total, through prefill and recorded the top-10 of each layer in token order.

  • Fixed general cold set (3 GiB per card): 8.08 host reads per token (Japanese 13.9, other 17.1, code 10.5).
  • With moves (every 8 steps, use counts decayed with a half-life of 512 steps, move when the difference exceeds 2, one pair per layer, prefill not counted): 1.71 reads + 0.36 moves in session order, 2.12 reads + 0.45 moves with the sessions shuffled.
  • Counting prefill tokens too gave 2.20 -> 2.18 reads, no difference. So only the decode router kernel counts.
  • Changing the threshold, half-life and interval, the total with one move counted as two reads stayed at 2.9-3.7 read equivalents, a small change. The defaults above are near the minimum.

4.4 Moving experts at run time: implementation

  • The router kernel of the fused decode path (decode-05 from last time) counts each selected row with atomic_add. It runs inside the CUDA graph.
  • model_state.prepare_inputs calls expert_table_step() every step, and every 8 steps it queues two kernels on the same stream as the forward. The decision kernel decays the use counts per layer and picks the most used host expert and the least used GPU expert. The swap kernel exchanges the four tensors and the gate row for all layers at once, through a pointer table.
  • Nothing is read back to the host. prepare_inputs is not called during CUDA graph capture (capture calls prepare_dummy_inputs).

To check correctness, I took a fingerprint of each row (the sum of the 5 tensors) right after startup, and confirmed that after thousands of rows had moved, the rows stayed together in all 48 layers (the check is not in the release). A build that runs only the decision kernel and never swaps did 86.6 tok/s, the same as the one-table-only build. The cost of the decision kernel is too small to measure.

4.5 Results

Three 3090s, general cold set at 3 GiB per card, chunk 2048, decode on the 16 new prompts:

Configurationtok/sHost reads/tokenMoves/token
No offload97.7
offload-02 (two calls)81.9
One table only (VLLM_EXPERT_SWAP_EVERY=0)86.612.10
One table + moves (offload-03 default)90.0-91.84.30.90

Counts are the sum over the 3 ranks, taken from counters in the router kernel. Working back from the ITL difference to no offload, one host read is about 0.108 ms (1.31 ms / 12.1 reads for the one-table-only row), and one move is about 0.2 ms. The cost of one read matches the experiment in 3.2 (111-136 µs).

The moves help more the longer a conversation stays on one subject. On a development build (three 3090s), I sent each of a code, an English story and a Japanese prompt twice, and repeated this round twice.

3 x 3090, 3 GiB per cardRound 1 (1st / 2nd)Round 2 (1st / 2nd)
offload-02 (no moves)83.7 / 84.2, 84.5 / 84.4, 80.8 / 80.9
One table + moves87.0 / 96.3, 90.6 / 94.3, 89.0 / 94.896.3 / 97.7, 94.9 / 95.4, 96.0 / 96.7

Each cell is code, English story, Japanese, in that order. Without moves (offload-02), the 1st and 2nd runs differ by less than 1 tok/s. With moves, the 2nd run onward is 94-98 tok/s, nearly the same as without the offload.

Final check with the release image (three 3090s, release cold set and chunk):

offload-03 (default)VLLM_EXPERT_ONE_TABLE=0 (offload-02 path)
16 new prompts91.5 tok/s81.8 tok/s
Code / English story / Japanese93-97 / 94-95 / 93-9583-84 / 85 / 81
Prefill 122k (needle correct)5,083 / 5,292 tok/s5,154 / 5,196 tok/s
Startup time335 s750 s
Worker host memory (RssAnon + RssShmem)21.6 / 19.6 / 20.8 GiB22.3 / 20.1 / 21.9 GiB

The logprob difference was 0.006-0.024, about the same as between two servers without the offload (0.004-0.021). The 90k needle was correct with 1 and 2 streams, and image input worked.

5. offload-02 held the host experts in RAM twice

On two 3090s (10.5 GiB per card in host RAM), startup took 27 minutes, and worker host memory (RssAnon + RssShmem) was 48.8 / 42.4 GiB.

Printing the pinned allocator’s statistics per layer showed that a copy of the host experts in their load-time format was not freed after the repack. When building the model, vLLM copies and keeps parameter attributes through record_metadata_for_reloading -> capture_layer_to_meta. These include offload-02’s weight_loader (a closure), and the closure held a reference to the load-time tensors through its argument _p=cold_p. With or without the one table, the host experts used twice their size in RAM (+6 GiB at the release settings, +21 GiB at 10.5 GiB per card).

The fix: the closure holds no tensor and looks it up by name (getattr(cold, _n)). It also calls torch._C._host_emptyCache() after each layer. I fixed the offload-02 patch itself and rebuilt offload-03 on top of it.

BeforeAfter
Startup, two 3090s (10.5 GiB per card)about 27 min275 s
Worker host memory, same48.8 / 42.4 GiB29.9 / 28.6 GiB
Worker host memory, three 3090s (3 GiB per card)about 25 GiB21.6 / 19.6 / 20.8 GiB

6. Measurements on two 3090s (PP=2)

To check the 2-card configuration, I ran PP=2 on two 3090s. To fit in 24 GB, the general cold set is 10.5 GiB per card, and the chunk is 512.

2 x 3090, 10.5 GiB per card, chunk 512offload-02 (before the fix)offload-03
Code55 tok/s68-83 tok/s
English story42-4479-87
Japanese71-86
16 new prompts (256 output tokens)66.1

The offload-02 column used --kv-cache-memory-bytes 360000000, the offload-03 column 550000000. For offload-03, code / English story / Japanese differ a lot between the 1st and 2nd run. As in 4.5, decode gets faster as experts move.

The 90k needle was correct with 1 and 2 streams. But the 89,716-token prompt with 1 stream took 92.3 s to answer. On three 3090s with 3 GiB per card and chunk 2048, the same prompt takes 16.5 s. With more experts in host RAM and a smaller chunk, prefill reads more from the host. A 32 GB card is estimated to run with 3 GiB per card and chunk 2048 (section 1), so this configuration is not a guide to 5090 performance. It checks that the 2-card configuration works correctly.

VLLM_EXPERT_ONE_TABLE=0 (the offload-02 path) took 1,831 s to start, longer than the measurement script waits, so there are no numbers for it.

7. Approaches tested and not adopted

7.1 Computing the host experts on the CPU

Instead of reading the host experts over PCIe, the CPU could compute them.

One expert (INT4 g128)Time
CPU (_weight_int4pack_mm_for_cpu, Threadripper 3945WX, AVX2), 1 token, 12 threads89 µs
Same, 6 threads137 µs
Same, 8 tokens338 µs
Read over PCIe (UVA), computed on the GPU110-136 µs (independent of token count)

For 1 token the CPU is slightly faster, but I did not adopt it, for these reasons.

  • A synchronous GPU -> CPU -> GPU round trip takes 32 µs. Decode runs as a FULL cudagraph, so each layer would need a host node that waits for the CPU, and every token would pay 48 round trips, including layers where no host expert was selected. That is about the same as the cost of host reads with the general 3 GiB per card (7-13 reads x 0.12 ms).
  • In prefill, PCIe reads once per chunk, but the CPU gets slower in proportion to the number of tokens.
  • The 5090 is on PCIe 5.0, which favors PCIe even more. The CPU’s DRAM bandwidth (79 GB/s measured) is only 3-4 times PCIe, and the computation uses that same bandwidth.

7.2 Moving experts asynchronously on a copy engine

Moving the experts asynchronously on a copy engine in a separate stream would hide the time of the moves. But it needs spare rows and a router that goes through an expert map, which brings an expert map into the prefill path. The gain would be 0.9 moves x 0.2 ms ≈ 0.18 ms/token (about 1.5 tok/s), so I did not implement it.

Overlapping host reads with GPU work already happens with the one table, inside the same kernel. The cost of one read is set by PCIe bandwidth, so it stays at about 0.11 ms.

8. Release settings and what is not verified

The main settings of vllm-patch/2x32GB/ in the release:

--pipeline-parallel-size 2 --tensor-parallel-size 1
--max-model-len 262144 --max-num-seqs 2
--max-num-batched-tokens 2048
--kv-cache-memory-bytes 360000000        # 262,144 x 1.04, 20.8 GiB pinned for QSA
--enable-prefix-caching --mamba-cache-mode align

VLLM_EXPERT_COLD_FILE=/cold/3GiB-per-card.json   # 3 GiB per card in host RAM, 6 GiB total

If the server runs out of GPU memory at start, use the next larger cold set (4 GiB). If more than about 1.5 GiB per card is free after start, a smaller cold set gives faster decode.

Not verified on 32 GB cards:

  • Whether the model actually fits in 32 GB. The estimate in section 1 leaves about 0.8 GiB with 3 GiB per card.
  • Decode speed on a 5090 without the offload. Memory bandwidth is about 1.9 times the 3090’s, and PP goes from 3 stages to 2.
  • The cost of one host read. I estimate about 0.11 ms on PCIe 5.0 x8, the same as on the 3090s, and about half that on PCIe 5.0 x16.
  • Whether the Triton patches from the earlier articles work on sm_120.

Summary

  • The 62.3 GiB of weights do not fit in two 32 GB cards. Of the routed experts, which are 93% of the weights, the least used stay in host RAM, 3 GiB per card, and are read over PCIe only when selected.
  • A cold set chosen from one kind of data increases reads on other subjects (chosen from code only, English story decode went from 95-96 to 57-61 tok/s). The release has a general cold set chosen from seven normalized data sets.
  • One host read costs about 0.11 ms (PCIe 4.0 x16). Putting whole layers in host RAM costs about 40% of decode speed with 6 layers, so it is not the default.
  • offload-03: CUDA VMM builds one expert table that spans VRAM and host RAM, and Marlin runs once (81.9 -> 86.6 tok/s). Experts move at run time based on the use counts the router keeps, which cuts host reads from 12.1 to 4.3 per token (-> 90.0-91.8 tok/s).
  • The first offload-02 held the host experts in RAM twice, through the weight_loader closure that vLLM copies and keeps. After the fix, the 2-card startup went from about 27 minutes to 275 s.
  • Everything was measured on RTX 3090s and is not verified on 32 GB cards. On three 3090s with the same cold set, decode was 91.5 tok/s against 97.7 tok/s without the offload.

The patches and compose files are in vllm-patch/2x32GB/ below. The target is vllm/vllm-openai 0.29.1rc1.dev47+gdc36fcce9. https://huggingface.co/Minachist/Qwen3.8-Flash-Next-INT4-Mixed-AutoRound

この記事の日本語版: 125B MoE を 32GB 2 枚で運用する vLLM 最適化編