[PyTorch] Enable NVFP4 row-scaled (per-token) backward for GroupedLinear

tests/pytorch/test_grouped_linear.py

@@ -97,6 +97,23 @@ def nvfp4_row_scaled():
     return nvfp4_recipe
 
 
+def nvfp4_row_scaled_quantized_backward():
+    # Same row-scaled activation recipe as nvfp4_row_scaled(), but with
+    # backward_override=None so the backward runs in NVFP4 instead of falling back
+    # to high precision.
+    nvfp4_recipe = recipe.NVFP4BlockScaling(
+        disable_rht=True,
+        disable_stochastic_rounding=True,
+        disable_2d_quantization=True,
+        row_scaled_activation=True,
+        backward_override=None,
+    )
+    nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams()
+    nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams()
+    nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams()
+    return nvfp4_recipe
+
+
 def nvfp4_4over6():
     nvfp4_recipe = recipe.NVFP4BlockScaling(
         disable_rht=True,
@@ -377,6 +394,84 @@ def test_grouped_linear_accuracy(
             torch.testing.assert_close(o, o_ref, rtol=0, atol=0)
 
 
+@pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4)
+@pytest.mark.parametrize("dtype", [torch.bfloat16], ids=str)
+@pytest.mark.parametrize("num_gemms", [1, 3])
+@pytest.mark.parametrize("bs", [2])
+@pytest.mark.parametrize("bias", all_boolean)
+def test_grouped_linear_row_scaled_quantized_backward(dtype, num_gemms, bs, bias, model="126m"):
+    """Row-scaled NVFP4 GroupedLinear with quantized (non-fallback) NVFP4 backward.
+
+    With ``backward_override=None`` the wgrad is computed in NVFP4: the row-scaled
+    activation becomes operand A of the ``NT`` grouped GEMM, which this PR routes
+    through the per-expert dense ``general_gemm`` loop. GroupedLinear must then
+    match a stack of independent dense ``Linear`` layers bit-for-bit, since both
+    execute the exact same per-expert quantize + GEMM kernels.
+    """
+    recipe_row_scaled = nvfp4_row_scaled_quantized_backward()
+    config = model_configs[model]
+    if dtype not in get_nvfp4_inp_supported_dtypes(recipe_row_scaled, dtype):
+        pytest.skip(f"Input dtype {dtype} not supported for row-scaled NVFP4.")
+
+    grouped_linear = (
+        GroupedLinear(
+            num_gemms,
+            config.hidden_size,
+            4 * config.hidden_size,
+            bias=bias,
+            params_dtype=dtype,
+            device="cuda",
+        )
+        .cuda()
+        .eval()
+    )
+    sequential_linear = torch.nn.ModuleList(
+        [
+            Linear(
+                config.hidden_size,
+                4 * config.hidden_size,
+                bias=bias,
+                params_dtype=dtype,
+                device="cuda",
+            ).eval()
+            for _ in range(num_gemms)
+        ]
+    )
+
+    # Share weights/biases so the two paths are numerically comparable.
+    with torch.no_grad():
+        for i in range(num_gemms):
+            sequential_linear[i].weight = Parameter(getattr(grouped_linear, f"weight{i}").clone())
+            if bias:
+                sequential_linear[i].bias = Parameter(getattr(grouped_linear, f"bias{i}").clone())
+
+    outputs_ref = _test_grouped_linear_accuracy(
+        sequential_linear,
+        num_gemms,
+        bs,
+        dtype,
+        config,
+        recipe_row_scaled,
+        fp8=True,
+        fuse_wgrad_accumulation=False,
+    )
+    outputs = _test_grouped_linear_accuracy(
+        grouped_linear,
+        num_gemms,
+        bs,
+        dtype,
+        config,
+        recipe_row_scaled,
+        fp8=True,
+        fuse_wgrad_accumulation=False,
+    )
+
+    # GroupedLinear is a per-expert loop over the same dense kernels, so the
+    # forward output, dgrad, and (row-scaled) wgrad must match bit-for-bit.
+    for o, o_ref in zip(outputs, outputs_ref):
+        torch.testing.assert_close(o, o_ref, rtol=0, atol=0)
+
+
 @pytest.mark.skipif(
     torch.cuda.get_device_capability() != (9, 0),
     reason="Only enable CUTLASS grouped gemm on Hopper",

transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh

@@ -99,18 +99,13 @@ __launch_bounds__(BLOCK_SIZE)
 #endif
 }
 
+// This is a general amax reduction, so it is not restricted to SM 10.0+.
 template <typename IType, int BLOCK_SIZE>
-__global__ void
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
-__launch_bounds__(BLOCK_SIZE)
-#endif
+__global__ void __launch_bounds__(BLOCK_SIZE)
     compute_columnwise_amax_kernel(const int num_rows, const int num_cols,
                                    const IType *__restrict__ input,
                                    float *__restrict__ output_columnwise_amax,
                                    const float *__restrict__ noop) {
-#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ < 1000)
-  NVTE_DEVICE_ERROR("SM 10.0+ is required.");
-#else
   if (noop != nullptr && noop[0] == 1.0f) {
     return;
   }
@@ -128,7 +123,6 @@ __launch_bounds__(BLOCK_SIZE)
   if (threadIdx.x == 0) {
     output_columnwise_amax[col_idx] = col_amax;
   }
-#endif
 }
 
 template <typename IType>

transformer_engine/pytorch/cpp_extensions/gemm.py

@@ -339,9 +339,12 @@ def general_grouped_gemm(
     else:
         bias_dtype = TE_DType[torch.bfloat16]
 
-    if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in A):
-        raise NotImplementedError("Row-scaled NVFP4 grouped GEMM does not support row-scaled A.")
-    if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in B):
+    # TODO: Row-scaled NVFP4 grouped GEMM is handled by looping over the per-expert
+    # dense general_gemm, fprop and wgrad without a dedicated grouped kernel.
+    # The per-expert loop is not CUDA-graph-safe and will be fixed in a future work
+    if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in A) or any(
+        _is_nvfp4_row_scaled_tensor(tensor) for tensor in B
+    ):
         assert D_dtype is None, "Row-scaled NVFP4 grouped GEMM currently does not support D_dtype."
         if single_output:
             assert (

transformer_engine/pytorch/csrc/extensions/cast.cpp

@@ -1001,8 +1001,9 @@ std::tuple<std::vector<py::object>, std::vector<TensorWrapper>, bool> bulk_alloc
   const auto columnwise_usage = quantizer_cpp_list[0]->columnwise_usage;
   if (row_scaled_nvfp4) {
     NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 bulk allocation requires rowwise usage.");
-    NVTE_CHECK(!columnwise_usage,
-               "Row-scaled NVFP4 bulk allocation does not support columnwise usage.");
+    // Columnwise (transpose) output is supported for row-scaled NVFP4: the
+    // columnwise amax is sized per-column below so the per-expert dense cast can
+    // emit the row-scaled transpose consumed by the wgrad GEMM.
   }
   const auto scaling_mode = quantizer_cpp_list[0]->get_scaling_mode();
   const auto fp4_dtype = quantizer_cpp_list[0]->dtype;
@@ -1146,7 +1147,10 @@ std::tuple<std::vector<py::object>, std::vector<TensorWrapper>, bool> bulk_alloc
     dtypes.insert(dtypes.end(), num_tensors, torch::kUInt8);
     alignments.insert(alignments.end(), num_tensors, 16);
     for (size_t i = 0; i < num_tensors; ++i) {
-      shapes.emplace_back(amax_shape(columnwise_data_shapes[i]));
+      // columnwise_data_shapes[i] is the transposed shape, so its leading dim is
+      // the original last dim (number of columns). For row-scaled NVFP4 this
+      // yields a per-column amax vector; otherwise it stays a scalar {1}.
+      shapes.emplace_back(amax_shape(columnwise_data_shapes[i], row_scaled_nvfp4));
     }
     dtypes.insert(dtypes.end(), num_tensors, torch::kFloat32);
     alignments.insert(alignments.end(), num_tensors, 16);
@@ -1206,7 +1210,7 @@ std::tuple<std::vector<py::object>, std::vector<TensorWrapper>, bool> bulk_alloc
       }
       if (columnwise_usage) {
         tensor_wrapper.set_columnwise_amax(amax_columnwise_list[i].data_ptr(), DType::kFloat32,
-                                           std::vector<size_t>{1});
+                                           getTensorShape(amax_columnwise_list[i]));
       }
 
       tensor_cpp_list.emplace_back(std::move(tensor_wrapper));