From de431affd634a3d9bafd4ece1423b8ba910cc6cc Mon Sep 17 00:00:00 2001 From: zhaoye <772971548@qq.com> Date: Thu, 29 Oct 2026 16:27:07 +0800 Subject: [PATCH 1/2] Fix grouped MXFP8 dequantize race on concurrent CUDA streams Grouped MXFP8 dequantization stored its per-tensor TMA descriptor storage in file-global `static __device__` arrays (`g_tensor_maps_input` / `g_tensor_maps_output`). Like the quantize path, that storage is shared by every kernel launch across all CUDA streams, so concurrent `group_dequantize` calls on different streams race on the descriptors: one stream's `update_tma_descriptors` overwrites tensor maps another stream is still consuming, producing incorrect output. Allocate per-invocation `CUtensorMap` arrays with `cudaMallocAsync` and pass them to both the descriptor-update and consumer kernels, then release them with `cudaFreeAsync`. The storage stays in global memory (required for TMA to consume the descriptors) but is disjoint across concurrent invocations; the stream-ordered allocate/free keeps the descriptors alive until the kernels that use them finish. This is the same fix as the quantize path in #3774, applied to the sibling dequantize path, which had the identical file-global buffer. Signed-off-by: zhaoye <772971548@qq.com> --- .../cast/mxfp8/group_dequantize_mxfp8.cuh | 53 ++++++++++++++----- 1 file changed, 41 insertions(+), 12 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh index fcf8819e5f8..d7dc2867c7a 100644 --- a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh @@ -28,8 +28,11 @@ namespace mxfp8 { namespace group_dequantize_kernel { constexpr int MAX_SUPPORTED_TENSOR_DESCRIPTORS = 64; -__device__ alignas(128) CUtensorMap g_tensor_maps_input[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; -__device__ alignas(128) CUtensorMap g_tensor_maps_output[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; + +// Per-invocation tensor-map storage (not file-global): a `static __device__` array here would be +// shared by every dequantize launch that includes this header, so calls on different CUDA streams +// would race on the descriptors. The launcher allocates private arrays with `cudaMallocAsync` and +// passes them to both kernels, keeping concurrent invocations disjoint. // Reuse helper types and functions from common namespace using common::fence_acquire_tensormap; @@ -88,7 +91,9 @@ __global__ void update_tma_descriptors(const __grid_constant__ CUtensorMap base_ const size_t last_logical_dim, const int64_t *const __restrict__ offsets_ptr, const int64_t *const __restrict__ first_dims_ptr, - const int64_t *const __restrict__ last_dims_ptr) { + const int64_t *const __restrict__ last_dims_ptr, + CUtensorMap *const tensor_maps_input, + CUtensorMap *const tensor_maps_output) { const size_t tensor_id = blockIdx.x; const size_t rows = @@ -107,12 +112,12 @@ __global__ void update_tma_descriptors(const __grid_constant__ CUtensorMap base_ 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], + modify_base_tensor_map(base_tensor_map_input, &tensor_maps_input[tensor_id], global_data_ptr, rows, cols, sizeof(IType)); } { const uintptr_t global_data_ptr = reinterpret_cast(output_data_ptr + offset_elts); - modify_base_tensor_map(base_tensor_map_output, &g_tensor_maps_output[tensor_id], + modify_base_tensor_map(base_tensor_map_output, &tensor_maps_output[tensor_id], global_data_ptr, rows, cols, sizeof(OType)); } } @@ -127,7 +132,9 @@ __global__ void __launch_bounds__(128) const int64_t *const __restrict__ offsets_ptr, const int64_t *const __restrict__ first_dims_ptr, const int64_t *const __restrict__ last_dims_ptr, - const e8m0_t *const __restrict__ scales_ptr) { + const e8m0_t *const __restrict__ scales_ptr, + const CUtensorMap *const tensor_maps_input, + const CUtensorMap *const tensor_maps_output) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) constexpr size_t CHUNK_DIM_Y = DequantizeConfig::CHUNK_DIM_Y; constexpr size_t CHUNK_DIM_X = DequantizeConfig::CHUNK_DIM_X; @@ -211,9 +218,9 @@ __global__ void __launch_bounds__(128) // Select TMA descriptors (static for single tensor, per-tensor for multi-tensor) const CUtensorMap &tensor_map_input = - is_single_tensor ? tensor_map_input_static : g_tensor_maps_input[tensor_id]; + is_single_tensor ? tensor_map_input_static : tensor_maps_input[tensor_id]; const CUtensorMap &tensor_map_output = - is_single_tensor ? tensor_map_output_static : g_tensor_maps_output[tensor_id]; + is_single_tensor ? tensor_map_output_static : tensor_maps_output[tensor_id]; if (!is_single_tensor) { fence_acquire_tensormap(&tensor_map_input); @@ -470,26 +477,48 @@ inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output, SHMEM_DIM_Y, SHMEM_DIM_X, last_logical_dim, 0, typeToNumBits(output->dtype())); + // Per-invocation tensor-map storage (see note above). Allocated with cudaMallocAsync so + // concurrent dequantize calls on different streams never share descriptor state. + CUtensorMap *tensor_maps_input = nullptr; + CUtensorMap *tensor_maps_output = nullptr; + // Update tensor descriptors before launching the kernel if (!is_single_tensor) { + const size_t tmap_bytes = MAX_SUPPORTED_TENSOR_DESCRIPTORS * sizeof(CUtensorMap); + NVTE_CHECK_CUDA(cudaMallocAsync( + reinterpret_cast(&tensor_maps_input), tmap_bytes, stream)); + NVTE_CHECK_CUDA(cudaMallocAsync( + reinterpret_cast(&tensor_maps_output), tmap_bytes, stream)); const IType *const input_dptr = reinterpret_cast(input_data.dptr); OType *const output_dptr = reinterpret_cast(output->data.dptr); 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); + last_dims_ptr, tensor_maps_input, tensor_maps_output); } if (use_rowwise_scaling) { group_dequantize_mxfp8_kernel<<>>( tensor_map_input, tensor_map_output, shape_rep, num_tensors, first_logical_dim, - last_logical_dim, offsets_ptr, first_dims_ptr, last_dims_ptr, scales_ptr); + last_logical_dim, offsets_ptr, first_dims_ptr, last_dims_ptr, scales_ptr, + tensor_maps_input, tensor_maps_output); } else { group_dequantize_mxfp8_kernel<<>>( tensor_map_input, tensor_map_output, shape_rep, num_tensors, first_logical_dim, - last_logical_dim, offsets_ptr, first_dims_ptr, last_dims_ptr, scales_ptr); - }); // NOLINT(*) + last_logical_dim, offsets_ptr, first_dims_ptr, last_dims_ptr, scales_ptr, + tensor_maps_input, tensor_maps_output); + } + + // Release the per-call tensor-map storage (stream-ordered; stays alive until the kernels + // above finish, and the pool won't reuse the block for another stream before then). + if (tensor_maps_input != nullptr) { + NVTE_CHECK_CUDA(cudaFreeAsync(tensor_maps_input, stream)); + } + if (tensor_maps_output != nullptr) { + NVTE_CHECK_CUDA(cudaFreeAsync(tensor_maps_output, stream)); + } + ); // NOLINT(*) ); // NOLINT(*) NVTE_CHECK_CUDA(cudaGetLastError()); } From 9e256abc137bdcf2fde61090c24f68cacb1a9364 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 29 Oct 2026 08:42:17 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../cast/mxfp8/group_dequantize_mxfp8.cuh | 39 ++++++++----------- 1 file changed, 17 insertions(+), 22 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh index d7dc2867c7a..145a66c75c4 100644 --- a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh @@ -82,18 +82,15 @@ struct DequantizeConfig { }; template -__global__ void update_tma_descriptors(const __grid_constant__ CUtensorMap base_tensor_map_input, - const __grid_constant__ CUtensorMap base_tensor_map_output, - const IType *const __restrict__ input_data_ptr, - const OType *const __restrict__ output_data_ptr, - const ShapeRepresentation shape_rep, - const size_t num_tensors, const size_t first_logical_dim, - const size_t last_logical_dim, - const int64_t *const __restrict__ offsets_ptr, - const int64_t *const __restrict__ first_dims_ptr, - const int64_t *const __restrict__ last_dims_ptr, - CUtensorMap *const tensor_maps_input, - CUtensorMap *const tensor_maps_output) { +__global__ void update_tma_descriptors( + const __grid_constant__ CUtensorMap base_tensor_map_input, + const __grid_constant__ CUtensorMap base_tensor_map_output, + const IType *const __restrict__ input_data_ptr, const OType *const __restrict__ output_data_ptr, + const ShapeRepresentation shape_rep, const size_t num_tensors, const size_t first_logical_dim, + const size_t last_logical_dim, const int64_t *const __restrict__ offsets_ptr, + const int64_t *const __restrict__ first_dims_ptr, + const int64_t *const __restrict__ last_dims_ptr, CUtensorMap *const tensor_maps_input, + CUtensorMap *const tensor_maps_output) { const size_t tensor_id = blockIdx.x; const size_t rows = @@ -112,8 +109,8 @@ __global__ void update_tma_descriptors(const __grid_constant__ CUtensorMap base_ 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, &tensor_maps_input[tensor_id], - global_data_ptr, rows, cols, sizeof(IType)); + modify_base_tensor_map(base_tensor_map_input, &tensor_maps_input[tensor_id], global_data_ptr, + rows, cols, sizeof(IType)); } { const uintptr_t global_data_ptr = reinterpret_cast(output_data_ptr + offset_elts); @@ -485,10 +482,10 @@ inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output, // Update tensor descriptors before launching the kernel if (!is_single_tensor) { const size_t tmap_bytes = MAX_SUPPORTED_TENSOR_DESCRIPTORS * sizeof(CUtensorMap); - NVTE_CHECK_CUDA(cudaMallocAsync( - reinterpret_cast(&tensor_maps_input), tmap_bytes, stream)); - NVTE_CHECK_CUDA(cudaMallocAsync( - reinterpret_cast(&tensor_maps_output), tmap_bytes, stream)); + NVTE_CHECK_CUDA( + cudaMallocAsync(reinterpret_cast(&tensor_maps_input), tmap_bytes, stream)); + NVTE_CHECK_CUDA(cudaMallocAsync(reinterpret_cast(&tensor_maps_output), + tmap_bytes, stream)); const IType *const input_dptr = reinterpret_cast(input_data.dptr); OType *const output_dptr = reinterpret_cast(output->data.dptr); @@ -514,11 +511,9 @@ inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output, // above finish, and the pool won't reuse the block for another stream before then). if (tensor_maps_input != nullptr) { NVTE_CHECK_CUDA(cudaFreeAsync(tensor_maps_input, stream)); - } - if (tensor_maps_output != nullptr) { + } if (tensor_maps_output != nullptr) { NVTE_CHECK_CUDA(cudaFreeAsync(tensor_maps_output, stream)); - } - ); // NOLINT(*) + }); // NOLINT(*) ); // NOLINT(*) NVTE_CHECK_CUDA(cudaGetLastError()); }