Running 125B MoE on Three 3090s - vLLM Edition

In the previous article, I explained the properties of Qwen Sparse Attention. Simply put, the KV read in a single decode becomes constant with indexer_budget and does not depend on the context length, so the KV cache can be offloaded to host RAM, and context length is no longer a bottleneck. In the previous article, I implemented it in ExLlamaV3 and reached 3.14M, so this time I reproduce it in vLLM.

What I usually run as a service is vLLM, structured so that clients make requests to an OpenAI-compatible API.

On paper, it looks like it would work properly.

In fact, as a result, I achieved the following.

1. Shape of This Model

From text_config.

hidden_size 2560   /  residual between blocks is 4x at 10240 (hc_count=4)
num_hidden_layers 48   layer_types = [linear x3, full] x 12
num_experts 512   num_experts_per_tok 10   moe_intermediate_size 640
num_attention_heads 24   num_key_value_heads 2   head_dim 256
indexer_budget 2048   indexer_compress_ratio 4   indexer_head_dim 128
ngram_size 3   split_ngram_parts 128   ple_layer_ids [2]
vocab_size 248320   max_position_embeddings 262144

As characteristics of the model,

Out of 48 layers, 36 layers do not have a KV cache. It is a linear attention called gated delta-net, and memory does not increase even if context grows. Only the 12 layers of full attention (QSA) have K/V.

The width of residual is 10240, not 2560. Hyper-connection has 4 residuals between blocks. Memory estimation per block must all be calculated with 10240.

A 95.37 GiB embedding table hangs on layer 2. It is not a vocab embedding, but a separate n-gram table with 128 shards of [2500012, 160]. It has a structure of looking up once per token.

2. Deciding the Bit Allocation

Subtracting simple calculations for KV, activations, and CUDA graphs from 69.2 GiB of VRAM leaves about 62 GiB that can be used for weights.

Because I did not want to write the allocation based on speculation, I referred to existing quantized models. EXL3’s 4.05bpw_h6_ng6 has bits_per_weight for 74,395 modules in quantization_config.json. Unsloth’s GGUF UD-IQ4_XS has types for 1,224 tensors in the header. What was found by comparing the two.

1. There is no implementation that sets dense layers to 4-bit. EXL3 uses 6, and Unsloth uses 8.5. My original plan had linear_attn / self_attn / shared_expert as INT5, which was a lower bit count than both. The difference is 0.31 GiB, and there is no reason to do something different from the existing two for this.
2. Neither quantizes the router. Making a mistake in selecting experts is scary. In fact, I have never seen anyone quantize here.
3. In hyper-connection, the degree of quantization differs between GGUF and EXL3. EXL3 uses FP16, and GGUF uses Q8_0.
4. There is also not a single implementation that keeps the n-gram table in BF16. EXL3 has 36.36 GiB in 6-bit trellis, and GGUF has 26.82 GiB in IQ4_NL.

Adopted allocation.

Experts Are INT4, and the Rest Is Mostly Rounding Error

Routed experts account for 120.8B out of the 125.75B resident on GPU, which is 96%, and this mostly determines the model size. INT8 experts are out of the question at 114 GiB, so INT4 is fixed, and the rest needs to be considered.

Making it INT4 g64 would raise it by +0.125 bpw, but increases by +1.76 GiB. group_size needs to be decided before starting the actual quantization, and if it turns out to be insufficient after running 48 blocks, redoing it will take several days, so margin was secured by keeping it small.

What Actually Worked Was Hyper-connection

input_mix_weight_down [320, 10240] and up [10240, 320] are low-rank matrices that compress residual to rank 320 once before mixing. Not only do they account for only 1.9% of the saved weights, but the quantization error on the low-rank side also affects the entire residual, and because 27B did not have this structure and intuition did not work, I had decided to leave them in BF16 untouched in the original plan.

Dropping to INT8 reduces VRAM from 1.193 to 0.622 GiB, and furthermore reduces read per token by 0.63 GB. If left in BF16, even though it is only 1.9% of the saved weights, it dominates 23% of the read per token. Raising dense from INT5 to INT6 increases by +0.36 GB/token, but well, it is faster than the original plan.

input_mix_weight_up has in_features=320, which becomes 2.5 groups with g128 and cannot be divided evenly.

Why PP Instead of TP

hidden_size 2560, linear_num_key_heads 16, num_experts 512, shared_expert_intermediate_size 640, and moe_intermediate_size 640 are none divisible by 3. Therefore, make it PP=3 / TP=1.

How Was the Quality

The expectation was mean KLD 0.010 to 0.014. AutoRound is not EXL3 trellis but a family of scaled integer quantization, so this figure was derived from a line connecting GGUF and NVFP4 on the referenced curve.

The measured value with the same trace and the same measurement script is 0.014888 nats, body 4.178 bpw.

It beat GGUF of the same bitwidth by 11%, and lost to EXL3 by 2.2 times. This difference is a difference in the model storage format itself, not a difference in bit allocation. It is not a difference that can be closed by moving bits around, and I proceeded knowing that. The merit here is not in KLD, but rather in latency and throughput.

3. Actually Quantizing

In the plan, it was DDP=3, 1 rank 1 GPU, iters=1000, but none was practical.

DDP does not work. One block is 2.58B, with 5.2 GB of BF16 weights, 10.3 GB of fp32 rounding parameters, and 10.3 GB of their gradients. Before putting activations on it, it is 25.8 GB, which does not fit in a 24 GB card. Therefore, it is necessary to load one process across 3 cards with device_map=auto.

iters is 200, not 1000. The measured time for 1 iter was 12.9 seconds. With 200, it is 34.6 hours for 48 blocks, and with 1000, it is 172 hours. Exceeding a week just for quantization is indeed unacceptable. In the actual production run, it took 56 hours and 37 minutes.

nsamples is 384, not 768. The activation cache is nsamples x seqlen x hidden x 2 bytes x 2 for input and output, and this hidden is the residual width 10240, not 2560. With 768, it is 64 GB by calculation, and about 161 GB at the peak extrapolated from actual measurements, which is impossible on a 128 GB server.

Other issues.

Qwen4ExpTextRMSNorm.group_size is deleted. apply_plan_to_model() of AutoRound traverses all modules and executes delattr on attributes that match field names of QuantizationScheme. In that field, there is group_size. On the other hand, RMSNorm of Qwen4Exp holds group_size as an architecture parameter, which is a value for dividing the 10240-width residual into four 2560 parts to normalize. If left deleted, AttributeError appears at block 0. If it becomes None, it changes to normalizing a single 10240 width, and the resulting model becomes an entirely different model. It is necessary to place a property in the class to escape the actual entity to another name, and count it with a checking script.

The memory allocation script does not recognize experts. AutoRound determines experts like this.

is_moe_expert = "expert" in name.lower() and isinstance(parent, nn.ModuleList)

However, the container that AutoRound itself creates when expanding experts one by one is _ExpertContainer, not nn.ModuleList. As a result, not a single expert is recognized, and output activations of 512 experts are added under the assumption that all occur simultaneously, resulting in “layer output 30.67 GB” and “card 0 requires 62.34 GB”, card 0 is entirely dropped from allocation, and blocks are placed on card 1 and 2, causing OOM. On top of that, the additional_memory factor 7 for placing remaining parts is only added to card 0, compounding the problem of unequal allocation like 10 to 45 to 45.

Reference forward accumulates outputs of all samples in GPU. The implementation stacks batch outputs in a list and moves them to CPU after exiting the loop. With residual 10240 width x 384 samples, 384 x 2048 x 10240 x 2 = 16.1 GiB is placed on card 0. Furthermore, the allocation when low_gpu_mem_usage=True sets this term to 0, so it is not reserved.

Block 1 carries a 95 GiB table. ple_layer_ids is [2], but the actual entity on the checkpoint is in layers.1.ple.*. Because forward of this block requires referencing n-grams, before production, 128 shards were expanded into a flat memmap on NVMe (5.7 minutes), and replaced with embeddings read directly from disk and a per-sample gather cache. Without the cache, more than 100 million random reads occur in a single run.

Post-processing After Export

• All 4 groups of config_groups came out with targets set to ["Linear"]. Because vLLM outputs to a dict as target_scheme_map[target] = {...}, the same key is overwritten 4 times and only the last group remains. When calling CompressedTensorsConfig inside the image directly and measuring, out of representative 26 layers including 73,728 INT4 experts, 16 layers were overwritten as INT8 g128. The weights themselves were correctly quantized, and only the metadata was wrong, so rewriting it was sufficient.

4. Stock vLLM Does Not Even Boot in the First Place

It took 5 fixes before a prompt passed once, and 4 more before it became practical.

A - Cannot Read Quantized embed_tokens / lm_head / Hyper-connection

Just as written in the article on INT5-7, qwen4_exp passes neither quant_config nor prefix when creating VocabParallelEmbedding. Therefore, quantized vocab tensors cannot be loaded.

Hyper-connection also has a similar problem. vLLM combines the low-rank mixing weights and block-inject weights into a single Linear called input_mix_weight_down_block_inject. In this checkpoint, the former is INT8 and the latter is BF16, and since only one quantization scheme can be tied to a single Linear, both cannot be handled, so it is necessary to split them into two with a patch.

C - Rejection of PP>1 Exists in Two Places

Since 62.14 GiB does not fit on a single 3090, PP=3 is required, meaning removing this rejection is also required.

vllm/model_executor/models/config.py          <- at engine startup
vllm/models/qwen4_exp/nvidia/model_state.py   <- at worker model runner construction

When I patched one, I was rejected by the other. The duplication is intentional, and the first comment states “Checked again in Qwen4ExpModelState”.

The reason for the guard itself is legitimate, whereas gpu/model_runner.py sets model_inputs["input_ids"] = None on non-first PP ranks, PLE lookup requires raw token IDs. PLE exists only in decoder layer 1, and get_pp_indices(48, 0, 3) == (0, 16), so layer 1 is always placed on the first rank. The first rank receives input_ids, and other ranks pass through with ple is None, so the correct fix is not distributing input_ids across ranks, but narrowing the guard to judge by looking at which rank the PLE layer is actually placed on, and configurations that send layer 1 to the back can simply be rejected as before.

B - PLE Table Cannot Be Allocated in the First Place

PLE offload in vLLM is not mmap. It is implemented to place the entire table in torch.empty(..., device="cpu", pin_memory=True) and read directly from GPU. Demanding 95.43 GiB of pinned memory with 125 GB RAM results in this.

torch.AcceleratorError: CUDA error: out of memory
  ngram_embedding.py:483 in allocate_embedding_weight
  -> torch.empty(320001536, 160, dtype=bfloat16, device="cpu", pin_memory=True)

The patch adds an mmap backend. If the file already exists, it opens it with copy-on-write and turns weight_loader into a no-op, so it does not even read 95 GiB from the checkpoint at startup. As a byproduct of that, startup from the second time onward becomes significantly faster.

A Single Line of madvise Reduces NVMe Read to 1/23

What gather really needs is 16 rows x 320 B = 5 KiB per token, and even rounding each row to a 4 KiB page is 64 KiB, so the initial measured value was 922 KiB/token.

The difference is readahead, where the kernel reads 128 KiB before and after upon every page fault, but because it is a completely random gather, almost all of what was read is discarded.

D - compressed-tensors Does Not Support PLE

The PLE table is plain BF16, and Qwen4ExpPLEEmbeddingMethod.from_quant_config() only accepts None / ModelOpt / Fp8. Therefore, with a compressed-tensors checkpoint, model construction halts with NotImplementedError. Because compressed-tensors only enumerates Linear there, it is not listed in ignore of the checkpoint either. It is necessary to look directly at the scheme map and determine it as Unquantized if no config group specifically names PLE.

E - Upstream Bug

Passing PP>1 causes ranks 1 and 2 to crash this time right before KV allocation.

File ".../vllm/v1/worker/utils.py", line 426, in allocate_kv_cache
    group_id, group = next(
StopIteration

_project_kv_cache_groups_to_worker() narrows down global KV groups to layers assigned to each rank, but it only reconstructs the inner spec dict “when that rank has 1 or more layers”, and empty groups keep layer_names=[] while the spec dict retains the entire model. Because tensor generation iterates over the spec dict rather than layer_names, it creates KVCacheTensor for layers that the rank does not have, and allocate_kv_cache cannot find the group to which that tensor belongs.

What becomes empty in this model is the short-conv state of PLE, which structurally does not exist on ranks other than rank 0. The patch merely narrows down the tensor generation side with group.layer_names. Fixing the filtering side is more dangerous, because creating an empty UniformTypeKVCacheSpecs causes next(iter(...)) or max() elsewhere to crash.

And It Worked

Loading weights took 31.86 seconds
Model loading took 21.71 / 21.67 / 21.67 GiB       -> Total 65.05 GiB
GPU KV cache size: 25,122 tokens
Application startup complete.                      (150 seconds)

It ran at context 8,192 and 8.5 tok/s, but usability is non-existent.

Please pay attention to the weights. By calculation it was supposed to be 62.14 GiB, but the actual measurement is 65.05 GiB. The difference of about 1 GiB per card is kernel workspace for Marlin / Humming and the prefetch buffer for PLE. vLLM is scary in this regard. If you do not deliberately leave a margin, it will end right here.

5. Bringing It Up to Practical Speed

5.1 Moving QSA K/V to Host RAM

The K/V row per token per layer of QSA is 2 kv heads x 256 dim x 2 (K and V) x 2 B = 2,048 B. For 262,144 tokens, it is 512 MiB per layer, and 6 GiB across 12 layers. If it were attention that reads the entire context, this would be read at every step. Because PCIe 4.0 x16 is about 25 GiB/s, placing the cache of such a model in host results in 4 tok/s.

Fortunately, QSA does not read the entire context. The indexer scores with pooled/compressed keys (indexer_head_dim=128 / indexer_compress_ratio=4), and selects up to indexer_budget=2048. It reads only those rows.

Selected 2048 x kv head 2 x dim 256 x 2 (K and V) x 2 B  =  4 MiB / layer
x 12 layers                                              =  48 MiB / token

At 80 tok/s it is 3.9 GB/s, which is about 1/6 of PCIe bandwidth, and can even be overlapped with computation. What remains on the GPU is only a 2 B slot number and a 64 B pooled key, which is 66 B per token per layer, and compared to the 2,048 B of K/V body, the KV placed in VRAM becomes 1/31. The process of selecting where to read does not pass through PCIe itself.

Decode does not drop at all from 3.6k to 248k. (This is the number before 5.6. Speed increases in 5.6.)

5.2 Waste in Block Size That Appears Only After Offloading

Even after offloading, the capacity that can be obtained is only about 140,000 tokens. What determines the size of the KV pool is another place.

vLLM creates the block pool aligned to the largest single page among all groups.

KVPROBE groups=39 bytes_per_block=3211264 available=720000000 num_blocks=224
  37 groups: mamba,    block=262144, 1 layer,  page 3,211,264   <- This determines it
   1 group : QSA,      block=1568,   8 names,  page [100352, 3136]
   1 group : PLE conv, block=4,      4 names,  page 1024

Although what the QSA group really needs is 4 layers x 1568 tokens x 66 B = 413,952 B, it allocates 3,211,264 B. Separately from this, _align_hybrid_block_size calculates page size per token of attention from model_config (2 heads x 256 x 2 x BF16 = 2048 B), so it does not recognize that the layer is currently using only 66 B. When not offloaded, the QSA page is 13.2 MB and exceeds mamba, so this waste does not occur.

The fix is to raise attention block size only during offloading up to the value where the QSA page matches the mamba page (12,144 tokens here) and early return.

At this point, 1,215,172 tokens in 400 MB. It was 7.7 times per the same number of bytes, matching the prediction of 7.8 times.

5.3 For Prefill, Swap the Loop Order

The read being constant every step is a story about 1 token. In prefill, thousands of tokens each arbitrarily choose 2048, so writing it naively means reading “number of tokens x number of selections” rows one by one randomly across PCIe.

As a fix, swap the order of loops. Split context by a fixed size, copy that section to a working buffer on VRAM only once, compute all tokens collectively for that section, and advance to the next section while merging intermediate results of softmax. The transfer volume comes to be determined by context length rather than number of tokens x number of selections.

You must not hold intermediate softmax results per section. The first version held number of sections x number of tokens x heads x dim x 4 B, resulting in the opposite behavior where smaller working buffers consumed more VRAM, shrinking card margin down to 2 MiB. Holding just one set of (acc, max, norm) and sequentially merging across sections is the correct form, and once fixed, the peak stopped moving.

5.4 Configuring Settings

If any one of the following is removed, long context prefill crashes.

--kv-cache-memory 550000000
--max-num-batched-tokens 512
--no-enable-prefix-caching
--gpu-memory-utilization 0.97
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0

When specifying --kv-cache-memory, --gpu-memory-utilization stops having an effect on KV cache size. This is because the memory profiling itself is skipped.

The recommended value for --kv-cache-memory produced by vLLM itself causes OOM with long contexts. There are two independent reasons. One is double counting of cudagraph, where peak_activation_memory already includes the prediction, but the actual measured value is added once more in the calculation of the recommended value. Therefore, simply setting VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 increases KV by +15%. The other is that the prediction of peak activation is only about half of actual long context prefill. Against the predicted 0.27 GiB, the actual measurement is about 0.53 GiB, and when starting with exactly the recommended value and sending 47,927 tokens, rank 0 fails to allocate 20 MiB and crashes.

With expandable_segments set to True, without optimization about 100 MiB disappears into “reserved unused” fragmentation, and the peak gradually rises with prompt length. When enabled, it stabilizes completely from 4k to 130k, and measured weights also decrease, albeit by only 0.07 GiB.

Disabling prefix caching changes how the number of tokens to process prefill at a time is determined. When enabled, the KV block size is the unit, but when disabled, --max-num-batched-tokens itself becomes the unit. Therefore, disabling prefix caching and leaving the default 2048 causes OOM even with a 4k prompt. With 512, the peak is 470 MiB lower than 1568, and speed is almost unchanged. With 256, prefill becomes 32% slower.

PyTorch’s caching host allocator rounds up allocations of pinned memory to powers of 2. With --kv-cache-memory 560000000, the host pool becomes 4.03 GiB per layer, rounding up to 8 GiB, requesting 96 GiB from the host for 12 QSA layers. With 550000000, it is 3.96 GiB per layer, rounding up to 4 GiB, and the same configuration suffices with 48 GiB. Total host RAM usage went from 111 GiB to 62 to 63 GiB.

The hardest one is not GPU0 but GPU1, that is, the middle PP rank. Although weights are the lightest, the peak is constantly 40 to 60 MiB higher. This is because it holds both transmit and receive buffers, and since --kv-cache-memory is uniform across all ranks, it is decided by looking only at the peak of GPU1.

5.5 4 Parallel Is Slower Than Single

The instant it reaches 4 requests, 1 step jumps from 25 ms to 110 ms, and the entire server becomes slower than a single request. The cause of this is misconfiguration of CUDA graph capture size. In PP, concurrency count and forward batch size become different numbers. Up to 3 requests can sometimes be arranged one per stage, but adding a 4th request inevitably puts 2 requests into a forward somewhere. If only size 1 is captured, such forward cannot use CUDA graphs and falls back to eager execution. Resolving it with cudagraph_capture_sizes set to [1,2,4], the total reached 139 to 146 tok/s.

When passing both --compilation-config '{...}' and --compilation-config.cudagraph_capture_sizes, this nightly silently discards the entire JSON and reverts mode to default. It is necessary to consolidate into a single JSON.

5.6 Making It Faster Than 45 tok/s

Single decode was stalled at 45 tok/s. Reading all safetensors headers and counting the bytes that 1 decode step actually touches gives 5.83 GB/token (dense 4.62 + 1.22 for experts 10/512), and with 936 GB/s of 3090, 160 tok/s is the upper limit. 45 tok/s is 29% in bandwidth utilization, clearly slow. The time the GPU spent executing kernels was 12.62 ms out of a 21.9 ms step, and 42% was doing nothing.

The cause is this single line in the PLE lookup,

ids = flat_ids.to(device="cpu", dtype=torch.int64).numpy()   # synchronous D2H

This synchronizes CPU to GPU every step. vLLM V2 + async scheduling assumes that CPU runs ahead of GPU to schedule tasks. When that disappears, the amount CPU lagged in the eager execution sections per layer (17 per rank) directly becomes GPU idle time.

I applied a 15-line patch switchable via environment variables and conducted 5 startup experiments.

That A and D are the same means that mmap reads, numpy gather, and H2D combined take only 0.16 ms. The value of making gather faster is zero. D -> B yielded +29 tok/s just by eliminating one synchronization. E means that --no-async-scheduling is a 15% loss so do not use it.

The fix is to move gather outside forward. Qwen4ExpModelState.prepare_inputs fills a device-side buffer whose address does not change prior to forward, and the PLE layer only reads it. Because CPU work disappears from forward, cudagraph_mode set to FULL_DECODE_ONLY becomes usable.

Measured 45.4 -> 80.0 tok/s (single), reaching 152 to 155 tok/s in 4 parallel. Prefill did not degrade (4 x 247,823 tokens went from 378 seconds to 368 seconds), GPU peak movement was 2 MiB, and the cost of capturing CUDA graphs on the decode side was 0.06 GiB per rank.

What is worth remembering is that moving gather outside forward itself did not make it faster by even 1 tok/s. If left at PIECEWISE, it was 45.1 tok/s, because eager execution sections per layer remain on the CPU-side critical path. What is making everything faster is entirely the full graph, and moving gather was only to make that usable.

5.7 A Bug That Had Nothing to Do with Anything Up to This Point

With everything working, Japanese output collapsed. From the middle of an answer, it turned into a mass of other languages and code snippets. It did not happen in English. This seems to have been caused by generation_config.json.

AutoRound export regenerates this file and drops sampling defaults. vLLM’s --generation-config defaults to auto, and get_diff_sampling_param() picks only 6 keys from it. If there is not a single one, it returns {} and falls back to defaults on the vLLM side. In other words, the entire 248,320 vocab ends up being sampled.

6. Current Configuration and What Cannot Be Done

--pipeline-parallel-size 3 --tensor-parallel-size 1
--max-model-len 262144 --max-num-seqs 4
--max-num-batched-tokens 512 --no-enable-prefix-caching
--gpu-memory-utilization 0.97 --kv-cache-memory 550000000
--compilation-config '{"mode":0,"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[1,2,4]}'

VLLM_PLE_MMAP_PATH=<table_file>   VLLM_QSA_KV_OFFLOAD=1
VLLM_QSA_KV_OFFLOAD_MAX_GIB=56          VLLM_QSA_KVO_ARENA=100663296
VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

Conclusion

There is only one structural factor, and what makes all of this possible in the first place is that QSA limits K/V reads per step with indexer_budget. Therefore, even if a 24 KiB cache per token is placed across PCIe, context length does not appear in bandwidth calculations. The rest of this article is nothing more than wiring work to actually make that property function.

The vLLM patch is placed here. The target is 0.29.1rc1.dev47+gdc36fcce9 of vllm/vllm-openai. https://huggingface.co/Minachist/Qwen3.8-Flash-Next-INT4-Mixed-AutoRound