From f47ffa0d7dfd5fe5d73d5ab64d6a3bd0bca431b8 Mon Sep 17 00:00:00 2001 From: Juice Date: Fri, 6 Nov 2026 15:30:00 +0000 Subject: [PATCH] [None][perf] CuTe DSL act-fusion GEMMs: key the early accumulator release on the loop index With a 256-wide MMA N tile the gather/act-fusion MoE FC1 kernel and the dense act-fusion GEMM overlap their two TMEM accumulator buffers (overlapping_accum). On every other tile the epilogue walks its subtiles in reverse, so that it reads the overlapped columns first, and it releases the accumulator after iter_acc_early_release_in_epilogue + 1 subtiles so the next tile's MMAs can start. Both kernels compared that count with the subtile index after the reverse remap, so on reversed tiles the accumulator was released only after the last subtile and the next tile's MMAs waited for the whole epilogue. Compare the subtile's position in processing order instead, as the other CuTe DSL grouped GEMMs already do. Loads, stores and outputs are unchanged. Signed-off-by: Juice Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_01VSV19DiCRbvQsMxdgAHrb8 --- .../kernels/blackwell/dense_blockscaled_gemm_act_fusion.py | 5 ++++- .../blockscaled_contiguous_gather_grouped_gemm_act_fusion.py | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/kernels/blackwell/dense_blockscaled_gemm_act_fusion.py b/tensorrt_llm/_torch/kernels/blackwell/dense_blockscaled_gemm_act_fusion.py index 4861940e9def..bc8baca6d673 100644 --- a/tensorrt_llm/_torch/kernels/blackwell/dense_blockscaled_gemm_act_fusion.py +++ b/tensorrt_llm/_torch/kernels/blackwell/dense_blockscaled_gemm_act_fusion.py @@ -1551,6 +1551,9 @@ def kernel( real_subtile_idx = subtile_idx // 2 else: real_subtile_idx = subtile_idx + # Position in processing order. The early accumulator release below counts + # processed subtiles, so it must not use the reversed index. + iter_subtile_idx = real_subtile_idx if cutlass.const_expr(self.overlapping_accum): if reverse_subtile: real_subtile_idx = ( @@ -1584,7 +1587,7 @@ def kernel( # Async arrive accumulator buffer empty earlier when overlapping_accum is enabled # if cutlass.const_expr(self.overlapping_accum): - if real_subtile_idx == self.iter_acc_early_release_in_epilogue: + if iter_subtile_idx == self.iter_acc_early_release_in_epilogue: # Fence for TMEM load cute.arch.fence_view_async_tmem_load() with cute.arch.elect_one(): diff --git a/tensorrt_llm/_torch/moe/kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py b/tensorrt_llm/_torch/moe/kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py index 7563b1f9483c..c41bbaefff38 100644 --- a/tensorrt_llm/_torch/moe/kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py +++ b/tensorrt_llm/_torch/moe/kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py @@ -2501,6 +2501,9 @@ def kernel( real_subtile_idx = subtile_idx // 2 else: real_subtile_idx = subtile_idx + # Position in processing order. The early accumulator release below counts + # processed subtiles, so it must not use the reversed index. + iter_subtile_idx = real_subtile_idx if cutlass.const_expr(self.overlapping_accum): if reverse_subtile: real_subtile_idx = ( @@ -2528,7 +2531,7 @@ def kernel( # Async arrive accumulator buffer empty earlier when overlapping_accum is enabled # if cutlass.const_expr(self.overlapping_accum): - if real_subtile_idx == self.iter_acc_early_release_in_epilogue: + if iter_subtile_idx == self.iter_acc_early_release_in_epilogue: # Fence for TMEM load cute.arch.fence_view_async_tmem_load() acc_pipeline.consumer_release(acc_consumer_state)