mirror fix for custom allreduce (#3124)

98522149 · yizhang2077 · GitHub · 5d9d15e7 · 98522149
Unverified Commit 98522149 authored Jan 25, 2025 by yizhang2077 Committed by GitHub Jan 25, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

sgl-kernel/src/sgl-kernel/csrc/trt_reduce_internal.cu sgl-kernel/src/sgl-kernel/csrc/trt_reduce_internal.cu +1 -1

No files found.
--- a/sgl-kernel/src/sgl-kernel/csrc/trt_reduce_internal.cu
+++ b/sgl-kernel/src/sgl-kernel/csrc/trt_reduce_internal.cu
@@ -160,7 +160,7 @@ __inline__ __device__ void block_barrier(uint32_t** signals, uint32_t const flag
 }

 template <typename T, int RANKS_PER_NODE, bool COPY_INPUT = true>
-static __global__ void oneShotAllReduceKernel(AllReduceParams params) {
+static __global__ void __launch_bounds__(512, 1) oneShotAllReduceKernel(AllReduceParams params) {
  // Suppose that two GPUs participate in the AR exchange, and we start four blocks.
  // The message is partitioned into chunks as detailed below:
  //               message