feat(vllm): refit-idempotent BF16 FlashInfer-TRTLLM MoE (vLLM 0.20.x stopgap) - #3296
feat(vllm): refit-idempotent BF16 FlashInfer-TRTLLM MoE (vLLM 0.20.x stopgap)#3296seonjinn wants to merge 5 commits into
Conversation
Unquantized vLLM with the FlashInfer TRTLLM MoE backend rebinds the fused-MoE weights to a 4D block layout on the first process_weights_after_loading pass, which breaks every subsequent refit load. Install idempotent replacements that keep checkpoint-layout params as permanent load targets, re-deriving the kernel layout from scratch buffers on every pass, and fail loudly if a layer is loaded without a full reprocessing pass. Kill switch: NRL_BF16_TRTLLM_REFIT=0. Also installs the patches for precision=fp8 runs that keep some MoE layers unquantized (num_first/last_layers_in_bf16, ignored-layer keywords). The engine-core compose hook chains to the fp8 patches only when they are present (guarded by hasattr on nrl_fp8_cfg), so the module builds and runs standalone without the fp8 refit patches. Squashed from 1cbf401b0, 4377fb223 (module/utils/worker wiring), and adb2b8c09; launcher hunks dropped. Signed-off-by: sna <sna@nvidia.com>
Validate the torch-only permutation/layout helpers in bf16_trtllm_moe_utils against the per-expert vLLM reference pipeline (bit equality), alias write-through, and two-cycle idempotency, all on CPU without vLLM or flashinfer. Squashed from 4377fb223 (test portion). Signed-off-by: sna <sna@nvidia.com>
The new torch-only module type-checks clean, so CI's zero-error whitelist check requires it in project-includes. Also add the missing docstring on swap_w13_to_w31_row_indices. Signed-off-by: sna <sna@nvidia.com>
|
Isolated-branch reproduction (follow-up promised in the PR body). Setup: Qwen3-30B-A3B GRPO (OpenMathInstruct-2), 2 nodes x 8 B200, BF16 recipe
*Generation is the causal saving; the total-step delta also absorbed |
Squash of NVIDIA-NeMo/RL PR #3545 (head 7180bff, 2026-08-11) onto main@d5fb8d04 with patches 0001-0003 applied. Replaces the carried PR #3296 module (bf16_trtllm_moe.py): the vllm_backend refit path now natively reloads unquantized FlashInfer TRTLLM MoE weights, detects the realized refit backend, constrains native refit scope, defaults injected reload state safely, and hardens layerwise-refit failures; the ModelOpt quant backend opts out explicitly. Retire when the upstream PR merges.
|
Contributor self-review status: this vLLM 0.20 compatibility stopgap is obsolete on the current vLLM stack, and the remaining BF16 TRTLLM refit requirement is covered by #3659. Recommendation: close this PR rather than merge current main into obsolete compatibility code; do not retrigger CI. |
…20260817 Signed-off-by: seonjinn <sna@nvidia.com>
|
/ok to test d5165a1 |
1 participant
What does this PR do?
vLLM 0.20's BF16 FlashInfer-TRTLLM MoE path rebinds
w13/w2to 4D block-layout Parameters on first processing, so the second RL refit'sload_weightsfails (ValueError: shard_dim=0 is not a valid data dimension for a 3D tensor) - the reason the shipped GRPO recipes pinmoe_backend: tritonfor BF16.This PR keeps checkpoint-layout parameters permanently and exposes the TRTLLM layout as 4D views aliasing the same storage (zero duplicate weight memory), recomputing the re-layout each refit as one batched gather.
NRL_BF16_TRTLLM_VERIFY=1: first-refit bit-equality assert vs the stock vLLM path.NRL_BF16_TRTLLM_REFIT=0: kill switch (patch fully inert).RuntimeErrorwith remediation) for load paths that skip reprocessing (ModelOpt real-quant, MTP draft reload).Version applicability
The rebind-on-reprocess issue this patch works around is handled natively in vLLM from v0.21.0 (
is_weight_update/prefer_copyin_setup_kernel; verified at tags v0.21.0-v0.24.0 and main - the upstream comment explicitly cites "RL weight updates that re-trigger process_weights_after_loading"). This patch is therefore a stopgap for the currently pinned vLLM 0.20.x only and should be removed at the next vLLM bump; the kill switch plus the isolated module make removal a two-line revert.Scope
Applies to SiLU-gated, 128-aligned MoE models on SM100 (B200/GB200). Validated on Qwen3-30B-A3B. The 128-alignment constraint applies to the per-partition intermediate size: Qwen3-235B-A22B at the recipe's vLLM TP=8 shards
moe_intermediate_size=1536to 192 per rank, which FlashInfer'sBf16MoeLauncher::check_moerejects (intermediate_size % 128 == 0fails) at engine build - before any refit, i.e. stock vLLM 0.20 withmoe_backend=flashinfer_trtllmfails identically on that shape. Qwen3-235B is therefore N/A at TP=8 and stays on triton. Dense models (Qwen3-32B, Llama) are N/A - the TRTLLM MoE backend is MoE-only. Nemotron Nano (RELU2 + non-aligned dims) excluded pending kernel support.Benchmark setup
Impact (BF16, triton -> flashinfer_trtllm)
Trade-offs
RuntimeError.Correctness
tests/unit/models/generation/test_bf16_trtllm_moe.py).fp8.pymodule in the same directory.