Demonstration of per-op targets

Summary:
For some new kernels where we are trying to add a lot of instances, it could cause code size bloat and we would encounter relocation issues on targets including fbgemm.

Although we should invest in ways to cut down our kernels code size (which I have some other ideas for later), inevitably over time as number of kernels increase, the chances of hitting this would increase a lot and is rather painful to deal with in fbcode. Currently we have a monolithic approach where all ops are pulled in when adding `:quantize_ops_gpu` as a dep. Instead we can go with a granular approach where kernels could be pulled in selectively.

This approach seems to work well, in this diff we give an example of it with adding only the bf16 grouped grad/wgrad kernels in `:quantize_bench` but not in `:quantize_ops_gpu`. The only minor down-side is the user of the ops would have to know to add the dep into their buck traget, but I think the upside is quite high as we increase the number of kernels in fbgemm and deal with more users of fbgemm, especially those with stricter requirements on code size (e.g. Sigrid predictor).

Reviewed By: jiawenliu64

Differential Revision: D83056686

Author: cthi

fbgemm_gpu/experimental/gen_ai/gen_ai/__init__.py

@@ -46,3 +46,15 @@
     torch.ops.load_library(
         "//deeplearning/fbgemm/fbgemm_gpu/experimental/gen_ai:gather_scatter_ops"
     )
+
+    gemm_ops = [
+        "//deeplearning/fbgemm/fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions:cutlass_bf16bf16bf16_grouped_grad",
+        "//deeplearning/fbgemm/fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions:cutlass_bf16bf16bf16_grouped_wgrad",
+    ]
+    for op in gemm_ops:
+        try:
+            torch.ops.load_library(
+                op,
+            )
+        except OSError:
+            pass

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

@@ -8,6 +8,7 @@
 
 #include <ATen/ATen.h>
 #include <ATen/cuda/CUDAContext.h>
+#include <torch/library.h>
 
 #include "bf16bf16bf16_grouped_grad/bf16bf16bf16_grouped_grad_manifest.cuh"
 #include "fbgemm_gpu/quantize/tuning_cache.hpp"
@@ -334,4 +335,28 @@ at::Tensor bf16bf16bf16_grouped_grad(at::Tensor, at::Tensor, at::Tensor) {
 
 #endif
 
+at::Tensor bf16bf16bf16_grouped_grad_meta(
+    at::Tensor X,
+    at::Tensor W,
+    at::Tensor /* M_sizes */) {
+  const at::SymInt total_M = X.sym_size(0);
+  const at::SymInt N = W.sym_size(1);
+  at::Tensor Y =
+      at::empty_symint({total_M, N}, X.options().dtype(at::kBFloat16));
+  return Y;
+}
+
+TORCH_LIBRARY_IMPL(fbgemm, CUDA, m) {
+  m.impl("bf16bf16bf16_grouped_grad", bf16bf16bf16_grouped_grad);
+}
+
+TORCH_LIBRARY_IMPL(fbgemm, Meta, m) {
+  m.impl("bf16bf16bf16_grouped_grad", bf16bf16bf16_grouped_grad_meta);
+}
+
+TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
+  m.def(
+      "bf16bf16bf16_grouped_grad(Tensor X, Tensor W, Tensor M_sizes) -> Tensor");
+}
+
 } // namespace fbgemm_gpu

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

@@ -8,6 +8,7 @@
 
 #include <ATen/ATen.h>
 #include <ATen/cuda/CUDAContext.h>
+#include <torch/library.h>
 
 #include "bf16bf16bf16_grouped_wgrad/bf16bf16bf16_grouped_wgrad_manifest.cuh"
 #include "fbgemm_gpu/quantize/tuning_cache.hpp"
@@ -1093,4 +1094,30 @@ at::Tensor bf16bf16bf16_grouped_wgrad(
 
 #endif
 
+at::Tensor bf16bf16bf16_grouped_wgrad_meta(
+    at::Tensor X,
+    at::Tensor W,
+    at::Tensor M_sizes,
+    std::optional<at::Tensor> /* output = std::nullopt */,
+    bool /* output_accum = false */) {
+  const at::SymInt G = M_sizes.size(0);
+  const at::SymInt N = X.sym_size(1);
+  const at::SymInt K = W.sym_size(1);
+  at::Tensor Y = at::empty_symint({G, N, K}, X.options().dtype(at::kBFloat16));
+  return Y;
+}
+
+TORCH_LIBRARY_IMPL(fbgemm, CUDA, m) {
+  m.impl("bf16bf16bf16_grouped_wgrad", bf16bf16bf16_grouped_wgrad);
+}
+
+TORCH_LIBRARY_IMPL(fbgemm, Meta, m) {
+  m.impl("bf16bf16bf16_grouped_wgrad", bf16bf16bf16_grouped_wgrad_meta);
+}
+
+TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
+  m.def(
+      "bf16bf16bf16_grouped_wgrad(Tensor X, Tensor W, Tensor M_sizes, Tensor(a!)? output=None, bool output_accum=False) -> Tensor");
+}
+
 } // namespace fbgemm_gpu

fbgemm_gpu/experimental/gen_ai/src/quantize/quantize.cpp

@@ -78,14 +78,6 @@ at::Tensor bf16bf16bf16_grouped_dynamic(
     at::Tensor zero_start_index_M);
 at::Tensor
 bf16bf16bf16_grouped_stacked(at::Tensor X, at::Tensor W, at::Tensor M_sizes);
-at::Tensor
-bf16bf16bf16_grouped_grad(at::Tensor X, at::Tensor W, at::Tensor M_sizes);
-at::Tensor bf16bf16bf16_grouped_wgrad(
-    at::Tensor X,
-    at::Tensor W,
-    at::Tensor M_sizes,
-    std::optional<at::Tensor> output = std::nullopt,
-    bool output_accum = false);
 at::Tensor f8f8bf16_rowwise(
     at::Tensor XQ,
     at::Tensor WQ,
@@ -325,8 +317,6 @@ TORCH_LIBRARY_IMPL(fbgemm, CUDA, m) {
   m.impl("bf16i4bf16_shuffled", bf16i4bf16_shuffled);
   m.impl("f8i4bf16_shuffled_grouped", f8i4bf16_shuffled_grouped);
   m.impl("bf16i4bf16_shuffled_grouped", bf16i4bf16_shuffled_grouped);
-  m.impl("bf16bf16bf16_grouped_grad", bf16bf16bf16_grouped_grad);
-  m.impl("bf16bf16bf16_grouped_wgrad", bf16bf16bf16_grouped_wgrad);
   m.impl("preshuffle_i4", preshuffle_i4);
   m.impl("bf16i4bf16_shuffled_batched", bf16i4bf16_shuffled_batched);
   m.impl("bf16i4bf16_rowwise_batched", bf16i4bf16_rowwise_batched);
@@ -382,7 +372,6 @@ TORCH_LIBRARY_IMPL(fbgemm, CPU, m) {
   m.impl("bf16i4bf16_shuffled", bf16i4bf16_shuffled);
   m.impl("f8i4bf16_shuffled_grouped", f8i4bf16_shuffled_grouped);
   m.impl("bf16i4bf16_shuffled_grouped", bf16i4bf16_shuffled_grouped);
-  m.impl("bf16bf16bf16_grouped_grad", bf16bf16bf16_grouped_grad);
   m.impl("preshuffle_i4", preshuffle_i4);
   m.impl("bf16i4bf16_shuffled_batched", bf16i4bf16_shuffled_batched);
   m.impl("bf16i4bf16_rowwise_batched", bf16i4bf16_rowwise_batched);
@@ -800,30 +789,6 @@ at::Tensor bf16bf16bf16_grouped_stacked_meta(
   return Y;
 }
 
-at::Tensor bf16bf16bf16_grouped_grad_meta(
-    at::Tensor X,
-    at::Tensor W,
-    at::Tensor /* M_sizes */) {
-  const at::SymInt total_M = X.sym_size(0);
-  const at::SymInt N = W.sym_size(1);
-  at::Tensor Y =
-      at::empty_symint({total_M, N}, X.options().dtype(at::kBFloat16));
-  return Y;
-}
-
-at::Tensor bf16bf16bf16_grouped_wgrad_meta(
-    at::Tensor X,
-    at::Tensor W,
-    at::Tensor M_sizes,
-    std::optional<at::Tensor> /* output = std::nullopt */,
-    bool /* output_accum = false */) {
-  const at::SymInt G = M_sizes.size(0);
-  const at::SymInt N = X.sym_size(1);
-  const at::SymInt K = W.sym_size(1);
-  at::Tensor Y = at::empty_symint({G, N, K}, X.options().dtype(at::kBFloat16));
-  return Y;
-}
-
 at::Tensor f8f8bf16_rowwise_grouped_stacked_meta(
     at::Tensor XQ,
     at::Tensor WQ,
@@ -864,8 +829,6 @@ TORCH_LIBRARY_IMPL(fbgemm, Meta, m) {
   m.impl("bf16i4bf16_rowwise", bf16i4bf16_rowwise_meta);
   m.impl("bf16i4bf16_shuffled_batched", bf16i4bf16_shuffled_batched_meta);
   m.impl("bf16i4bf16_rowwise_batched", bf16i4bf16_rowwise_batched_meta);
-  m.impl("bf16bf16bf16_grouped_grad", bf16bf16bf16_grouped_grad_meta);
-  m.impl("bf16bf16bf16_grouped_wgrad", bf16bf16bf16_grouped_wgrad_meta);
   m.impl("f8f8bf16_lite", f8f8bf16_lite_meta);
   m.impl("scaled_fp4_quant", scaled_fp4_quant_meta);
   m.impl("preshuffle_i4", preshuffle_i4_meta);

fbgemm_gpu/experimental/gen_ai/src/quantize/quantize_defs.cpp

@@ -64,10 +64,6 @@ TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
       "bf16bf16bf16_grouped_dynamic(Tensor X, Tensor W, Tensor zero_start_index_M) -> Tensor");
   m.def(
       "bf16bf16bf16_grouped_stacked(Tensor X, Tensor W, Tensor M_sizes) -> Tensor");
-  m.def(
-      "bf16bf16bf16_grouped_grad(Tensor X, Tensor W, Tensor M_sizes) -> Tensor");
-  m.def(
-      "bf16bf16bf16_grouped_wgrad(Tensor X, Tensor W, Tensor M_sizes, Tensor(a!)? output=None, bool output_accum=False) -> Tensor");
   m.def(
       "f8f8bf16_blockwise(Tensor XQ, Tensor WQ, Tensor x_scale, Tensor w_scale, int block_m=128, int block_n=128, int block_k=128) -> Tensor");
   m.def(