Add better assertions for MXFP8 group gemm (#5203)

Summary:
Pull Request resolved: #5203

X-link: https://github.com/facebookresearch/FBGEMM/pull/2200

Currently only supporst BF16 out, but would silently fail if someone tries FP32 out. Add explicit assertion for this. And some small code refactor.

Reviewed By: q10, summerdengfb, jianyuh

Differential Revision: D88747914

fbshipit-source-id: da8337816d7792c3abb64ce67e768bd8ef539873

Author: cthi

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/mx8mx8bf16_grouped.cu

@@ -165,8 +165,9 @@ at::Tensor mx8mx8bf16_grouped_mm(
     TORCH_CHECK(false, "Invalid input shapes. Must be one of 2D-2D, 2D-3D.");
   }
 
+  TORCH_CHECK(output_actual.dtype() == at::kBFloat16, "output must be bf16.");
   // Early exit for empty inputs.
-  if (M == 0) {
+  if (output_actual.numel() == 0) {
     return output_actual;
   }
 

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/mx8mx8bf16_grouped/mx8mx8bf16_grouped_common.cuh

@@ -65,16 +65,11 @@ at::Tensor mx8mx8bf16_grouped_impl(
     at::Tensor offsets) {
   c10::cuda::CUDAGuard deviceGuard(XQ.device());
   // The number of groups the kernel uses may vary.
-  int kernel_groups = G;
+  const int kernel_groups = G;
 
   at::TensorOptions options;
   options = XQ.options();
 
-  // Return early if there are no elements in the output.
-  if (output.numel() == 0) {
-    return output;
-  }
-
   // WQ is shape (K,N) or (E,K,N) in column major layout, to align with
   // torch._scaled_grouped_mm. We transpose here to match cutlass kernel
   // requirements.