fix fbgemm build issues after upgrading CK (#4517)

Summary:
Pull Request resolved: #4517

X-link: facebookresearch/FBGEMM#1565

Some files are copied over from CK so need to update them.

Differential Revision: D78455742

Author: 842974287

fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fused_moe/fused_moe.hpp

@@ -58,6 +58,8 @@ struct fused_moe_traits {
   bool local_expert_masking; // if mask experts as local expert
 };
 
+// if return zero, no ws needed
+int fused_moe_get_workspace_size(int tokens, int num_experts, int topk);
 float fused_moe(
     fused_moe_traits,
     fused_moe_args,

fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fused_moe/fused_moe_kernel.hip

@@ -77,14 +77,14 @@ at::Tensor fused_moe_impl(
   auto prec_o = get_prec_str(output);
   auto prec_tkw = get_prec_str(topk_weights);
 
-  int workspace_size = ck_tile::moe_sorting_get_workspace_size(tokens, experts);
+  int workspace_size = ck_tile::moe_sorting_get_workspace_size(tokens, experts, topk);
   void *ws_ptr = nullptr;
   if (workspace_size > 0)
   {
       auto ws = at::zeros({workspace_size}, at::TensorOptions().dtype(topk_ids.dtype()).device(device_of(topk_ids)));
       ws_ptr = ws.data_ptr();
   }
-  
+
 
   // Set up traits structure
   fused_moe_traits traits{

fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fused_moe/fused_moesorting.hpp

@@ -15,13 +15,8 @@ struct fused_moesorting_trait {
 
 struct fused_moesorting_args : public ck_tile::MoeSortingHostArgs {};
 
+int fused_moe_get_workspace_size(int tokens, int num_experts, int topk);
 float fused_moesorting(
     fused_moesorting_trait t,
     fused_moesorting_args a,
     ck_tile::stream_config s);
-
-int moe_sorting_get_workspace_size(int tokens, int num_experts);
-float moe_sorting_mp(
-    fused_moesorting_trait t,
-    fused_moesorting_args a,
-    ck_tile::stream_config s);

fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fused_moe/instances/fused_moe_api.hip

@@ -3,6 +3,11 @@
 
 #include "fused_moe.hpp"
 
+int fused_moe_get_workspace_size(int tokens, int num_experts, int topk)
+{
+    return ck_tile::moe_sorting_get_workspace_size(tokens, num_experts, topk);
+}
+
 float fused_moe(fused_moe_traits t, fused_moe_args a, const ck_tile::stream_config& s)
 {
     auto s_sub = ck_tile::stream_config{s.stream_id_, false, s.log_level_, 0, 1};
@@ -19,20 +24,21 @@ float fused_moe(fused_moe_traits t, fused_moe_args a, const ck_tile::stream_conf
 
     auto t0 = fused_moesorting_trait{"int32", "fp32", t.local_expert_masking};
     auto a0 = fused_moesorting_args{
-        a.topk_ids_ptr,                              // const void* p_topk_ids;
-        a.topk_weight_ptr,                           // const void* p_weights;
-        a.local_expert_mask_ptr,                     // const void* p_local_expert_mask;
-        a.sorted_token_ids_ptr,                      // void* p_sorted_token_ids;
-        a.sorted_weight_ptr,                         // void* p_sorted_weights;
-        a.sorted_expert_ids_ptr,                     // void* p_sorted_expert_ids;
-        a.num_sorted_tiles_ptr,                      // void* p_total_tokens_post_pad;
-        a.o_ptr,                                     // void* p_moe_buf;
-        a.ws_ptr,                                    // moe_sorting_ws
-        a.num_tokens,                                // index_t tokens;
-        a.block_m,                                   // index_t unit_size;
-        a.num_experts,                               // index_t num_experts;
-        a.topk,                                      // index_t topk;
-        a.num_tokens * a.stride_token * o_data_bytes // index_t moe_buf_bytes;
+        a.topk_ids_ptr,          // const void* p_topk_ids;
+        a.topk_weight_ptr,       // const void* p_weights;
+        a.local_expert_mask_ptr, // const void* p_local_expert_mask;
+        a.sorted_token_ids_ptr,  // void* p_sorted_token_ids;
+        a.sorted_weight_ptr,     // void* p_sorted_weights;
+        a.sorted_expert_ids_ptr, // void* p_sorted_expert_ids;
+        a.num_sorted_tiles_ptr,  // void* p_total_tokens_post_pad;
+        a.o_ptr,                 // void* p_moe_buf;
+        a.ws_ptr,                // void* p_ws;
+        a.num_tokens,            // index_t tokens;
+        a.block_m,               // index_t unit_size;
+        a.num_experts,           // index_t num_experts;
+        a.topk,                  // index_t topk;
+        static_cast<ck_tile::long_index_t>(a.num_tokens) * a.stride_token *
+            o_data_bytes // index_t moe_buf_bytes;
     };
 
     auto t1 = fused_moegemm_traits{t.prec_i,