Support row-only MXFP8 distributed master-weight casts

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")

tests/pytorch/test_partial_cast.py

@@ -128,10 +128,47 @@ 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)
     run_one_case(32768, 256, 128, 0)
     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)])
+@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)
+    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")
+    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
+        )

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.
+ *                                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.
@@ -303,7 +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.
+ *  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).

transformer_engine/common/recipe/mxfp8_scaling.cu

@@ -23,7 +23,7 @@ constexpr int kColsPerTile = 128;  // Columns each block processes
 
 constexpr int kThreadsPerBlock = 128;
 
-template <typename IType>
+template <typename IType, bool kWithColumnwise>
 __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<float>(input_minus_offset[idx]));
       row_amax = fmaxf(row_amax, abs_input);
-      col_amax = fmaxf(col_amax, abs_input);
+      if constexpr (kWithColumnwise) 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<IType>(col_amax);
+  if constexpr (kWithColumnwise) {
+    amax_colwise[blockIdx.y * amax_colwise_stride + c] = static_cast<IType>(col_amax);
+  }
 
   __syncthreads();
 
@@ -74,7 +76,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock)
   amax_rowwise[r * amax_rowwise_stride + c] = static_cast<IType>(smem_amax_rowwise[r_][c_]);
 }
 
-template <typename IType, typename OType>
+template <typename IType, typename OType, bool kWithColumnwise>
 __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,
@@ -95,7 +97,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock)
   }
 
   // Load scales_colwise
-  {
+  if constexpr (kWithColumnwise) {
     int c_ = threadIdx.x;
     int r = blockIdx.y * kRowsPerTile / 32;
     int c = blockIdx.x * kColsPerTile + c_;
@@ -108,7 +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 - start_offset;
+  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;
@@ -121,9 +123,11 @@ __global__ void __launch_bounds__(kThreadsPerBlock)
     if (r < rows && c < cols && idx >= start_offset && idx < end_offset) {
       float inp = static_cast<float>(input_minus_offset[idx]);
       OType out_rowwise = static_cast<OType>(inp * smem_scales_rowwise[i][warp_idx]);
-      OType out_colwise = static_cast<OType>(inp * smem_scales_colwise[threadIdx.x]);
       output_rowwise_minus_offset[idx] = out_rowwise;
-      output_colwise_minus_offset[idx] = out_colwise;
+      if constexpr (kWithColumnwise) {
+        OType out_colwise = static_cast<OType>(inp * smem_scales_colwise[threadIdx.x]);
+        output_colwise_minus_offset[idx] = out_colwise;
+      }
     }
 
     r++;
@@ -149,26 +153,33 @@ 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");
+  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,
+               "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;
   dim3 grid(blocks_x, blocks_y);
 
   TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(
       input.dtype(), IType,
-      mxfp8_scaling_compute_partial_amax_kernel<IType><<<grid, kColsPerTile, 0, stream>>>(
-          reinterpret_cast<const IType *>(input.data.dptr),
-          reinterpret_cast<IType *>(amax_rowwise.data.dptr),
-          reinterpret_cast<IType *>(amax_colwise.data.dptr), amax_rowwise.data.shape[1],
-          amax_colwise.data.shape[1], rows, cols, start_offset, input.data.shape[0]);)
+      TRANSFORMER_ENGINE_SWITCH_CONDITION(
+          with_columnwise, kWithColumnwise,
+          mxfp8_scaling_compute_partial_amax_kernel<IType, kWithColumnwise>
+          <<<grid, kColsPerTile, 0, stream>>>(
+              reinterpret_cast<const IType *>(input.data.dptr),
+              reinterpret_cast<IType *>(amax_rowwise.data.dptr),
+              with_columnwise ? reinterpret_cast<IType *>(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 +193,23 @@ 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 = 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 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");
+  }
 
   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,29 +220,36 @@ 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;
   dim3 grid(blocks_x, blocks_y);
 
   TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(
       input.dtype(), IType,
-      mxfp8_scaling_partial_cast_kernel<IType, fp8e4m3><<<grid, kColsPerTile, 0, stream>>>(
-          reinterpret_cast<const IType *>(input.data.dptr),
-          reinterpret_cast<fp8e4m3 *>(output_rowwise.data.dptr),
-          reinterpret_cast<fp8e4m3 *>(output_colwise.data.dptr),
-          reinterpret_cast<const e8m0_t *>(scale_inv_rowwise.data.dptr),
-          reinterpret_cast<const e8m0_t *>(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]);)
+      TRANSFORMER_ENGINE_SWITCH_CONDITION(
+          with_columnwise, kWithColumnwise,
+          mxfp8_scaling_partial_cast_kernel<IType, fp8e4m3, kWithColumnwise>
+          <<<grid, kColsPerTile, 0, stream>>>(
+              reinterpret_cast<const IType *>(input.data.dptr),
+              reinterpret_cast<fp8e4m3 *>(output_rowwise.data.dptr),
+              with_columnwise ? reinterpret_cast<fp8e4m3 *>(output_colwise.data.dptr) : nullptr,
+              reinterpret_cast<const e8m0_t *>(scale_inv_rowwise.data.dptr),
+              with_columnwise ? reinterpret_cast<const e8m0_t *>(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