fix moe_align_kernel shm init not sync bug (#3534)

2025-02-13 16:47:00 +08:00
parent bf2a70872e
commit f076328bb7
1 changed files with 2 additions and 0 deletions
--- a/sgl-kernel/src/sgl-kernel/csrc/moe_align_kernel.cu
+++ b/sgl-kernel/src/sgl-kernel/csrc/moe_align_kernel.cu
@@ -53,6 +53,8 @@ __global__ void moe_align_block_size_kernel(scalar_t* __restrict__ topk_ids, int
    }
  }

+  __syncthreads();
+
  const size_t tokens_per_thread = CEILDIV(numel, blockDim.x);
  const size_t start_idx = threadIdx.x * tokens_per_thread;