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
Additional context
Add any other context about the problem here.
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
Expected behavior
Not to crash with groups of size zero and simply perform no quantization work.
Environment overview (please complete the following information)
docker pull&docker runcommands usedEnvironment details
If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:
Device details
Additional context
Add any other context about the problem here.