A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Bugfix][Rocm] fix qr error when different inp shape (#25892)
Signed-off-by: Haoyang Li <lihaoyang0109@gmail.com> Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: ilmarkov <markovilya197@gmail.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>
haoyangli-amd committed
Oct 14, 2025 at 01:04 UTC
134f70b3eddf05f01f55ecee9c2a14ec0732e8b6
3 files changed
+96
-11
csrc/quickreduce/quick_reduce.h
+6
-5
@@ -22,13 +22,14 @@ 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 flag_color,
26
+ int64_t data_size_per_phase) {
27
int block = blockIdx.x;
28
int grid = gridDim.x;
29
30
while (block < num_blocks) {
31
AllReduceKernel::run(A, B, N, block, rank, dbuffer_list, data_offset,
31
- flag_color);
32
+ flag_color, data_size_per_phase);
33
block += grid;
34
flag_color++;
35
}
@@ -41,21 +42,21 @@ allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks,
42
hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
43
dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
44
num_blocks, rank, dbuffer_list, data_offset, \
44
- flag_color); \
45
+ flag_color, this->kMaxProblemSize); \
46
} else if (world_size == 4) { \
47
using LineCodec = __codec<T, 4>; \
48
using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
49
hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
50
dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
51
num_blocks, rank, dbuffer_list, data_offset, \
51
- flag_color); \
52
+ flag_color, this->kMaxProblemSize); \
53
} else if (world_size == 8) { \
54
using LineCodec = __codec<T, 8>; \
55
using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
56
hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
57
dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
58
num_blocks, rank, dbuffer_list, data_offset, \
58
- flag_color); \
59
+ flag_color, this->kMaxProblemSize); \
60
}
61
62
enum QuickReduceQuantLevel {
csrc/quickreduce/quick_reduce_impl.cuh
+3
-6
@@ -553,13 +553,12 @@ struct AllReduceTwoshot {
553
int const rank, // rank index
554
uint8_t** __restrict__ buffer_list, // communication buffers
555
uint32_t const data_offset, // offset to start of the data buffer
556
- uint32_t flag_color) {
556
+ uint32_t flag_color, int64_t data_size_per_phase) {
557
// Topology
558
int thread = threadIdx.x + threadIdx.y * kWavefront;
559
uint8_t* rank_buffer = buffer_list[rank];
560
Codec codec(thread, rank);
561
int block_id = blockIdx.x;
562
- int grid_size = gridDim.x;
562
// --------------------------------------------------------
563
// Read input into registers
564
int32x4_t tA[kAtoms];
@@ -588,12 +587,10 @@ struct AllReduceTwoshot {
587
// rank responsible for this segment.
588
uint32_t comm_data0_offset =
589
data_offset + block_id * Codec::kTransmittedTileSize;
591
- uint32_t comm_data1_offset =
592
- grid_size * Codec::kTransmittedTileSize + comm_data0_offset;
590
+ uint32_t comm_data1_offset = data_size_per_phase + comm_data0_offset;
591
592
uint32_t comm_flags0_offset = block_id * (kWorldSize * sizeof(uint32_t));
595
- uint32_t comm_flags1_offset =
596
- grid_size * (kWorldSize * sizeof(uint32_t)) + comm_flags0_offset;
593
+ uint32_t comm_flags1_offset = (data_offset / 2) + comm_flags0_offset;
594
595
for (int r = 0; r < kWorldSize; r++) {
596
int32x4_t* send_buffer =
tests/distributed/test_quick_all_reduce.py
+87
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import multiprocessing
5
import random
6
7
import pytest
@@ -8,6 +9,7 @@ import ray
9
import torch
10
import torch.distributed as dist
11
12
+from vllm import _custom_ops as ops
13
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa
14
from vllm.distributed.parallel_state import get_tp_group, graph_capture
15
from vllm.platforms import current_platform
@@ -134,3 +136,88 @@ def test_custom_quick_allreduce(
136
monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quant_mode)
137
138
multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target)
139
+
140
+
141
+def qr_variable_input(rank, world_size):
142
+ """
143
+ When the tensor parallelism is set to 4 or 8, frequent changes
144
+ in the input shape can cause QuickReduce to hang (this issue
145
+ has been observed with the gpt_oss model).
146
+ """
147
+ device = torch.device(f"cuda:{rank}")
148
+ torch.cuda.set_device(device)
149
+ qr_max_size = None # MB
150
+ _ptr = ops.init_custom_qr(rank, world_size, qr_max_size)
151
+ ranks = []
152
+ for i in range(world_size):
153
+ ranks.append(i)
154
+ dist.init_process_group(
155
+ backend="nccl",
156
+ init_method="tcp://127.0.0.1:29500",
157
+ rank=rank,
158
+ world_size=world_size,
159
+ )
160
+ cpu_group = torch.distributed.new_group(ranks, backend="nccl")
161
+
162
+ handle = ops.qr_get_handle(_ptr)
163
+ world_size = dist.get_world_size(group=cpu_group)
164
+ handles = [None] * world_size
165
+ dist.all_gather_object(handles, handle, group=cpu_group)
166
+ ops.qr_open_handles(_ptr, handles)
167
+
168
+ num = 1
169
+ s1 = 1024
170
+ while num < 50000: # 50000 is sufficient to identify issues.
171
+ dtype = torch.float16
172
+ if num % 2 == 0:
173
+ s2 = 1024
174
+ inp1 = torch.zeros(
175
+ (s1, s2), dtype=dtype, device=torch.cuda.current_device()
176
+ )
177
+ else:
178
+ s2 = 2048
179
+ inp1 = torch.ones((s1, s2), dtype=dtype, device=torch.cuda.current_device())
180
+ result = torch.empty_like(inp1)
181
+ # FP = 0 INT8 = 1 INT6 = 2 INT4 = 3 NONE = 4
182
+ ops.qr_all_reduce(_ptr, inp1, result, 3, cast_bf2half=True)
183
+ try:
184
+ if inp1[0, 0] == 0:
185
+ assert torch.all(result == 0)
186
+ else:
187
+ assert torch.all(result == world_size)
188
+ except AssertionError:
189
+ print("Assertion failed! Allreduce results are incorrect.")
190
+ raise
191
+ num += 1
192
+
193
+
194
+@pytest.mark.skipif(
195
+ not current_platform.is_rocm(), reason="only test quick allreduce for rocm"
196
+)
197
+@pytest.mark.parametrize("tp_size", [4, 8])
198
+@pytest.mark.parametrize("pipeline_parallel_size", [1])
199
+def test_custom_quick_allreduce_variable_input(tp_size, pipeline_parallel_size):
200
+ world_size = tp_size * pipeline_parallel_size
201
+ if world_size > torch.cuda.device_count():
202
+ pytest.skip("Not enough GPUs to run the test.")
203
+
204
+ multiprocessing.set_start_method("spawn", force=True)
205
+ # 60s is enough
206
+ timeout = 60
207
+ processes = []
208
+ for rank in range(tp_size):
209
+ p = multiprocessing.Process(target=qr_variable_input, args=(rank, tp_size))
210
+ p.start()
211
+ processes.append((rank, p))
212
+ for rank, p in processes:
213
+ p.join(timeout=timeout)
214
+ if p.is_alive():
215
+ for r, proc in processes:
216
+ if proc.is_alive():
217
+ proc.terminate()
218
+ proc.join()
219
+ raise RuntimeError(f"QuickReduce hang detected after {timeout} seconds!")
220
+
221
+
222
+if __name__ == "__main__":
223
+ test_custom_quick_allreduce_variable_input(tp_size=4, pipeline_parallel_size=1)