diff --git a/sgl-kernel/csrc/allreduce/quick_all_reduce.cuh b/sgl-kernel/csrc/allreduce/quick_all_reduce.cuh index bd9e7b10f..4032e6666 100644 --- a/sgl-kernel/csrc/allreduce/quick_all_reduce.cuh +++ b/sgl-kernel/csrc/allreduce/quick_all_reduce.cuh @@ -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 = diff --git a/sgl-kernel/csrc/allreduce/quick_all_reduce.h b/sgl-kernel/csrc/allreduce/quick_all_reduce.h index 1d629e018..5cf961b86 100644 --- a/sgl-kernel/csrc/allreduce/quick_all_reduce.h +++ b/sgl-kernel/csrc/allreduce/quick_all_reduce.h @@ -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; \ using AllReduceKernel = AllReduceTwoshot; \ @@ -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; \ using AllReduceKernel = AllReduceTwoshot; \ @@ -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 { diff --git a/test/srt/test_quick_allreduce.py b/test/srt/test_quick_allreduce.py index ed081255f..8789426a3 100644 --- a/test/srt/test_quick_allreduce.py +++ b/test/srt/test_quick_allreduce.py @@ -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()