NVIDIA / NVIDIA/TensorRT-LLM

[Performance]: Fuse local NVFP4 tail after NCCL fallback

Open
#17,900 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Model optimization
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.