diff --git a/sgl-kernel/csrc/allreduce/quick_all_reduce.h b/sgl-kernel/csrc/allreduce/quick_all_reduce.h index 5cf961b86..0e6a68fa3 100644 --- a/sgl-kernel/csrc/allreduce/quick_all_reduce.h +++ b/sgl-kernel/csrc/allreduce/quick_all_reduce.h @@ -28,16 +28,25 @@ __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* d_flag_counters, int64_t data_size_per_phase) { int block = blockIdx.x; int grid = gridDim.x; + // Read this block's counter from device memory and bump it here in the + // kernel. Keeping the value in device memory (instead of a host scalar + // baked into the launch) lets every CUDA-graph replay see a fresh color. + uint32_t flag_color = d_flag_counters[blockIdx.x]; + while (block < num_blocks) { AllReduceKernel::run(A, B, N, block, rank, dbuffer_list, data_offset, flag_color, data_size_per_phase); block += grid; flag_color++; } + // The whole block ends up with the same value, so a single writer suffices. + if (threadIdx.x == 0 && threadIdx.y == 0) { + d_flag_counters[blockIdx.x] = flag_color; + } } #define TWOSHOT_DISPATCH(__codec) \ @@ -57,7 +66,7 @@ __global__ __quickreduce_launch_bounds_two_shot__ static void allreduce_prototyp rank, \ dbuffer_list, \ data_offset, \ - flag_color, \ + d_flag_counters, \ this->kMaxProblemSize); \ } else if (world_size == 4) { \ using LineCodec = __codec; \ @@ -75,7 +84,7 @@ __global__ __quickreduce_launch_bounds_two_shot__ static void allreduce_prototyp rank, \ dbuffer_list, \ data_offset, \ - flag_color, \ + d_flag_counters, \ this->kMaxProblemSize); \ } else if (world_size == 8) { \ using LineCodec = __codec; \ @@ -93,7 +102,7 @@ __global__ __quickreduce_launch_bounds_two_shot__ static void allreduce_prototyp rank, \ dbuffer_list, \ data_offset, \ - flag_color, \ + d_flag_counters, \ this->kMaxProblemSize); \ } @@ -112,7 +121,7 @@ struct DeviceComms { static int constexpr kMaxWorldSize = 8; bool initialized = false; - uint32_t flag_color = 1; + uint32_t* d_flag_counters = nullptr; int world_size; int rank; @@ -145,6 +154,14 @@ struct DeviceComms { // Clear the flags buffer. HIP_CHECK(hipMemset(dbuffer, 0, flags_buffer_size)); + // A per-block color counter that the kernel advances itself. Seed it with + // 1 rather than 0 so it never matches the freshly zeroed flags buffer. + HIP_CHECK(hipMalloc(&d_flag_counters, kMaxNumBlocks * sizeof(uint32_t))); + { + std::vector init_color(kMaxNumBlocks, 1u); + HIP_CHECK(hipMemcpy(d_flag_counters, init_color.data(), kMaxNumBlocks * sizeof(uint32_t), hipMemcpyHostToDevice)); + } + // Device-side list of IPC buffers. buffer_list.resize(world_size); HIP_CHECK(hipMalloc(&dbuffer_list, world_size * sizeof(uint8_t*))); @@ -169,6 +186,12 @@ struct DeviceComms { } void destroy() { + // This buffer is created before `initialized` becomes true, so release it + // on its own check to keep a half-finished init from leaking it. + if (d_flag_counters) { + HIP_CHECK(hipFree(d_flag_counters)); + d_flag_counters = nullptr; + } if (initialized) { for (int i = 0; i < world_size; i++) { if (i != rank) { @@ -229,8 +252,7 @@ struct DeviceComms { break; } HIP_CHECK(cudaGetLastError()); - // Rotate the flag color. - flag_color += divceil(N, grid); + // The color now advances on-device inside the kernel; no host-side bump. } }; diff --git a/test/manual/test_quick_allreduce.py b/test/manual/test_quick_allreduce.py index 42bd4c9c0..f88049b03 100644 --- a/test/manual/test_quick_allreduce.py +++ b/test/manual/test_quick_allreduce.py @@ -15,6 +15,7 @@ from sglang.srt.distributed.communication_op import ( # noqa tensor_model_parallel_all_reduce, ) from sglang.srt.distributed.device_communicators.quick_all_reduce import ( + QuickAllReduce, qr_rocm_arch_available, ) from sglang.srt.distributed.parallel_state import ( @@ -258,6 +259,117 @@ def qr_variable_input(rank, world_size): num += 1 +def qr_graph_replay(rank, world_size, quant_mode="FP", num_replays=10): + """Capture ONE CUDA graph with a single quick-reduce and replay it many + times with changing input. Every rank contributes the same value v in a + round, so the true all-reduce sum is world_size * v; the FP regime is + lossless, so the comparison is bit-exact. + + The pre-fix kernel bakes the per-block flag color into the graph launch and + reuses it on every replay -- the waiting peer is satisfied by the previous + round's residual flag and reads stale data, giving wrong results on some + replays. The fixed kernel advances the color on-device each replay. + """ + os.environ["ROCM_QUICK_REDUCE_QUANTIZATION"] = quant_mode + os.environ["ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16"] = "0" + device = torch.device(f"cuda:{rank}") + torch.cuda.set_device(device) + + # gloo (CPU) group: QuickAllReduce must attach to a non-NCCL group; it is + # used only for the one-time IPC-handle exchange. + dist.init_process_group( + backend="gloo", + init_method="tcp://127.0.0.1:29500", + rank=rank, + world_size=world_size, + ) + qr = QuickAllReduce(group=dist.group.WORLD, device=device) + assert not qr.disabled, ( + "quick-reduce unavailable on this arch/env " + "(needs ROCm MI300 gfx94/gfx95, even GPU count, same node, " + "and a non-NONE ROCM_QUICK_REDUCE_QUANTIZATION)." + ) + + N = 1 << 21 # 4 MB fp16, above the QR size threshold for the direct path + inp = torch.empty(N, dtype=torch.float16, device=device) + out = torch.empty(N, dtype=torch.float16, device=device) + + # Warmup, then capture a graph with EXACTLY ONE quick-reduce. + inp.fill_(1.0) + qr.quick_all_reduce(inp, out=out) + torch.cuda.synchronize() + dist.barrier() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + qr.quick_all_reduce(inp, out=out) + torch.cuda.synchronize() + dist.barrier() + + try: + for v in range(1, num_replays + 1): + inp.fill_(float(v)) # in-place: same value on every rank + dist.barrier() + graph.replay() + torch.cuda.synchronize() + dist.barrier() + expected = float(v * world_size) + assert torch.all(out.float() == expected), ( + f"[rank {rank}] round {v}: got {out.float().flatten()[0].item()}, " + f"expected {expected} (stale-flag corruption across replays)" + ) + finally: + dist.destroy_process_group() + + +class TestQuickreduceGraphReplay(CustomTestCase): + """Regression test for the QuickReduce CUDA-graph stale-flag bug. + + Unlike test_graph_allreduce (which captures a fresh graph each iteration + and replays it once), this captures a single graph and replays it many + times -- the exact scenario the on-device flag-color fix addresses. + """ + + TP_SIZES = [4, 8] + + @unittest.skipIf( + not qr_rocm_arch_available(), + "Only test Quick AllReduce on ROCm architectures >= gfx94*", + ) + def test_quick_allreduce_graph_replay(self): + for tp_size in self.TP_SIZES: + world_size = tp_size + if world_size > torch.cuda.device_count(): + continue + + multiprocessing.set_start_method("spawn", force=True) + timeout = 120 + processes = [] + for rank in range(tp_size): + p = multiprocessing.Process( + target=qr_graph_replay, 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 graph-replay hang detected after {timeout}s!" + ) + for rank, p in processes: + self.assertEqual( + p.exitcode, + 0, + f"QuickReduce graph-replay (tp={tp_size}, rank={rank}) " + f"produced wrong results -- stale-flag bug not fixed.", + ) + + class TestQuickreduceVariableInput(CustomTestCase): """ When the tensor parallelism is set to 4 or 8, frequent changes