Fix grouped MXFP8 quantization race on concurrent CUDA streams

tests/pytorch/mxfp8/test_mxfp8_group_quantize_multistream.py

@@ -0,0 +1,74 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Regression test for NVIDIA/TransformerEngine#3630.
+
+Grouped MXFP8 quantization stored its per-tensor TMA descriptor storage in a
+file-global ``static __device__`` buffer. Because that buffer is shared by every launch of the
+kernel, two ``group_quantize`` calls running on different CUDA streams raced on the same
+descriptors: one stream's ``update_tma_descriptors`` overwrote tensor maps the other stream was
+still consuming, producing incorrect output.
+
+The fix allocates a private ``TensorMapStorage`` per invocation with ``cudaMallocAsync`` and
+passes it to both the descriptor-update and consumer kernels, so concurrent invocations never
+share descriptor state. This test exercises concurrent multi-stream grouped quantization and
+checks every result against a single-stream reference.
+"""
+
+import pytest
+import torch
+
+import transformer_engine.pytorch as te
+import transformer_engine_torch as tex
+from transformer_engine.pytorch import MXFP8Quantizer
+
+mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True)
+
+# VARYING_LAST_DIM layout: 2D input is [rows, sum(last_dims)] and each group is a column slice.
+NUM_GROUPS = 8
+ROWS = 4096
+COLS_PER_GROUP = 1024
+
+
+def _make_input(seed):
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    return torch.randn(
+        ROWS, NUM_GROUPS * COLS_PER_GROUP, device="cuda", dtype=torch.bfloat16, generator=g
+    )
+
+
+def _last_dims():
+    return torch.full((NUM_GROUPS,), COLS_PER_GROUP, dtype=torch.int64, device="cuda")
+
+
+def _quantize(x):
+    quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False)
+    return tex.group_quantize(x, quantizer, NUM_GROUPS, None, _last_dims())
+
+
+@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
+@pytest.mark.parametrize("iteration", range(8))
+def test_group_quantize_multistream_matches_reference(iteration):
+    """Concurrent multi-stream grouped quantize must equal the single-stream reference."""
+    x1 = _make_input(iteration)
+    x2 = _make_input(1_000_000 + iteration)
+
+    # Deterministic single-stream references.
+    ref1 = _quantize(x1)
+    ref2 = _quantize(x2)
+
+    s1 = torch.cuda.Stream()
+    s2 = torch.cuda.Stream()
+    with torch.cuda.stream(s1):
+        out1 = _quantize(x1)
+    with torch.cuda.stream(s2):
+        out2 = _quantize(x2)
+
+    torch.cuda.current_stream().wait_stream(s1)
+    torch.cuda.current_stream().wait_stream(s2)
+
+    assert torch.equal(out1.rowwise_data, ref1.rowwise_data)
+    assert torch.equal(out1.scale_inv, ref1.scale_inv)
+    assert torch.equal(out2.rowwise_data, ref2.rowwise_data)
+    assert torch.equal(out2.scale_inv, ref2.scale_inv)

transformer_engine/common/cast/core/grouped_tma.cuh

@@ -43,8 +43,12 @@ struct alignas(128) TensorMapStorage {
   size_t offsets[MAX_SUPPORTED_TENSOR_DESCRIPTORS];
 };
 
-// Internal linkage avoids device-link ODR issues when this header is included by multiple .cu TUs.
-static __device__ TensorMapStorage g_tensor_maps;
+// `TensorMapStorage` is per-invocation, not file-global. A `static __device__` instance would be
+// shared by every launch that includes this header, so two `group_quantize` calls running on
+// different CUDA streams would race on the same descriptors (one stream's `update_tma_descriptors`
+// overwrites the tensor maps another stream is still consuming). Instead the launcher allocates a
+// private `TensorMapStorage` with `cudaMallocAsync` and passes it to both kernels, which keeps the
+// storage disjoint across concurrent invocations while remaining stream-ordered and graph-safe.
 
 inline bool dimensions_supported_by_TMA(const Tensor *const t) {
   const size_t cols = t->flat_last_dim();
@@ -119,17 +123,18 @@ __global__ void __launch_bounds__(THREADS_PER_WARP)
                            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 rowwise,
-                           const bool colwise, const bool compute_dactivations) {
+                           const bool colwise, const bool compute_dactivations,
+                           TensorMapStorage *const tensor_maps) {
   const size_t tensor_id = blockIdx.x;
   const size_t rows =
       get_tensor_rows_num(tensor_id, shape_rep, first_logical_dim, first_dims_ptr, num_tensors);
   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];
   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;
+    tensor_maps->rows[tensor_id] = rows;
+    tensor_maps->cols[tensor_id] = cols;
+    tensor_maps->offsets[tensor_id] = offset_elts;
   }
 
   // Zero-sized groups: skip TMA descriptor update. The main kernel already returns
@@ -141,27 +146,27 @@ __global__ void __launch_bounds__(THREADS_PER_WARP)
 
   if (tensor_id < num_tensors) {
     {
-      CUtensorMap *modified_tensor_map_input = &g_tensor_maps.input[tensor_id];
+      CUtensorMap *modified_tensor_map_input = &tensor_maps->input[tensor_id];
       const uintptr_t global_data_ptr = reinterpret_cast<uintptr_t>(input_data_ptr + offset_elts);
       modify_base_tensor_map(base_tensor_map_input, modified_tensor_map_input, global_data_ptr,
                              rows, cols, sizeof(IType));
     }
     if (compute_dactivations) {
-      CUtensorMap *modified_tensor_map_act_input = &g_tensor_maps.act_input[tensor_id];
+      CUtensorMap *modified_tensor_map_act_input = &tensor_maps->act_input[tensor_id];
       const uintptr_t global_data_ptr =
           reinterpret_cast<uintptr_t>(act_input_data_ptr + offset_elts);
       modify_base_tensor_map(base_tensor_map_act_input, modified_tensor_map_act_input,
                              global_data_ptr, rows, cols, sizeof(IType));
     }
     if (rowwise) {
-      CUtensorMap *modified_tensor_map_output_rowwise = &g_tensor_maps.output_rowwise[tensor_id];
+      CUtensorMap *modified_tensor_map_output_rowwise = &tensor_maps->output_rowwise[tensor_id];
       const uintptr_t global_data_ptr =
           reinterpret_cast<uintptr_t>(output_rowwise_data_ptr + offset_elts);
       modify_base_tensor_map(base_tensor_map_output_rowwise, modified_tensor_map_output_rowwise,
                              global_data_ptr, rows, cols, sizeof(OType));
     }
     if (colwise) {
-      CUtensorMap *modified_tensor_map_output_colwise = &g_tensor_maps.output_colwise[tensor_id];
+      CUtensorMap *modified_tensor_map_output_colwise = &tensor_maps->output_colwise[tensor_id];
       const uintptr_t global_data_ptr =
           reinterpret_cast<uintptr_t>(output_colwise_data_ptr + offset_elts);
       modify_base_tensor_map(base_tensor_map_output_colwise, modified_tensor_map_output_colwise,

transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh

@@ -625,7 +625,8 @@ __global__ void __launch_bounds__(CastTraits::THREADS_PER_CHUNK) group_quantize_
     const int64_t *const __restrict__ last_dims_ptr, e8m0_t *const __restrict__ scales_rowwise_ptr,
     e8m0_t *const __restrict__ scales_colwise_ptr, const float *__restrict__ noop,
     float *const __restrict__ dbias_workspace, float *const __restrict__ amax_ptr,
-    const size_t work_blocks_X, const size_t work_blocks_Y) {
+    const size_t work_blocks_X, const size_t work_blocks_Y,
+    const TensorMapStorage *const tensor_maps) {
 #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
   constexpr uint SCALE_DIM_Y = CastTraits::SCALE_DIM_Y;
   constexpr uint SCALE_DIM_X = CastTraits::SCALE_DIM_X;
@@ -785,9 +786,9 @@ __global__ void __launch_bounds__(CastTraits::THREADS_PER_CHUNK) group_quantize_
       if (fixed_tensor_id >= num_tensors) {
         return;
       }
-      fixed_rows = g_tensor_maps.rows[fixed_tensor_id];
-      fixed_cols = g_tensor_maps.cols[fixed_tensor_id];
-      fixed_tensor_base = g_tensor_maps.offsets[fixed_tensor_id];
+      fixed_rows = tensor_maps->rows[fixed_tensor_id];
+      fixed_cols = tensor_maps->cols[fixed_tensor_id];
+      fixed_tensor_base = tensor_maps->offsets[fixed_tensor_id];
       if (fixed_rows == 0 || fixed_cols == 0) {
         return;
       }
@@ -907,15 +908,15 @@ __global__ void __launch_bounds__(CastTraits::THREADS_PER_CHUNK) group_quantize_
     const size_t dbias_offset_Y = block_offset_Y / TILE_DIM_Y;
 
     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_act_input =
-        is_single_tensor ? tensor_map_act_input_static : g_tensor_maps.act_input[tensor_id];
+        is_single_tensor ? tensor_map_act_input_static : tensor_maps->act_input[tensor_id];
     const CUtensorMap &tensor_map_output_rowwise = is_single_tensor
                                                        ? tensor_map_output_rowwise_static
-                                                       : g_tensor_maps.output_rowwise[tensor_id];
+                                                       : tensor_maps->output_rowwise[tensor_id];
     const CUtensorMap &tensor_map_output_colwise = is_single_tensor
                                                        ? tensor_map_output_colwise_static
-                                                       : g_tensor_maps.output_colwise[tensor_id];
+                                                       : tensor_maps->output_colwise[tensor_id];
 
     if (leading_thread && (!is_single_tensor) && (last_acquired_tensor_id != tensor_id)) {
       fence_acquire_tensormap(&tensor_map_input);
@@ -1304,8 +1305,14 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations
                         const size_t out_mem = out_rowwise_mem + out_colwise_mem;
                         const size_t dshmem_size = in_mem + out_mem + TMA_SHMEM_ALIGNMENT;
 
-                        // Update tensor descriptors before launching the kernel
+                        // Per-invocation tensor-map storage. Allocated with cudaMallocAsync so the
+                        // descriptors live in stream-ordered, per-call storage instead of a file-global
+                        // `static __device__` buffer that concurrent streams would clobber.
+                        TensorMapStorage *tensor_maps_storage = nullptr;
                         if (!is_single_tensor) {
+                          NVTE_CHECK_CUDA(
+                              cudaMallocAsync(reinterpret_cast<void **>(&tensor_maps_storage),
+                                              sizeof(TensorMapStorage), stream));
                           const IType *const input_dptr =
                               reinterpret_cast<const IType *>(input->data.dptr);
                           const IType *const act_input_dptr =
@@ -1324,7 +1331,8 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations
                                   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);
+                                  last_dims_ptr, use_rowwise_scaling, use_colwise_scaling, IS_DACT,
+                                  tensor_maps_storage);
                         }
 
                         TRANSFORMER_ENGINE_SWITCH_CONDITION(
@@ -1344,7 +1352,8 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations
                                   last_logical_dim, launch_config.same_both_rows, offsets_ptr,
                                   first_dims_ptr, last_dims_ptr, scales_rowwise_ptr,
                                   scales_colwise_ptr, noop_ptr, workspace_ptr, amax_ptr,
-                                  launch_config.work_blocks_X, launch_config.work_blocks_Y);
+                                  launch_config.work_blocks_X, launch_config.work_blocks_Y,
+                                  tensor_maps_storage);
                             });
 
                         if constexpr (IS_DBIAS) {
@@ -1355,6 +1364,13 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations
                         }
 
                         NVTE_CHECK_CUDA(cudaGetLastError());
+
+                        // Release the per-call tensor-map storage. Stream-ordered free: the block
+                        // stays alive until the update and consumer kernels above have finished, and
+                        // the pool will not hand it to another stream in the meantime.
+                        if (tensor_maps_storage != nullptr) {
+                          NVTE_CHECK_CUDA(cudaFreeAsync(tensor_maps_storage, stream));
+                        }
                       });  // NOLINT(*)
               );           // NOLINT(*)
           );               // NOLINT(*)