[communication] [bugfix] fix quickreduce acc error in cudagraph mode (#46913)
Signed-off-by: Haoyang Li <lihaoyang0109@gmail.com> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
haoyangli0109 committed
Jul 27, 2026 at 16:38 UTC
7f599d78546819948c32f2b23d913507bbb38875
2 files changed
+154
-8
csrc/quickreduce/quick_reduce.h
+30
-8
@@ -22,17 +22,25 @@ template <typename AllReduceKernel, typename T>
22
__global__ __quickreduce_launch_bounds_two_shot__ static void
23
allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks,
24
int rank, uint8_t** dbuffer_list,
25
- uint32_t data_offset, uint32_t flag_color,
25
+ uint32_t data_offset, uint32_t* d_flag_counters,
26
int64_t data_size_per_phase) {
27
int block = blockIdx.x;
28
int grid = gridDim.x;
29
30
+ // Load this block's counter from device memory and advance it on-device,
31
+ // so the color keeps changing across graph replays instead of being frozen.
32
+ uint32_t flag_color = d_flag_counters[blockIdx.x];
33
+
34
while (block < num_blocks) {
35
AllReduceKernel::run(A, B, N, block, rank, dbuffer_list, data_offset,
36
flag_color, data_size_per_phase);
37
block += grid;
38
flag_color++;
39
}
40
+ // All threads compute the same final value; one writer per block is enough.
41
+ if (threadIdx.x == 0 && threadIdx.y == 0) {
42
+ d_flag_counters[blockIdx.x] = flag_color;
43
+ }
44
}
45
46
#define TWOSHOT_DISPATCH(__codec) \
@@ -42,21 +50,21 @@ allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks,
50
hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
51
dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
52
num_blocks, rank, dbuffer_list, data_offset, \
45
- flag_color, this->kMaxProblemSize); \
53
+ d_flag_counters, this->kMaxProblemSize); \
54
} else if (world_size == 4) { \
55
using LineCodec = __codec<T, 4>; \
56
using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
57
hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
58
dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
59
num_blocks, rank, dbuffer_list, data_offset, \
52
- flag_color, this->kMaxProblemSize); \
60
+ d_flag_counters, this->kMaxProblemSize); \
61
} else if (world_size == 8) { \
62
using LineCodec = __codec<T, 8>; \
63
using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
64
hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
65
dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
66
num_blocks, rank, dbuffer_list, data_offset, \
59
- flag_color, this->kMaxProblemSize); \
67
+ d_flag_counters, this->kMaxProblemSize); \
68
}
69
70
// INT3 only retains good performance on TP2 (world_size == 2). On TP4/TP8
@@ -69,7 +77,7 @@ allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks,
77
hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
78
dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
79
num_blocks, rank, dbuffer_list, data_offset, \
72
- flag_color, this->kMaxProblemSize); \
80
+ d_flag_counters, this->kMaxProblemSize); \
81
} else { \
82
throw std::runtime_error( \
83
"INT3 quick all-reduce is only supported for world_size == 2 " \
@@ -94,7 +102,7 @@ struct DeviceComms {
102
static int constexpr kMaxWorldSize = 8;
103
104
bool initialized = false;
97
- uint32_t flag_color = 1;
105
+ uint32_t* d_flag_counters = nullptr;
106
int world_size;
107
int rank;
108
@@ -128,6 +136,16 @@ struct DeviceComms {
136
// Clear the flags buffer.
137
HIP_CHECK(hipMemset(dbuffer, 0, flags_buffer_size));
138
139
+ // One flag-color counter per block, advanced by the kernel. Start at 1
140
+ // to stay clear of the flags buffer we just zeroed.
141
+ HIP_CHECK(hipMalloc(&d_flag_counters, kMaxNumBlocks * sizeof(uint32_t)));
142
+ {
143
+ std::vector<uint32_t> init_color(kMaxNumBlocks, 1u);
144
+ HIP_CHECK(hipMemcpy(d_flag_counters, init_color.data(),
145
+ kMaxNumBlocks * sizeof(uint32_t),
146
+ hipMemcpyHostToDevice));
147
+ }
148
+
149
// Device-side list of IPC buffers.
150
buffer_list.resize(world_size);
151
HIP_CHECK(hipMalloc(&dbuffer_list, world_size * sizeof(uint8_t*)));
@@ -144,6 +162,12 @@ struct DeviceComms {
162
hipIpcMemHandle_t const get_handle() { return buffer_ipc_handle; }
163
164
void destroy() {
165
+ // Allocated before `initialized` flips true, so free it on its own guard
166
+ // to avoid a leak if init fails partway through.
167
+ if (d_flag_counters) {
168
+ HIP_CHECK(hipFree(d_flag_counters));
169
+ d_flag_counters = nullptr;
170
+ }
171
if (initialized) {
172
for (int i = 0; i < world_size; i++) {
173
if (i != rank) {
@@ -211,8 +235,6 @@ struct DeviceComms {
235
break;
236
}
237
HIP_CHECK(cudaGetLastError());
214
- // Rotate the flag color.
215
- flag_color += divceil(N, grid);
238
}
239
};
240
tests/distributed/test_rocm_quick_reduce.py
+124
@@ -23,6 +23,8 @@ from vllm.platforms import current_platform
23
from vllm.transformers_utils.repo_utils import hf_api
24
from vllm.utils.network_utils import get_open_port
25
26
+from ..utils import multi_gpu_test
27
+
28
pytestmark = pytest.mark.skipif(
29
not current_platform.is_rocm(),
30
reason="ROCm-only quick-reduce tests",
@@ -151,6 +153,116 @@ def _run_two_gpu_quick_allreduce_test(
153
)
154
155
156
+CUDAGRAPH_WORLD_SIZE = 2
157
+CUDAGRAPH_ROUNDS = 10
158
+CUDAGRAPH_NUM_ELEMENTS = 1 << 21 # 2M fp16 = 4 MB, above the quick-reduce thresholds
159
+
160
+
161
+def _quick_allreduce_cudagraph_worker(
162
+ rank: int,
163
+ world_size: int,
164
+ port: int,
165
+ quant_level: str,
166
+):
167
+ # FP keeps the all-reduce bit-exact for small fp16 integers, so every
168
+ # replayed round can be checked exactly.
169
+ os.environ["VLLM_ROCM_QUICK_REDUCE_QUANTIZATION"] = quant_level
170
+ os.environ["VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16"] = "0"
171
+ _log(f"cudagraph worker start: rank={rank} quant={quant_level}")
172
+
173
+ device = torch.device(f"cuda:{rank}")
174
+ torch.accelerator.set_device_index(device)
175
+ dist.init_process_group(
176
+ backend="gloo",
177
+ init_method=f"tcp://127.0.0.1:{port}",
178
+ rank=rank,
179
+ world_size=world_size,
180
+ )
181
+
182
+ qar = None
183
+ try:
184
+ from vllm.distributed.device_communicators.quick_all_reduce import (
185
+ QuickAllReduce,
186
+ )
187
+
188
+ qar = QuickAllReduce(group=dist.GroupMember.WORLD, device=rank)
189
+ assert not qar.disabled
190
+
191
+ N = CUDAGRAPH_NUM_ELEMENTS
192
+ inp = torch.empty(N, dtype=torch.float16, device=device)
193
+ out = torch.empty(N, dtype=torch.float16, device=device)
194
+ assert qar.should_quick_allreduce(inp)
195
+
196
+ # Every rank contributes the same value v in a round, so the true
197
+ # cross-rank all-reduce sum is simply world_size * v.
198
+ def expected(v):
199
+ return float(world_size * v)
200
+
201
+ if rank == 0:
202
+ print(
203
+ f"[repro] world_size={world_size} elems={N} regime={quant_level} fp16",
204
+ flush=True,
205
+ )
206
+
207
+ # Warmup, then capture a graph with EXACTLY ONE quick-reduce (isolated
208
+ # qr, so it is the sole writer of its flag slot -- the condition that
209
+ # triggers the stale-flag bug).
210
+ inp.fill_(1.0)
211
+ qar.quick_all_reduce(inp, out=out)
212
+ torch.accelerator.synchronize()
213
+ dist.barrier()
214
+
215
+ g = torch.cuda.CUDAGraph()
216
+ with torch.cuda.graph(g):
217
+ qar.quick_all_reduce(inp, out=out)
218
+ torch.accelerator.synchronize()
219
+ dist.barrier()
220
+
221
+ for v in range(CUDAGRAPH_ROUNDS):
222
+ inp.fill_(float(v)) # in-place: same value on every rank
223
+ dist.barrier()
224
+ g.replay()
225
+ torch.accelerator.synchronize()
226
+ dist.barrier()
227
+ got = out.float()
228
+ expect = expected(v)
229
+ if rank == 0:
230
+ print(f"round {v}: got={got[:10]}, expected={expect}", flush=True)
231
+ mismatch = ~torch.isclose(got, torch.full_like(got, expect))
232
+ num_mismatch = int(mismatch.sum().item())
233
+ assert num_mismatch == 0, (
234
+ f"rank={rank} round={v} expected={expect} "
235
+ f"mismatched {num_mismatch}/{got.numel()} elements; "
236
+ f"unique wrong values={torch.unique(got[mismatch])[:8].tolist()}"
237
+ )
238
+ _log(f"cudagraph worker complete: rank={rank} rounds={CUDAGRAPH_ROUNDS}")
239
+ finally:
240
+ if qar is not None:
241
+ qar.close()
242
+ if dist.is_initialized():
243
+ dist.destroy_process_group()
244
+
245
+
246
+def _run_cudagraph_replay_test(*, world_size: int, quant_level: str):
247
+ _log(f"launch {world_size}-GPU cudagraph replay case: quant={quant_level}")
248
+ ctx = mp.get_context("spawn")
249
+ port = get_open_port()
250
+ procs = []
251
+
252
+ for rank in range(world_size):
253
+ proc = ctx.Process(
254
+ target=_quick_allreduce_cudagraph_worker,
255
+ args=(rank, world_size, port, quant_level),
256
+ )
257
+ proc.start()
258
+ procs.append(proc)
259
+
260
+ for proc in procs:
261
+ proc.join(timeout=120)
262
+ assert proc.exitcode == 0, f"worker exited with code {proc.exitcode}"
263
+ _log(f"finished {world_size}-GPU cudagraph replay case: quant={quant_level}")
264
+
265
+
266
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
267
E2E_PREFILL_TOKENS = 1024
268
E2E_MAX_MODEL_LEN = 1536
@@ -729,6 +841,18 @@ def test_quick_allreduce_two_gpu_correctness(quant_level):
841
)
842
843
844
+@multi_gpu_test(num_gpus=CUDAGRAPH_WORLD_SIZE)
845
+def test_quick_allreduce_cudagraph_replay():
846
+ # Regression test for the stale flag_color bug: a quick-reduce captured in a
847
+ # CUDA graph must return correct results on every replay, not the data from
848
+ # a previous round.
849
+ _log("cudagraph replay case")
850
+ _run_cudagraph_replay_test(
851
+ world_size=CUDAGRAPH_WORLD_SIZE,
852
+ quant_level="FP",
853
+ )
854
+
855
+
856
@pytest.mark.skipif(
857
current_platform.device_count() < WORLD_SIZE,
858
reason="requires 2 ROCm GPUs",