feat(distributed): wire TensorRT-LLM all-reduce backend

[Xiangyi1996]

Summary

• add an explicit --enable-trtllm-allreduce opt-in for eligible single-node NVIDIA tensor-parallel collectives
• configure deduplicated attention, dense, and pure-TP MoE groups before CUDA graph capture; groups containing expert parallelism remain on NCCL
• fail closed on unsupported topology, dtype, tensor rank, shape, operation, or workspace capacity and fall back to NCCL for ineligible calls
• size the one-shot workspace limit from the selected Lamport buffer dtype, including the FP32 path
• expose workspace cleanup through the tokenspeed-kernel boundary
• add backend-selection, registered 2-GPU correctness/fallback/CUDA Graph replay coverage, manual TP8 coverage, and server configuration documentation

Root cause

TrtllmAllReduceBackend existed, but no runtime call configured its per-group IPC workspace. Consequently AutoBackend.has_trtllm_ar() remained false and DeepSeek-V4-Pro TP8 decode collectives fell back to NCCL Ring LL.

This PR wires the workspace explicitly after process-group initialization and before model CUDA Graph capture. The feature remains disabled by default.

Selection and fallback

The Lamport path is selected only for configured, contiguous, 2-D CUDA BF16 SUM tensors that fit the workspace shape and the 2 MiB one-shot byte limit. Large prefill tensors, non-2-D layouts, unsupported dtypes/ops, and unconfigured groups continue to use NCCL. MoE groups containing expert parallelism are intentionally out of scope for this change.

Decode trace evidence

Matched exact p4 decode traces show the backend switch directly:

Kernel	Off launches/step	On launches/step	Off duration/step	On duration/step
ncclDevKernel_AllReduce_Sum_bf16_RING_LL	130.11	0 (not observed)	4.848 ms	0
allreduce_fusion_kernel_oneshot_lamport	0 (not observed)	130.68	0	1.976 ms

The on-side decode window therefore executed the Lamport kernel rather than silently falling back. The nearly identical launch counts show that the comparison switches the backend for the same collective workload.

Validation

• pre-commit run --all-files passed on the final local head
• the new CUDA test is registered in the per-commit runtime-2gpu suite; it covers eager correctness, oversized and 1-D NCCL fallback, and 10 changing-input CUDA Graph replays
• TP8 BF16 eager correctness matched NCCL in the manual 8-GPU run
• TP8 CUDA Graph replay passed for 10 changing per-rank inputs
• full DeepSeek-V4-Pro CUDA Graph capture passed for batch sizes [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 79, 80]
• matched non-profile p1→p2→p4→p8 curves completed with all requests successful on both sides

At p4, enabling the backend improved CI TPS/user by 8.99%, CI TPS/GPU by 10.29%, and steady completion throughput by 11.69%; TPOT decreased by 8.25%. Cache hit and speculative acceptance were unchanged. Across p1, p2, p4, and p8, TPS/user improved by 8.99–12.69% and TPOT decreased by 8.25–11.26%.

The endpoint curve currently has one matched off/on run per concurrency point, so it does not claim a variance estimate.

Stacking

This draft is intentionally based on the head branch of #563 so that its diff contains only the all-reduce wiring changes.

[github-actions]

This PR has been inactive for 14 days and is marked as stale. It will be closed in 3 days if there is no further activity.