minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[MUSA][8/N] Port CUDA kernels that are compatible with MUSA (#17946)

Browse Source 
Signed-off-by: yafeng.li <yafeng.li@mthreads.com>
Co-authored-by: Alex Nails <alex.nails@radixark.ai>
This commit is contained in:
MARATRIX
2026-04-23 18:04:58 -07:00
committed by GitHub
co-authored by Alex Nails
parent c0166355ae
commit 74c2e5bacd
15 changed files with 1905 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/allreduce/custom_all_reduce.cuh
+206 -25
View File
@@ -17,7 +17,24 @@
namespace sglang {
#ifndef USE_MUSA
constexpr int kMaxBlocks = 36;
constexpr int kDefaultThreads = 512;
constexpr int kDefaultBlockLimit = 36;
constexpr int kMaxThreadsPerBlock = 512;
#else
constexpr int kMaxBlocks = 60;
constexpr int kDefaultThreads = 1024;
constexpr int kDefaultBlockLimit = 60;
constexpr int kMaxThreadsPerBlock = 1024;
#endif
// Allreduce algorithm selection thresholds
constexpr int kAllReduceGPUSmall = 4;
constexpr int kAllReduceGPULarge = 8;
constexpr size_t kAllReduceSmallThreshold = 512 * 1024; // 512KB
constexpr size_t kAllReduceLargeThreshold = 256 * 1024; // 256KB
// Counter may overflow, but it's fine since unsigned int overflow is
// well-defined behavior.
using FlagType = uint32_t;
@@ -134,7 +151,9 @@ DINLINE O downcast(array_t<float, O::size> val) {
}
static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
#ifdef USE_MUSA
volatile_store((uint32_t)flag, (uint32_t*)flag_addr);
#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
asm volatile("st.release.sys.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr));
#else
asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr));
@@ -142,6 +161,11 @@ static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) {
}
static DINLINE FlagType ld_flag_acquire(FlagType* flag_addr) {
#ifdef USE_MUSA
flushInv_byp();
return (uint32_t)volatile_load((uint32_t*)flag_addr);
#endif
FlagType flag;
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(flag) : "l"(flag_addr));
@@ -167,7 +191,12 @@ static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) {
// barrier.
template <int ngpus, bool is_start, bool need_fence = false>
DINLINE void multi_gpu_barrier(const RankSignals& sg, Signal* self_sg, int rank) {
if constexpr (!is_start) __syncthreads();
if constexpr (!is_start)
#ifndef USE_MUSA
__syncthreads();
#else
__syncthreads_lm();
#endif
static_assert(!(is_start && need_fence)); // Start barrier shouldn't need fence.
if (threadIdx.x < ngpus) {
// Increment the counter. Technically we only need one counter, but we use
@@ -187,7 +216,12 @@ DINLINE void multi_gpu_barrier(const RankSignals& sg, Signal* self_sg, int rank)
;
}
}
if constexpr (is_start || need_fence) __syncthreads();
if constexpr (is_start || need_fence)
#ifndef USE_MUSA
__syncthreads();
#else
__syncthreads_lm();
#endif
}
template <typename P, int ngpus, typename A>
@@ -201,7 +235,7 @@ DINLINE P packed_reduce(const P* ptrs[], int idx) {
}
template <typename T, int ngpus>
__global__ void __launch_bounds__(512, 1) cross_device_reduce_1stage(
__global__ void __launch_bounds__(kMaxThreadsPerBlock, 1) cross_device_reduce_1stage(
RankData* _dp, RankSignals sg, Signal* self_sg, T* __restrict__ result, int rank, int size) {
using P = typename packed_t<T>::P;
using A = typename packed_t<T>::A;
@@ -221,8 +255,132 @@ DINLINE P* get_tmp_buf(Signal* sg) {
return (P*)(((Signal*)sg) + 1);
}
#ifdef USE_MUSA
template <typename T, int32_t nranks, int32_t vlen = 8>
DINLINE void shfl_reduce(float* res) {
if constexpr (nranks >= 4) {
#pragma unroll
for (int32_t i = 0; i < vlen; i++) {
res[i] += __shfl_xor_sync(0xffffffff, res[i], 16);
}
}
#pragma unroll
for (int32_t i = 0; i < vlen; i++) {
res[i] += __shfl_xor_sync(0xffffffff, res[i], 8);
}
}
template <typename T, int32_t nranks, int32_t vlen = 8>
__global__ void __launch_bounds__(kMaxThreadsPerBlock, 1) custom_all_reduce_2shot(
RankData* _dp, RankSignals sg, Signal* self_sg, T* __restrict__ result, int32_t local_rank, int32_t size) {
constexpr int32_t nranks_sft = (nranks >> 1) - (nranks >> 3); // 8->3, 4->2, 2->1
constexpr int32_t coalesce_num = 8;
constexpr int32_t coalesce_sft = 3; // 8 threads per rank in group
constexpr int32_t group_size = nranks << coalesce_sft; // tp 8 -> 64 threads, tp 4 -> 32 threads, tp 2 -> 16 threads
constexpr int32_t group_stride_sft = nranks_sft + coalesce_sft;
const int32_t tidx = threadIdx.x;
const int32_t bidx = blockIdx.x;
const int32_t thread_num = blockDim.x;
const int32_t lane_idx = tidx & 31;
const int32_t warp_idx = tidx >> 5;
const int32_t group_num = thread_num >> group_stride_sft;
const int32_t target_rank = (tidx >> coalesce_sft) & (nranks - 1);
const int32_t group_id = tidx >> group_stride_sft;
const int32_t coalesce_tid = tidx & (coalesce_num - 1);
typedef int16_t Vec __attribute__((vector_size(16)));
const int32_t stride = gridDim.x * thread_num;
int32_t idx_base = bidx * thread_num;
int32_t idx_in_blk = coalesce_tid + (local_rank << coalesce_sft) + (group_id << group_stride_sft);
// first sync barrier
FlagType* target_barrier = nullptr;
FlagType* local_barrier = nullptr;
FlagType flag;
if (tidx < nranks) {
flag = atomicAdd(&(self_sg->self_counter[bidx][tidx]), 1);
target_barrier = &sg.signals[tidx]->peer_counter[flag & 1][bidx][local_rank];
local_barrier = &self_sg->peer_counter[flag & 1][bidx][tidx];
atomicExch(target_barrier, flag);
while (atomicAdd(local_barrier, 0) != flag) {
}
}
__syncthreads_lm();
// reduce scatter
Vec* target_ptr = (Vec*)_dp->ptrs[target_rank];
Vec* buffer_ptr = get_tmp_buf<Vec>(sg.signals[local_rank]);
do {
int32_t idx = idx_in_blk + idx_base;
float temp_res[vlen] = {0};
if (idx < size) {
T* data = reinterpret_cast<T*>(&(target_ptr[idx]));
#pragma unroll
for (int32_t i = 0; i < vlen; i++) {
temp_res[i] = upcast_s(data[i]);
}
}
shfl_reduce<T, nranks, vlen>(temp_res);
// reduce cross warp, only trigger when tp 8
if constexpr (nranks == 8) {
__shared__ float smem[kMaxThreadsPerBlock << 1];
if (lane_idx < coalesce_num) {
#pragma unroll
for (int32_t i = 0; i < vlen; i++) {
smem[warp_idx * vlen * coalesce_num + coalesce_tid * vlen + i] = temp_res[i];
}
}
__syncthreads_lm();
#pragma unroll
for (int32_t i = 0; i < vlen; i++) {
temp_res[i] += smem[(warp_idx ^ 1) * vlen * coalesce_num + coalesce_tid * vlen + i];
}
}
if (local_rank == target_rank && idx < size) {
Vec res;
#pragma unroll
for (int32_t i = 0; i < vlen; i++) {
reinterpret_cast<T*>(&res)[i] = downcast_s<T>(temp_res[i]);
}
buffer_ptr[idx] = res;
}
idx_base += stride;
} while (idx_base < size);
// make sure buffer_ptr data ready
__musa_barrier_slc();
__syncthreads_lm();
if (tidx == 0) {
__threadfence_system_noflush();
}
buffer_ptr = get_tmp_buf<Vec>(sg.signals[target_rank]);
// second sync barrier
if (tidx < nranks) {
flag = atomicAdd(&(self_sg->self_counter[bidx][tidx]), 1);
target_barrier = &sg.signals[tidx]->peer_counter[flag & 1][bidx][local_rank];
local_barrier = &self_sg->peer_counter[flag & 1][bidx][tidx];
atomicExch(target_barrier, flag);
while (atomicAdd(local_barrier, 0) != flag) {
}
}
__syncthreads_lm();
// all gather
idx_in_blk = coalesce_tid + (target_rank << coalesce_sft) + (group_id << group_stride_sft);
idx_base = bidx * thread_num;
do {
int32_t idx = idx_in_blk + idx_base;
if (idx < size) {
reinterpret_cast<Vec*>(result)[idx] = buffer_ptr[idx];
}
idx_base += stride;
} while (idx_base < size);
}
#endif // USE_MUSA
template <typename T, int ngpus>
__global__ void __launch_bounds__(512, 1) cross_device_reduce_2stage(
__global__ void __launch_bounds__(kMaxThreadsPerBlock, 1) cross_device_reduce_2stage(
RankData* _dp, RankSignals sg, Signal* self_sg, T* __restrict__ result, int rank, int size) {
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int stride = gridDim.x * blockDim.x;
@@ -414,7 +572,13 @@ class CustomAllreduce {
* guess is that too many SMs will cause contention on NVLink bus.
*/
template <typename T>
void allreduce(cudaStream_t stream, T* input, T* output, int size, int threads = 512, int block_limit = 36) {
void allreduce(
cudaStream_t stream,
T* input,
T* output,
int size,
int threads = kDefaultThreads,
int block_limit = kDefaultBlockLimit) {
auto d = packed_t<T>::P::size;
if (size % d != 0)
throw std::runtime_error(
@@ -462,26 +626,43 @@ class CustomAllreduce {
#define KL(ngpus, name) name<T, ngpus><<<blocks, threads, 0, stream>>>(ptrs, sg_, self_sg_, output, rank_, size);
// TODO(hanzhi713): Threshold is different for A100 and H100.
// Add per device threshold.
#define REDUCE_CASE(ngpus) \
case ngpus: { \
if (force_1stage) { \
KL(ngpus, cross_device_reduce_1stage); \
} else if (force_2stage) { \
KL(ngpus, cross_device_reduce_2stage); \
} else { \
if (world_size_ == 2) { \
KL(ngpus, cross_device_reduce_1stage); \
} else if (full_nvlink_) { \
if ((world_size_ <= 4 && bytes < 512 * 1024) || (world_size_ <= 8 && bytes < 256 * 1024)) { \
KL(ngpus, cross_device_reduce_1stage); \
} else { \
KL(ngpus, cross_device_reduce_2stage); \
} \
} \
} \
break; \
#ifndef USE_MUSA
#define REDUCE_CASE(ngpus) \
case ngpus: { \
if (force_1stage) { \
KL(ngpus, cross_device_reduce_1stage); \
} else if (force_2stage) { \
KL(ngpus, cross_device_reduce_2stage); \
} else { \
if (world_size_ == 2) { \
KL(ngpus, cross_device_reduce_1stage); \
} else if (full_nvlink_) { \
if ((world_size_ <= kAllReduceGPUSmall && bytes < kAllReduceSmallThreshold) || \
(world_size_ <= kAllReduceGPULarge && bytes < kAllReduceLargeThreshold)) { \
KL(ngpus, cross_device_reduce_1stage); \
} else { \
KL(ngpus, cross_device_reduce_2stage); \
} \
} \
} \
break; \
}
#else
#define REDUCE_CASE(ngpus) \
case ngpus: { \
if constexpr (!std::is_same<T, float>::value) { \
custom_all_reduce_2shot<T, ngpus><<<blocks, threads, 0, stream>>>(ptrs, sg_, self_sg_, output, rank_, size); \
} else { \
if ((world_size_ <= kAllReduceGPUSmall && bytes < kAllReduceSmallThreshold) || \
(world_size_ <= kAllReduceGPULarge && bytes < kAllReduceLargeThreshold)) { \
KL(ngpus, cross_device_reduce_1stage); \
} else { \
KL(ngpus, cross_device_reduce_2stage); \
} \
} \
break; \
}
#endif
switch (world_size_) {
REDUCE_CASE(2)
REDUCE_CASE(4)