[Bugfix] fix quickreduce acc error in cudagraph mode (#29508)

Signed-off-by: Haoyang Li <lihaoyang0109@gmail.com>

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<T, 4>;                                      \
@@ -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<T, 8>;                                      \
@@ -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<uint32_t> 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.
+    // The color now advances on-device inside the kernel; no host-side bump.
-    flag_color += divceil(N, grid);
   }
 };
 

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