Use cudaMemsetAsync to setup IndexShuffling kernel. (#4016)

Summary:
Pull Request resolved: #4016

X-link: facebookresearch/FBGEMM#1104

It is too expensive to launch a ATen kernel to do setup. Use cudaMemsetAsync instead.
hipMemsetAsync is somehow more expensive than launching a kernel. Avoid doing so for now.

Reviewed By: Alkaid-Benetnash

Differential Revision: D73602755

fbshipit-source-id: 6cdb2b5d489a7fd9c1a2791fcc1c26f2e6dd63ec

Author: levendlee

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

@@ -276,7 +276,23 @@ 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);
 
+#ifdef USE_ROCM
   counts.zero_();
+  // TODO(shikaili): hipMetsetAsync is more expensive than ATen set zero.
+  /*
+  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) {
@@ -298,7 +314,6 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> index_shuffling_torch(
   kernel = (void*)index_shuffling_kernel<DataType, IndexType, E, B>; \
   smem_size = sizeof(SharedStorage<DataType, IndexType, E, B>);
 
-  int num_tokens_per_tile;
   if (num_experts == 16) {
     DISPATCH(16, kNumTokensPerTile);
   } else {