[Common] Fix Grouped MXFP8 work mapping and TMA synchronization

transformer_engine/common/cast/core/grouped_tma.cuh

@@ -53,45 +53,59 @@ inline bool dimensions_supported_by_TMA(const Tensor *const t) {
   return cols % alignment_requirement == 0;
 }
 
-// Copies the base tensor map to shmem, modifies the copy, stores the modified tensor map at index
+// Copy the base tensor map to shared memory, modify it, and publish it to global memory. This
+// function must be called convergently by every thread in a one-warp CTA.
 __device__ __forceinline__ void modify_base_tensor_map(const CUtensorMap base_tensor_map,
                                                        CUtensorMap *global_tensor_map,
                                                        const uintptr_t global_data_ptr,
                                                        const size_t global_dim_Y,
                                                        const size_t global_dim_X,
                                                        const size_t data_type_size_bytes) {
-  __shared__ CUtensorMap shared_tensor_map;
-  shared_tensor_map = base_tensor_map;  // Copy the base tensor map into shmem
+  __shared__ alignas(128) CUtensorMap shared_tensor_map;
   constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY;
   if constexpr (is_blackwell) {
-    const size_t global_stride_bytes = global_dim_X * data_type_size_bytes;
-    if (global_stride_bytes % TMA_GMEM_ALIGNMENT != 0) {
-      NVTE_DEVICE_ERROR("Shape not supported. Data stride must be 16B aligned.");
-    }
-    if (global_data_ptr % TMA_GMEM_ALIGNMENT != 0) {
-      NVTE_DEVICE_ERROR("Tensor data pointer must be 16B aligned");
+    const uint32_t shared_tensor_map_ptr = __cvta_generic_to_shared(&shared_tensor_map);
+    if (threadIdx.x == 0) {
+      shared_tensor_map = base_tensor_map;
+
+      const size_t global_stride_bytes = global_dim_X * data_type_size_bytes;
+      if (global_stride_bytes % TMA_GMEM_ALIGNMENT != 0) {
+        NVTE_DEVICE_ERROR("Shape not supported. Data stride must be 16B aligned.");
+      }
+      if (global_data_ptr % TMA_GMEM_ALIGNMENT != 0) {
+        NVTE_DEVICE_ERROR("Tensor data pointer must be 16B aligned");
+      }
+
+      asm volatile(
+          "tensormap.replace.tile.global_address.shared::cta.b1024.b64 [%0], %1;\n\t"
+          "tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 1, %2;\n\t"
+          "tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 0, %3;\n\t"
+          "tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 0, %4;\n"
+          :
+          : "r"(shared_tensor_map_ptr), "l"(global_data_ptr),
+            "r"(static_cast<uint32_t>(global_dim_Y)), "r"(static_cast<uint32_t>(global_dim_X)),
+            "l"(static_cast<uint64_t>(global_stride_bytes))
+          : "memory");
     }
 
+    // tensormap.cp_fenceproxy is a warp-collective instruction. Besides copying the complete
+    // 128-byte descriptor, its GPU-scope release makes the update visible to tensor-map proxy
+    // accesses that perform a matching acquire in a consumer CTA.
+    __syncwarp();
+    const uintptr_t global_tensor_map_ptr = reinterpret_cast<uintptr_t>(global_tensor_map);
     asm volatile(
-        "{\n\t"
-        ".reg.b64 tensor_map_ptr; \n\t"
-        "mov.b64 tensor_map_ptr, %0; \n\t"
-        "tensormap.replace.tile.global_address.b1024.b64  [tensor_map_ptr], %1; \n\t"
-        "tensormap.replace.tile.global_dim.b1024.b32  [tensor_map_ptr], 1, %2; \n\t"  // DIM Y
-        "tensormap.replace.tile.global_dim.b1024.b32  [tensor_map_ptr], 0, %3; \n\t"  // DIM X
-        "tensormap.replace.tile.global_stride.b1024.b64  [tensor_map_ptr], 0, %4; \n"
-        "}\n" ::"l"(reinterpret_cast<uintptr_t>(&shared_tensor_map)),
-        "l"(global_data_ptr), "r"(static_cast<uint32_t>(global_dim_Y)),
-        "r"(static_cast<uint32_t>(global_dim_X)), "l"(static_cast<uint64_t>(global_stride_bytes))
+        "tensormap.cp_fenceproxy.global.shared::cta.tensormap::generic.release.gpu.sync.aligned "
+        "[%0], [%1], 128;"
+        :
+        : "l"(global_tensor_map_ptr), "r"(shared_tensor_map_ptr)
         : "memory");
-    *global_tensor_map = shared_tensor_map;
   } else {
     NVTE_DEVICE_ERROR("tensormap.replace is architecture-specific. ");
   }
 }
 
 template <typename IType, typename OType>
-__global__ void __launch_bounds__(1)
+__global__ void __launch_bounds__(THREADS_PER_WARP)
     update_tma_descriptors(const __grid_constant__ CUtensorMap base_tensor_map_input,
                            const __grid_constant__ CUtensorMap base_tensor_map_act_input,
                            const __grid_constant__ CUtensorMap base_tensor_map_output_rowwise,
@@ -112,9 +126,11 @@ __global__ void __launch_bounds__(1)
   const size_t cols = get_tensor_cols_num(tensor_id, shape_rep, last_logical_dim, last_dims_ptr);
 
   const size_t offset_elts = offsets_ptr[tensor_id];
-  g_tensor_maps.rows[tensor_id] = rows;
-  g_tensor_maps.cols[tensor_id] = cols;
-  g_tensor_maps.offsets[tensor_id] = offset_elts;
+  if (threadIdx.x == 0) {
+    g_tensor_maps.rows[tensor_id] = rows;
+    g_tensor_maps.cols[tensor_id] = cols;
+    g_tensor_maps.offsets[tensor_id] = offset_elts;
+  }
 
   // Zero-sized groups: skip TMA descriptor update. The main kernel already returns
   // early for rows==0 or cols==0, but creating a TMA descriptor with a zero dimension
@@ -156,7 +172,8 @@ __global__ void __launch_bounds__(1)
 
 __device__ __forceinline__ void fence_acquire_tensormap(const CUtensorMap *tensor_map) {
 #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-  asm volatile("fence.proxy.tensormap::generic.acquire.cta [%0], 128;" ::"l"(tensor_map));
+  // 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));
 #else
   NVTE_DEVICE_ERROR("fence_acquire_tensormap is only supported on SM 9.0+.");
 #endif  // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)

transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh

@@ -89,7 +89,6 @@ __global__ void update_tma_descriptors(const __grid_constant__ CUtensorMap base_
                                        const int64_t *const __restrict__ offsets_ptr,
                                        const int64_t *const __restrict__ first_dims_ptr,
                                        const int64_t *const __restrict__ last_dims_ptr) {
-  const bool leading_thread = (threadIdx.x == 0);
   const size_t tensor_id = blockIdx.x;
 
   const size_t rows =
@@ -105,7 +104,7 @@ __global__ void update_tma_descriptors(const __grid_constant__ CUtensorMap base_
     return;
   }
 
-  if (leading_thread && (tensor_id < num_tensors)) {
+  if (tensor_id < num_tensors) {
     {
       const uintptr_t global_data_ptr = reinterpret_cast<uintptr_t>(input_data_ptr + offset_elts);
       modify_base_tensor_map(base_tensor_map_input, &g_tensor_maps_input[tensor_id],
@@ -469,7 +468,7 @@ inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output,
             const IType *const input_dptr = reinterpret_cast<const IType *>(input_data.dptr);
             OType *const output_dptr = reinterpret_cast<OType *>(output->data.dptr);
 
-            update_tma_descriptors<IType, OType><<<num_tensors, 32, 0, stream>>>(
+            update_tma_descriptors<IType, OType><<<num_tensors, THREADS_PER_WARP, 0, stream>>>(
                 tensor_map_input, tensor_map_output, input_dptr, output_dptr, shape_rep,
                 num_tensors, first_logical_dim, last_logical_dim, offsets_ptr, first_dims_ptr,
                 last_dims_ptr);

transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh

@@ -133,6 +133,13 @@ struct LaunchConfig {
   dim3 grid;
 };
 
+struct alignas(8) DirectVaryingFirstMapperStorage {
+  size_t active_elements;
+  size_t tensor_id;
+  size_t rows;
+  size_t tensor_start_offset;
+};
+
 template <typename CastTraits>
 LaunchConfig get_launch_config(const size_t first_logical_dim, const size_t last_logical_dim,
                                const size_t elts_total, const size_t num_tensors) {
@@ -659,6 +666,22 @@ __global__ void __launch_bounds__(CastTraits::THREADS_PER_CHUNK) group_quantize_
   constexpr bool is_single_tensor = (shape_rep == SAME_BOTH_DIMS || shape_rep == VARYING_FIRST_DIM);
 
   const bool leading_thread = (threadIdx.x == 0);
+  // Keep mapper metadata separate from dynamic shared memory. The latter is the destination of
+  // asynchronous TMA loads and may be overwritten before every warp has consumed the metadata.
+  __shared__ DirectVaryingFirstMapperStorage direct_mapper_storage;
+
+  if constexpr (use_direct_varying_first_mapper) {
+    static_assert(CastTraits::THREADS_PER_CHUNK >= MAX_SUPPORTED_TENSOR_DESCRIPTORS,
+                  "The first CTA must have enough threads to validate every tensor.");
+    // The direct mapper relies on every tensor boundary being TILE_DIM_Y-aligned. The legacy
+    // mapper validated this while decoding each job, but the direct path no longer calls it.
+    // Validate all per-tensor row counts once, using the first CTA. num_tensors is bounded by
+    // MAX_SUPPORTED_TENSOR_DESCRIPTORS (64), which is smaller than the CTA size.
+    if (blockIdx.x == 0 && threadIdx.x < num_tensors) {
+      get_tensor_rows_num<ShapeRepresentation::VARYING_FIRST_DIM>(threadIdx.x, first_logical_dim,
+                                                                  first_dims_ptr, num_tensors);
+    }
+  }
 
   // Decode the linear direct-mapper grid once per CTA. Valid CUDA grid extents fit in uint, which
   // also keeps this one-time coordinate calculation in 32-bit arithmetic.
@@ -719,26 +742,26 @@ __global__ void __launch_bounds__(CastTraits::THREADS_PER_CHUNK) group_quantize_
 
   if constexpr (use_direct_varying_first_mapper) {
     // logical_shape may describe graph-safe capacity beyond the active tensors. Resolve the
-    // active tail once per CTA and reject it before initializing TMA barriers. Reuse the same
-    // temporary storage for the exceptional colwise-swizzled tensor metadata.
-    size_t *const mapper_storage = reinterpret_cast<size_t *>(dshmem);
+    // active tail once per CTA and reject it before initializing TMA barriers. Cache the
+    // exceptional colwise-swizzled tensor metadata in dedicated static shared memory.
     const size_t block_offset_Y = direct_block_id_Y * CHUNK_DIM_Y;
     const size_t tensor_offset = block_offset_Y * last_logical_dim;
     if (leading_thread) {
       const size_t active_elements = static_cast<size_t>(offsets_ptr[num_tensors]);
-      mapper_storage[0] = active_elements;
+      direct_mapper_storage.active_elements = active_elements;
       if constexpr (WITH_GEMM_SWIZZLED_SCALES && COLWISE_SCALING) {
         if (tensor_offset < active_elements) {
           const size_t mapped_tensor_id =
               find_tensor_from_offsets(offsets_ptr, num_tensors, tensor_offset);
-          mapper_storage[1] = mapped_tensor_id;
-          mapper_storage[2] = static_cast<size_t>(first_dims_ptr[mapped_tensor_id]);
-          mapper_storage[3] = static_cast<size_t>(offsets_ptr[mapped_tensor_id]);
+          direct_mapper_storage.tensor_id = mapped_tensor_id;
+          direct_mapper_storage.rows = static_cast<size_t>(first_dims_ptr[mapped_tensor_id]);
+          direct_mapper_storage.tensor_start_offset =
+              static_cast<size_t>(offsets_ptr[mapped_tensor_id]);
         }
       }
     }
     __syncthreads();
-    if (tensor_offset >= mapper_storage[0]) {
+    if (tensor_offset >= direct_mapper_storage.active_elements) {
       return;
     }
   }
@@ -839,10 +862,10 @@ __global__ void __launch_bounds__(CastTraits::THREADS_PER_CHUNK) group_quantize_
       if constexpr (WITH_GEMM_SWIZZLED_SCALES && COLWISE_SCALING) {
         // Colwise GEMM-swizzled scale indices restart at each tensor and depend on M_i.
         // The leading thread decoded this exceptional metadata before barrier initialization.
-        size_t *const mapper_storage = reinterpret_cast<size_t *>(dshmem);
-        tensor_id = mapper_storage[1];
-        rows = mapper_storage[2];
-        tensor_start_offset = mapper_storage[3];
+        tensor_id = direct_mapper_storage.tensor_id;
+        rows = direct_mapper_storage.rows;
+        tensor_start_offset = direct_mapper_storage.tensor_start_offset;
+        tensor_offset_Y = block_offset_Y - tensor_start_offset / cols;
       }
     } else {
       block_id_Y = current_block_id / fixed_blocks_X;
@@ -1289,12 +1312,13 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations
                               use_colwise_scaling
                                   ? reinterpret_cast<OType *>(output->columnwise_data.dptr)
                                   : nullptr;
-                          update_tma_descriptors<IType, OType><<<num_tensors, 1, 0, stream>>>(
-                              tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise,
-                              tensor_map_output_colwise, input_dptr, act_input_dptr,
-                              output_rowwise_dptr, output_colwise_dptr, shape_rep, num_tensors,
-                              first_logical_dim, last_logical_dim, offsets_ptr, first_dims_ptr,
-                              last_dims_ptr, use_rowwise_scaling, use_colwise_scaling, IS_DACT);
+                          update_tma_descriptors<IType, OType>
+                              <<<num_tensors, THREADS_PER_WARP, 0, stream>>>(
+                                  tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise,
+                                  tensor_map_output_colwise, input_dptr, act_input_dptr,
+                                  output_rowwise_dptr, output_colwise_dptr, shape_rep, num_tensors,
+                                  first_logical_dim, last_logical_dim, offsets_ptr, first_dims_ptr,
+                                  last_dims_ptr, use_rowwise_scaling, use_colwise_scaling, IS_DACT);
                         }
 
                         TRANSFORMER_ENGINE_SWITCH_CONDITION(