fix: expose trigger_completion_at_end through unified API - #2894
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request refines the Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
📝 WalkthroughWalkthrough
Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ 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 a new trigger_completion_at_end parameter to the allreduce_fusion function in flashinfer/comm/allreduce.py. This parameter provides more granular control over Programmatic Dependent Launch (PDL) completion signaling, enabling potential kernel overlap and performance improvements. The review feedback suggests enhancing the docstring to explicitly state that this parameter is specific to the trtllm backend, ensuring clarity for users regarding its applicability.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@flashinfer/comm/allreduce.py`:
- Around line 499-505: The parameter trigger_completion_at_end is documented as
part of the unified contract but is not forwarded/validated in the MNNVL code
path (specifically where MNNVLAllReduceFusionWorkspace is constructed/used),
causing silent no-op behavior; update the code to explicitly reject/raise on any
call that passes trigger_completion_at_end when the workspace is an
MNNVLAllReduceFusionWorkspace (or when backend == "MNNVL") and adjust the
docstring to state that trigger_completion_at_end is TRTLLM-only until MNNVL
implements equivalent support so callers get an immediate error rather than
silent behavior.
- Around line 455-456: The allreduce_fusion signature must make
trigger_completion_at_end keyword-only to avoid breaking positional args: change
the signature in allreduce_fusion so trigger_completion_at_end is declared after
a * (e.g., ..., launch_with_pdl: bool = False, *, trigger_completion_at_end:
Optional[bool] = None, output: Tensor, ...) so existing positional calls keep
binding correctly. In the MNNVL-specific branches inside allreduce_fusion that
currently ignore this parameter, either forward trigger_completion_at_end to the
underlying MNNVL call if you implement support, or add validation that raises
NotImplementedError when trigger_completion_at_end is not None with a clear
message indicating MNNVL does not support it (the TRTLLM path already correctly
forwards the param). Ensure the change references the allreduce_fusion signature
and the MNNVL handling blocks so callers get consistent behavior or a clear
error.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 10857c3a-03b6-4d15-8716-417bb6a348ad
📒 Files selected for processing (1)
flashinfer/comm/allreduce.py
There was a problem hiding this comment.
♻️ Duplicate comments (2)
flashinfer/comm/allreduce.py (2)
498-505:⚠️ Potential issue | 🟠 MajorReject
trigger_completion_at_endon MNNVL until backend support exists.Lines 498-505 document this as part of the unified API, but the MNNVL branch (Lines 667-699) neither forwards nor validates it, causing a silent no-op. Add explicit validation for MNNVL and clarify docs as TRTLLM-only for now.
Proposed fix
elif isinstance(workspace, MNNVLAllReduceFusionWorkspace): + if trigger_completion_at_end is not None: + raise ValueError( + "trigger_completion_at_end is only supported with TRTLLMAllReduceFusionWorkspace" + ) if ( pattern != AllReduceFusionPattern.kARResidualRMSNorm and pattern != AllReduceFusionPattern.kAllReduce ):Also applies to: 648-699
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@flashinfer/comm/allreduce.py` around lines 498 - 505, The MNNVL code path silently ignores the new trigger_completion_at_end flag; update the MNNVL branch in allreduce.py to explicitly reject/validate trigger_completion_at_end when backend == "MNNVL" by raising a clear error (or ValueError) if it is not None/unsupported, and update the docstring/comment near launch_with_pdl/trigger_completion_at_end to state that trigger_completion_at_end is TRTLLM-only for now; look for the backend dispatch / MNNVL-specific block (the code handling MNNVL in the allreduce implementation) and add the validation there referencing trigger_completion_at_end and launch_with_pdl.
451-459:⚠️ Potential issue | 🟠 MajorMake
trigger_completion_at_endkeyword-only to avoid positional API breakage.Line 456 inserts a new positional parameter before
output, so existing positional calls can bind tensors to the wrong argument. Move this new arg to keyword-only at the end of the signature.Proposed fix
def allreduce_fusion( input: torch.Tensor, workspace: AllReduceFusionWorkspace, pattern: int, launch_with_pdl: bool = False, - trigger_completion_at_end: Optional[bool] = None, # ===== OUTPUT tensors (pre-allocated, will be filled) ===== output: Optional[torch.Tensor] = None, residual_out: Optional[torch.Tensor] = None, norm_out: Optional[torch.Tensor] = None, @@ # ===== Control parameters ===== use_oneshot: Optional[bool] = None, fp32_acc: bool = False, + *, + trigger_completion_at_end: Optional[bool] = None, ) -> torch.Tensor:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@flashinfer/comm/allreduce.py` around lines 451 - 459, The new parameter trigger_completion_at_end should be made keyword-only to avoid breaking existing positional calls to allreduce_fusion; move trigger_completion_at_end after the output and residual_out parameters and introduce a bare * before it (e.g. ..., residual_out: Optional[torch.Tensor] = None, *, trigger_completion_at_end: Optional[bool] = None) so only trigger_completion_at_end is keyword-only, then update any internal references and the function signature in allreduce_fusion accordingly.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@flashinfer/comm/allreduce.py`:
- Around line 498-505: The MNNVL code path silently ignores the new
trigger_completion_at_end flag; update the MNNVL branch in allreduce.py to
explicitly reject/validate trigger_completion_at_end when backend == "MNNVL" by
raising a clear error (or ValueError) if it is not None/unsupported, and update
the docstring/comment near launch_with_pdl/trigger_completion_at_end to state
that trigger_completion_at_end is TRTLLM-only for now; look for the backend
dispatch / MNNVL-specific block (the code handling MNNVL in the allreduce
implementation) and add the validation there referencing
trigger_completion_at_end and launch_with_pdl.
- Around line 451-459: The new parameter trigger_completion_at_end should be
made keyword-only to avoid breaking existing positional calls to
allreduce_fusion; move trigger_completion_at_end after the output and
residual_out parameters and introduce a bare * before it (e.g. ...,
residual_out: Optional[torch.Tensor] = None, *, trigger_completion_at_end:
Optional[bool] = None) so only trigger_completion_at_end is keyword-only, then
update any internal references and the function signature in allreduce_fusion
accordingly.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: b3adf7e1-2e15-41e9-abc5-632e31b4efc2
📒 Files selected for processing (1)
flashinfer/comm/allreduce.py
|
/bot run |
nv-yunzheq
changed the title
|
[FAILED] Pipeline #47004057: 13/20 passed |
4 participants
📌 Description
Fix issue #2887, original logic of
trigger_completion_at_endequals toenable_pdlis wrong. Expose as a separate parameter in the unified API.🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
New Features
Documentation