[Bugfix] Fix MoE model load OOM in FlashInfer_TRTLLM backend with sleep mode

[andakai]

Summary

This PR fixes the load-time OOM reported in #43951 for large BF16 MoE models using --enable-sleep-mode with the FlashInfer TRTLLM MoE backend.

Instead of bypassing the max_split_size_mb=20 allocator policy introduced by #41268, this PR reduces the transient GPU memory created by FlashInfer TRTLLM BF16 MoE weight conversion itself. The conversion now preallocates the final block-layout output tensors and writes each expert directly into the target storage, avoiding the previous per-expert clone/list/stack pattern.

Issue

#43951 reports that loading Qwen/Qwen3-235B-A22B-Instruct-2507 with sleep mode OOMs during weight loading, before CUDA graph profiling starts. The traceback points to convert_moe_weights_to_flashinfer_trtllm_block_layout, and the failing allocation can be as small as 20 MiB.

I reproduced the load path on H20-141GB by forcing the FlashInfer TRTLLM backend to pass the SM100 capability gate. The problem happens in model load stage, so this H20-forcing-to-run method is enough to reproduce this issue.

Repro command:

CUDA_VISIBLE_DEVICES=0,1,2,3 \
.venv/bin/python -m vllm.entrypoints.cli.main serve \
  Qwen/Qwen3-235B-A22B-Instruct-2507 \
  --tensor-parallel-size 4 \
  --enable-sleep-mode \
  --load-format dummy \
  --dtype bfloat16 \
  --moe-backend flashinfer_trtllm \
  --gpu-memory-utilization 0.93 \
  --max-num-seqs 1024 \
  --skip-tokenizer-init

The issue is reproduced. Only Flashinfer_trtllm backend has the error, Trtion, flashinfer_cutlass backend can succeed because they have simple postprocess.

Root Cause

For Qwen/Qwen3-235B-A22B-Instruct-2507:

• num_hidden_layers = 94
• num_experts = 128
• hidden_size = 4096
• moe_intermediate_size = 1536
• TP = 4
• dtype = BF16

With TP=4:

• intermediate_per_partition = 1536 / 4 = 384
• w13 = [128, 768, 4096]
• w2 = [128, 4096, 384]

The sensitive path is FlashInfer TRTLLM BF16 MoE conversion. The original implementation creates per-expert temporary tensors such as cloned weights, permute indices, reordered tensors, contiguous copies, block-layout outputs, and finally stacked outputs.

Per expert:

• w13[i]: [768, 4096] BF16 = 6 MiB
• w2[i] : [4096, 384] BF16 = 3 MiB

The permute index tensors are larger:

• w13 uint8 view size = 768 * 4096 * 2 = 6,291,456 elements, so int64 permute indices are about 48 MiB
• w2 uint8 view size = 4096 * 384 * 2 = 3,145,728 elements, so int64 permute indices are about 24 MiB

That 24 MiB size is important because it is larger than the 20 MiB split threshold from #41268.

Observed allocator behavior for this model:

• max_split_size_mb=20 -> OOM
• max_split_size_mb=25 -> succeeds
• max_split_size_mb=32 -> succeeds
• max_split_size_mb=64 -> succeeds
• max_split_size_mb=128 -> succeeds

This shows a conflict between #41268 and this workload. #41268 uses max_split_size_mb=20 to protect large cached blocks and avoid fragmentation during model loading. That behavior is useful for the original fragmentation problem. However, FlashInfer TRTLLM MoE conversion repeatedly creates medium-sized temporary tensors, especially in the 20-50 MiB range, across 128 experts and 94 layers. With sleep mode enabled, these allocations go through vLLM's CuMem weights pool, and the pool does not fully reclaim transient backing segments mid-load.

FlashInfer CUTLASS and Triton did not reproduce this issue in my tests because their MoE post-processing does not create the same repeated transient allocation pattern.

Why not bypass max_split_size_mb=20

I considered bypassing max_split_size_mb=20 for the FlashInfer TRTLLM MoE conversion. That does fix this issue locally, but it also weakens the protection added by #41268.

I also tested the behavior without #41268's allocator policy and observed OOM later during warmup. This suggests #41268 is still necessary: during model loading, vLLM still needs to protect large tensors from fragmentation.

Changing the global max_split_size_mb value is also not robust. The right threshold depends on model shape and tensor parallel size. Larger MoE models can produce larger transient tensors, so values like 32, 64, or 128 MiB are only case-specific workarounds. Raising the global threshold also moves away from #41268's original goal.

Fix

The original FlashInfer TRTLLM BF16 MoE conversion did this for each expert:

clone -> index reorder -> contiguous -> convert_to_block_layout -> append to list
torch.stack(all_experts)

This PR changes the conversion to:

1. Preallocate the final w13 and w2 block-layout output tensors.
2. View each expert as uint8.
3. Inline the equivalent of convert_to_block_layout by reshaping to blocks and permuting dimensions.
4. Use torch.index_select(..., out=...) to write the permuted expert directly into the preallocated output tensor.
5. Fold the swap_w13_to_w31 operation into the W13 permute indices, avoiding a full materialized swapped W13 tensor.

The block-layout transformation is still performed. It is just fused with the row permutation and written directly into the destination tensor, which avoids the previous clone/list/stack intermediates.

Validation

Time

Qwen3-235B TP=4 dummy load, forced H20 FlashInfer TRTLLM backend:

• Old conversion + sleep mode: OOM
• New conversion + sleep mode:
  • load completed
  • memory: 109.61 GiB
  • load time: 3.001113s
  • later expected H20 failure: No supported CUDA architectures found for major versions [10, 12]

No sleep mode comparison:

• Old conversion + no sleep mode:
  • run 1: 10.723114s
  • run 2: 10.214007s
• New conversion + no sleep mode:
  • 2.385791s

So in this dummy-load, conversion-heavy setup, the new path is roughly 4x faster than the original conversion.

Correctness check

I compared the old conversion and the new conversion using random BF16 tensors and bitwise uint16 comparison.

Covered cases:

• Qwen-like TP4 full shape:
  • w13=[128,768,4096]
  • w2=[128,4096,384]
• DeepSeek-like TP4 shape
• DeepSeek-like TP2 shape

Result:

ALL_BITWISE_EQUAL

Limitations

The load/conversion path was validated on H20 by forcing the FlashInfer TRTLLM backend to pass the capability gate. Since FlashInfer TRTLLM BF16 MoE kernels are SM100/B-series only, H20 cannot validate final runtime kernel correctness. This should still be tested on B-series hardware.

[aoshen02]

GB200 Validation Report

Validated PR #45589 on NVIDIA GB200 (SM100) with real weights. All tests use vllm/vllm-openai:v0.23.0 as the base image with the PR patch applied on top.

1. OOM Fix (Qwen3-235B-A22B-Instruct-2507, TP4)

Test	Result	Detail
Baseline (no patch, --enable-sleep-mode)	OOM	torch.OutOfMemoryError: Tried to allocate 20.00 MiB at flashinfer_utils.py:176, all 4 GPUs
PR #45589 (patch + --enable-sleep-mode)	PASS	Model loaded 109.61 GiB, served completions successfully

2. Output Consistency (Qwen3-30B-A3B, TP2, greedy)

Single-request greedy completion (temperature=0, top_p=1.0) with --enforce-eager:

Group	Config	Output
A (baseline)	No patch, no sleep mode	" Paris. The capital of the United Kingdom is London..."
B (PR #45589)	Patch + --enable-sleep-mode	" Paris. The capital of the United Kingdom is London..."

Token-level exact match.

3. GSM8K Accuracy (Qwen3-30B-A3B, TP1, lm_eval --model vllm)

Standard vLLM CI evaluation: --tasks gsm8k --num_fewshot 5 --limit 250 --batch_size auto, CUDA graphs enabled (default).

Group	Config	strict-match	flexible-extract
A	No patch, no sleep mode	0.908 ± 0.0183	0.864 ± 0.0217
B	PR #45589 + --enable-sleep-mode	0.908 ± 0.0183	0.864 ± 0.0217
C	No patch + --enable-sleep-mode	0.908 ± 0.0183	0.864 ± 0.0217

All three groups produce identical scores — the PR's preallocated inplace conversion is numerically equivalent to the original path.

Environment

• GPU: NVIDIA GB200 NVL72 (184 GiB, SM100)
• Image: vllm/vllm-openai:v0.23.0 (PyTorch 2.11.0+cu130, FlashInfer 0.6.12)
• Node: GCP GB200 rack

[MatthewBonanni]

LGTM, thanks for fixing this!