Use cudaMemsetAsync/hipMemsetAsync to setup IndexShuffling kernel. (pytorch#4016)

Summary:
Pull Request resolved: pytorch#4016

X-link: facebookresearch/FBGEMM#1104

It is too expensive to launch a ATen kernel to do setup. Use cudaMemsetAsync/hipMemsetAsync instead.
Note we need to use cudaMemsetAsync/hipMemsetAsync on current stream to be compatible with CUDAGraph capture.

Reviewed By: Alkaid-Benetnash

Differential Revision: D73602755

Author: levendlee

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

@@ -276,7 +276,19 @@ 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
+  hipMemsetAsync(
+      counts.data_ptr(),
+      0,
+      counts.numel() * counts.dtype().itemsize(),
+      at::cuda::getCurrentCUDAStream());
+#else
+  cudaMemsetAsync(
+      counts.data_ptr(),
+      0,
+      counts.numel() * counts.dtype().itemsize(),
+      at::cuda::getCurrentCUDAStream());
+#endif
 
   // Avoid expensive `cudaGetDeviceProperties` call.
   if (num_sms < 0) {