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)