chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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 17: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

7
sgl-kernel/csrc/allreduce/quick_all_reduce.cuh
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 =

14
sgl-kernel/csrc/allreduce/quick_all_reduce.h
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 {

91
test/srt/test_quick_allreduce.py
View File

@@ -1,3 +1,4 @@
import multiprocessing
import os
import random
import socket
@@ -8,6 +9,7 @@ import ray
import torch
import torch.distributed as dist
from sglang.srt import _custom_ops as ops
from sglang.srt.distributed import init_distributed_environment
from sglang.srt.distributed.communication_op import ( # noqa
tensor_model_parallel_all_reduce,
@@ -208,5 +210,94 @@ class TestQuickAllReduce(CustomTestCase):
# print("Max rel diff:", ((out1 - inp1).abs() / inp1.abs().clamp(min=1e-5)).max())
def qr_variable_input(rank, world_size):
device = torch.device(f"cuda:{rank}")
torch.cuda.set_device(device)
qr_max_size = None # MB
_ptr = ops.init_custom_qr(rank, world_size, qr_max_size)
ranks = []
for i in range(world_size):
ranks.append(i)
dist.init_process_group(
backend="nccl",
init_method="tcp://127.0.0.1:29500",
rank=rank,
world_size=world_size,
)
cpu_group = torch.distributed.new_group(ranks, backend="nccl")
handle = ops.qr_get_handle(_ptr)
world_size = dist.get_world_size(group=cpu_group)
handles = [None] * world_size
dist.all_gather_object(handles, handle, group=cpu_group)
ops.qr_open_handles(_ptr, handles)
num = 1
s1 = 1024
while num < 50000: # 50000 is sufficient to identify issues.
dtype = torch.float16
if num % 2 == 0:
s2 = 1024
inp1 = torch.zeros(
(s1, s2), dtype=dtype, device=torch.cuda.current_device()
)
else:
s2 = 2048
inp1 = torch.ones((s1, s2), dtype=dtype, device=torch.cuda.current_device())
result = torch.empty_like(inp1)
# FP = 0 INT8 = 1 INT6 = 2 INT4 = 3 NONE = 4
ops.qr_all_reduce(_ptr, inp1, result, 3, cast_bf2half=True)
try:
if inp1[0, 0] == 0:
assert torch.all(result == 0)
else:
assert torch.all(result == world_size)
except AssertionError:
print("Assertion failed! Allreduce results are incorrect.")
raise
num += 1
class TestQuickreduceVariableInput(CustomTestCase):
"""
When the tensor parallelism is set to 4 or 8, frequent changes
in the input shape can cause QuickReduce to hang (this issue
has been observed with the gpt_oss model).
"""
TP_SIZES = [4, 8]
@unittest.skipIf(
not qr_rocm_arch_available(),
"Only test Quick AllReduce on ROCm architectures >= gfx94*",
)
def test_custom_quick_allreduce_variable_input(self):
for tp_size in self.TP_SIZES:
world_size = tp_size
if world_size > torch.cuda.device_count():
return
multiprocessing.set_start_method("spawn", force=True)
# 90s is enough
timeout = 90
processes = []
for rank in range(tp_size):
p = multiprocessing.Process(
target=qr_variable_input, args=(rank, tp_size)
)
p.start()
processes.append((rank, p))
for rank, p in processes:
p.join(timeout=timeout)
if p.is_alive():
for r, proc in processes:
if proc.is_alive():
proc.terminate()
proc.join()
raise RuntimeError(
f"QuickReduce hang detected after {timeout} seconds!"
)
if __name__ == "__main__":
unittest.main()