[Core] MXFP8 grouped quantization kernel crashes with groups of size zero

**Describe the bug** In MXFP8 grouped quantization, TMA descriptors are created with size zero when a group exists with size zero. **Steps/Code to reproduce bug** Try running the kernel with input with a group of size zero. This diff fixes it ``` diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index 129d6724..d7eaf028 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -192,6 +192,10 @@ __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<uintptr_t>(input_data_ptr + offset_elts); modify_base_tensor_map(base_tensor_map_input, &g_tensor_maps_input[tensor_id], ``` **Expected behavior** Not to crash with groups of size zero and simply perform no quantization work. **Environment overview (please complete the following information)** - Environment location: [Bare-metal, Docker, Cloud(specify cloud provider - AWS, Azure, GCP, Collab)] - Method of Transformer Engine install: [pip install or from source]. Please specify exact commands you used to install. - If method of install is [Docker], provide `docker pull` & `docker run` commands used **Environment details** If NVIDIA docker image is used you don't need to specify these. Otherwise, please provide: - OS version - PyTorch version - Python version - Transformer Engine version - CUDA version - CUDNN version **Device details** - GPU model **Additional context** Add any other context about the problem here.