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

[bug][rocm]fix qr when variable inp (#11609)

Browse Source 
Signed-off-by: Haoyang Li <lihaoyang0109@gmail.com>
This commit is contained in:
haoyangli-amd
2025-11-11 01:43:48 -08:00
committed by GitHub
parent 71aea45c41
commit ea10a9d165
3 changed files with 104 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/allreduce/quick_all_reduce.cuh
+4 -3
View File
@@ -501,7 +501,8 @@ struct AllReduceTwoshot {
int const rank, // rank index
uint8_t** __restrict__ buffer_list, // communication buffers
uint32_t const data_offset, // offset to start of the data buffer
uint32_t flag_color) {
uint32_t flag_color,
int64_t data_size_per_phase) {
// Topology
int thread = threadIdx.x + threadIdx.y * kWavefront;
uint8_t* rank_buffer = buffer_list[rank];
@@ -534,10 +535,10 @@ struct AllReduceTwoshot {
// Phase-1A: Write segment data into the communication buffer of the target
// rank responsible for this segment.
uint32_t comm_data0_offset = data_offset + block_id * Codec::kTransmittedTileSize;
uint32_t comm_data1_offset = grid_size * Codec::kTransmittedTileSize + comm_data0_offset;
uint32_t comm_data1_offset = data_size_per_phase + comm_data0_offset;
uint32_t comm_flags0_offset = block_id * (kWorldSize * sizeof(uint32_t));
uint32_t comm_flags1_offset = grid_size * (kWorldSize * sizeof(uint32_t)) + comm_flags0_offset;
uint32_t comm_flags1_offset = (data_offset / 2) + comm_flags0_offset;
for (int r = 0; r < kWorldSize; r++) {
int32x4_t* send_buffer =
sgl-kernel/csrc/allreduce/quick_all_reduce.h
+9 -5
View File
@@ -28,12 +28,13 @@ __global__ __quickreduce_launch_bounds_two_shot__ static void allreduce_prototyp
int rank,
uint8_t** dbuffer_list,
uint32_t data_offset,
uint32_t flag_color) {
uint32_t flag_color,
int64_t data_size_per_phase) {
int block = blockIdx.x;
int grid = gridDim.x;
while (block < num_blocks) {
AllReduceKernel::run(A, B, N, block, rank, dbuffer_list, data_offset, flag_color);
AllReduceKernel::run(A, B, N, block, rank, dbuffer_list, data_offset, flag_color, data_size_per_phase);
block += grid;
flag_color++;
}
@@ -56,7 +57,8 @@ __global__ __quickreduce_launch_bounds_two_shot__ static void allreduce_prototyp
rank, \
dbuffer_list, \
data_offset, \
flag_color); \
flag_color, \
this->kMaxProblemSize); \
} else if (world_size == 4) { \
using LineCodec = __codec<T, 4>; \
using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
@@ -73,7 +75,8 @@ __global__ __quickreduce_launch_bounds_two_shot__ static void allreduce_prototyp
rank, \
dbuffer_list, \
data_offset, \
flag_color); \
flag_color, \
this->kMaxProblemSize); \
} else if (world_size == 8) { \
using LineCodec = __codec<T, 8>; \
using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
@@ -90,7 +93,8 @@ __global__ __quickreduce_launch_bounds_two_shot__ static void allreduce_prototyp
rank, \
dbuffer_list, \
data_offset, \
flag_color); \
flag_color, \
this->kMaxProblemSize); \
}
enum QuickReduceQuantLevel {