From 5bf61bb1e393631dca23a43ad6b7f5ade519a04a Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Wed, 18 Mar 2026 16:27:23 -0700 Subject: [PATCH 1/4] Fix zero-sized groups in update_tma_descriptors Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> --- .../common/cast/mxfp8/group_quantize_mxfp8.cuh | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index 129d6724ac..5d6db58d8d 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -191,6 +191,13 @@ __global__ void update_tma_descriptors( const size_t offset_elts = offsets_ptr[tensor_id]; + // 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 + // is invalid and causes CUDA_ERROR_ILLEGAL_ADDRESS. + if (rows == 0 || cols == 0) { + return; + } + if (leading_thread && (tensor_id < num_tensors)) { { const uintptr_t global_data_ptr = reinterpret_cast(input_data_ptr + offset_elts); From 526c852aa81e395c8f936f1099b70f2105240743 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Wed, 18 Mar 2026 16:32:06 -0700 Subject: [PATCH 2/4] Update test_cast_mxfp8_grouped.cu Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> --- tests/cpp/operator/test_cast_mxfp8_grouped.cu | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/cpp/operator/test_cast_mxfp8_grouped.cu b/tests/cpp/operator/test_cast_mxfp8_grouped.cu index e469ad0845..09bd21657a 100644 --- a/tests/cpp/operator/test_cast_mxfp8_grouped.cu +++ b/tests/cpp/operator/test_cast_mxfp8_grouped.cu @@ -649,6 +649,7 @@ std::vector> input_config = { {SAME_BOTH_DIMS, 2, 256,128}, {VARYING_FIRST_DIM, 2, 512,128, 128,384}, {VARYING_FIRST_DIM, 3, 1024,144, 128,384,512}, + {VARYING_FIRST_DIM, 4, 1024,144, 128,384,0,512}, {VARYING_FIRST_DIM, 4, 1536,160, 128,384,512,512}, {VARYING_FIRST_DIM, 5, 4096,512, 128,256,384,1024,2304}, {VARYING_LAST_DIM, 3, 256,896, 128,256,512}, From ea30dbd019bbe357b33363cd43f4ec2d1ec7f6cd Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Wed, 18 Mar 2026 16:32:53 -0700 Subject: [PATCH 3/4] Apply suggestions from code review Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> --- .../common/cast/mxfp8/group_quantize_mxfp8.cuh | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index 5d6db58d8d..f8508f9e39 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -189,8 +189,6 @@ __global__ void update_tma_descriptors( 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]; - // 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 // is invalid and causes CUDA_ERROR_ILLEGAL_ADDRESS. @@ -198,6 +196,9 @@ __global__ void update_tma_descriptors( return; } + const size_t offset_elts = offsets_ptr[tensor_id]; + + if (leading_thread && (tensor_id < num_tensors)) { { const uintptr_t global_data_ptr = reinterpret_cast(input_data_ptr + offset_elts); From d9d1da49f160e14a93414488a0b99e6a6dfd1947 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 18 Mar 2026 23:33:36 +0000 Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh | 1 - 1 file changed, 1 deletion(-) diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index f8508f9e39..d0d15d8d6c 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -198,7 +198,6 @@ __global__ void update_tma_descriptors( const size_t offset_elts = offsets_ptr[tensor_id]; - if (leading_thread && (tensor_id < num_tensors)) { { const uintptr_t global_data_ptr = reinterpret_cast(input_data_ptr + offset_elts);