FP4 grouped refactor

Summary:
X-link: facebookresearch/FBGEMM#1957

Split some clean-up/refactors from the core FP4 Torch API support to make the next diff more focused.

- Removed `zero_start_index_M` as it's unused
- Removed passing `G` into the kernel directly as it can be inferred
- Rename `ElementComputeEpilogue` -> `ElementScale`
- Add `namespace fbgemm_gpu` in `f4f4bf16_grouped_common.cuh`
- Removed `num_x_scale_per_group` and `num_w_scale_per_group` as they are both unused
- Removed un-neccesary cutlass headers in `f4f4bf16_grouped.cu`

Differential Revision: D83166227

Author: cthi

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

@@ -8,16 +8,6 @@
 
 #include <ATen/ATen.h>
 #include <ATen/cuda/CUDAContext.h>
-#include <cutlass/util/device_memory.h>
-#include <cutlass/util/packed_stride.hpp>
-
-// clang-format off
-// The fixed ordering of the headers is required for CUTLASS 3.2+
-#include <cute/tensor.hpp>
-#include <cutlass/gemm/collective/collective_builder.hpp>     // @manual
-#include <cutlass/gemm/device/gemm_universal_adapter.h>       // @manual
-#include <cutlass/epilogue/collective/collective_builder.hpp> // @manual
-// clang-format on
 
 #if defined(CUDA_VERSION) && (CUDA_VERSION >= 12080)
 #include "f4f4bf16_grouped/f4f4bf16_grouped_manifest.cuh"
@@ -160,14 +150,10 @@ at::Tensor dispatch_fp4_grouped_kernel(
     at::Tensor x_scale,
     at::Tensor w_scale,
     at::Tensor output,
-    std::optional<at::Tensor> zero_start_index_M = std::nullopt,
     std::optional<at::Tensor> M_sizes = std::nullopt,
     std::optional<at::Tensor> global_scale = std::nullopt,
     std::optional<at::Tensor> starting_row_after_padding = std::nullopt,
     bool use_mx = true) {
-  TORCH_CHECK(
-      zero_start_index_M.has_value() != M_sizes.has_value(),
-      "One of zero_start_index_M or M_sizes must be provided.");
   TORCH_CHECK(M_sizes.has_value(), "M_sizes is assumed to be provided.");
   TORCH_CHECK(
       starting_row_after_padding.has_value(),
@@ -187,8 +173,6 @@ at::Tensor dispatch_fp4_grouped_kernel(
       x_scale,
       w_scale,
       output,
-      G,
-      zero_start_index_M,
       M_sizes,
       global_scale,
       starting_row_after_padding);
@@ -228,7 +212,6 @@ at::Tensor f4f4bf16_grouped_stacked(
       x_scale,
       w_scale,
       Y,
-      std::nullopt,
       M_sizes,
       global_scale,
       starting_row_after_padding,

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f4f4bf16_grouped/f4f4bf16_grouped_128_128_256_1_1_1_f.cu

@@ -18,8 +18,6 @@ at::Tensor f4f4bf16_grouped_128_128_256_1_1_1_f(
     at::Tensor x_scale,
     at::Tensor w_scale,
     at::Tensor output,
-    int64_t G,
-    std::optional<at::Tensor> zero_start_index_M,
     std::optional<at::Tensor> M_sizes,
     std::optional<at::Tensor> global_scale,
     std::optional<at::Tensor> starting_row_after_padding) {
@@ -36,8 +34,6 @@ at::Tensor f4f4bf16_grouped_128_128_256_1_1_1_f(
       x_scale,
       w_scale,
       output,
-      G,
-      zero_start_index_M,
       M_sizes,
       global_scale,
       starting_row_after_padding);

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f4f4bf16_grouped/f4f4bf16_grouped_128_128_256_1_1_1_t.cu

@@ -18,8 +18,6 @@ at::Tensor f4f4bf16_grouped_128_128_256_1_1_1_t(
     at::Tensor x_scale,
     at::Tensor w_scale,
     at::Tensor output,
-    int64_t G,
-    std::optional<at::Tensor> zero_start_index_M,
     std::optional<at::Tensor> M_sizes,
     std::optional<at::Tensor> global_scale,
     std::optional<at::Tensor> starting_row_after_padding) {
@@ -36,8 +34,6 @@ at::Tensor f4f4bf16_grouped_128_128_256_1_1_1_t(
       x_scale,
       w_scale,
       output,
-      G,
-      zero_start_index_M,
       M_sizes,
       global_scale,
       starting_row_after_padding);

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f4f4bf16_grouped/f4f4bf16_grouped_128_64_256_1_1_1_f.cu

@@ -18,8 +18,6 @@ at::Tensor f4f4bf16_grouped_128_64_256_1_1_1_f(
     at::Tensor x_scale,
     at::Tensor w_scale,
     at::Tensor output,
-    int64_t G,
-    std::optional<at::Tensor> zero_start_index_M,
     std::optional<at::Tensor> M_sizes,
     std::optional<at::Tensor> global_scale,
     std::optional<at::Tensor> starting_row_after_padding) {
@@ -36,8 +34,6 @@ at::Tensor f4f4bf16_grouped_128_64_256_1_1_1_f(
       x_scale,
       w_scale,
       output,
-      G,
-      zero_start_index_M,
       M_sizes,
       global_scale,
       starting_row_after_padding);

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f4f4bf16_grouped/f4f4bf16_grouped_128_64_256_1_1_1_t.cu

@@ -18,8 +18,6 @@ at::Tensor f4f4bf16_grouped_128_64_256_1_1_1_t(
     at::Tensor x_scale,
     at::Tensor w_scale,
     at::Tensor output,
-    int64_t G,
-    std::optional<at::Tensor> zero_start_index_M,
     std::optional<at::Tensor> M_sizes,
     std::optional<at::Tensor> global_scale,
     std::optional<at::Tensor> starting_row_after_padding) {
@@ -36,8 +34,6 @@ at::Tensor f4f4bf16_grouped_128_64_256_1_1_1_t(
       x_scale,
       w_scale,
       output,
-      G,
-      zero_start_index_M,
       M_sizes,
       global_scale,
       starting_row_after_padding);

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f4f4bf16_grouped/f4f4bf16_grouped_256_128_256_2_1_1_f.cu

@@ -18,8 +18,6 @@ at::Tensor f4f4bf16_grouped_256_128_256_2_1_1_f(
     at::Tensor x_scale,
     at::Tensor w_scale,
     at::Tensor output,
-    int64_t G,
-    std::optional<at::Tensor> zero_start_index_M,
     std::optional<at::Tensor> M_sizes,
     std::optional<at::Tensor> global_scale,
     std::optional<at::Tensor> starting_row_after_padding) {
@@ -36,8 +34,6 @@ at::Tensor f4f4bf16_grouped_256_128_256_2_1_1_f(
       x_scale,
       w_scale,
       output,
-      G,
-      zero_start_index_M,
       M_sizes,
       global_scale,
       starting_row_after_padding);

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f4f4bf16_grouped/f4f4bf16_grouped_256_128_256_2_1_1_t.cu

@@ -18,8 +18,6 @@ at::Tensor f4f4bf16_grouped_256_128_256_2_1_1_t(
     at::Tensor x_scale,
     at::Tensor w_scale,
     at::Tensor output,
-    int64_t G,
-    std::optional<at::Tensor> zero_start_index_M,
     std::optional<at::Tensor> M_sizes,
     std::optional<at::Tensor> global_scale,
     std::optional<at::Tensor> starting_row_after_padding) {
@@ -36,8 +34,6 @@ at::Tensor f4f4bf16_grouped_256_128_256_2_1_1_t(
       x_scale,
       w_scale,
       output,
-      G,
-      zero_start_index_M,
       M_sizes,
       global_scale,
       starting_row_after_padding);

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f4f4bf16_grouped/f4f4bf16_grouped_256_256_256_2_1_1_f.cu

@@ -18,8 +18,6 @@ at::Tensor f4f4bf16_grouped_256_256_256_2_1_1_f(
     at::Tensor x_scale,
     at::Tensor w_scale,
     at::Tensor output,
-    int64_t G,
-    std::optional<at::Tensor> zero_start_index_M,
     std::optional<at::Tensor> M_sizes,
     std::optional<at::Tensor> global_scale,
     std::optional<at::Tensor> starting_row_after_padding) {
@@ -36,8 +34,6 @@ at::Tensor f4f4bf16_grouped_256_256_256_2_1_1_f(
       x_scale,
       w_scale,
       output,
-      G,
-      zero_start_index_M,
       M_sizes,
       global_scale,
       starting_row_after_padding);

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f4f4bf16_grouped/f4f4bf16_grouped_256_256_256_2_1_1_t.cu

@@ -18,8 +18,6 @@ at::Tensor f4f4bf16_grouped_256_256_256_2_1_1_t(
     at::Tensor x_scale,
     at::Tensor w_scale,
     at::Tensor output,
-    int64_t G,
-    std::optional<at::Tensor> zero_start_index_M,
     std::optional<at::Tensor> M_sizes,
     std::optional<at::Tensor> global_scale,
     std::optional<at::Tensor> starting_row_after_padding) {
@@ -36,8 +34,6 @@ at::Tensor f4f4bf16_grouped_256_256_256_2_1_1_t(
       x_scale,
       w_scale,
       output,
-      G,
-      zero_start_index_M,
       M_sizes,
       global_scale,
       starting_row_after_padding);

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f4f4bf16_grouped/f4f4bf16_grouped_256_64_256_2_1_1_f.cu

@@ -18,8 +18,6 @@ at::Tensor f4f4bf16_grouped_256_64_256_2_1_1_f(
     at::Tensor x_scale,
     at::Tensor w_scale,
     at::Tensor output,
-    int64_t G,
-    std::optional<at::Tensor> zero_start_index_M,
     std::optional<at::Tensor> M_sizes,
     std::optional<at::Tensor> global_scale,
     std::optional<at::Tensor> starting_row_after_padding) {
@@ -36,8 +34,6 @@ at::Tensor f4f4bf16_grouped_256_64_256_2_1_1_f(
       x_scale,
       w_scale,
       output,
-      G,
-      zero_start_index_M,
       M_sizes,
       global_scale,
       starting_row_after_padding);