comm: multicast-free (IPC) workspace for TRT-LLM AR fusion under Confidential Computing - #3993
Conversation
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for NVIDIA Confidential Computing (CC) in FlashInfer. Key changes include auto-detection of CC via NVML, a CC-safe autotuner timing path using the GPU %globaltimer register (since cudaEventElapsedTime is unreliable under CC), and multicast-free TRT-LLM AllReduce fusion workspace allocation (disabling symmetric device memory and NVLink multicast-based mnnvl backend under CC). The review feedback points out a potential ZeroDivisionError in test_mnnvl_raises_under_cc if CUDA is not available or device_count() is 0, and suggests adding a guard to skip the test in such environments.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
…idential Computing(CC) create_allreduce_fusion_workspace / TRTLLMAllReduceFusionWorkspace allocate a multicast-free IPC workspace instead of symmetric device memory whenever NVIDIA Confidential Computing is detected (is_confidential_compute(), overridable via FLASHINFER_CONFIDENTIAL_COMPUTE) -- the symmetric-memory setup does a cuMulticast probe that fails under CC. The trtllm one-shot Lamport and two-shot sync fusion kernels are both multicast-free (0 multimem in trtllm_allreduce_fusion.cuh) and the IPC workspace is sized identically to symm, so both strategies run on it -- only the workspace allocator differs. The mnnvl backend requires NVLink multicast, which is unavailable under CC, so create_allreduce_fusion_workspace raises there rather than failing deep inside workspace creation. Adds tests/comm test_allreduce_trtllm_multicast_free (forcing CC via FLASHINFER_CONFIDENTIAL_COMPUTE=1, both kernels via use_oneshot True/False, both AR patterns, fp16/bf16) and test_mnnvl_raises_under_cc. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
724f6a8 to
a4a17a2
Compare
|
Caution Review failedAn error occurred during the review process. Please try again later. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
/bot run tests/comm |
|
[FAILED] Pipeline #58370904: 13/20 passed |
|
@leejnau could you help us to put this PR into P0 priority list? This is for CC. Thanks! |
|
@elvischenv please rebase and fix ci failure. thanks! |
6 participants
Motivation
Under NVIDIA Confidential Computing (CC), the symmetric-memory allocator's
cuMulticastsetup fails (the bounce-buffer path can't complete the fabric/multicast rendezvous), so the TRT-LLM AllReduce-fusion workspace can't be created the usual way.What this changes
create_allreduce_fusion_workspace/TRTLLMAllReduceFusionWorkspacenow allocate a multicast-free IPC workspace instead of symmetric device memory whenever CC is detected — auto-detected viais_confidential_compute()(overridable withFLASHINFER_CONFIDENTIAL_COMPUTE).multimemintrtllm_allreduce_fusion.cuh), and the IPC workspace is sized identically to the symmetric one, so both strategies run on it — only the allocator differs.(ipc_handles, workspace_tensor, metadata)return tuple is handled, andmem_handlesis set to[]so the "handles attached?" guard inallreduce_fusionanddestroy()iterate a no-op.mnnvlbackend requires NVLink multicast (unavailable under CC), socreate_allreduce_fusion_workspaceraises there rather than failing deep inside workspace creation.Tests
tests/comm/test_allreduce_unified_api.py:test_allreduce_trtllm_multicast_free— forces CC viaFLASHINFER_CONFIDENTIAL_COMPUTE=1, exercising both kernels (use_oneshotTrue/False), both AR patterns, and fp16/bf16.test_mnnvl_raises_under_cc— asserts the mnnvl+CC guard.Verified 17/17 on 2×B200.
🤖 Generated with Claude Code