torch.ops.fbgemm.gather_scale_dense_tokens for oss. (#3855)

Summary:
Pull Request resolved: #3855

X-link: facebookresearch/FBGEMM#946

Move internal kernel to FBGEMM.

Reviewed By: jingsh, jianyuh, xw285cornell

Differential Revision: D71559646

Author: levendlee

fbgemm_gpu/experimental/gemm/triton_gemm/grouped_gemm.py

@@ -128,7 +128,7 @@ def early_config_prune(configs, named_args, dtsize=None, dtype=None, **kwargs):
     prune_configs_by={"early_config_prune": early_config_prune},
 )
 @triton.jit
-def _kernel_grouped_gemm(
+def _fbgemm_grouped_gemm(
     a_desc_ptr,
     b_desc_ptr,
     c_ptr,
@@ -271,7 +271,7 @@ def _kernel_grouped_gemm(
     },
 )
 @triton.jit
-def _kernel_grouped_gemm_fp8_rowwise(
+def _fbgemm_grouped_gemm_fp8_rowwise(
     a_desc_ptr,
     a_scale_ptr,
     b_desc_ptr,
@@ -487,7 +487,7 @@ def grid(META):
     if x_scale is not None and w_scale is not None:
         assert x_scale.is_contiguous()
         assert w_scale.is_contiguous()
-        _kernel_grouped_gemm_fp8_rowwise[grid](
+        _fbgemm_grouped_gemm_fp8_rowwise[grid](
             desc_x,
             x_scale,
             desc_w,
@@ -507,7 +507,7 @@ def grid(META):
     else:
         assert x_scale is None
         assert w_scale is None
-        _kernel_grouped_gemm[grid](
+        _fbgemm_grouped_gemm[grid](
             desc_x,
             desc_w,
             y,