From 47074caa2a5b6b3f395f2ccc0a9b4188918fd402 Mon Sep 17 00:00:00 2001
From: Oleg Goncharov <ogoncharov@nvidia.com>
Date: Sat, 22 Aug 2026 15:46:40 +0000
Subject: [PATCH 1/3] Fix TMA shared-memory reuse synchronization

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>
---
 .../common/cast/fp8/gated_fp8.cuh             |  9 ++++-
 .../common/cast/fp8/quantize_fp8.cuh          | 22 +++++++++---
 .../common/cast/fused_group_requantize.cu     |  2 +-
 .../common/cast/mxfp8/dequantize_mxfp8.cuh    | 11 ++++--
 .../common/cast/mxfp8/gated_mxfp8.cuh         | 13 ++++---
 .../cast/mxfp8/group_dequantize_mxfp8.cuh     | 11 ++++--
 .../common/cast/mxfp8/quantize_mxfp8.cuh      | 20 ++++++++---
 .../cast/mxfp8/specialized/quantize_mxfp8.cuh | 22 ++++++++----
 .../nvfp4/group_quantize_transpose_nvfp4.cuh  | 18 +++++++---
 .../cast/nvfp4/quantize_transpose_nvfp4.cuh   | 36 ++++++++++++++-----
 .../quantize_transpose_nvfp4_tuned_1D.cuh     | 18 ++++++++--
 .../quantize_transpose_square_blockwise.cu    |  4 +--
 transformer_engine/common/util/ptx.cuh        | 18 +++++-----
 13 files changed, 155 insertions(+), 49 deletions(-)

diff --git a/transformer_engine/common/cast/fp8/gated_fp8.cuh b/transformer_engine/common/cast/fp8/gated_fp8.cuh
index 631143c6aef..9e7f844e13b 100644
--- a/transformer_engine/common/cast/fp8/gated_fp8.cuh
+++ b/transformer_engine/common/cast/fp8/gated_fp8.cuh
@@ -247,8 +247,15 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
       // Wait for TMA transfer to have finished reading shared memory.
       ptx::cp_async_bulk_wait_group_read<BUFFERS_NUM - 1>();
     }
+    // The bulk async-group is owned by the issuing thread. Hand its completion
+    // off to all cooperative writers before the output ring can be reused.
+    if (next_it < ITERATIONS && next_it >= BUFFERS_NUM) {
+      __syncthreads();
+    }
+  }
+  if (is_master_thread) {
+    ptx::cp_async_bulk_wait_group();
   }
-  ptx::cp_async_bulk_wait_group_read<0>();
   __syncthreads();
 
   if (amax_ptr != nullptr) {
diff --git a/transformer_engine/common/cast/fp8/quantize_fp8.cuh b/transformer_engine/common/cast/fp8/quantize_fp8.cuh
index 9b1ebf5c040..1f6e33e441a 100644
--- a/transformer_engine/common/cast/fp8/quantize_fp8.cuh
+++ b/transformer_engine/common/cast/fp8/quantize_fp8.cuh
@@ -198,10 +198,17 @@ __global__ void __launch_bounds__(FP8_THREADS_PER_CHUNK)
       ptx::cp_async_bulk_commit_group();
 
       // Wait for TMA transfer to have finished reading shared memory.
-      ptx::cp_async_bulk_wait_group_read<FP8_PREFETCH_BUFFERS_NUM>();
+      ptx::cp_async_bulk_wait_group_read<FP8_BUFFERS_NUM - 1>();
     }
+    // The bulk async-group is owned by the issuing thread. Hand its completion
+    // off to all cooperative writers before the output ring can be reused.
+    if (iter + 1 < FP8_ITERATIONS && iter + 1 >= FP8_BUFFERS_NUM) {
+      __syncthreads();
+    }
+  }
+  if (is_master_thread) {
+    ptx::cp_async_bulk_wait_group();
   }
-  ptx::cp_async_bulk_wait_group_read<0>();
   __syncthreads();
 
   parity ^= 1;
@@ -331,10 +338,17 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK)
       ptx::cp_async_bulk_commit_group();
 
       // Wait for TMA transfer to have finished reading shared memory.
-      ptx::cp_async_bulk_wait_group_read<1>();
+      ptx::cp_async_bulk_wait_group_read<SHMEM_BUFFERS - 1>();
     }
+    // The bulk async-group is owned by the issuing thread. Hand its completion
+    // off to all cooperative writers before the output ring can be reused.
+    if (next_iter < ITERATIONS && next_iter >= SHMEM_BUFFERS) {
+      __syncthreads();
+    }
+  }
+  if (is_master_thread) {
+    ptx::cp_async_bulk_wait_group();
   }
-  ptx::cp_async_bulk_wait_group_read<0>();
   __syncthreads();
 
   if (amax_ptr != nullptr) {
diff --git a/transformer_engine/common/cast/fused_group_requantize.cu b/transformer_engine/common/cast/fused_group_requantize.cu
index a2156f369a9..7efe629d112 100644
--- a/transformer_engine/common/cast/fused_group_requantize.cu
+++ b/transformer_engine/common/cast/fused_group_requantize.cu
@@ -380,7 +380,7 @@ __global__ void __launch_bounds__(kThreads)
         reinterpret_cast<const uint64_t *>(&output_tensor_map), col_base, row_base,
         reinterpret_cast<uint64_t *>(quantized));
     ptx::cp_async_bulk_commit_group();
-    ptx::cp_async_bulk_wait_group_read<0>();
+    ptx::cp_async_bulk_wait_group();
   }
   __syncthreads();
 #else
diff --git a/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh
index 1face261bd6..91af8b139d9 100644
--- a/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh
@@ -213,10 +213,17 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
       ptx::cp_async_bulk_commit_group();
 
       // Wait for TMA transfer to have finished reading shared memory.
-      ptx::cp_async_bulk_wait_group_read<1>();
+      ptx::cp_async_bulk_wait_group_read<BUFFERS_NUM - 1>();
     }
+    // The bulk async-group is owned by the issuing thread. Hand its completion
+    // off to all cooperative writers before the output ring can be reused.
+    if (next_iter < ITERATIONS && next_iter >= BUFFERS_NUM) {
+      __syncthreads();
+    }
+  }
+  if (is_master_thread) {
+    ptx::cp_async_bulk_wait_group();
   }
-  ptx::cp_async_bulk_wait_group_read<0>();
   __syncthreads();
 
   parity ^= 1;
diff --git a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh
index 44dea1a7c76..31e445d4fb4 100644
--- a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh
@@ -74,6 +74,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
 
   constexpr size_t STAGES = CHUNK_DIM_Y / BUFF_DIM_Y;
   static_assert(STAGES >= 1);
+  static_assert(STAGES <= BUFFS_NUM,
+                "Output-buffer reuse requires a TMA completion handoff to the CTA");
 
   constexpr bool IS_CACHED_ACT_OP = ROWWISE_SCALING && COLWISE_SCALING;
   constexpr bool ONLY_COLWISE_SCALING = COLWISE_SCALING && (!ROWWISE_SCALING);
@@ -195,10 +197,6 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
     const size_t stage_offset_Y = stage * BUFF_DIM_Y;
 
     if (next_stage < STAGES) {
-      // Wait for TMA transfer to have finished reading shared memory.
-      // I.e. the buffer is ready to be written to
-      ptx::cp_async_bulk_wait_group_read<1>();
-
       const size_t next_buff = next_stage % BUFFS_NUM;
       const size_t next_stage_offset_Y = next_stage * BUFF_DIM_Y;
       const size_t global_offset_Y = block_offset_Y + next_stage_offset_Y;
@@ -704,6 +702,13 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
   }
 
   parity ^= 1;
+
+  // Ensure all S2G operations issued by the master have completed before the CTA exits.
+  if (is_master_thread) {
+    ptx::cp_async_bulk_wait_group();
+  }
+  __syncthreads();
+
   destroy_barriers<STAGES>(mbar, is_master_thread);
 #endif  // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
 }  // NOLINT(readability/fn_size)
diff --git a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh
index dad8d18d6f8..f4ae221601d 100644
--- a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh
@@ -360,10 +360,17 @@ __global__ void __launch_bounds__(128)
           chunk_it_offset_y, reinterpret_cast<uint64_t *>(&out_sh[buff]));
 
       ptx::cp_async_bulk_commit_group();
-      ptx::cp_async_bulk_wait_group_read<1>();
+      ptx::cp_async_bulk_wait_group_read<BUFFERS_NUM - 1>();
     }
+    // The bulk async-group is owned by the issuing thread. Hand its completion
+    // off to all cooperative writers before the output ring can be reused.
+    if (next_iter < ITERATIONS && next_iter >= BUFFERS_NUM) {
+      __syncthreads();
+    }
+  }
+  if (is_master_thread) {
+    ptx::cp_async_bulk_wait_group();
   }
-  ptx::cp_async_bulk_wait_group_read<0>();
   __syncthreads();
 
   destroy_barriers<ITERATIONS>(mbar, is_master_thread);
diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh
index 9f312ac3f56..b286d80697c 100644
--- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh
@@ -180,11 +180,17 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
     const size_t next_stage = stage + 1;
     const size_t stage_offset_Y = stage * BUFF_DIM_Y;
 
-    if (next_stage < STAGES) {
-      // Wait for TMA transfer to have finished reading shared memory.
-      // I.e. the buffer is ready to be written to
-      ptx::cp_async_bulk_wait_group_read<1>();
+    // S2G bulk async-groups are thread-local, so only the issuing thread can wait for the
+    // previous TMA read of this output buffer to complete. Hand that completion off to the
+    // rest of the CTA before any thread overwrites the reused buffer.
+    if (stage >= BUFFS_NUM) {
+      if (is_master_thread) {
+        ptx::cp_async_bulk_wait_group_read<BUFFS_NUM - 1>();
+      }
+      __syncthreads();
+    }
 
+    if (next_stage < STAGES) {
       const size_t next_buff = next_stage % BUFFS_NUM;
       const size_t next_stage_offset_Y = next_stage * BUFF_DIM_Y;
       const size_t global_offset_Y = block_offset_Y + next_stage_offset_Y;
@@ -578,6 +584,12 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
     atomicMaxFloat(amax_ptr, block_amax);
   }
 
+  // Ensure all S2G operations issued by the master have completed before the CTA exits.
+  if (is_master_thread) {
+    ptx::cp_async_bulk_wait_group();
+  }
+  __syncthreads();
+
   destroy_barriers<STAGES>(mbar, is_master_thread);
 #endif  // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
 }
diff --git a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh
index 24aa154393e..9a58a88c18f 100644
--- a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh
@@ -1027,8 +1027,8 @@ __global__ void quantize_mxfp8_kernel_cast_only(
         ptx::cp_async_bulk_wait_group_read<CastTraits::numStages - 1>();
         ptx::mbarrier_arrive_expect_tx(&stg_consumer[read_state.index()], 0u);
       }
+      ptx::cp_async_bulk_wait_group();
     }
-    ptx::cp_async_bulk_wait_group_read<0>();
   } else {
     PipeState<CastTraits::numStages> read_state;
 
@@ -1510,6 +1510,17 @@ __global__ void quantize_mxfp8_kernel_cast_only(
   }
 #pragma unroll 1
   for (int32_t iter = 0; iter < CastTraits::iterLayout::num; iter++) {
+    // S2G bulk async-groups are thread-local. Once the output ring wraps, the issuing
+    // thread must observe completion and hand it off to the CTA before cooperative STS.
+    if constexpr (CastTraits::_need_wait_group) {
+      if (iter >= CastTraits::numStages) {
+        if (warpId == 0 && leader) {
+          ptx::cp_async_bulk_wait_group_read<CastTraits::numStages - 1>();
+        }
+        __syncthreads();
+      }
+    }
+
     {
       int32_t next = iter + (CastTraits::numStages - 1);
       int32_t next_stage = next % CastTraits::numStages;
@@ -1520,10 +1531,6 @@ __global__ void quantize_mxfp8_kernel_cast_only(
       coords.x = block_coords.x + iter_n * CastTraits::blockIterDim::N;
       if (coords.x < cols && coords.y < rows) {
         if (warpId == 0 && leader) {
-          if constexpr (CastTraits::_need_wait_group) {
-            ptx::cp_async_bulk_wait_group_read<CastTraits::numStages - 1>();
-          }
-
           ptx::cp_async_bulk_tensor_2d_global_to_shared(
               reinterpret_cast<uint64_t *>(sInput + next_stage * CastTraits::blockIterDim::num),
               reinterpret_cast<const uint64_t *>(&tensor_map_input),
@@ -1984,7 +1991,10 @@ __global__ void quantize_mxfp8_kernel_cast_only(
     }
   }
 
-  ptx::cp_async_bulk_wait_group_read<0>();
+  if (warpId == 0 && leader) {
+    ptx::cp_async_bulk_wait_group();
+  }
+  __syncthreads();
 
 #endif  // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
 }  // NOLINT(readability/fn_size)
diff --git a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh
index 3c6d9585e41..7347336a32c 100644
--- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh
+++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh
@@ -331,6 +331,15 @@ __global__ void __launch_bounds__(THREADS_NUM)
     const size_t buff_offset_out = buff * BUFF_OUT_SIZE;
     const size_t buff_offset_out_t = buff * BUFF_OUT_T_SIZE;
 
+    if (stage >= BUFFS_NUM) {
+      if (is_master_thread) {
+        ptx::cp_async_bulk_wait_group_read<BUFFS_NUM - 1>();
+      }
+      // Bulk async-groups are thread-local. Publish the master's completion before any thread
+      // overwrites this reused output buffer.
+      __syncthreads();
+    }
+
     // for stages from 1 to STAGES - 1, we need to update the tensor id
     // skip updating tensor id if it's the last CTA, and some stages will be out of bounds
     if (need_update_tensor_id && stage > 0 && (block_offset_Y + stage_offset_Y < rows)) {
@@ -349,10 +358,6 @@ __global__ void __launch_bounds__(THREADS_NUM)
     }
 
     if (next_stage < STAGES) {
-      // Wait for TMA transfer to have finished reading shared memory.
-      // I.e. the buffer is ready to be written to
-      ptx::cp_async_bulk_wait_group_read<1>();
-
       const size_t next_buff = next_stage % BUFFS_NUM;
       const size_t next_stage_offset_Y = next_stage * BUFF_DIM_Y;
       const size_t global_offset_Y = block_offset_Y + next_stage_offset_Y;
@@ -726,6 +731,11 @@ __global__ void __launch_bounds__(THREADS_NUM)
   //   }
   // }
 
+  if (is_master_thread) {
+    ptx::cp_async_bulk_wait_group();
+  }
+  __syncthreads();
+
   destroy_barriers<STAGES>(mbar, is_master_thread);
 #else
   NVTE_DEVICE_ERROR("sm_100 or higher is required.");
diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh
index 2734cf3fefb..1edc69de5f4 100644
--- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh
+++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh
@@ -455,11 +455,16 @@ __global__ void __launch_bounds__(THREADS_NUM)
     const size_t buff_offset_out = buff * BUFF_OUT_SIZE;
     const size_t buff_offset_out_t = buff * BUFF_OUT_T_SIZE;
 
-    if (next_stage < STAGES) {
-      // Wait for TMA transfer to have finished reading shared memory.
-      // I.e. the buffer is ready to be written to
-      ptx::cp_async_bulk_wait_group_read<1>();
+    if (stage >= BUFFS_NUM) {
+      if (is_master_thread) {
+        ptx::cp_async_bulk_wait_group_read<BUFFS_NUM - 1>();
+      }
+      // Bulk async-groups are thread-local. Publish the master's completion before any thread
+      // overwrites this reused output buffer.
+      __syncthreads();
+    }
 
+    if (next_stage < STAGES) {
       const size_t next_buff = next_stage % BUFFS_NUM;
       const size_t next_stage_offset_Y = next_stage * BUFF_DIM_Y;
       const size_t global_offset_Y = block_offset_Y + next_stage_offset_Y;
@@ -843,6 +848,11 @@ __global__ void __launch_bounds__(THREADS_NUM)
     }
   }
 
+  if (is_master_thread) {
+    ptx::cp_async_bulk_wait_group();
+  }
+  __syncthreads();
+
   destroy_barriers<STAGES>(mbar, is_master_thread);
 #else
   NVTE_DEVICE_ERROR("sm_100 or higher is required.");
@@ -1010,11 +1020,16 @@ __global__ void __launch_bounds__(THREADS_NUM)
     const size_t buff_offset_out = buff * BUFF_OUT_SIZE;
     const size_t buff_offset_out_t = buff * BUFF_OUT_T_SIZE;
 
-    if (next_stage < STAGES) {
-      // Wait for TMA transfer to have finished reading shared memory.
-      // I.e. the buffer is ready to be written to
-      ptx::cp_async_bulk_wait_group_read<1>();
+    if (stage >= BUFFS_NUM) {
+      if (is_master_thread) {
+        ptx::cp_async_bulk_wait_group_read<BUFFS_NUM - 1>();
+      }
+      // Bulk async-groups are thread-local. Publish the master's completion before any thread
+      // overwrites this reused output buffer.
+      __syncthreads();
+    }
 
+    if (next_stage < STAGES) {
       const size_t next_buff = next_stage % BUFFS_NUM;
       const size_t next_stage_offset_Y = next_stage * BUFF_DIM_Y;
       const size_t global_offset_Y = block_offset_Y + next_stage_offset_Y;
@@ -1401,6 +1416,11 @@ __global__ void __launch_bounds__(THREADS_NUM)
     }
   }
 
+  if (is_master_thread) {
+    ptx::cp_async_bulk_wait_group();
+  }
+  __syncthreads();
+
   destroy_barriers<STAGES>(mbar, is_master_thread);
 #endif  // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
 }
diff --git a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh
index cdd0d4916af..a038ed34e72 100644
--- a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh
+++ b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh
@@ -80,7 +80,11 @@ constexpr int STAGES = STAGES_Y * STAGES_X;
 constexpr int BUFFS_NUM = TunableConfig::PREFETCH_STAGES + 1;
 constexpr int BUFFS_NUM_IN = BUFFS_NUM;
 constexpr int BUFFS_NUM_OUT = BUFFS_NUM;
-constexpr int BUFFS_NUM_OUT_TR = 2;
+constexpr int BUFFS_NUM_OUT_TR = TunableConfig::PREFETCH_STAGES + 1;
+static_assert(BUFFS_NUM_OUT >= TunableConfig::PREFETCH_STAGES + 1,
+              "The output buffer ring must be larger than the TMA store wait depth");
+static_assert(BUFFS_NUM_OUT_TR >= TunableConfig::PREFETCH_STAGES + 1,
+              "The transposed output buffer ring must be larger than the TMA store wait depth");
 constexpr int BUFF_DIM_Y = TILE_DIM_Y;
 constexpr int BUFF_DIM_X = TILE_DIM_X;
 constexpr int BUFF_SIZE = BUFF_DIM_Y * BUFF_DIM_X;
@@ -585,7 +589,12 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D
 
       // Wait for TMA transfer to have finished reading shared memory
       // I.e. the OUT buffer is ready to be written to
-      ptx::cp_async_bulk_wait_group_read<TunableConfig::PREFETCH_STAGES>();
+      if (leading_thread) {
+        ptx::cp_async_bulk_wait_group_read<TunableConfig::PREFETCH_STAGES>();
+      }
+      // Bulk async-groups are thread-local. Publish the leading thread's completion to all
+      // threads before they cooperatively overwrite a reused output buffer.
+      __syncthreads();
 
       // NVFP4 Quantization
       rowwise_scaling<USE_STOCHASTIC_ROUNDING, USE_FAST_MATH, ROW_SCALED_NVFP4>(
@@ -673,6 +682,11 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D
     }
   }
 
+  if (leading_thread) {
+    ptx::cp_async_bulk_wait_group();
+  }
+  __syncthreads();
+
   if (leading_thread) {
 #pragma unroll
     for (int buff = 0; buff < BUFFS_NUM; ++buff) {
diff --git a/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu b/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu
index 02d64bcfffa..2d1a62e0fae 100644
--- a/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu
+++ b/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu
@@ -230,8 +230,8 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK)
       // Wait for TMA transfer to have finished reading shared memory.
       // Create a "bulk async-group" out of the previous bulk copy operation.
       ptx::cp_async_bulk_commit_group();
-      // Wait for the group to have completed reading from shared memory.
-      ptx::cp_async_bulk_wait_group_read<0>();
+      // Wait for the group to have completed the shared-to-global transfer.
+      ptx::cp_async_bulk_wait_group();
     }
 #else
     // Step 4 Alternative (when TMA is not available, skip writing to shared memory)
diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh
index f48224c3651..a84de8681ee 100644
--- a/transformer_engine/common/util/ptx.cuh
+++ b/transformer_engine/common/util/ptx.cuh
@@ -374,7 +374,7 @@ __device__ __forceinline__ void cp_async_bulk_tensor_2d_shared_to_global(
 // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-wait-group
 __device__ __forceinline__ void cp_async_bulk_wait_group() {
 #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-  asm volatile("cp.async.bulk.wait_group 0;");
+  asm volatile("cp.async.bulk.wait_group 0;" ::: "memory");
 #else
   NVTE_DEVICE_ERROR("cp_async_bulk_wait_group is only supported on SM 9.0+.");
 #endif  // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
@@ -384,7 +384,7 @@ __device__ __forceinline__ void cp_async_bulk_wait_group() {
 template <size_t W>
 __device__ __forceinline__ void cp_async_bulk_wait_group_read() {
 #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-  asm volatile("cp.async.bulk.wait_group.read 0;");
+  asm volatile("cp.async.bulk.wait_group.read 0;" ::: "memory");
 #else
   NVTE_DEVICE_ERROR("cp_async_bulk_wait_group_read is only supported on SM 9.0+.");
 #endif  // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
@@ -393,7 +393,7 @@ __device__ __forceinline__ void cp_async_bulk_wait_group_read() {
 template <>
 __device__ __forceinline__ void cp_async_bulk_wait_group_read<0>() {
 #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-  asm volatile("cp.async.bulk.wait_group.read 0;");
+  asm volatile("cp.async.bulk.wait_group.read 0;" ::: "memory");
 #else
   NVTE_DEVICE_ERROR("cp_async_bulk_wait_group_read is only supported on SM 9.0+.");
 #endif  // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
@@ -401,7 +401,7 @@ __device__ __forceinline__ void cp_async_bulk_wait_group_read<0>() {
 template <>
 __device__ __forceinline__ void cp_async_bulk_wait_group_read<1>() {
 #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-  asm volatile("cp.async.bulk.wait_group.read 1;");
+  asm volatile("cp.async.bulk.wait_group.read 1;" ::: "memory");
 #else
   NVTE_DEVICE_ERROR("cp_async_bulk_wait_group_read is only supported on SM 9.0+.");
 #endif  // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
@@ -409,7 +409,7 @@ __device__ __forceinline__ void cp_async_bulk_wait_group_read<1>() {
 template <>
 __device__ __forceinline__ void cp_async_bulk_wait_group_read<2>() {
 #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-  asm volatile("cp.async.bulk.wait_group.read 2;");
+  asm volatile("cp.async.bulk.wait_group.read 2;" ::: "memory");
 #else
   NVTE_DEVICE_ERROR("cp_async_bulk_wait_group_read is only supported on SM 9.0+.");
 #endif  // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
@@ -417,7 +417,7 @@ __device__ __forceinline__ void cp_async_bulk_wait_group_read<2>() {
 template <>
 __device__ __forceinline__ void cp_async_bulk_wait_group_read<4>() {
 #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-  asm volatile("cp.async.bulk.wait_group.read 4;");
+  asm volatile("cp.async.bulk.wait_group.read 4;" ::: "memory");
 #else
   NVTE_DEVICE_ERROR("cp_async_bulk_wait_group_read is only supported on SM 9.0+.");
 #endif  // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
@@ -426,7 +426,7 @@ __device__ __forceinline__ void cp_async_bulk_wait_group_read<4>() {
 // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-commit-group
 __device__ __forceinline__ void cp_async_bulk_commit_group() {
 #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-  asm volatile("cp.async.bulk.commit_group;");
+  asm volatile("cp.async.bulk.commit_group;" ::: "memory");
 #else
   NVTE_DEVICE_ERROR("cp_async_bulk_commit_group is only supported on SM 9.0+.");
 #endif  // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
@@ -435,7 +435,7 @@ __device__ __forceinline__ void cp_async_bulk_commit_group() {
 // Proxy fence (bi-directional):
 __device__ __forceinline__ void fence_proxy_async() {
 #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-  asm volatile("fence.proxy.async;");
+  asm volatile("fence.proxy.async;" ::: "memory");
 #else
   NVTE_DEVICE_ERROR("fence_proxy_async is only supported on SM 9.0+.");
 #endif  // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
@@ -443,7 +443,7 @@ __device__ __forceinline__ void fence_proxy_async() {
 
 __device__ __forceinline__ void fence_proxy_async_shared_cta() {
 #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-  asm volatile("fence.proxy.async.shared::cta;");
+  asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
 #else
   NVTE_DEVICE_ERROR("fence_proxy_async_shared_cta is only supported on SM 9.0+.");
 #endif  // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)

From 7fc6459603c66cc949d5c82e88ad8fc697e894a7 Mon Sep 17 00:00:00 2001
From: Oleg Goncharov <ogoncharov@nvidia.com>
Date: Fri, 11 Sep 2026 16:54:47 +0000
Subject: [PATCH 2/3] Fix remaining TMA synchronization gaps

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>
---
 .../common/cast/core/grouped_tma.cuh             |  2 +-
 .../group_quantize_fp8_blockwise.cuh             |  6 ++++--
 .../cast/mxfp8/group_scaled_swiglu_mxfp8.cuh     | 16 +++++++++++-----
 ...ow_cast_col_hadamard_transform_cast_fusion.cu |  2 ++
 .../group_hadamard_transform_cast_fusion.cu      |  2 ++
 ...ow_cast_col_hadamard_transform_cast_fusion.cu |  2 ++
 .../hadamard_transform_cast_fusion.cu            |  2 ++
 ...ow_cast_col_hadamard_transform_cast_fusion.cu |  2 ++
 transformer_engine/common/util/ptx.cuh           |  4 ++--
 9 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/transformer_engine/common/cast/core/grouped_tma.cuh b/transformer_engine/common/cast/core/grouped_tma.cuh
index a919b0a5eec..d0f0c181799 100644
--- a/transformer_engine/common/cast/core/grouped_tma.cuh
+++ b/transformer_engine/common/cast/core/grouped_tma.cuh
@@ -173,7 +173,7 @@ __global__ void __launch_bounds__(THREADS_PER_WARP)
 __device__ __forceinline__ void fence_acquire_tensormap(const CUtensorMap *tensor_map) {
 #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
   // The descriptor updater and consumer execute in different CTAs, so CTA scope is insufficient.
-  asm volatile("fence.proxy.tensormap::generic.acquire.gpu [%0], 128;" ::"l"(tensor_map));
+  asm volatile("fence.proxy.tensormap::generic.acquire.gpu [%0], 128;" ::"l"(tensor_map) : "memory");
 #else
   NVTE_DEVICE_ERROR("fence_acquire_tensormap is only supported on SM 9.0+.");
 #endif  // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
diff --git a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh
index 31feaf833dc..bf73e935014 100644
--- a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh
+++ b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh
@@ -341,8 +341,9 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 4) group_block_scaled_2d_tma
         static_cast<uint32_t>(global_col_base), static_cast<uint32_t>(global_row_base), &tma_mbar);
   }
   ptx::mbarrier_wait_parity(&tma_mbar, 0);
-  if (leading_thread) ptx::mbarrier_invalid(&tma_mbar);
+  // Every thread must finish waiting before the leader invalidates the barrier.
   __syncthreads();
+  if (leading_thread) ptx::mbarrier_invalid(&tma_mbar);
 
   // ---- Optional dbias: per-tile column sum of the high-precision input (smem-resident) ----
   if (dbias_workspace != nullptr) {
@@ -628,8 +629,9 @@ __global__ void __launch_bounds__(kThreadsPerBlock) group_block_scaled_1d_tma_ke
         static_cast<uint32_t>(global_col_base), static_cast<uint32_t>(global_row_base), &tma_mbar);
   }
   ptx::mbarrier_wait_parity(&tma_mbar, 0);
-  if (leading_thread) ptx::mbarrier_invalid(&tma_mbar);
+  // Every thread must finish waiting before the leader invalidates the barrier.
   __syncthreads();
+  if (leading_thread) ptx::mbarrier_invalid(&tma_mbar);
 
   // ---- Optional dbias: per-tile column sum of the high-precision input (smem-resident) ----
   if (dbias_workspace != nullptr) {
diff --git a/transformer_engine/common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh
index 878fc93107c..0dab09f163e 100644
--- a/transformer_engine/common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh
@@ -344,11 +344,11 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_scaled_swiglu_mxfp8_k
         ptx::mbarrier_wait_parity_acquire_cta_shared_cta(&IN_buff_readable_mbar[buff_in],
                                                          IN_buff_readable_parity[buff_in]);
         IN_buff_readable_parity[buff_in] ^= 1;
-        // Wait until the store groups still holding an output slice have drained. Only
-        // the leading thread commits those groups, so the wait is a no-op on the other
-        // threads and the barrier is what stops them from overwriting a slice the TMA
-        // unit has not finished reading.
-        ptx::cp_async_bulk_wait_group_read<OUT_BUFFS_NUM - 1>();
+        // Only the issuing thread can wait for the TMA read of the output slice.
+        // Hand its completion off to all cooperative writers before reusing the slice.
+        if (leading_thread) {
+          ptx::cp_async_bulk_wait_group_read<OUT_BUFFS_NUM - 1>();
+        }
         __syncthreads();
 
         const size_t buff = buff_in;
@@ -379,6 +379,12 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_scaled_swiglu_mxfp8_k
                           total_work_blocks, work_blocks_X);
     }
 
+    // The last stage has no following iteration to wait for its TMA store.
+    if (leading_thread) {
+      ptx::cp_async_bulk_wait_group();
+    }
+    __syncthreads();
+
     destroy_barriers<BUFFS_NUM>(IN_buff_readable_mbar, leading_thread);
   }  // if constexpr (is_single_tensor)
 #endif  // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
diff --git a/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu
index 0f2456c975c..171eaaa857a 100644
--- a/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu
+++ b/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu
@@ -499,6 +499,8 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device_g
 
     if (warp_idx == 2 && elect_one_sync()) {
       cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1);
+      // Publish this separately initialized barrier before TMA can complete against it.
+      cutlass::arch::fence_barrier_init();
     }
     __syncthreads();
 
diff --git a/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu
index 4b1435f9eb2..60ba2ba45e3 100644
--- a/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu
+++ b/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu
@@ -361,6 +361,8 @@ __global__ static void group_rht_gemm_device(
 
     if (warp_idx == 2 && elect_one_sync()) {
       cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1);
+      // Publish this separately initialized barrier before TMA can complete against it.
+      cutlass::arch::fence_barrier_init();
     }
     __syncthreads();
     using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x;
diff --git a/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu
index 2e6d383ce16..82863117807 100644
--- a/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu
+++ b/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu
@@ -487,6 +487,8 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device(
 
     if (warp_idx == 2 && elect_one_sync()) {
       cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1);
+      // Publish this separately initialized barrier before TMA can complete against it.
+      cutlass::arch::fence_barrier_init();
     }
     __syncthreads();
 
diff --git a/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu
index 433da1f0f0d..1f62191759a 100644
--- a/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu
+++ b/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu
@@ -326,6 +326,8 @@ rht_gemm_device(MShape M, NShape N, KShape K, ClusterTileShape cluster_tile,
 
   if (warp_idx == 2 && elect_one_sync()) {
     cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1);
+    // Publish this separately initialized barrier before TMA can complete against it.
+    cutlass::arch::fence_barrier_init();
   }
   __syncthreads();
   using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x;
diff --git a/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu
index 8d8ab201653..37311750204 100644
--- a/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu
+++ b/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu
@@ -491,6 +491,8 @@ __global__ static void row_col_rht_gemm_device(
 
   if (warp_idx == 2 && elect_one_sync()) {
     cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1);
+    // Publish this separately initialized barrier before TMA can complete against it.
+    cutlass::arch::fence_barrier_init();
   }
   __syncthreads();
 
diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh
index a84de8681ee..1cbd6d51e19 100644
--- a/transformer_engine/common/util/ptx.cuh
+++ b/transformer_engine/common/util/ptx.cuh
@@ -181,7 +181,7 @@ __device__ __forceinline__ void mbarrier_arrive_expect_tx_cta_relaxed_shared_cta
 
 __device__ __forceinline__ void fence_mbarrier_init_release_cluster() {
 #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
-  asm volatile("fence.mbarrier_init.release.cluster;");
+  asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
 #else
   NVTE_DEVICE_ERROR("fence_mbarrier_init_release_cluster is only supported on SM 10.0+.");
 #endif  // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
@@ -735,7 +735,7 @@ __device__ __forceinline__ int32_t elect_one_sync(uint32_t mask = 0xFFFFFFFFu) {
 
 __device__ __forceinline__ void numbered_barrier_sync(uint32_t num_threads,
                                                       uint32_t barrier_id = 1u) {
-  asm volatile("bar.sync %0, %1;\n" ::"r"(barrier_id), "r"(num_threads));
+  asm volatile("bar.sync %0, %1;\n" ::"r"(barrier_id), "r"(num_threads) : "memory");
 }
 
 __device__ __forceinline__ void fma_f32_f16(float &out, uint16_t const &a, uint16_t const &b,

From aee64b99cdb6224e65347b9dc678573562a439f2 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 11 Sep 2026 16:56:14 +0000
Subject: [PATCH 3/3] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 transformer_engine/common/cast/core/grouped_tma.cuh | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/transformer_engine/common/cast/core/grouped_tma.cuh b/transformer_engine/common/cast/core/grouped_tma.cuh
index d0f0c181799..6b3fd43f001 100644
--- a/transformer_engine/common/cast/core/grouped_tma.cuh
+++ b/transformer_engine/common/cast/core/grouped_tma.cuh
@@ -173,7 +173,8 @@ __global__ void __launch_bounds__(THREADS_PER_WARP)
 __device__ __forceinline__ void fence_acquire_tensormap(const CUtensorMap *tensor_map) {
 #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
   // The descriptor updater and consumer execute in different CTAs, so CTA scope is insufficient.
-  asm volatile("fence.proxy.tensormap::generic.acquire.gpu [%0], 128;" ::"l"(tensor_map) : "memory");
+  asm volatile("fence.proxy.tensormap::generic.acquire.gpu [%0], 128;" ::"l"(tensor_map)
+               : "memory");
 #else
   NVTE_DEVICE_ERROR("fence_acquire_tensormap is only supported on SM 9.0+.");
 #endif  // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)