[Common/PyTorch] Enable FP8BS for Grouped GEMM on Blackwell via MXFP8 Broadcast

tests/cpp/operator/CMakeLists.txt

@@ -42,6 +42,7 @@ add_executable(test_operator
                test_causal_softmax.cu
                test_swizzle.cu
                test_multi_swizzle.cu
+               test_swizzle_block_scaling.cu
                test_swap_first_dims.cu
                test_grouped_gemm.cu
                ../test_common.cu)

tests/cpp/operator/test_cast_float8blockwise_grouped.cu

@@ -74,10 +74,13 @@ template <typename InputType, typename OutputType>
 void perform_test(ShapeRep shape_rep, BlockDim block_dim, ScalingDir dir,
                   const std::vector<size_t>& first_dims_h, size_t K,
                   bool force_pow_2_scales, float epsilon) {
-  if (getDeviceComputeCapability() < hopperComputeCapability ||
-      getDeviceComputeCapability() >= blackwellComputeCapability) {
+  if (getDeviceComputeCapability() < hopperComputeCapability) {
     GTEST_SKIP();
   }
+  if (!force_pow_2_scales && getDeviceComputeCapability() >= blackwellComputeCapability) {
+    GTEST_SKIP() << "Non-power-of-2 FP8 block scales are rejected on Blackwell and newer "
+                    "(covered by GroupedFP8BlockwiseBlackwell.RejectsNonPow2Scales).";
+  }
 
   DType itype = TypeInfo<InputType>::dtype;
   DType otype = TypeInfo<OutputType>::dtype;
@@ -418,4 +421,48 @@ std::string make_name(const ::testing::TestParamInfo<TestConfig>& info) {
 INSTANTIATE_TEST_SUITE_P(GroupedFP8Blockwise, GroupedFP8BlockwiseTestSuite,
                          ::testing::ValuesIn(make_configs()), make_name);
 
+// On Blackwell and newer, FP8 block scaling is emulated with MXFP8 GEMM, whose E8M0 scales can
+// only represent powers of 2, so grouped quantize must reject non-power-of-2 scales.
+TEST(GroupedFP8BlockwiseBlackwell, RejectsNonPow2Scales) {
+  if (getDeviceComputeCapability() < blackwellComputeCapability) {
+    GTEST_SKIP();
+  }
+  constexpr size_t num_tensors = 2, M = 128, K = 128;
+  const size_t numel = num_tensors * M * K;
+  bf16* input_d = nullptr;
+  fp8e4m3* output_d = nullptr;
+  float* scale_inv_d = nullptr;
+  const size_t scale_floats = num_tensors * per_expert_scale_floats(BlockDim::ONE_D, false, M, K);
+  cudaMalloc(&input_d, numel * sizeof(bf16));
+  cudaMemset(input_d, 0, numel * sizeof(bf16));
+  cudaMalloc(&output_d, numel);
+  cudaMalloc(&scale_inv_d, scale_floats * sizeof(float));
+
+  std::vector<size_t> logical_shape_vec = {num_tensors * M, K};
+  const NVTEShape logical_shape =
+      nvte_make_shape(logical_shape_vec.data(), logical_shape_vec.size());
+  NVTEGroupedTensor in_gt =
+      nvte_create_grouped_tensor(NVTE_DELAYED_TENSOR_SCALING, num_tensors, logical_shape);
+  NVTEGroupedTensor out_gt =
+      nvte_create_grouped_tensor(NVTE_BLOCK_SCALING_1D, num_tensors, logical_shape);
+  NVTEBasicTensor in_data = {input_d, kNVTEBFloat16, logical_shape};
+  nvte_set_grouped_tensor_param(in_gt, kNVTEGroupedRowwiseData, &in_data, sizeof(in_data));
+  NVTEBasicTensor out_data = {output_d, kNVTEFloat8E4M3, logical_shape};
+  nvte_set_grouped_tensor_param(out_gt, kNVTEGroupedRowwiseData, &out_data, sizeof(out_data));
+  std::vector<size_t> scale_shape_vec = {scale_floats};
+  NVTEBasicTensor scale_bt = {scale_inv_d, kNVTEFloat32,
+                              nvte_make_shape(scale_shape_vec.data(), scale_shape_vec.size())};
+  nvte_set_grouped_tensor_param(out_gt, kNVTEGroupedRowwiseScaleInv, &scale_bt, sizeof(scale_bt));
+
+  QuantizationConfigWrapper quant_config;
+  quant_config.set_force_pow_2_scales(false);
+  EXPECT_THROW(nvte_group_quantize(in_gt, out_gt, quant_config, 0), std::runtime_error);
+
+  nvte_destroy_grouped_tensor(in_gt);
+  nvte_destroy_grouped_tensor(out_gt);
+  cudaFree(input_d);
+  cudaFree(output_d);
+  cudaFree(scale_inv_d);
+}
+
 }  // namespace

tests/cpp/operator/test_dequantize_float8blockwise_grouped.cu

@@ -78,9 +78,8 @@ inline size_t per_expert_scale_floats(BlockDim block_dim, bool columnwise, size_
 template <typename InputType, typename OutputType>
 void performTest(ShapeRep shape_rep, BlockDim block_dim, bool rowwise,
                  const std::vector<size_t>& first_dims_h, size_t K) {
-  // FP8 block-scaling grouped kernels are Hopper-only (SM90-SM99).
-  if (getDeviceComputeCapability() < hopperComputeCapability ||
-      getDeviceComputeCapability() >= blackwellComputeCapability) {
+  // Grouped FP8 block-scaling kernels need TMA (SM90+).
+  if (getDeviceComputeCapability() < hopperComputeCapability) {
     GTEST_SKIP();
   }
 

tests/cpp/operator/test_grouped_gemm.cu

@@ -269,7 +269,9 @@ inline std::string grouped_gemm_skip_reason(const TestParams& params) {
       }
     }
     if (is_blackwell_plus && fp8_block) {
-      return "FP8 block scaling grouped GEMM is only supported on Hopper (SM90), " + cc_suffix;
+      return "FP8 block scaling grouped GEMM is native only on Hopper (SM90); Blackwell+ emulates "
+             "it with MXFP8 in the framework binding, " +
+             cc_suffix;
     }
     if (params.recipe == InputRecipe::kNVFP4 && params.output_dtype == DType::kFloat16) {
       return "NVFP4 grouped GEMM does not support FP16 output.";