Reverse MXFP8 quantization row raster

tests/pytorch/mxfp8/test_mxfp8_quantize_raster.py

@@ -0,0 +1,72 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+import pytest
+import torch
+import transformer_engine.pytorch as te
+import transformer_engine_torch as tex
+from transformer_engine.pytorch import MXFP8Quantizer
+
+recipe_available, reason_for_no_recipe = te.is_mxfp8_available(return_reason=True)
+
+
+def _rowwise_quantize(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    quantizer = MXFP8Quantizer(
+        fp8_dtype=te.DType.kFloat8E4M3,
+        rowwise=True,
+        columnwise=False,
+    )
+    y = quantizer(x)
+    return y._rowwise_data.view(dtype=torch.uint8), y._rowwise_scale_inv
+
+
+def _rowwise_quantize_padded_reference(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    padded_cols = ((x.size(1) + 127) // 128) * 128
+    x_padded = torch.zeros((x.size(0), padded_cols), dtype=x.dtype, device=x.device)
+    x_padded[:, : x.size(1)] = x
+
+    q_ref, s_ref = _rowwise_quantize(x_padded)
+    valid_scale_cols = x.size(1) // 32
+    return q_ref[:, : x.size(1)].contiguous(), s_ref[: x.size(0), :valid_scale_cols].contiguous()
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+def test_mxfp8_generic_quantize_reverse_raster_preserves_values() -> None:
+    torch.manual_seed(0)
+    torch.cuda.manual_seed(0)
+
+    # N=96 avoids the specialized rowwise cast-only path, exercising the generic TMA path.
+    x = torch.randn((320, 96), dtype=torch.bfloat16, device="cuda")
+
+    q, s = _rowwise_quantize(x)
+    q_ref, s_ref = _rowwise_quantize_padded_reference(x)
+
+    torch.testing.assert_close(q, q_ref, atol=0.0, rtol=0.0)
+    torch.testing.assert_close(s[: x.size(0), : x.size(1) // 32], s_ref, atol=0.0, rtol=0.0)
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+def test_mxfp8_grouped_quantize_reverse_raster_preserves_values() -> None:
+    torch.manual_seed(1)
+    torch.cuda.manual_seed(1)
+
+    split_sections = [128, 256, 128]
+    x = torch.randn((sum(split_sections), 96), dtype=torch.bfloat16, device="cuda")
+    split_section_tensor = torch.tensor(split_sections, dtype=torch.int64, device="cuda")
+    quantizer = MXFP8Quantizer(
+        fp8_dtype=te.DType.kFloat8E4M3,
+        rowwise=True,
+        columnwise=False,
+    )
+
+    grouped_output = tex.group_quantize(x, quantizer, len(split_sections), split_section_tensor)
+    outputs = grouped_output.split_into_quantized_tensors()
+
+    for x_chunk, output in zip(torch.split(x, split_sections), outputs):
+        q_ref, s_ref = _rowwise_quantize_padded_reference(x_chunk)
+        q = output._rowwise_data.view(dtype=torch.uint8)
+        s = output._rowwise_scale_inv
+
+        torch.testing.assert_close(q, q_ref, atol=0.0, rtol=0.0)
+        torch.testing.assert_close(s, s_ref, atol=0.0, rtol=0.0)

transformer_engine/common/cast/core/grouped_layout.cuh

@@ -374,14 +374,24 @@ __device__ __forceinline__ bool job_has_work(const JobDescriptor &job) {
   return job.rows != 0 && job.cols != 0;
 }
 
+__device__ __forceinline__ void linear_block_id_to_reverse_y_cta_coords(const size_t block_id,
+                                                                        const size_t work_blocks_X,
+                                                                        const size_t work_blocks_Y,
+                                                                        int32_t &ctaid_X,
+                                                                        int32_t &ctaid_Y) {
+  ctaid_X = static_cast<int32_t>(block_id % work_blocks_X);
+  ctaid_Y = static_cast<int32_t>(work_blocks_Y - 1 - block_id / work_blocks_X);
+}
+
 __device__ __forceinline__ void advance_to_next_job(bool &job_finished, int32_t &ctaid_X,
                                                     int32_t &ctaid_Y, size_t &static_next_block_id,
                                                     const size_t static_block_stride,
                                                     const size_t total_work_blocks,
-                                                    const size_t work_blocks_X) {
+                                                    const size_t work_blocks_X,
+                                                    const size_t work_blocks_Y) {
   if (static_next_block_id < total_work_blocks) {
-    ctaid_X = static_cast<int32_t>(static_next_block_id % work_blocks_X);
-    ctaid_Y = static_cast<int32_t>(static_next_block_id / work_blocks_X);
+    linear_block_id_to_reverse_y_cta_coords(static_next_block_id, work_blocks_X, work_blocks_Y,
+                                            ctaid_X, ctaid_Y);
     static_next_block_id += static_block_stride;
   } else {
     job_finished = true;

transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh

@@ -528,8 +528,10 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel
   if (launch_block_id >= total_work_blocks) {
     return;
   }
-  int32_t ctaid_X = static_cast<int32_t>(launch_block_id % work_blocks_X);
-  int32_t ctaid_Y = static_cast<int32_t>(launch_block_id / work_blocks_X);
+  int32_t ctaid_X = 0;
+  int32_t ctaid_Y = 0;
+  linear_block_id_to_reverse_y_cta_coords(launch_block_id, work_blocks_X, work_blocks_Y, ctaid_X,
+                                          ctaid_Y);
   size_t static_block_stride = gridDim.x * gridDim.y;
   size_t static_next_block_id = launch_block_id + static_block_stride;
 
@@ -555,7 +557,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel
     if (!job_has_work(current_job)) {
       // Zero-sized tensors are valid grouped-tensor entries; skip them and keep scheduling work.
       advance_to_next_job(job_finished, ctaid_X, ctaid_Y, static_next_block_id, static_block_stride,
-                          total_work_blocks, work_blocks_X);
+                          total_work_blocks, work_blocks_X, work_blocks_Y);
       continue;
     }
 
@@ -632,7 +634,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel
 #pragma unroll
     for (int stage = 0; stage < PREFETCH_STAGES; ++stage) {
       const size_t buff = stage;
-      const size_t stage_offset_Y = stage * BUFF_DIM_Y;
+      const int logical_stage = static_cast<int>(STAGES) - 1 - stage;
+      const size_t stage_offset_Y = logical_stage * BUFF_DIM_Y;
       const size_t global_offset_Y = block_offset_Y + stage_offset_Y;
       const size_t global_offset_X = block_offset_X;
       const size_t buff_offset = buff * BUFF_DIM;
@@ -654,11 +657,13 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel
 // Process one [CHUNK_DIM_Y x CHUNK_DIM_X] block in STAGES slices (32 rows each).
 #pragma unroll
     for (int stage = 0; stage < STAGES; ++stage) {
-      const size_t stage_offset_Y = stage * BUFF_DIM_Y;
+      const int logical_stage = static_cast<int>(STAGES) - 1 - stage;
+      const size_t stage_offset_Y = logical_stage * BUFF_DIM_Y;
       if (stage < STAGES - PREFETCH_STAGES) {
         const size_t next_prefetch_buff = (buff_in + PREFETCH_STAGES) % BUFFS_NUM;
         const size_t next_prefetch_stage = stage + PREFETCH_STAGES;
-        const size_t next_prefetch_stage_offset_Y = next_prefetch_stage * BUFF_DIM_Y;
+        const int next_logical_stage = static_cast<int>(STAGES) - 1 - next_prefetch_stage;
+        const size_t next_prefetch_stage_offset_Y = next_logical_stage * BUFF_DIM_Y;
 
         const size_t global_offset_Y = block_offset_Y + next_prefetch_stage_offset_Y;
         const size_t global_offset_X = block_offset_X;
@@ -679,7 +684,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel
       if constexpr (COLWISE_SCALING) {
         process_colwise_stage<IS_DBIAS, IS_DACT, IS_ACT, ParamOP, OP, IType, OType, ROWWISE_SCALING,
                               WITH_GEMM_SWIZZLED_SCALES>(
-            buff, stage, tid_X_colwise, scales_offset_Y_colwise, scales_offset_X_colwise,
+            buff, logical_stage, tid_X_colwise, scales_offset_Y_colwise, scales_offset_X_colwise,
             scale_stride_colwise, tensor_base_for_scales, rows, cols, sIn_ptr, sActIn_ptr,
             sCachedAct_ptr, sOutColwise_ptr, scales_colwise, partial_dbias_colwise);
       }
@@ -750,7 +755,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel
     }
 
     advance_to_next_job(job_finished, ctaid_X, ctaid_Y, static_next_block_id, static_block_stride,
-                        total_work_blocks, work_blocks_X);
+                        total_work_blocks, work_blocks_X, work_blocks_Y);
   }
 
   destroy_barriers<BUFFS_NUM>(IN_buff_readable_mbar, leading_thread);

transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh

@@ -82,11 +82,12 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
 
   constexpr bool IS_CACHED_ACT_OP = COMPUTE_ACTIVATIONS && ROWWISE_SCALING && COLWISE_SCALING;
 
-  const size_t block_offset_Y = blockIdx.y * CHUNK_DIM_Y;
+  const size_t block_id_Y = gridDim.y - 1 - blockIdx.y;
+  const size_t block_offset_Y = block_id_Y * CHUNK_DIM_Y;
   const size_t block_offset_X = blockIdx.x * CHUNK_DIM_X;
-  const size_t scales_block_offset_Y_rowwise = blockIdx.y * CHUNK_DIM_Y;
+  const size_t scales_block_offset_Y_rowwise = block_id_Y * CHUNK_DIM_Y;
   const size_t scales_block_offset_X_rowwise = blockIdx.x * CHUNK_DIM_X / SCALE_DIM_X;
-  const size_t scales_block_offset_Y_colwise = blockIdx.y * CHUNK_DIM_Y / SCALE_DIM_Y;
+  const size_t scales_block_offset_Y_colwise = block_id_Y * CHUNK_DIM_Y / SCALE_DIM_Y;
   const size_t scales_block_offset_X_colwise = blockIdx.x * CHUNK_DIM_X;
 
   const size_t tid_Y_rowwise = threadIdx.x / THREADS_X;
@@ -168,27 +169,33 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
   int parity = 0;
 
   if constexpr (IS_DACT) {
-    copy_2d_to_sharedx2(&in_sh[0], &tensor_map_input, block_offset_X, block_offset_Y, &act_in_sh[0],
-                        &tensor_map_act_input, block_offset_X, block_offset_Y, shmem_buff_size,
-                        &mbar[0], is_master_thread);
+    const size_t first_stage_offset_Y = (STAGES - 1) * BUFF_DIM_Y;
+    const size_t global_offset_Y = block_offset_Y + first_stage_offset_Y;
+    copy_2d_to_sharedx2(&in_sh[0], &tensor_map_input, block_offset_X, global_offset_Y,
+                        &act_in_sh[0], &tensor_map_act_input, block_offset_X, global_offset_Y,
+                        shmem_buff_size, &mbar[0], is_master_thread);
   } else {
-    copy_2d_to_shared(&in_sh[0], &tensor_map_input, block_offset_X, block_offset_Y, shmem_buff_size,
-                      &mbar[0], is_master_thread);
+    const size_t first_stage_offset_Y = (STAGES - 1) * BUFF_DIM_Y;
+    const size_t global_offset_Y = block_offset_Y + first_stage_offset_Y;
+    copy_2d_to_shared(&in_sh[0], &tensor_map_input, block_offset_X, global_offset_Y,
+                      shmem_buff_size, &mbar[0], is_master_thread);
   }
 
 #pragma unroll
   for (int stage = 0; stage < STAGES; ++stage) {
     const size_t buff = stage % BUFFS_NUM;
     const size_t next_stage = stage + 1;
-    const size_t stage_offset_Y = stage * BUFF_DIM_Y;
+    const int logical_stage = static_cast<int>(STAGES) - 1 - stage;
+    const size_t stage_offset_Y = logical_stage * BUFF_DIM_Y;
 
     if (next_stage < STAGES) {
       // Wait for TMA transfer to have finished reading shared memory.
       // I.e. the buffer is ready to be written to
       ptx::cp_async_bulk_wait_group_read<1>();
 
       const size_t next_buff = next_stage % BUFFS_NUM;
-      const size_t next_stage_offset_Y = next_stage * BUFF_DIM_Y;
+      const int next_logical_stage = static_cast<int>(STAGES) - 1 - next_stage;
+      const size_t next_stage_offset_Y = next_logical_stage * BUFF_DIM_Y;
       const size_t global_offset_Y = block_offset_Y + next_stage_offset_Y;
       const size_t global_offset_X = block_offset_X;
       const size_t next_buff_offset = next_buff * BUFF_DIM;
@@ -267,7 +274,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
       // 2. Compute E8M0 scaling factor
       const e8m0_t biased_exponent =
           ptx::float_to_e8m0(thread_amax * Quantized_Limits<OType>::max_norm_rcp);
-      const size_t global_scales_offset_Y = scales_offset_Y_colwise + stage;
+      const size_t global_scales_offset_Y = scales_offset_Y_colwise + logical_stage;
       const size_t global_scales_offset_X = scales_offset_X_colwise;
       size_t scale_idx;
       if constexpr (WITH_GEMM_SWIZZLED_SCALES) {
@@ -531,7 +538,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
       }
     }
     const int dbias_stride = cols;
-    const int dbias_offset_Y = blockIdx.y;
+    const int dbias_offset_Y = block_id_Y;
     const int dbias_offset_X = blockIdx.x * CHUNK_DIM_X + threadIdx.x;
     const int dbias_idx = dbias_offset_Y * dbias_stride + dbias_offset_X;
     const bool col_out_of_bounds_dbias = (dbias_offset_X >= cols);