fix(vllm): support native BF16 FlashInfer TRTLLM refit

[seonjinn]

TL;DR for reviewers

FlashInfer TRTLLM repacks BF16 MoE expert weights into a backend-private layout at load time, so NeMo-RL's generic refit (write canonical tensors into live params, run generic post-processing) silently corrupts the runtime weights on this backend. This PR makes the affected refit paths follow vLLM's native layerwise reload lifecycle instead. The three things to review: (1) the lifecycle context manager and its fail-closed semantics in vllm_backend.py; (2) the transport-buffer detach that keeps deferred reload tensors alive after NeMo-RL reuses NCCL/IPC buffers; (3) the single-buffer collective path that removes a cross-stream read-after-write hazard.

Problem

FlashInfer TRTLLM does not execute incoming BF16 expert weights in their original layout — vLLM repacks them into backend-specific runtime storage at model load. Two failure modes follow from the generic refit path:

1. Layout bypass: loading new weights without the native reload protocol never rebuilds the repacked runtime storage, so the model keeps executing stale expert weights even though the transfer succeeded.
2. Buffer aliasing: the native protocol defers some weight_loader calls until finalization, and those deferred calls hold references into NCCL/IPC transport buffers that NeMo-RL reuses for the next transfer — so by the time they replay, the data may have been overwritten.

How the fix works

For this backend only, each refit follows vLLM's native lifecycle:

1. Initialize vLLM's layerwise reload state.
2. Load the new BF16 weights through the existing collective or IPC/ZMQ transport.
3. Detach any deferred reload tensors from reusable transport buffers (clone-on-alias, keyed by storage pointer).
4. Finalize so vLLM rebuilds the FlashInfer TRTLLM runtime layout (plus the HpcModule refresh that vLLM's own finalizer omits — upstream gap, to be filed).
5. Only then reuse the transport buffers.

Backend detection inspects the realized model, so auto recipes that resolve to FlashInfer TRTLLM at runtime are covered.

Stream safety: with the native lifecycle active, vLLM replays buffered weight_loaders eagerly when a layer's last shard arrives — potentially on a different CUDA stream than the one that produced the clones. The collective consumer therefore runs single-buffer on this path (trading broadcast/load overlap for a closed race window, this backend only).

Fail-closed: any failure after reload initialization marks the worker permanently unusable (_refit_unusable_reason); subsequent refit attempts return False instead of running against a partially updated model.

Scope

Path	Behavior
collective, IPC/ZMQ (unquantized FlashInfer TRTLLM)	native lifecycle (this PR)
nccl_reshard	keeps its direct-buffer protocol (setup-time mapping needs stable live storage)
checkpoint-engine, sparse-delta (unquantized FlashInfer TRTLLM)	rejected loudly — they bypass the native lifecycle and would corrupt silently
FP8 / ModelOpt-quantized / Triton / other backends	unchanged
static MTP draft weights	same lifecycle during their one-time disk load
co-trained MTP refits	unsupported, fail early with an explicit error

Runtime support only; recipe changes follow separately. Independent of #3478 (MXFP8 refit performance). Sibling ordering: this PR → #3659 (extends the lifecycle to nccl_reshard) → #3669 (perf); #3651 (opt-in reload_weights API) is complementary and rebases on top.

Validation

• Unit tests: repeated refits, runtime auto backend selection, transport-buffer ownership, failure handling (including public-method fatal-raise and worker poisoning), unaffected backend paths, static MTP draft loading.
• 3-step Qwen3-30B-A3B smoke on 4 nodes x 4 GB200: all 48 worker refits completed without fatal error (job 5950903). Note: the smoke predates the single-buffer stream fix and validates the lifecycle, not stream safety — the race window was narrow enough that a clean smoke was consistent with it.
• CI: /ok to test issued on the current head; earlier CI:Lfast runs targeted a pre-review SHA.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[seonjinn]

/ok to test a73421e

[seonjinn]

/ok to test 7180bff

[yuki-97]

hi @youngeunkwon0405 , could you help review this PR? thanks!

[seonjinn]

/ok to test 95452e9

[copy-pr-bot]

/ok to test 95452e91062b3a13f08db82ae6c8219513782f70

@seonjinn, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[seonjinn]

/ok to test 95452e9

[seonjinn]

Self-review completed with /review-pr-team on head 10e40a591.

• Reviewed the native layerwise reload lifecycle, repeated refits, MTP handling, collective and IPC transports, failure handling, ModelOpt gating, and compatibility with the changes now on main.
• Found and fixed one correctness issue: the vLLM layerwise finalizer does not run the separate HpcModule post-load pass. The updated code now refreshes derived HPC state for both the policy model and a disk-loaded MTP drafter without re-running quantization or attention post-processing.
• Added tests for the HPC-only pass and both native layerwise reload call sites.
• No other blocking correctness, compatibility, security, or deadlock issues were found.

Local validation completed:

• git diff --check
• Python syntax compilation for the changed implementation and test files

The full vLLM unit suite is left to CI because the local host cannot create the repository Linux-only locked vLLM environment. Follow-up PRs #3659 and #3669 should be restacked after this PR lands.

[youngeunkwon0405]

Team review (5 agents: RL-code, vLLM-API, tests, bug-finder, devil's advocate; every vLLM claim verified against the pinned v0.25.1 wheel source).

This is a well-built change — the lifecycle design, the realized-backend gating (including auto), the failure latching, and especially the buffer-overwrite regression test all held up under adversarial verification. The detach-necessity claim, the FlashInfer TRTLLM finalize semantics, and the HPC-pass gap were each confirmed in vLLM source. The earlier FP8-KV-cache concern was re-checked and is indeed covered: the layerwise finalizer re-runs layer.process_weights_after_loading(model_config.dtype) and reloads streamed k/v scales, equivalent to the skipped generic pass.

6 inline comments: 1 confirmed concurrency bug on the collective path (cross-stream RAW race introduced by the deferral+detach), 1 fail-loud gap (checkpoint-engine/sparse-delta transports), 1 upstream-issue suggestion, 1 test gap, 1 lint failure (repo-pinned ruff 0.9.9), 1 docstring nit.

Lower-confidence notes, not requesting changes:

• The detach helper relies on private reload.layerwise internals (LayerReloadingInfo.loaded_weights shape, the loaded_weight arg name); imports are function-local so older vLLM fails loudly only when the feature runs. A version-pinning comment would help future upgrades.
• except Exception in _weight_update_lifecycle misses BaseException (e.g. a cancelled refit): the worker isn't marked unusable and the next refit would silently merge stale buffered clones (vLLM skips still-initialized layers). Low reachability; except BaseException closes it if you think it's worth it.
• Question: does BF16-TRTLLM + nccl_reshard fail loudly today, or run the generic finalize silently until #3659 lands? Our read is the reshard mapping most plausibly errors at prepare time on the 4-D block layout, but a confirmation would help.
• Given the per-refit init/detach/finalize overhead, one before/after refit-latency number in the description would be a nice-to-have (perf work itself is understood to be #3669).

Generated by Claude Code

[youngeunkwon0405]

I did a thorough review of this PR. The logic is hard to understand, but I think the problem itself is just complex. I think the implementation functionally makes sense and is reasonable. Please address the comments. Thank you!

[youngeunkwon0405]

@guyueh1 I studied this PR today and learned that this is also using the layer-wise reload feature of vLLM, which is I didn't realize before. I am adding you because you have been working on the reload PR. Maybe you might also want to check an overlap with your PR? I will also review your PR soon.

[seonjinn]

I did a thorough review of this PR. The logic is hard to understand, but I think the problem itself is just complex. I think the implementation functionally makes sense and is reasonable. Please address the comments. Thank you!

Thank you so much for your thorugh review Youngeun, I'll try to fix codes for clear understanding.

[seonjinn]

/ok to test 86c94f1

[seonjinn]

Superseding my 08-18 comment ("close in favor of #3659"): that had the dependency direction backwards — #3659 explicitly extends this PR's native layerwise refit to the nccl_reshard transport, and #3669 is a perf follow-up on top of that. The intended order is #3545 → #3659 → #3669.

On the overlap question with @guyueh1's #3651 (native reload_weights API): the two are complementary layers, not duplicates. This PR fixes the default NeMo-RL loader paths (IPC + collective) by wrapping them in vLLM's layerwise reload lifecycle, streaming chunk-by-chunk; #3651 adds an opt-in refit_with_reload_api flag that bypasses the NeMo-RL loader and hands a whole weight iterator to model_runner.reload_weights (collective-only, non-colocated, also covers MXFP8), while keeping this PR's path as the default else-branch. They touch the same regions of update_weights_from_collective and packed_tensor.py, so whichever lands second needs a rebase — proposal: land this first (default-path correctness + fail-closed lifecycle + IPC coverage), rebase #3651 on top for the opt-in native API.

[seonjinn]

Self-review summary (supplemental to youngeunkwon0405's review, whose 8 comments were addressed in e81aba7 with per-thread replies).

Method: five review agents over the post-fix diff, then an adversarial verification pass that re-checked every citation, verified vLLM internals against a real vLLM tree, and re-ran CI queries live. No inline code findings survived adversarial verification -- the lifecycle state machine (fail-closed on every failure window including finalize), the transport-buffer detach, the single-buffer race fix, and the mock fidelity of the new tests were all confirmed sound. Two follow-up fixes were pushed as 86c94f1: a version-guard on the fused_moe.oracle import (on vLLM builds without that module layout, every unquantized-model refit would have raised ImportError instead of skipping the TRTLLM check) and a missing num_buffers kwarg on a pre-existing test fake that would have failed the vllm-gated test lane.

Merge gates (process, not code):

1. CI has not run on any recent head. The last /ok to test targeted f0c23ae (08-18), fourteen commits and a 169-file upstream merge ago. A fresh /ok to test on the final head is required; posted above.
2. PR body needs a refresh. Three statements are stale relative to the code: "Checkpoint-engine transports are unchanged" (they are now rejected for the unquantized TRTLLM backend, by design); the CI claim references a stale SHA; and the 48-refit smoke predates the cross-stream race fix -- a clean smoke was consistent with the race being present (narrow-window silent corruption), so it should be cited as evidence for the lifecycle, not for stream-safety. Worth adding while editing: the native path now runs the collective consumer single-buffer, trading broadcast/load overlap for stream-safety on this backend only.
3. Sibling-PR order is clarified in the comment above: #3545 -> #3659 (nccl_reshard extension) -> #3669 (perf), with #3651 (opt-in reload_weights API) complementary and rebased on top.

[seonjinn]

/ok to test e450256

[seonjinn]

/ok to test 87e9e98