diff --git a/transformer_engine/common/cast/core/grouped_tma.cuh b/transformer_engine/common/cast/core/grouped_tma.cuh index 8603fd1fd2..a919b0a5ee 100644 --- a/transformer_engine/common/cast/core/grouped_tma.cuh +++ b/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(global_dim_Y)), "r"(static_cast(global_dim_X)), + "l"(static_cast(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(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(&shared_tensor_map)), - "l"(global_data_ptr), "r"(static_cast(global_dim_Y)), - "r"(static_cast(global_dim_X)), "l"(static_cast(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 -__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) diff --git a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh index dad8d18d6f..141efabbaf 100644 --- a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh +++ b/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(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(input_data.dptr); OType *const output_dptr = reinterpret_cast(output->data.dptr); - update_tma_descriptors<<>>( + update_tma_descriptors<<>>( 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); diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index b0383d95f3..bd8c3052f6 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/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 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(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(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(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(first_dims_ptr[mapped_tensor_id]); - mapper_storage[3] = static_cast(offsets_ptr[mapped_tensor_id]); + direct_mapper_storage.tensor_id = mapped_tensor_id; + direct_mapper_storage.rows = static_cast(first_dims_ptr[mapped_tensor_id]); + direct_mapper_storage.tensor_start_offset = + static_cast(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(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(output->columnwise_data.dptr) : nullptr; - update_tma_descriptors<<>>( - 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 + <<>>( + 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(