[common] Improved performance of Group MXFP8 kernels

tests/cpp/operator/test_cast_mxfp8_grouped.cu

@@ -677,6 +677,42 @@ std::vector<std::vector<size_t>> input_config_small = {
     {VARYING_FIRST_DIM,     4,      512,160,                    128,0,0,256},
 };
 
+// Focused coverage for chunks that span multiple X and Y stages. Keep this separate from the
+// general shape matrix to avoid multiplying these larger cases by every operation and dtype.
+std::vector<std::vector<size_t>> input_config_multichunk = {
+    {SAME_BOTH_DIMS,        2,      768,384},
+    {VARYING_FIRST_DIM,     4,      1024,384,                   128,256,384,256},
+    {VARYING_LAST_DIM,      3,      384,1152,                   128,384,640},
+    {VARYING_BOTH_DIMS,     3,      1,(128*128)+(256*384)+(384*640),
+                                                               128,256,384,     128,384,640},
+};
+
+std::vector<std::vector<size_t>> input_config_same_both_multichunk = {
+    {SAME_BOTH_DIMS,        2,      768,384},
+};
+
+std::vector<size_t> make_max_descriptor_input_config() {
+    constexpr size_t num_tensors = 64;
+    std::vector<size_t> first_dims(num_tensors, 128);
+    std::vector<size_t> last_dims(num_tensors, 128);
+    first_dims.front() = 256;
+    last_dims.back() = 256;
+
+    size_t elts_num = 0;
+    for (size_t t = 0; t < num_tensors; ++t) {
+        elts_num += first_dims[t] * last_dims[t];
+    }
+
+    std::vector<size_t> config = {VARYING_BOTH_DIMS, num_tensors, 1, elts_num};
+    config.insert(config.end(), first_dims.begin(), first_dims.end());
+    config.insert(config.end(), last_dims.begin(), last_dims.end());
+    return config;
+}
+
+std::vector<std::vector<size_t>> input_config_max_descriptors = {
+    make_max_descriptor_input_config(),
+};
+
 class GroupedFusedCastMXFP8TestSuite : public ::testing::TestWithParam
     <std::tuple<ProcessingMethod,
                 ActivationKind,
@@ -894,3 +930,51 @@ INSTANTIATE_TEST_SUITE_P(
         ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16),
         ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2)),
     MakeGroupedFusedCastMXFP8TestName);
+
+INSTANTIATE_TEST_SUITE_P(
+    OperatorTest_GroupedFusedCastMXFP8_MultiChunk,
+    GroupedFusedCastMXFP8TestSuite,
+    ::testing::Combine(
+        ::testing::Values(ProcessingMethod::CAST_ONLY),
+        ::testing::Values(ActivationKind::Identity),
+        ::testing::ValuesIn(scaling_directions),
+        ::testing::ValuesIn(input_config_multichunk),
+        ::testing::Values(DType::kBFloat16),
+        ::testing::Values(DType::kFloat8E4M3)),
+    MakeGroupedFusedCastMXFP8TestName);
+
+INSTANTIATE_TEST_SUITE_P(
+    OperatorTest_GroupedFusedCastMXFP8_MultiChunkDBias,
+    GroupedFusedCastMXFP8TestSuite,
+    ::testing::Combine(
+        ::testing::Values(ProcessingMethod::CAST_DBIAS),
+        ::testing::Values(ActivationKind::Identity),
+        ::testing::ValuesIn(scaling_directions),
+        ::testing::ValuesIn(input_config_same_both_multichunk),
+        ::testing::Values(DType::kBFloat16),
+        ::testing::Values(DType::kFloat8E4M3)),
+    MakeGroupedFusedCastMXFP8TestName);
+
+INSTANTIATE_TEST_SUITE_P(
+    OperatorTest_GroupedFusedCastMXFP8_MultiChunkDBiasDAct,
+    GroupedFusedCastMXFP8TestSuite,
+    ::testing::Combine(
+        ::testing::Values(ProcessingMethod::CAST_DBIAS_DACT),
+        ::testing::Values(ActivationKind::GeLU),
+        ::testing::ValuesIn(scaling_directions),
+        ::testing::ValuesIn(input_config_same_both_multichunk),
+        ::testing::Values(DType::kBFloat16),
+        ::testing::Values(DType::kFloat8E4M3)),
+    MakeGroupedFusedCastMXFP8TestName);
+
+INSTANTIATE_TEST_SUITE_P(
+    OperatorTest_GroupedFusedCastMXFP8_MaxDescriptors,
+    GroupedFusedCastMXFP8TestSuite,
+    ::testing::Combine(
+        ::testing::Values(ProcessingMethod::CAST_ONLY),
+        ::testing::Values(ActivationKind::Identity),
+        ::testing::Values(ScalingDirection::BOTH),
+        ::testing::ValuesIn(input_config_max_descriptors),
+        ::testing::Values(DType::kBFloat16),
+        ::testing::Values(DType::kFloat8E4M3)),
+    MakeGroupedFusedCastMXFP8TestName);

transformer_engine/common/cast/core/grouped_tma.cuh

@@ -38,6 +38,9 @@ struct alignas(128) TensorMapStorage {
   alignas(128) CUtensorMap act_input[MAX_SUPPORTED_TENSOR_DESCRIPTORS];
   alignas(128) CUtensorMap output_rowwise[MAX_SUPPORTED_TENSOR_DESCRIPTORS];
   alignas(128) CUtensorMap output_colwise[MAX_SUPPORTED_TENSOR_DESCRIPTORS];
+  size_t rows[MAX_SUPPORTED_TENSOR_DESCRIPTORS];
+  size_t cols[MAX_SUPPORTED_TENSOR_DESCRIPTORS];
+  size_t offsets[MAX_SUPPORTED_TENSOR_DESCRIPTORS];
 };
 
 // Internal linkage avoids device-link ODR issues when this header is included by multiple .cu TUs.
@@ -115,6 +118,9 @@ __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;
 
   // 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