Synchronize bulk-async reductions before sC reuse in Rubin fused MoE finalize
@peaceh-nv is already working on this.
Since Sep 17, 2026.
- Dominant language
- Python
- Stars
- 14.7k
- Forks
- 2.8k
- Avg merge
- 2d 23h
- Merged PRs (30d)
- 489
Description
Summary
Add bulk-async group synchronization after the fused MoE finalize block reductions in the Rubin SM107 contiguous grouped GEMM kernel.
Rationale
blk_reduce_bf16, blk_reduce_fp32, and blk_reduce_fp16 issue cp.reduce.async.bulk...bulk_group operations. The current epilogue synchronization does not wait for completion of that bulk-async group. A later reuse of the single-stage sC buffer can overwrite data while a reduction is still reading it.
This follow-up is intentionally separate from the merge-back because the current implementation matches the validated source branch.
Affected area
tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_grouped_gemm_finalize_fusion.py- Fused finalize epilogue after the dtype-specific
blk_reduce_bf16,blk_reduce_fp32, andblk_reduce_fp16calls
Required change
After the dtype-specific reduction calls, commit and wait for the bulk-async group before sC can be reused and before tmem.free:
cute.arch.cp_async_bulk_commit_group()cute.arch.cp_async_bulk_wait_group(0, read=True)
Keep the existing dtype dispatch unchanged.
Acceptance criteria
- The epilogue commits the issued
cp.reduce.async.bulkoperations. - The epilogue waits for bulk-async completion before
sCreuse and before tensor-memory release. - The BF16, FP32, and FP16 reduction paths retain their existing dispatch behavior.
- Kernel validation covers repeated/multi-tile execution that reuses
sC.
Backlinks
- Pull request: https://github.com/NVIDIA/TensorRT-LLM/pull/18311
- Review comment: https://github.com/NVIDIA/TensorRT-LLM/pull/18311#discussion_r3872091317
- Requested by: @farazkh80
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.