Fix potential bugs in internode dispatch kernel (#25)

When I review the code, I think there are two issues in current
internode dispatch kernel:
1. The size of shared memory allocation `sharedMemoryRecv` during kernel
launch is set to `sizeof(uint32_t) * expertsPerBlock`. However during
recv in the kernel, both `sharedExpert` and `sharedToken` uses shared
memory, with each an `uint32_t` array of `expertsPerBlock`. So we need
to double the allocation size to `sizeof(uint32_t) * expertsPerBlock *
2`
2. When synchronizing the number of tokens received, the for loop should
increment with `blockDim.x`, instead of `gridDim.x * expertsPerBlock`

Author: libertyeagle

csrc/all_to_all/internode_dispatch.cu

@@ -183,8 +183,7 @@ __global__ __launch_bounds__(NUM_WARPS * 32, 1) void dispatchKernel(
     unsigned firstGroup = blockIdx.x * expertsPerBlock;
     unsigned lastGroup = std::min(firstGroup + expertsPerBlock, numExpertsAndGroups);
 
-    for (unsigned group = firstGroup + threadIdx.x; group < lastGroup;
-         group += gridDim.x * expertsPerBlock) {
+    for (unsigned group = firstGroup + threadIdx.x; group < lastGroup; group += blockDim.x) {
       const uint32_t expert = group / numDPGroups;
 
       // Fetch the token count per DP, which is non-zero to indicate receipt.
@@ -276,7 +275,7 @@ void AllToAllInterNode::dispatch(
 
   const size_t expertsPerBlock = ceil_div<size_t>(numLocalExperts * numDPGroups, numBlocks);
   const size_t sharedMemorySend = sizeof(uint32_t) * numExperts;
-  const size_t sharedMemoryRecv = sizeof(uint32_t) * expertsPerBlock;
+  const size_t sharedMemoryRecv = sizeof(uint32_t) * expertsPerBlock * 2;
 
   void *args[] = {
       const_cast<int32_t **>(&outNumTokensPerExpert.data),