Skip to content

feat: update flashinfer ar oneshot params #22108

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

feat: update flashinfer ar oneshot params #22108

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1aac329
init
yyihuang Aug 1, 2025
49a5a78
Merge branch 'auto-oneshot' of github.com:yyihuang/vllm; branch 'main…
yyihuang Aug 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion vllm/compilation/collective_fusion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -457,7 +457,6 @@ def call_trtllm_fused_allreduce_norm(
hidden_dim=allreduce_in.shape[-1],
workspace_ptrs=_FI_WORKSPACE_TENSOR,
launch_with_pdl=launch_with_pdl,
use_oneshot=True,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you have more context over this? At the very least the comment above might need to be updated. "For the sizes that are smaller than the max size,
# we only use flashinfer one shot allreduce". Is there a test we can add for this?

Copy link
Contributor Author

@yyihuang yyihuang Aug 7, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We deduce the strategy by token num: https://github.com/flashinfer-ai/flashinfer/blob/main/flashinfer/comm/trtllm_ar.py#L826.

Different kernels would be called by this strategy: https://github.com/flashinfer-ai/flashinfer/blob/main/include/flashinfer/comm/trtllm_allreduce_fusion.cuh#L1388-L1400

We can add a unit test of token_num > 128 if needed. And using some general model tests would also be okay.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

And you might be interested in the first of the PR series on flashinfer's allreduce_fusion.
#20691

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@yyihuang Do you have results of benchmarking for oneshot vs twoshot? Firstly, usage of two shot should not only depend on token_num but world_size, similarly what is done in custom_all_reduce.cuh. Secondly, In my benchmarking using two shot only made sense on Hopper, whereas one shot was better across all workloads on Blackwell.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @ilmarkov I benched on h200. Let me do more benchmarks on blackwell.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The benchmark results are updated above. thx for review @mgoin @ilmarkov

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@yyihuang I don't see any speedups with this PR. Also, in this PR, when I do isolated benchmarking against non-fused version, FI has different performance results on H100 and B200 but in none of them two shot is the one we'd want to use.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for you benchmark @ilmarkov! Let's keep this as a draft PR since we did not get speedup by this auto, until we figure out the problem shape and use case of each strategy cross-DLFW. In tllm we're taking this for min-latency case (https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/allreduceOp.cpp#L453), which might not be the target case or there might be some framework diffs.

trigger_completion_at_end=trigger_completion_at_end,
fp32_acc=fp32_acc,
pattern_code=pattern_code,
Expand Down