diff --git a/tests/pytorch/mxfp8/test_mxfp8_quantize_raster.py b/tests/pytorch/mxfp8/test_mxfp8_quantize_raster.py new file mode 100644 index 0000000000..6e765da97d --- /dev/null +++ b/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) diff --git a/transformer_engine/common/cast/core/grouped_layout.cuh b/transformer_engine/common/cast/core/grouped_layout.cuh index 8337e44815..98df9370e3 100644 --- a/transformer_engine/common/cast/core/grouped_layout.cuh +++ b/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(block_id % work_blocks_X); + ctaid_Y = static_cast(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(static_next_block_id % work_blocks_X); - ctaid_Y = static_cast(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; diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index 14832573d7..afc399799c 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/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(launch_block_id % work_blocks_X); - int32_t ctaid_Y = static_cast(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(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(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(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( - 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(IN_buff_readable_mbar, leading_thread); diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index a9da143c3c..62ca502e89 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/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,19 +169,24 @@ __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(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. @@ -188,7 +194,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) 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(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::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);