NVIDIA / NVIDIA/TensorRT-LLM

Synchronize bulk-async reductions before sC reuse in Rubin fused MoE finalize

Open
#18,335 2 comments 0 reactions 1 assignee View on GitHub

@peaceh-nv is already working on this.

Since Sep 17, 2026.

Customized kernels
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, and blk_reduce_fp16 calls

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.bulk operations.
  • The epilogue waits for bulk-async completion before sC reuse 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

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.