Skip to content

perf(vllm): batch BF16 TRTLLM expert layout conversion - #3669

Draft
seonjinn wants to merge 83 commits into
NVIDIA-NeMo:mainfrom
seonjinn:sna/bf16-trtllm-two-gather-refit
Draft

perf(vllm): batch BF16 TRTLLM expert layout conversion#3669
seonjinn wants to merge 83 commits into
NVIDIA-NeMo:mainfrom
seonjinn:sna/bf16-trtllm-two-gather-refit

Conversation

@seonjinn

@seonjinn seonjinn commented Aug 16, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Summary

  • Convert all BF16 W13 experts with one indexed gather and all W2 experts with a second gather, instead of running two gathers per expert.
  • Limit the batched converter to the vLLM reload lifecycle and keep the output layout and runtime parameter addresses unchanged.

Validation

  • GB200, Qwen3-30B-A3B, 20-step BF16 training and BF16 rollout: refit bubble decreased from 3.637 s to 3.124 s (-14.1%).
  • Generation KL error remained equivalent: 0.001887 baseline vs. 0.001886 batched.
  • End-to-end throughput stayed within the run-to-run difference and is not claimed as improved.
  • 82 targeted GB200 BF16 TRTLLM refit tests pass.

Depends on #3659.

seonjinn added 24 commits August 7, 2026 02:01
@seonjinn
test(vllm): cover native BF16 MoE refit lifecycle
7db8244 
Signed-off-by: seonjinn <sna@nvidia.com>
(cherry picked from commit 707968a)
@seonjinn
fix(vllm): reload unquantized TRTLLM MoE weights
e919d26 
Signed-off-by: seonjinn <sna@nvidia.com>
(cherry picked from commit cbd4138)
@seonjinn
fix(vllm): preserve refit error fallback
2ce743e 
Signed-off-by: seonjinn <sna@nvidia.com>
(cherry picked from commit bda8858)
@seonjinn
fix(vllm): default injected reload state safely
b79256b 
Signed-off-by: seonjinn <sna@nvidia.com>
(cherry picked from commit 2aa1570)
@seonjinn
fix(vllm): constrain native TRTLLM refit scope
e3d561d 
Signed-off-by: seonjinn <sna@nvidia.com>
(cherry picked from commit 4aece4b)
@seonjinn
test(vllm): harden native refit failure coverage
a81e440 
Signed-off-by: seonjinn <sna@nvidia.com>
(cherry picked from commit 1292dbe)
@seonjinn
test(vllm): isolate refit capability contract
d791753 
Signed-off-by: seonjinn <sna@nvidia.com>
(cherry picked from commit 56e27a8)
@seonjinn
fix(vllm): harden native layerwise refit failures
6009452 
Signed-off-by: seonjinn <sna@nvidia.com>
(cherry picked from commit 0c33c30)
@seonjinn
Merge remote-tracking branch 'origin/main' into sna/bf16-trtllm-nativ…
8b8d658 
…e-refit-pr-20260807

Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
Merge origin/main into sna/bf16-trtllm-native-refit-pr-20260807
3ea0ce5 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
Merge remote-tracking branch 'origin/main' into sna/bf16-trtllm-nativ…
7cc62b4 
…e-refit-pr-20260807

Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
fix(vllm): detect realized TRTLLM refit backend
a73421e 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(vllm): use module-shaped refit fixtures
7180bff 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
Merge remote-tracking branch 'origin/main' into sna/bf16-trtllm-nativ…
809ddd6 
…e-refit-pr-20260807
@seonjinn
test(vllm): address native refit self-review
95452e9 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(vllm): cover refit buffer reuse
8881a6d 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(vllm): satisfy import ordering
cfd6cab 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
fix(vllm): support native BF16 TRTLLM refit
87af15c 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(vllm): cover TRTLLM NCCL reshard refit
4edab68 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
fix(vllm): support TRTLLM NCCL reshard refit
06b8967 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(vllm): preserve TRTLLM staging dtype
2cadb92 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
docs(vllm): clarify TRTLLM reshard staging
64cbbdc 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
refactor(refit): expose local shard slices
9d5a10a 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(vllm): cover unsupported TRTLLM reshard layout
2844542 
Signed-off-by: seonjinn <sna@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Aug 16, 2026

Copy link
Copy Markdown

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

seonjinn commented Aug 16, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 1dc3cf1

@seonjinn
seonjinn marked this pull request as ready for review August 16, 2026 21:29
@seonjinn
seonjinn requested review from a team as code owners August 16, 2026 21:29
@seonjinn
test(vllm): cover nested layer count config
7f7b8dd 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
fix(vllm): resolve nested text layer counts
2de2d65 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(vllm): use normalized layer names
b5d52ed 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(vllm): cover static MTP ownership
299a357 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
fix(vllm): preserve static MTP drafters
52900b8 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(vllm): accept remote-code model construction
347b24d 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
fix(mxfp8): validate mixed layer boundaries
830a6ca 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
fix(mxfp8): validate scope before patching workers
5b14cad 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
docs(fp8): explain mixed MXFP8 rollout scope
96101df 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
chore: minimize Qwen3.5 TRTLLM recipe
8eddc45 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(vllm): update FP8 refit loader stub
f20709b 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(refit): reproduce mixed TRTLLM init ordering
d8e8088 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
fix(refit): build TRTLLM maps after communicator init
9326f41 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(refit): cover both reshard setup orders
d3f3f06 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
fix(refit): support both reshard setup orders
c2f8c1f 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
Merge main into BF16 TRTLLM NCCL reshard refit
593e4eb 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(vllm): cover FP8 layerwise generator loading
ad6d0dc 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
docs(refit): document BF16 TRTLLM limits
78d066d 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
fix(vllm): preserve FP8 reload iterator contents
08d8687 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
chore(vllm): type tracked reload storage pointers
38deaed 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(qwen3.5): exempt long model setup from idle reaper
7ecfc32 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
test(qwen3.5): align async launch topology
eb7529f 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
fix(vllm): satisfy refit lint checks
90bbc30 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
Merge remote-tracking branch 'upstream/main' into sna/pr3659-review-v…
3d2d769 
…erify-20260906

Signed-off-by: seonjinn <sna@nvidia.com>

# Conflicts:
#	tests/test_suites/disabled.txt
@seonjinn
test(vllm): cover batched BF16 TRTLLM layout
7dfc93a 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
perf(vllm): batch BF16 TRTLLM expert layout conversion
7c1856a 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
fix(vllm): cover full TRTLLM reload lifecycle
1256eaa 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
fix(vllm): isolate TRTLLM reload converter by context
a62e3e0 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
style: format TRTLLM refit tests
2c2b6cd 
Signed-off-by: seonjinn <sna@nvidia.com>
@seonjinn
seonjinn force-pushed the sna/bf16-trtllm-two-gather-refit branch from d2f77c9 to 2c2b6cd Compare September 7, 2026 05:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

CI:L1 Run doctests, unit tests, and functional tests Documentation Improvements or additions to documentation

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@seonjinn