From 1f306a100841b179b969d30ee217974ce0ae8b34 Mon Sep 17 00:00:00 2001 From: xiuhu17 Date: Sun, 6 Sep 2026 01:27:52 -0500 Subject: [PATCH 1/5] Support row-only MXFP8 distributed master-weight casts Signed-off-by: xiuhu17 --- .../pytorch/mxfp8/test_rowwise_master_cast.py | 132 ++++++++++++++++++ .../include/transformer_engine/recipe.h | 4 +- .../common/recipe/mxfp8_scaling.cu | 91 +++++++----- transformer_engine/pytorch/tensor/utils.py | 41 +++++- 4 files changed, 227 insertions(+), 41 deletions(-) create mode 100644 tests/pytorch/mxfp8/test_rowwise_master_cast.py diff --git a/tests/pytorch/mxfp8/test_rowwise_master_cast.py b/tests/pytorch/mxfp8/test_rowwise_master_cast.py new file mode 100644 index 00000000000..9b2a952884d --- /dev/null +++ b/tests/pytorch/mxfp8/test_rowwise_master_cast.py @@ -0,0 +1,132 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# See LICENSE for license information. + +"""Run with pytest on one GPU or torchrun -m pytest on multiple GPUs.""" + +import os + +import pytest +import torch +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.pytorch.tensor.utils import ( + cast_master_weights_to_fp8, + quantize_master_weights, +) + +available, reason = te.is_mxfp8_available(return_reason=True) +pytestmark = pytest.mark.skipif(not available, reason=reason) + + +@pytest.fixture +def group(): + owned = not torch.distributed.is_initialized() + if owned: + torch.cuda.set_device(int(os.getenv("LOCAL_RANK", "0"))) + if "RANK" in os.environ: + torch.distributed.init_process_group("nccl") + else: + torch.distributed.init_process_group( + "nccl", store=torch.distributed.HashStore(), rank=0, world_size=1 + ) + yield torch.distributed.group.WORLD + if owned: + torch.distributed.destroy_process_group() + + +@pytest.mark.parametrize("mixed", [False, True]) +@pytest.mark.parametrize("fragments", [False, True]) +def test_rowwise_master_cast(group, monkeypatch, mixed, fragments): + rank = torch.distributed.get_rank(group) + world = torch.distributed.get_world_size(group) + shape = (160, 96) # Both scale padding and non-tile-aligned logical shapes. + full = torch.linspace(-4, 7, 160 * 96, device="cuda").reshape(shape) + quantizers = [te.MXFP8Quantizer(te.DType.kFloat8E4M3, rowwise=True, columnwise=False)] + if mixed: + quantizers.append(te.MXFP8Quantizer(te.DType.kFloat8E4M3, rowwise=True, columnwise=True)) + weights = [q(full.to(torch.bfloat16)) for q in quantizers] + pointers = [(w._rowwise_data.data_ptr(), w._rowwise_scale_inv.data_ptr()) for w in weights] + if world == 1: + lo, hi = 0, full.numel() + elif rank == 0: + lo, hi = 0, 17 # Split inside a 32-value block. + elif rank == 1: + lo, hi = 17, full.numel() + else: + lo, hi = 0, 0 # No master shard on tail ranks. + real_reduce = torch.distributed.all_reduce + reduced_sizes = [] + + def record_reduce(tensor, *args, **kwargs): + reduced_sizes.append(tensor.numel()) + return real_reduce(tensor, *args, **kwargs) + + monkeypatch.setattr(torch.distributed, "all_reduce", record_reduce) + for step in range(2): + master = full + step * 0.25 + shard = master.flatten()[lo:hi] if lo < hi else None + outputs = [ + ( + torch.empty(hi - lo, dtype=torch.uint8, device="cuda"), + torch.empty(hi - lo, dtype=torch.uint8, device="cuda") if i else None, + ) + for i in range(len(weights)) + ] + reduced_sizes.clear() + cast = quantize_master_weights if step == 0 else cast_master_weights_to_fp8 + cast( + weights, + [shard] * len(weights), + [lo if shard is not None else None] * len(weights), + group, + fsdp_shard_model_weights=outputs if fragments else None, + ) + expected_amax = sum(w._rowwise_scale_inv.numel() for w in weights) + if mixed: + expected_amax += weights[1]._columnwise_scale_inv.numel() + assert reduced_sizes == [expected_amax] + for idx, (weight, quantizer) in enumerate(zip(weights, quantizers)): + expected = quantizer(master.to(torch.bfloat16)) + for direction in ["rowwise", "columnwise"] if idx else ["rowwise"]: + data = getattr(weight, f"_{direction}_data") + actual = torch.zeros_like(data).flatten() + if lo < hi: + src = ( + outputs[idx][direction == "columnwise"] + if fragments + else data.flatten()[lo:hi] + ) + actual[lo:hi].copy_(src) + real_reduce(actual, op=torch.distributed.ReduceOp.MAX, group=group) + torch.testing.assert_close( + actual.view(shape), getattr(expected, f"_{direction}_data"), rtol=0, atol=0 + ) + scale = getattr(weight, f"_{direction}_scale_inv") + ref = getattr(expected, f"_{direction}_scale_inv") + rows, cols = (160, 3) if direction == "rowwise" else (5, 96) + torch.testing.assert_close(scale[:rows, :cols], ref[:rows, :cols], rtol=0, atol=0) + assert pointers[idx] == ( + weight._rowwise_data.data_ptr(), + weight._rowwise_scale_inv.data_ptr(), + ) + assert weights[0]._columnwise_data is None + assert weights[0]._columnwise_scale_inv is None + + +@pytest.mark.parametrize("offset,length", [(0, 0), (31, 3), (1, 4094)]) +def test_rowwise_partial_kernels(offset, length): + inp = torch.randn(length, device="cuda", dtype=torch.bfloat16) + row = torch.zeros((128, 4), device="cuda", dtype=inp.dtype) + ref = torch.zeros_like(row) + col = torch.zeros((4, 128), device="cuda", dtype=inp.dtype) + # Empty view may still have backing storage. Omission is shape-based. + omitted = col.flatten()[:0].view(0, 0) + tex.mxfp8_scaling_compute_partial_amax(inp, row, omitted, 64, 64, offset) + tex.mxfp8_scaling_compute_partial_amax(inp, ref, col, 64, 64, offset) + torch.testing.assert_close(row, ref, rtol=0, atol=0) + scales = torch.full((128, 4), 127, dtype=torch.uint8, device="cuda") + out = torch.empty(length, dtype=torch.uint8, device="cuda") + tex.mxfp8_scaling_partial_cast( + inp, out, out[:0], scales, scales.flatten()[:0].view(0, 0), 64, 64, offset + ) + torch.testing.assert_close(out, inp.to(torch.float8_e4m3fn).view(torch.uint8), rtol=0, atol=0) diff --git a/transformer_engine/common/include/transformer_engine/recipe.h b/transformer_engine/common/include/transformer_engine/recipe.h index 47539a89a17..1ce8923a960 100644 --- a/transformer_engine/common/include/transformer_engine/recipe.h +++ b/transformer_engine/common/include/transformer_engine/recipe.h @@ -290,6 +290,7 @@ void nvte_fp8_block_scaling_partial_cast(const NVTETensor inp, NVTETensor out, * \param[in] input Input tensor (continuous segment of flattened original tensor). * \param[in,out] amax_rowwise Output tensor for row-wise maximum absolute values. * \param[in,out] amax_colwise Output tensor for column-wise maximum absolute values. + * Shape [0, 0] skips column-wise computation. * \param[in] rows Number of rows in the logical tensor. * \param[in] cols Number of columns in the logical tensor. * \param[in] start_offset Starting offset in the flattened tensor. @@ -303,7 +304,8 @@ void nvte_mxfp8_scaling_compute_partial_amax(const NVTETensor input, NVTETensor * * This function casts the input tensor to MXFP8 format, producing both row-wise and * column-wise scaled outputs. input contains a continuous segment from the flattened - * original tensor. + * original tensor. To skip column-wise computation, pass output_colwise with shape [0] + * and scale_inv_colwise with shape [0, 0]. Row-wise outputs remain required. * * \param[in] input Input (continuous segment of flattened original tensor). * \param[out] output_rowwise Output tensor with row-wise scaling (MXFP8 format). diff --git a/transformer_engine/common/recipe/mxfp8_scaling.cu b/transformer_engine/common/recipe/mxfp8_scaling.cu index be692d45633..5d40d5475e0 100644 --- a/transformer_engine/common/recipe/mxfp8_scaling.cu +++ b/transformer_engine/common/recipe/mxfp8_scaling.cu @@ -47,7 +47,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) if (r < rows && c < cols && idx >= start_offset && idx < end_offset) { float abs_input = fabs(static_cast(input_minus_offset[idx])); row_amax = fmaxf(row_amax, abs_input); - col_amax = fmaxf(col_amax, abs_input); + if (amax_colwise != nullptr) col_amax = fmaxf(col_amax, abs_input); } #pragma unroll @@ -63,7 +63,9 @@ __global__ void __launch_bounds__(kThreadsPerBlock) r++; } - amax_colwise[blockIdx.y * amax_colwise_stride + c] = static_cast(col_amax); + if (amax_colwise != nullptr) { + amax_colwise[blockIdx.y * amax_colwise_stride + c] = static_cast(col_amax); + } __syncthreads(); @@ -95,7 +97,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) } // Load scales_colwise - { + if (output_colwise != nullptr) { int c_ = threadIdx.x; int r = blockIdx.y * kRowsPerTile / 32; int c = blockIdx.x * kColsPerTile + c_; @@ -108,7 +110,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock) size_t end_offset = start_offset + len; const IType *input_minus_offset = input - start_offset; OType *output_rowwise_minus_offset = output_rowwise - start_offset; - OType *output_colwise_minus_offset = output_colwise - start_offset; + OType *output_colwise_minus_offset = + output_colwise != nullptr ? output_colwise - start_offset : nullptr; int warp_idx = threadIdx.x / 32; // int lane_idx = threadIdx.x % 32; int c = blockIdx.x * kColsPerTile + threadIdx.x; @@ -121,9 +124,11 @@ __global__ void __launch_bounds__(kThreadsPerBlock) if (r < rows && c < cols && idx >= start_offset && idx < end_offset) { float inp = static_cast(input_minus_offset[idx]); OType out_rowwise = static_cast(inp * smem_scales_rowwise[i][warp_idx]); - OType out_colwise = static_cast(inp * smem_scales_colwise[threadIdx.x]); output_rowwise_minus_offset[idx] = out_rowwise; - output_colwise_minus_offset[idx] = out_colwise; + if (output_colwise != nullptr) { + OType out_colwise = static_cast(inp * smem_scales_colwise[threadIdx.x]); + output_colwise_minus_offset[idx] = out_colwise; + } } r++; @@ -149,14 +154,20 @@ void mxfp8_scaling_compute_partial_amax(const Tensor input, Tensor amax_rowwise, NVTE_CHECK(amax_rowwise.data.shape[1] >= cols / 32, "Invalid cols"); NVTE_CHECK(amax_rowwise.dtype() == input.dtype(), "Wrong dtype of amax_rowwise"); - NVTE_CHECK(amax_colwise.data.shape.size() == 2, "amax_colwise must be a 2D tensor"); - NVTE_CHECK(amax_colwise.data.shape[0] % colwise_row_padding == 0, - "Wrong padding of amax_colwise's rows"); - NVTE_CHECK(amax_colwise.data.shape[0] >= rows / 32, "Invalid rows"); - NVTE_CHECK(amax_colwise.data.shape[1] % colwise_col_padding == 0, - "Wrong padding of amax_colwise's cols"); - NVTE_CHECK(amax_colwise.data.shape[1] >= cols, "Invalid cols"); - NVTE_CHECK(amax_colwise.dtype() == input.dtype(), "Wrong dtype of amax_colwise"); + // Empty views can have a non-null data pointer: use shape to detect omission. + const bool with_columnwise = + !(amax_colwise.data.shape.size() == 2 && amax_colwise.data.shape[0] == 0 && + amax_colwise.data.shape[1] == 0); + if (with_columnwise) { + NVTE_CHECK(amax_colwise.data.shape.size() == 2, "amax_colwise must be a 2D tensor"); + NVTE_CHECK(amax_colwise.data.shape[0] % colwise_row_padding == 0, + "Wrong padding of amax_colwise's rows"); + NVTE_CHECK(amax_colwise.data.shape[0] >= rows / 32, "Invalid rows"); + NVTE_CHECK(amax_colwise.data.shape[1] % colwise_col_padding == 0, + "Wrong padding of amax_colwise's cols"); + NVTE_CHECK(amax_colwise.data.shape[1] >= cols, "Invalid cols"); + NVTE_CHECK(amax_colwise.dtype() == input.dtype(), "Wrong dtype of amax_colwise"); + } int blocks_x = (cols + kColsPerTile - 1) / kColsPerTile; int blocks_y = (rows + kRowsPerTile - 1) / kRowsPerTile; @@ -167,8 +178,9 @@ void mxfp8_scaling_compute_partial_amax(const Tensor input, Tensor amax_rowwise, mxfp8_scaling_compute_partial_amax_kernel<<>>( reinterpret_cast(input.data.dptr), reinterpret_cast(amax_rowwise.data.dptr), - reinterpret_cast(amax_colwise.data.dptr), amax_rowwise.data.shape[1], - amax_colwise.data.shape[1], rows, cols, start_offset, input.data.shape[0]);) + with_columnwise ? reinterpret_cast(amax_colwise.data.dptr) : nullptr, + amax_rowwise.data.shape[1], with_columnwise ? amax_colwise.data.shape[1] : 0, rows, cols, + start_offset, input.data.shape[0]);) } void mxfp8_scaling_partial_cast(const Tensor input, Tensor output_rowwise, Tensor output_colwise, @@ -182,16 +194,27 @@ void mxfp8_scaling_partial_cast(const Tensor input, Tensor output_rowwise, Tenso "Invalid start_offset"); NVTE_CHECK(output_rowwise.data.shape.size() == 1, "output_rowwise must be a 1D tensor"); - NVTE_CHECK(output_colwise.data.shape.size() == 1, "output_colwise must be a 1D tensor"); NVTE_CHECK(output_rowwise.data.shape[0] == input.data.shape[0], "Size of input and output_rowwise mismatch"); - NVTE_CHECK(output_colwise.data.shape[0] == input.data.shape[0], - "Size of input and output_colwise mismatch"); NVTE_CHECK(output_rowwise.dtype() == DType::kFloat8E4M3 || output_rowwise.dtype() == DType::kByte, "output_rowwise should be e4m3 or uint8"); - NVTE_CHECK(output_colwise.dtype() == DType::kFloat8E4M3 || output_colwise.dtype() == DType::kByte, - "output_colwise should be e4m3 or uint8"); + const bool with_columnwise = + !(scale_inv_colwise.data.shape.size() == 2 && scale_inv_colwise.data.shape[0] == 0 && + scale_inv_colwise.data.shape[1] == 0); + NVTE_CHECK(with_columnwise || + (output_colwise.data.shape.size() == 1 && output_colwise.data.shape[0] == 0), + "Omitted columnwise scales require an empty columnwise output"); + if (with_columnwise) { + NVTE_CHECK(output_colwise.data.shape.size() == 1, "output_colwise must be a 1D tensor"); + NVTE_CHECK(output_colwise.data.shape[0] == input.data.shape[0], + "Size of input and output_colwise mismatch"); + NVTE_CHECK(input.data.shape[0] == 0 || output_colwise.data.dptr != nullptr, + "Columnwise scales require columnwise output"); + NVTE_CHECK( + output_colwise.dtype() == DType::kFloat8E4M3 || output_colwise.dtype() == DType::kByte, + "output_colwise should be e4m3 or uint8"); + } NVTE_CHECK(scale_inv_rowwise.data.shape.size() == 2, "scale_inv_rowwise must be a 2D tensor"); NVTE_CHECK(scale_inv_rowwise.data.shape[0] % rowwise_row_padding == 0, @@ -202,14 +225,16 @@ void mxfp8_scaling_partial_cast(const Tensor input, Tensor output_rowwise, Tenso NVTE_CHECK(scale_inv_rowwise.data.shape[1] >= cols / 32, "Invalid cols"); NVTE_CHECK(scale_inv_rowwise.dtype() == DType::kByte, "Wrong dtype of scale_inv_rowwise"); - NVTE_CHECK(scale_inv_colwise.data.shape.size() == 2, "scale_inv_colwise must be a 2D tensor"); - NVTE_CHECK(scale_inv_colwise.data.shape[0] % colwise_row_padding == 0, - "Wrong padding of scale_inv_colwise's rows"); - NVTE_CHECK(scale_inv_colwise.data.shape[0] >= rows / 32, "Invalid rows"); - NVTE_CHECK(scale_inv_colwise.data.shape[1] % colwise_col_padding == 0, - "Wrong padding of scale_inv_colwise's cols"); - NVTE_CHECK(scale_inv_colwise.data.shape[1] >= cols, "Invalid cols"); - NVTE_CHECK(scale_inv_colwise.dtype() == DType::kByte, "Wrong dtype of scale_inv_colwise"); + if (with_columnwise) { + NVTE_CHECK(scale_inv_colwise.data.shape.size() == 2, "scale_inv_colwise must be a 2D tensor"); + NVTE_CHECK(scale_inv_colwise.data.shape[0] % colwise_row_padding == 0, + "Wrong padding of scale_inv_colwise's rows"); + NVTE_CHECK(scale_inv_colwise.data.shape[0] >= rows / 32, "Invalid rows"); + NVTE_CHECK(scale_inv_colwise.data.shape[1] % colwise_col_padding == 0, + "Wrong padding of scale_inv_colwise's cols"); + NVTE_CHECK(scale_inv_colwise.data.shape[1] >= cols, "Invalid cols"); + NVTE_CHECK(scale_inv_colwise.dtype() == DType::kByte, "Wrong dtype of scale_inv_colwise"); + } int blocks_x = (cols + kColsPerTile - 1) / kColsPerTile; int blocks_y = (rows + kRowsPerTile - 1) / kRowsPerTile; @@ -220,11 +245,11 @@ void mxfp8_scaling_partial_cast(const Tensor input, Tensor output_rowwise, Tenso mxfp8_scaling_partial_cast_kernel<<>>( reinterpret_cast(input.data.dptr), reinterpret_cast(output_rowwise.data.dptr), - reinterpret_cast(output_colwise.data.dptr), + with_columnwise ? reinterpret_cast(output_colwise.data.dptr) : nullptr, reinterpret_cast(scale_inv_rowwise.data.dptr), - reinterpret_cast(scale_inv_colwise.data.dptr), - scale_inv_rowwise.data.shape[1], scale_inv_colwise.data.shape[1], rows, cols, - start_offset, input.data.shape[0]);) + with_columnwise ? reinterpret_cast(scale_inv_colwise.data.dptr) : nullptr, + scale_inv_rowwise.data.shape[1], with_columnwise ? scale_inv_colwise.data.shape[1] : 0, + rows, cols, start_offset, input.data.shape[0]);) } } // namespace mxfp8_scaling_recipe diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index cef45c0223f..18e2ae0c2dd 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -1017,6 +1017,8 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( ): # pylint: disable=unused-argument r"""Helper function to cast master weights to FP8 primary weights for mxfp8 scaling. + Row-only primaries omit columnwise amax, scale updates, and output writes. + Parameters ---------- params : List of tuple, each tuple contains a model weight, a master weight, and an offset @@ -1046,7 +1048,15 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( raise ValueError( f"rowwise_scale_inv must be 2D, got {len(rowwise_shape)}D shape {rowwise_shape}" ) - colwise_shape = model_weight._columnwise_scale_inv.shape + if (model_weight._columnwise_data is None) != ( + model_weight._columnwise_scale_inv is None + ): + raise ValueError("MXFP8 columnwise data and scales must both be present or absent") + colwise_shape = ( + model_weight._columnwise_scale_inv.shape + if model_weight._columnwise_scale_inv is not None + else (0, 0) + ) if len(colwise_shape) != 2: raise ValueError( f"columnwise_scale_inv must be 2D, got {len(colwise_shape)}D shape {colwise_shape}" @@ -1075,7 +1085,11 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( amaxes_colwise, scale_invs_colwise = [], [] for i, (model_weight, master_weight, start_offset, _) in enumerate(params): rowwise_shape = model_weight._rowwise_scale_inv.shape - colwise_shape = model_weight._columnwise_scale_inv.shape + colwise_shape = ( + model_weight._columnwise_scale_inv.shape + if model_weight._columnwise_scale_inv is not None + else (0, 0) + ) rowwise_start = cu_rowwise_amax_sizes[i] rowwise_end = cu_rowwise_amax_sizes[i + 1] colwise_start = cu_rowwise_amax_sizes[-1] + cu_colwise_amax_sizes[i] @@ -1085,7 +1099,12 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( amaxes_rowwise.append(amax_rowwise) amaxes_colwise.append(amax_colwise) scale_invs_rowwise.append(model_weight._rowwise_scale_inv) - scale_invs_colwise.append(model_weight._columnwise_scale_inv) + # Empty tensors mark an omitted direction without allocating a GPU payload. + scale_invs_colwise.append( + model_weight._columnwise_scale_inv + if model_weight._columnwise_scale_inv is not None + else torch.empty((0, 0), dtype=torch.uint8, device=device) + ) # Compute amax of the master weight and store it in packed_amaxes. if master_weight is not None: @@ -1111,8 +1130,8 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( multi_tensor_compute_scale_inv_e8m0, None, # dummy_overflow_buf [ - amaxes_rowwise + amaxes_colwise, - scale_invs_rowwise + scale_invs_colwise, + amaxes_rowwise + [amax for amax in amaxes_colwise if amax.numel()], + scale_invs_rowwise + [scale for scale in scale_invs_colwise if scale.numel()], ], ) @@ -1133,10 +1152,18 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( end_offset = start_offset + master_weight.numel() if use_fsdp_shard_model_weights: rowwise_fragment = model_weight_fragment[0] - colwise_fragment = model_weight_fragment[1] + colwise_fragment = ( + model_weight_fragment[1] + if scale_inv_colwise.numel() + else torch.empty(0, dtype=torch.uint8, device=device) + ) else: rowwise_fragment = model_weight._rowwise_data.reshape(-1)[start_offset:end_offset] - colwise_fragment = model_weight._columnwise_data.reshape(-1)[start_offset:end_offset] + colwise_fragment = ( + model_weight._columnwise_data.reshape(-1)[start_offset:end_offset] + if scale_inv_colwise.numel() + else torch.empty(0, dtype=torch.uint8, device=device) + ) if len(model_weight.shape) != 2: raise ValueError( "model_weight must be 2D for MXFP8 scaling partial cast, " From 9c04c1a1d27cda4923084f843cdf87996e404174 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 6 Sep 2026 06:30:05 +0000 Subject: [PATCH 2/5] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/tensor/utils.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index 18e2ae0c2dd..6da03679a1f 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -1048,9 +1048,7 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( raise ValueError( f"rowwise_scale_inv must be 2D, got {len(rowwise_shape)}D shape {rowwise_shape}" ) - if (model_weight._columnwise_data is None) != ( - model_weight._columnwise_scale_inv is None - ): + if (model_weight._columnwise_data is None) != (model_weight._columnwise_scale_inv is None): raise ValueError("MXFP8 columnwise data and scales must both be present or absent") colwise_shape = ( model_weight._columnwise_scale_inv.shape From 497dd39adbca7bc8feeacabf56036781ece96e6e Mon Sep 17 00:00:00 2001 From: xiuhu17 Date: Sun, 6 Sep 2026 01:34:40 -0500 Subject: [PATCH 3/5] Share the NCCL process group across row-only cast tests Signed-off-by: xiuhu17 --- tests/pytorch/mxfp8/test_rowwise_master_cast.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pytorch/mxfp8/test_rowwise_master_cast.py b/tests/pytorch/mxfp8/test_rowwise_master_cast.py index 9b2a952884d..a876a8fb236 100644 --- a/tests/pytorch/mxfp8/test_rowwise_master_cast.py +++ b/tests/pytorch/mxfp8/test_rowwise_master_cast.py @@ -18,7 +18,7 @@ pytestmark = pytest.mark.skipif(not available, reason=reason) -@pytest.fixture +@pytest.fixture(scope="module") def group(): owned = not torch.distributed.is_initialized() if owned: From 704b1707da6f67c716ea3ec8e295cb258082a14b Mon Sep 17 00:00:00 2001 From: Zhihao Wang Date: Wed, 7 Oct 2026 22:58:13 -0700 Subject: [PATCH 4/5] Address row-only MXFP8 review and preserve bidirectional performance Signed-off-by: Zhihao Wang --- .../test_cast_master_weights_to_fp8.py | 84 +++++++++++ .../pytorch/mxfp8/test_rowwise_master_cast.py | 132 ------------------ tests/pytorch/test_partial_cast.py | 21 +++ .../include/transformer_engine/recipe.h | 6 +- .../common/recipe/mxfp8_scaling.cu | 63 +++++---- .../csrc/extensions/fp8_partial_cast.cpp | 13 +- 6 files changed, 150 insertions(+), 169 deletions(-) delete mode 100644 tests/pytorch/mxfp8/test_rowwise_master_cast.py diff --git a/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py b/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py index ab047dc32ff..7bf9b3699dd 100644 --- a/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py +++ b/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py @@ -9,6 +9,7 @@ import subprocess import sys import pathlib +from unittest.mock import patch sys.path.append(str(pathlib.Path(__file__).resolve().parent.parent)) from utils import run_distributed @@ -958,6 +959,86 @@ def _test_mxfp8_empty_master_shard(dp_group): ) +def _test_rowwise_master_cast(group, mixed, fragments): + rank = torch.distributed.get_rank(group) + world = torch.distributed.get_world_size(group) + shape = (160, 96) # Both scale padding and non-tile-aligned logical shapes. + full = torch.linspace(-4, 7, 160 * 96, device="cuda").reshape(shape) + quantizers = [te.MXFP8Quantizer(te.DType.kFloat8E4M3, rowwise=True, columnwise=False)] + if mixed: + quantizers.append(te.MXFP8Quantizer(te.DType.kFloat8E4M3, rowwise=True, columnwise=True)) + weights = [q(full.to(torch.bfloat16)) for q in quantizers] + pointers = [(w._rowwise_data.data_ptr(), w._rowwise_scale_inv.data_ptr()) for w in weights] + if world == 1: + lo, hi = 0, full.numel() + elif rank == 0: + lo, hi = 0, 17 # Split inside a 32-value block. + elif rank == 1: + lo, hi = 17, full.numel() + else: + lo, hi = 0, 0 # No master shard on tail ranks. + real_reduce = torch.distributed.all_reduce + reduced_sizes = [] + + def record_reduce(tensor, *args, **kwargs): + reduced_sizes.append(tensor.numel()) + return real_reduce(tensor, *args, **kwargs) + + for step in range(2): + # Also exercise empty ranks in the existing two-rank CI launch. + if step == 1 and world > 1: + lo, hi = (0, full.numel()) if rank == 0 else (0, 0) + master = full + step * 0.25 + shard = master.flatten()[lo:hi] if lo < hi else None + outputs = [ + ( + torch.empty(hi - lo, dtype=torch.uint8, device="cuda"), + torch.empty(hi - lo, dtype=torch.uint8, device="cuda") if i else None, + ) + for i in range(len(weights)) + ] + reduced_sizes.clear() + cast = quantize_master_weights if step == 0 else cast_master_weights_to_fp8 + with patch("torch.distributed.all_reduce", side_effect=record_reduce): + cast( + weights, + [shard] * len(weights), + [lo if shard is not None else None] * len(weights), + group, + fsdp_shard_model_weights=outputs if fragments else None, + ) + expected_amax = sum(w._rowwise_scale_inv.numel() for w in weights) + if mixed: + expected_amax += weights[1]._columnwise_scale_inv.numel() + assert reduced_sizes == [expected_amax] + for idx, (weight, quantizer) in enumerate(zip(weights, quantizers)): + expected = quantizer(master.to(torch.bfloat16)) + for direction in ["rowwise", "columnwise"] if idx else ["rowwise"]: + data = getattr(weight, f"_{direction}_data") + actual = torch.zeros_like(data).flatten() + if lo < hi: + src = ( + outputs[idx][direction == "columnwise"] + if fragments + else data.flatten()[lo:hi] + ) + actual[lo:hi].copy_(src) + real_reduce(actual, op=torch.distributed.ReduceOp.MAX, group=group) + torch.testing.assert_close( + actual.view(shape), getattr(expected, f"_{direction}_data"), rtol=0, atol=0 + ) + scale = getattr(weight, f"_{direction}_scale_inv") + ref = getattr(expected, f"_{direction}_scale_inv") + rows, cols = (160, 3) if direction == "rowwise" else (5, 96) + torch.testing.assert_close(scale[:rows, :cols], ref[:rows, :cols], rtol=0, atol=0) + assert pointers[idx] == ( + weight._rowwise_data.data_ptr(), + weight._rowwise_scale_inv.data_ptr(), + ) + assert weights[0]._columnwise_data is None + assert weights[0]._columnwise_scale_inv is None + + def run_parallel_tests() -> None: """Run parallel tests""" @@ -1000,6 +1081,9 @@ def run_parallel_tests() -> None: if is_mxfp8_available(): print("starting mxfp8 empty master shard test") _test_mxfp8_empty_master_shard(dp_group) + for mixed in (False, True): + for fragments in (False, True): + _test_rowwise_master_cast(dp_group, mixed, fragments) nvfp4_available, _ = is_nvfp4_available(return_reason=True) if nvfp4_available: print("starting cast master weights to nvfp4 test") diff --git a/tests/pytorch/mxfp8/test_rowwise_master_cast.py b/tests/pytorch/mxfp8/test_rowwise_master_cast.py deleted file mode 100644 index a876a8fb236..00000000000 --- a/tests/pytorch/mxfp8/test_rowwise_master_cast.py +++ /dev/null @@ -1,132 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# See LICENSE for license information. - -"""Run with pytest on one GPU or torchrun -m pytest on multiple GPUs.""" - -import os - -import pytest -import torch -import transformer_engine.pytorch as te -import transformer_engine_torch as tex -from transformer_engine.pytorch.tensor.utils import ( - cast_master_weights_to_fp8, - quantize_master_weights, -) - -available, reason = te.is_mxfp8_available(return_reason=True) -pytestmark = pytest.mark.skipif(not available, reason=reason) - - -@pytest.fixture(scope="module") -def group(): - owned = not torch.distributed.is_initialized() - if owned: - torch.cuda.set_device(int(os.getenv("LOCAL_RANK", "0"))) - if "RANK" in os.environ: - torch.distributed.init_process_group("nccl") - else: - torch.distributed.init_process_group( - "nccl", store=torch.distributed.HashStore(), rank=0, world_size=1 - ) - yield torch.distributed.group.WORLD - if owned: - torch.distributed.destroy_process_group() - - -@pytest.mark.parametrize("mixed", [False, True]) -@pytest.mark.parametrize("fragments", [False, True]) -def test_rowwise_master_cast(group, monkeypatch, mixed, fragments): - rank = torch.distributed.get_rank(group) - world = torch.distributed.get_world_size(group) - shape = (160, 96) # Both scale padding and non-tile-aligned logical shapes. - full = torch.linspace(-4, 7, 160 * 96, device="cuda").reshape(shape) - quantizers = [te.MXFP8Quantizer(te.DType.kFloat8E4M3, rowwise=True, columnwise=False)] - if mixed: - quantizers.append(te.MXFP8Quantizer(te.DType.kFloat8E4M3, rowwise=True, columnwise=True)) - weights = [q(full.to(torch.bfloat16)) for q in quantizers] - pointers = [(w._rowwise_data.data_ptr(), w._rowwise_scale_inv.data_ptr()) for w in weights] - if world == 1: - lo, hi = 0, full.numel() - elif rank == 0: - lo, hi = 0, 17 # Split inside a 32-value block. - elif rank == 1: - lo, hi = 17, full.numel() - else: - lo, hi = 0, 0 # No master shard on tail ranks. - real_reduce = torch.distributed.all_reduce - reduced_sizes = [] - - def record_reduce(tensor, *args, **kwargs): - reduced_sizes.append(tensor.numel()) - return real_reduce(tensor, *args, **kwargs) - - monkeypatch.setattr(torch.distributed, "all_reduce", record_reduce) - for step in range(2): - master = full + step * 0.25 - shard = master.flatten()[lo:hi] if lo < hi else None - outputs = [ - ( - torch.empty(hi - lo, dtype=torch.uint8, device="cuda"), - torch.empty(hi - lo, dtype=torch.uint8, device="cuda") if i else None, - ) - for i in range(len(weights)) - ] - reduced_sizes.clear() - cast = quantize_master_weights if step == 0 else cast_master_weights_to_fp8 - cast( - weights, - [shard] * len(weights), - [lo if shard is not None else None] * len(weights), - group, - fsdp_shard_model_weights=outputs if fragments else None, - ) - expected_amax = sum(w._rowwise_scale_inv.numel() for w in weights) - if mixed: - expected_amax += weights[1]._columnwise_scale_inv.numel() - assert reduced_sizes == [expected_amax] - for idx, (weight, quantizer) in enumerate(zip(weights, quantizers)): - expected = quantizer(master.to(torch.bfloat16)) - for direction in ["rowwise", "columnwise"] if idx else ["rowwise"]: - data = getattr(weight, f"_{direction}_data") - actual = torch.zeros_like(data).flatten() - if lo < hi: - src = ( - outputs[idx][direction == "columnwise"] - if fragments - else data.flatten()[lo:hi] - ) - actual[lo:hi].copy_(src) - real_reduce(actual, op=torch.distributed.ReduceOp.MAX, group=group) - torch.testing.assert_close( - actual.view(shape), getattr(expected, f"_{direction}_data"), rtol=0, atol=0 - ) - scale = getattr(weight, f"_{direction}_scale_inv") - ref = getattr(expected, f"_{direction}_scale_inv") - rows, cols = (160, 3) if direction == "rowwise" else (5, 96) - torch.testing.assert_close(scale[:rows, :cols], ref[:rows, :cols], rtol=0, atol=0) - assert pointers[idx] == ( - weight._rowwise_data.data_ptr(), - weight._rowwise_scale_inv.data_ptr(), - ) - assert weights[0]._columnwise_data is None - assert weights[0]._columnwise_scale_inv is None - - -@pytest.mark.parametrize("offset,length", [(0, 0), (31, 3), (1, 4094)]) -def test_rowwise_partial_kernels(offset, length): - inp = torch.randn(length, device="cuda", dtype=torch.bfloat16) - row = torch.zeros((128, 4), device="cuda", dtype=inp.dtype) - ref = torch.zeros_like(row) - col = torch.zeros((4, 128), device="cuda", dtype=inp.dtype) - # Empty view may still have backing storage. Omission is shape-based. - omitted = col.flatten()[:0].view(0, 0) - tex.mxfp8_scaling_compute_partial_amax(inp, row, omitted, 64, 64, offset) - tex.mxfp8_scaling_compute_partial_amax(inp, ref, col, 64, 64, offset) - torch.testing.assert_close(row, ref, rtol=0, atol=0) - scales = torch.full((128, 4), 127, dtype=torch.uint8, device="cuda") - out = torch.empty(length, dtype=torch.uint8, device="cuda") - tex.mxfp8_scaling_partial_cast( - inp, out, out[:0], scales, scales.flatten()[:0].view(0, 0), 64, 64, offset - ) - torch.testing.assert_close(out, inp.to(torch.float8_e4m3fn).view(torch.uint8), rtol=0, atol=0) diff --git a/tests/pytorch/test_partial_cast.py b/tests/pytorch/test_partial_cast.py index bbb18503b10..cb1d5335e6b 100644 --- a/tests/pytorch/test_partial_cast.py +++ b/tests/pytorch/test_partial_cast.py @@ -128,6 +128,7 @@ def _pad(x, padding): def test_mxfp8_scaling_partial_cast(): torch.cuda.manual_seed(1234) + run_one_case(0, 64, 64, 0) run_one_case(3, 32, 64, 31) run_one_case(64 * 64 - 2, 64, 64, 1) run_one_case(16384 * 6144, 16384, 6144, 0) @@ -135,3 +136,23 @@ def test_mxfp8_scaling_partial_cast(): run_one_case(131072, 768, 256, 0) run_one_case(65536, 768, 256, 131072) run_one_case(98304, 128, 768, 0) + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.parametrize("offset,length", [(0, 0), (31, 3), (1, 4094)]) +def test_rowwise_partial_kernels(offset, length): + inp = torch.randn(length, device="cuda", dtype=torch.bfloat16) + row = torch.zeros((128, 4), device="cuda", dtype=inp.dtype) + ref = torch.zeros_like(row) + col = torch.zeros((4, 128), device="cuda", dtype=inp.dtype) + # The PyTorch extension must normalize empty views with backing storage. + omitted = col.flatten()[:0].view(0, 0) + tex.mxfp8_scaling_compute_partial_amax(inp, row, omitted, 64, 64, offset) + tex.mxfp8_scaling_compute_partial_amax(inp, ref, col, 64, 64, offset) + torch.testing.assert_close(row, ref, rtol=0, atol=0) + scales = torch.full((128, 4), 127, dtype=torch.uint8, device="cuda") + out = torch.empty(length, dtype=torch.uint8, device="cuda") + tex.mxfp8_scaling_partial_cast( + inp, out, out[:0], scales, scales.flatten()[:0].view(0, 0), 64, 64, offset + ) + torch.testing.assert_close(out, inp.to(torch.float8_e4m3fn).view(torch.uint8), rtol=0, atol=0) diff --git a/transformer_engine/common/include/transformer_engine/recipe.h b/transformer_engine/common/include/transformer_engine/recipe.h index 1ce8923a960..e7a9b0e22fa 100644 --- a/transformer_engine/common/include/transformer_engine/recipe.h +++ b/transformer_engine/common/include/transformer_engine/recipe.h @@ -290,7 +290,7 @@ void nvte_fp8_block_scaling_partial_cast(const NVTETensor inp, NVTETensor out, * \param[in] input Input tensor (continuous segment of flattened original tensor). * \param[in,out] amax_rowwise Output tensor for row-wise maximum absolute values. * \param[in,out] amax_colwise Output tensor for column-wise maximum absolute values. - * Shape [0, 0] skips column-wise computation. + * An empty NVTETensor skips column-wise computation. * \param[in] rows Number of rows in the logical tensor. * \param[in] cols Number of columns in the logical tensor. * \param[in] start_offset Starting offset in the flattened tensor. @@ -304,8 +304,8 @@ void nvte_mxfp8_scaling_compute_partial_amax(const NVTETensor input, NVTETensor * * This function casts the input tensor to MXFP8 format, producing both row-wise and * column-wise scaled outputs. input contains a continuous segment from the flattened - * original tensor. To skip column-wise computation, pass output_colwise with shape [0] - * and scale_inv_colwise with shape [0, 0]. Row-wise outputs remain required. + * original tensor. To skip column-wise computation, pass empty NVTETensors for + * output_colwise and scale_inv_colwise. Row-wise outputs remain required. * * \param[in] input Input (continuous segment of flattened original tensor). * \param[out] output_rowwise Output tensor with row-wise scaling (MXFP8 format). diff --git a/transformer_engine/common/recipe/mxfp8_scaling.cu b/transformer_engine/common/recipe/mxfp8_scaling.cu index 5d40d5475e0..5ae5bf791aa 100644 --- a/transformer_engine/common/recipe/mxfp8_scaling.cu +++ b/transformer_engine/common/recipe/mxfp8_scaling.cu @@ -23,7 +23,7 @@ constexpr int kColsPerTile = 128; // Columns each block processes constexpr int kThreadsPerBlock = 128; -template +template __global__ void __launch_bounds__(kThreadsPerBlock) mxfp8_scaling_compute_partial_amax_kernel(const IType *input, IType *amax_rowwise, IType *amax_colwise, int amax_rowwise_stride, @@ -47,7 +47,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) if (r < rows && c < cols && idx >= start_offset && idx < end_offset) { float abs_input = fabs(static_cast(input_minus_offset[idx])); row_amax = fmaxf(row_amax, abs_input); - if (amax_colwise != nullptr) col_amax = fmaxf(col_amax, abs_input); + if constexpr (kWithColumnwise) col_amax = fmaxf(col_amax, abs_input); } #pragma unroll @@ -63,7 +63,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) r++; } - if (amax_colwise != nullptr) { + if constexpr (kWithColumnwise) { amax_colwise[blockIdx.y * amax_colwise_stride + c] = static_cast(col_amax); } @@ -76,7 +76,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) amax_rowwise[r * amax_rowwise_stride + c] = static_cast(smem_amax_rowwise[r_][c_]); } -template +template __global__ void __launch_bounds__(kThreadsPerBlock) mxfp8_scaling_partial_cast_kernel(const IType *input, OType *output_rowwise, OType *output_colwise, const e8m0_t *scale_inv_rowwise, @@ -97,7 +97,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) } // Load scales_colwise - if (output_colwise != nullptr) { + if constexpr (kWithColumnwise) { int c_ = threadIdx.x; int r = blockIdx.y * kRowsPerTile / 32; int c = blockIdx.x * kColsPerTile + c_; @@ -110,8 +110,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) size_t end_offset = start_offset + len; const IType *input_minus_offset = input - start_offset; OType *output_rowwise_minus_offset = output_rowwise - start_offset; - OType *output_colwise_minus_offset = - output_colwise != nullptr ? output_colwise - start_offset : nullptr; + OType *output_colwise_minus_offset = kWithColumnwise ? output_colwise - start_offset : nullptr; int warp_idx = threadIdx.x / 32; // int lane_idx = threadIdx.x % 32; int c = blockIdx.x * kColsPerTile + threadIdx.x; @@ -125,7 +124,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) float inp = static_cast(input_minus_offset[idx]); OType out_rowwise = static_cast(inp * smem_scales_rowwise[i][warp_idx]); output_rowwise_minus_offset[idx] = out_rowwise; - if (output_colwise != nullptr) { + if constexpr (kWithColumnwise) { OType out_colwise = static_cast(inp * smem_scales_colwise[threadIdx.x]); output_colwise_minus_offset[idx] = out_colwise; } @@ -154,10 +153,7 @@ void mxfp8_scaling_compute_partial_amax(const Tensor input, Tensor amax_rowwise, NVTE_CHECK(amax_rowwise.data.shape[1] >= cols / 32, "Invalid cols"); NVTE_CHECK(amax_rowwise.dtype() == input.dtype(), "Wrong dtype of amax_rowwise"); - // Empty views can have a non-null data pointer: use shape to detect omission. - const bool with_columnwise = - !(amax_colwise.data.shape.size() == 2 && amax_colwise.data.shape[0] == 0 && - amax_colwise.data.shape[1] == 0); + const bool with_columnwise = amax_colwise.has_data(); if (with_columnwise) { NVTE_CHECK(amax_colwise.data.shape.size() == 2, "amax_colwise must be a 2D tensor"); NVTE_CHECK(amax_colwise.data.shape[0] % colwise_row_padding == 0, @@ -175,12 +171,15 @@ void mxfp8_scaling_compute_partial_amax(const Tensor input, Tensor amax_rowwise, TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( input.dtype(), IType, - mxfp8_scaling_compute_partial_amax_kernel<<>>( - reinterpret_cast(input.data.dptr), - reinterpret_cast(amax_rowwise.data.dptr), - with_columnwise ? reinterpret_cast(amax_colwise.data.dptr) : nullptr, - amax_rowwise.data.shape[1], with_columnwise ? amax_colwise.data.shape[1] : 0, rows, cols, - start_offset, input.data.shape[0]);) + TRANSFORMER_ENGINE_SWITCH_CONDITION( + with_columnwise, kWithColumnwise, + mxfp8_scaling_compute_partial_amax_kernel + <<>>( + reinterpret_cast(input.data.dptr), + reinterpret_cast(amax_rowwise.data.dptr), + with_columnwise ? reinterpret_cast(amax_colwise.data.dptr) : nullptr, + amax_rowwise.data.shape[1], with_columnwise ? amax_colwise.data.shape[1] : 0, rows, + cols, start_offset, input.data.shape[0]);)) } void mxfp8_scaling_partial_cast(const Tensor input, Tensor output_rowwise, Tensor output_colwise, @@ -199,11 +198,8 @@ void mxfp8_scaling_partial_cast(const Tensor input, Tensor output_rowwise, Tenso NVTE_CHECK(output_rowwise.dtype() == DType::kFloat8E4M3 || output_rowwise.dtype() == DType::kByte, "output_rowwise should be e4m3 or uint8"); - const bool with_columnwise = - !(scale_inv_colwise.data.shape.size() == 2 && scale_inv_colwise.data.shape[0] == 0 && - scale_inv_colwise.data.shape[1] == 0); - NVTE_CHECK(with_columnwise || - (output_colwise.data.shape.size() == 1 && output_colwise.data.shape[0] == 0), + const bool with_columnwise = scale_inv_colwise.has_data(); + NVTE_CHECK(with_columnwise || !output_colwise.has_data(), "Omitted columnwise scales require an empty columnwise output"); if (with_columnwise) { NVTE_CHECK(output_colwise.data.shape.size() == 1, "output_colwise must be a 1D tensor"); @@ -242,14 +238,19 @@ void mxfp8_scaling_partial_cast(const Tensor input, Tensor output_rowwise, Tenso TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( input.dtype(), IType, - mxfp8_scaling_partial_cast_kernel<<>>( - reinterpret_cast(input.data.dptr), - reinterpret_cast(output_rowwise.data.dptr), - with_columnwise ? reinterpret_cast(output_colwise.data.dptr) : nullptr, - reinterpret_cast(scale_inv_rowwise.data.dptr), - with_columnwise ? reinterpret_cast(scale_inv_colwise.data.dptr) : nullptr, - scale_inv_rowwise.data.shape[1], with_columnwise ? scale_inv_colwise.data.shape[1] : 0, - rows, cols, start_offset, input.data.shape[0]);) + TRANSFORMER_ENGINE_SWITCH_CONDITION( + with_columnwise, kWithColumnwise, + mxfp8_scaling_partial_cast_kernel + <<>>( + reinterpret_cast(input.data.dptr), + reinterpret_cast(output_rowwise.data.dptr), + with_columnwise ? reinterpret_cast(output_colwise.data.dptr) : nullptr, + reinterpret_cast(scale_inv_rowwise.data.dptr), + with_columnwise ? reinterpret_cast(scale_inv_colwise.data.dptr) + : nullptr, + scale_inv_rowwise.data.shape[1], + with_columnwise ? scale_inv_colwise.data.shape[1] : 0, rows, cols, start_offset, + input.data.shape[0]);)) } } // namespace mxfp8_scaling_recipe diff --git a/transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp b/transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp index d6693a485ee..1c6fb20ce43 100644 --- a/transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp @@ -57,7 +57,9 @@ void mxfp8_scaling_compute_partial_amax(const at::Tensor &input, at::Tensor amax const TensorWrapper input_cu = makeTransformerEngineTensor(input); TensorWrapper amax_rowwise_cu = makeTransformerEngineTensor(amax_rowwise); - TensorWrapper amax_colwise_cu = makeTransformerEngineTensor(amax_colwise); + // Empty PyTorch views may have a backing pointer. Normalize omission at this boundary. + TensorWrapper amax_colwise_cu = + amax_colwise.numel() == 0 ? TensorWrapper() : makeTransformerEngineTensor(amax_colwise); nvte_mxfp8_scaling_compute_partial_amax(input_cu.data(), amax_rowwise_cu.data(), amax_colwise_cu.data(), rows, cols, start_offset, @@ -76,9 +78,14 @@ void mxfp8_scaling_partial_cast(const at::Tensor &input, at::Tensor output_rowwi const TensorWrapper input_cu = makeTransformerEngineTensor(input); TensorWrapper output_rowwise_cu = makeTransformerEngineTensor(output_rowwise); - TensorWrapper output_colwise_cu = makeTransformerEngineTensor(output_colwise); + const bool with_columnwise = scale_inv_colwise.numel() != 0; + TORCH_CHECK(with_columnwise || output_colwise.numel() == 0, + "Omitted columnwise scales require an empty columnwise output"); + TensorWrapper output_colwise_cu = + with_columnwise ? makeTransformerEngineTensor(output_colwise) : TensorWrapper(); const TensorWrapper scale_inv_rowwise_cu = makeTransformerEngineTensor(scale_inv_rowwise); - const TensorWrapper scale_inv_colwise_cu = makeTransformerEngineTensor(scale_inv_colwise); + const TensorWrapper scale_inv_colwise_cu = + with_columnwise ? makeTransformerEngineTensor(scale_inv_colwise) : TensorWrapper(); nvte_mxfp8_scaling_partial_cast(input_cu.data(), output_rowwise_cu.data(), output_colwise_cu.data(), scale_inv_rowwise_cu.data(), From 034a4bb43bbc780c78cf9d2deb2b97af03a60146 Mon Sep 17 00:00:00 2001 From: Codex Date: Mon, 26 Oct 2026 14:44:44 -0700 Subject: [PATCH 5/5] Select MXFP8 partial-cast directions from output storage Use columnwise output presence to select partial-cast computation in the common API and PyTorch binding. Require scales for requested columnwise output and ignore them when that output is omitted. Document the contract and cover unused scales, empty shards, and missing required scales in regression tests. Signed-off-by: Codex --- tests/pytorch/test_partial_cast.py | 24 +++++++++++++++---- .../include/transformer_engine/recipe.h | 5 ++-- .../common/recipe/mxfp8_scaling.cu | 7 +++--- .../csrc/extensions/fp8_partial_cast.cpp | 9 +++---- 4 files changed, 31 insertions(+), 14 deletions(-) diff --git a/tests/pytorch/test_partial_cast.py b/tests/pytorch/test_partial_cast.py index cb1d5335e6b..b7f683b222c 100644 --- a/tests/pytorch/test_partial_cast.py +++ b/tests/pytorch/test_partial_cast.py @@ -140,7 +140,8 @@ def test_mxfp8_scaling_partial_cast(): @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) @pytest.mark.parametrize("offset,length", [(0, 0), (31, 3), (1, 4094)]) -def test_rowwise_partial_kernels(offset, length): +@pytest.mark.parametrize("with_columnwise_scales", [False, True]) +def test_rowwise_partial_kernels(offset, length, with_columnwise_scales): inp = torch.randn(length, device="cuda", dtype=torch.bfloat16) row = torch.zeros((128, 4), device="cuda", dtype=inp.dtype) ref = torch.zeros_like(row) @@ -151,8 +152,23 @@ def test_rowwise_partial_kernels(offset, length): tex.mxfp8_scaling_compute_partial_amax(inp, ref, col, 64, 64, offset) torch.testing.assert_close(row, ref, rtol=0, atol=0) scales = torch.full((128, 4), 127, dtype=torch.uint8, device="cuda") - out = torch.empty(length, dtype=torch.uint8, device="cuda") - tex.mxfp8_scaling_partial_cast( - inp, out, out[:0], scales, scales.flatten()[:0].view(0, 0), 64, 64, offset + columnwise_scales = ( + torch.full((4, 128), 126, dtype=torch.uint8, device="cuda") + if with_columnwise_scales + else scales.flatten()[:0].view(0, 0) ) + out = torch.empty(length, dtype=torch.uint8, device="cuda") + tex.mxfp8_scaling_partial_cast(inp, out, out[:0], scales, columnwise_scales, 64, 64, offset) torch.testing.assert_close(out, inp.to(torch.float8_e4m3fn).view(torch.uint8), rtol=0, atol=0) + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +def test_mxfp8_partial_cast_requires_columnwise_scales(): + inp = torch.randn(3, device="cuda", dtype=torch.bfloat16) + row = torch.empty(3, device="cuda", dtype=torch.uint8) + col = torch.empty_like(row) + scales = torch.full((128, 4), 127, device="cuda", dtype=torch.uint8) + with pytest.raises(RuntimeError, match="Columnwise output requires columnwise scales"): + tex.mxfp8_scaling_partial_cast( + inp, row, col, scales, scales.flatten()[:0].view(0, 0), 64, 64, 31 + ) diff --git a/transformer_engine/common/include/transformer_engine/recipe.h b/transformer_engine/common/include/transformer_engine/recipe.h index dffd93677c3..98c5eed0985 100644 --- a/transformer_engine/common/include/transformer_engine/recipe.h +++ b/transformer_engine/common/include/transformer_engine/recipe.h @@ -304,8 +304,9 @@ void nvte_mxfp8_scaling_compute_partial_amax(const NVTETensor input, NVTETensor * * This function casts the input tensor to MXFP8 format, producing both row-wise and * column-wise scaled outputs. input contains a continuous segment from the flattened - * original tensor. To skip column-wise computation, pass empty NVTETensors for - * output_colwise and scale_inv_colwise. Row-wise outputs remain required. + * original tensor. To skip column-wise computation, pass an empty NVTETensor for + * output_colwise; scale_inv_colwise is ignored in this case. Column-wise output + * requires column-wise scales. Row-wise outputs remain required. * * \param[in] input Input (continuous segment of flattened original tensor). * \param[out] output_rowwise Output tensor with row-wise scaling (MXFP8 format). diff --git a/transformer_engine/common/recipe/mxfp8_scaling.cu b/transformer_engine/common/recipe/mxfp8_scaling.cu index 5ae5bf791aa..b25e3d4b00b 100644 --- a/transformer_engine/common/recipe/mxfp8_scaling.cu +++ b/transformer_engine/common/recipe/mxfp8_scaling.cu @@ -198,15 +198,14 @@ void mxfp8_scaling_partial_cast(const Tensor input, Tensor output_rowwise, Tenso NVTE_CHECK(output_rowwise.dtype() == DType::kFloat8E4M3 || output_rowwise.dtype() == DType::kByte, "output_rowwise should be e4m3 or uint8"); - const bool with_columnwise = scale_inv_colwise.has_data(); - NVTE_CHECK(with_columnwise || !output_colwise.has_data(), - "Omitted columnwise scales require an empty columnwise output"); + const bool with_columnwise = output_colwise.has_data(); if (with_columnwise) { + NVTE_CHECK(scale_inv_colwise.has_data(), "Columnwise output requires columnwise scales"); NVTE_CHECK(output_colwise.data.shape.size() == 1, "output_colwise must be a 1D tensor"); NVTE_CHECK(output_colwise.data.shape[0] == input.data.shape[0], "Size of input and output_colwise mismatch"); NVTE_CHECK(input.data.shape[0] == 0 || output_colwise.data.dptr != nullptr, - "Columnwise scales require columnwise output"); + "Columnwise output must have a data pointer"); NVTE_CHECK( output_colwise.dtype() == DType::kFloat8E4M3 || output_colwise.dtype() == DType::kByte, "output_colwise should be e4m3 or uint8"); diff --git a/transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp b/transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp index 1c6fb20ce43..dc02e67e5fb 100644 --- a/transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp @@ -74,13 +74,14 @@ void mxfp8_scaling_partial_cast(const at::Tensor &input, at::Tensor output_rowwi TORCH_CHECK(output_rowwise.is_contiguous(), "output_rowwise must be contiguous"); TORCH_CHECK(output_colwise.is_contiguous(), "output_colwise must be contiguous"); TORCH_CHECK(scale_inv_rowwise.is_contiguous(), "scale_inv_rowwise must be contiguous"); - TORCH_CHECK(scale_inv_colwise.is_contiguous(), "scale_inv_colwise must be contiguous"); const TensorWrapper input_cu = makeTransformerEngineTensor(input); TensorWrapper output_rowwise_cu = makeTransformerEngineTensor(output_rowwise); - const bool with_columnwise = scale_inv_colwise.numel() != 0; - TORCH_CHECK(with_columnwise || output_colwise.numel() == 0, - "Omitted columnwise scales require an empty columnwise output"); + const bool with_columnwise = output_colwise.numel() != 0; + if (with_columnwise) { + TORCH_CHECK(scale_inv_colwise.numel() != 0, "Columnwise output requires columnwise scales"); + TORCH_CHECK(scale_inv_colwise.is_contiguous(), "scale_inv_colwise must be contiguous"); + } TensorWrapper output_colwise_cu = with_columnwise ? makeTransformerEngineTensor(output_colwise) : TensorWrapper(); const TensorWrapper scale_inv_rowwise_cu = makeTransformerEngineTensor(scale_inv_rowwise);