Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (3)
💤 Files with no reviewable changes (3)
📝 WalkthroughWalkthroughThis PR adds ChangesMNNVL Trigger Completion End-to-End
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 the trigger_completion_at_end parameter to the MNNVL all-reduce implementation, allowing control over Programmatic Dependent Launch (PDL) trigger timing for improved kernel overlap. The changes span the C++ dispatch logic, CUDA kernels (one-shot, two-shot, and RMSNorm), and the Python API. Feedback highlights a critical issue in the rmsNormLamport kernel where cudaGridDependencySynchronize() is placed at the end instead of the beginning, which breaks dependency tracking. Additionally, there is a recommendation to ensure that out-of-bounds thread returns do not inadvertently skip PDL triggers in the one-shot kernel.
| #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) | ||
| cudaTriggerProgrammaticLaunchCompletion(); | ||
| if constexpr (!TriggerCompletionAtEnd) { | ||
| cudaTriggerProgrammaticLaunchCompletion(); | ||
| } | ||
| #endif |
There was a problem hiding this comment.
For PDL-aware kernels, cudaGridDependencySynchronize() should be called at the very beginning of the kernel to establish the dependency on the previous kernel's trigger. In rmsNormLamport, this call is currently placed at the end of the kernel (line 1101), which is incorrect as it allows the kernel to execute its work before the dependency is satisfied.
Please move the cudaGridDependencySynchronize() call to the start of the kernel, before the early trigger check.
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
cudaGridDependencySynchronize();
if constexpr (!TriggerCompletionAtEnd) {
cudaTriggerProgrammaticLaunchCompletion();
}
#endif
There was a problem hiding this comment.
This is the intention of this PR
| // Invariant: every block in the cluster must contain at least one in-bounds | ||
| // thread, otherwise an entirely-OOB block would skip the end PDL trigger and | ||
| // delay launch dependents until the block "completes normally". The | ||
| // adjustGridConfig logic sets blockSize = ceil_div(threadsNeeded, | ||
| // clusterSize) (or larger) so OOB count < clusterSize <= blockSize; this | ||
| // check guards against future tuning regressions. | ||
| int const threadsNeeded = ceil_div(tokenDim, eltsPerThread); | ||
| FLASHINFER_CHECK(threadsNeeded > (clusterSize - 1) * blockSize, | ||
| "[MNNVL AllReduceOneShot] grid config has at least one fully-OOB " | ||
| "block (threadsNeeded=%d, blockSize=%d, clusterSize=%d); PDL end " | ||
| "trigger would not be issued by that block", | ||
| threadsNeeded, blockSize, clusterSize); |
There was a problem hiding this comment.
This check is excellent for ensuring that the adjustGridConfig logic doesn't produce entirely out-of-bounds blocks that would skip the PDL trigger. However, for the TriggerCompletionAtEnd == true case, individual threads that return early due to the OOB check at line 555 will still skip the trigger at line 667. While the cluster-wide check ensures at least one thread per block reaches the trigger, it is generally safer to ensure all threads in a block reach the trigger or that the trigger is called by a specific thread (e.g., threadIdx.x == 0) that is guaranteed to be in-bounds by this check.
There was a problem hiding this comment.
The grid-wide invariant we already added (no entirely-OOB block) is the actual correct condition
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/comm/test_allreduce_unified_api.py`:
- Around line 313-344: The new test test_mnnvl_allreduce_trigger_completion
invokes the MNNVL trigger-path but isn't gated for unsupported GPU SM
capability; update the test to check the appropriate flashinfer.utils capability
helper (e.g., is_sm90a_supported() or the correct SM check for MNNVL) at the
start of test_mnnvl_allreduce_trigger_completion and skip the test when the
capability is absent (use pytest.skip or pytest.mark.skipif) before calling
run_allreduce_test so CI won't run unsupported paths.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: fd6daa1f-3684-4fd0-85b0-59f0335ce3fd
📒 Files selected for processing (5)
csrc/trtllm_mnnvl_allreduce.cuflashinfer/comm/allreduce.pyflashinfer/comm/trtllm_mnnvl_ar.pyinclude/flashinfer/comm/trtllm_mnnvl_allreduce.cuhtests/comm/test_allreduce_unified_api.py
|
@aleozlx Could you assign someone to review this? Thanks! |
…etion-mnnvl # Conflicts: # csrc/trtllm_mnnvl_allreduce.cu # flashinfer/comm/allreduce.py # flashinfer/comm/trtllm_mnnvl_ar.py # include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh
This branch has not been deployed
3 participants
To achieve parity with TRTLLM
SGLang before:
After:
TRTLLM:
Summary by CodeRabbit
trigger_completion_at_endparameter to MNNVL and unified all-reduce/fused all-reduce APIs to control when launch completion is signaled (defaults toTrue), with PDL/S M90+ behavior respected.