Use cudaMemset/hipMemset to setup IndexShuffling kernel. (pytorch#4016)

Summary:

X-link: facebookresearch/FBGEMM#1104

It is too expensive to launch a ATen kernel to do setup. Use cudaMemset/hipMemset instead.

Differential Revision: D73602755

Author: levendlee

fbgemm_gpu/experimental/gen_ai/src/moe/index_shuffling.cu

@@ -276,7 +276,11 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> index_shuffling_torch(
   at::Tensor shuffled_expert_indices = allocate_index_tensor(num_tokens);
   at::Tensor shuffled_token_indices = allocate_index_tensor(num_tokens);
 
-  counts.zero_();
+#ifdef USE_ROCM
+  hipMemset(counts.data_ptr(), 0, counts.numel() * counts.dtype().itemsize());
+#else
+  cudaMemset(counts.data_ptr(), 0, counts.numel() * counts.dtype().itemsize());
+#endif
 
   // Avoid expensive `cudaGetDeviceProperties` call.
   if (num_sms < 0) {