[Performance]: Fuse local NVFP4 tail after NCCL fallback
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 14.7k
- Forks
- 2.8k
- Avg merge
- 2d 23h
- Merged PRs (30d)
- 489
Description
Proposal to improve performance
On topologies without CUDA P2P support, AllReduce::fallbackRunSubsequentOps completes the collective with NCCL and then launches separate residual/RMSNorm and NVFP4 quantization kernels. For the NVFP4 residual-RMSNorm patterns on Blackwell, the existing residualRmsNormFp4Quant kernel can perform that local tail in one launch.
The proposed change would use that fused local kernel after the NCCL collective when all of the following are true:
- the fusion pattern is NVFP4 residual + RMSNorm, with or without a returned high-precision norm output;
- the GPU is SM 10.x or SM 12.x;
- tensors are contiguous FP16/BF16 with matching dtypes;
- the per-tensor scale is a contiguous FP32 scalar;
- there is no bias and the hidden dimension is divisible by 16.
Every unsupported case keeps the current separated fallback path. The change does not alter NCCL itself or the P2P/custom-all-reduce paths.
Performance result
On four RTX 5060 Ti GPUs (SM120, no CUDA P2P), the local tail at FP16 shape [64, 8192] improved from 0.103 ms for the separated sequence to 0.039 ms for the fused kernel: 62.1% lower latency, or 2.64x throughput for this component.
A TensorRT-LLM PyTorch-backend E2E run with NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4, TP4/EP1 and NCCL all-reduce completed 25/25 non-empty requests. This E2E result validates the candidate path but is not presented as a model-level before/after speedup.
Correctness and build validation
- SM120 source build completed with the candidate files compiled and linked.
- Upstream multi-GPU pytest selection: 6 passed across hidden sizes 128/7168 and sequence lengths 16/256/8192. The node reported no P2P support, exercising the NCCL fallback topology.
- Four-rank directed test of
RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4: residual max error 0, RMSNorm max error 0.00390, expected packed NVFP4 output shape. - Pre-commit passed for both changed C++ files.
An unrelated SM120 paired-FP8 PTX build compatibility workaround was needed elsewhere in the checkout to complete the full source build; that workaround is intentionally not part of the proposed change.
Environment
- GPU: 4x NVIDIA GeForce RTX 5060 Ti (SM120)
- Python: 3.12
- CUDA: 13.0
- TensorRT: 11.2.1.2
- PyTorch: 2.12
- NCCL: 2.29.7
- TensorRT-LLM base commit:
0eeda343df8cf44e743d00d1b666b08a43fde83e
I searched existing and closed issues for NCCL fallback, NVFP4, and RMSNorm combinations and did not find an issue covering this local tail fusion.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start at AllReduce::fallbackRunSubsequentOps and inspect the existing residualRmsNormFp4Quant kernel and the separated NCCL fallback tail. Run the upstream multi-GPU pytest selection described in the issue, including hidden sizes 128/7168 and sequence lengths 16/256/8192. Done means supported NVFP4 residual-RMSNorm cases use the fused local tail while unsupported cases retain the current path and correctness checks pass.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, pytorch
- Domain
- backend, machine-learning, performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 52/100