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 3, 05:30 UTC
[ROCm] Add QuickReduce min-size override and codec threshold (#41675)
Signed-off-by: <>
akii96 committed
May 21, 2026 at 01:46 UTC
bde560ed6e1dc889debf68410ccbcb00b749513b
3 files changed
+258
-8
tests/distributed/test_quick_all_reduce.py
+178
@@ -11,7 +11,14 @@ 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.device_communicators.quick_all_reduce import (
15
+ KB,
16
+ MB,
17
+ QuickAllReduce,
18
+ QuickReduceRegime,
19
+)
20
from vllm.distributed.parallel_state import get_tp_group, graph_capture
21
+from vllm.envs import disable_envs_cache
22
from vllm.platforms import current_platform
23
24
from ..utils import (
@@ -28,6 +35,177 @@ for i, v in enumerate(test_sizes):
35
test_sizes[i] -= v % 8
36
37
38
+@pytest.fixture
39
+def envs_cache_disabled():
40
+ disable_envs_cache()
41
+ yield
42
+ disable_envs_cache()
43
+
44
+
45
+def _make_quick_allreduce_for_test(
46
+ min_size_mb: int | None = None,
47
+ quantization_min_size: int | None = None,
48
+) -> QuickAllReduce:
49
+ quick_reduce = QuickAllReduce.__new__(QuickAllReduce)
50
+ quick_reduce.disabled = False
51
+ quick_reduce.qr_max_size = 16 * MB
52
+ quick_reduce.qr_min_size = min_size_mb * MB if min_size_mb is not None else None
53
+ quick_reduce.qr_quant_level = QuickReduceRegime.INT4
54
+ quick_reduce.qr_quantization_min_size = quantization_min_size
55
+ quick_reduce.use_fp16_kernels = False
56
+ quick_reduce.world_size = 2
57
+ return quick_reduce
58
+
59
+
60
+def test_should_quick_allreduce_uses_builtin_min_size_when_unset():
61
+ quick_reduce = _make_quick_allreduce_for_test(min_size_mb=None)
62
+
63
+ below_builtin_min = torch.empty(MB // 4, dtype=torch.float16)
64
+ at_builtin_min = torch.empty(MB // 2, dtype=torch.float16)
65
+
66
+ assert not quick_reduce.should_quick_allreduce(below_builtin_min)
67
+ assert quick_reduce.should_quick_allreduce(at_builtin_min)
68
+
69
+
70
+def test_should_quick_allreduce_uses_min_size_override():
71
+ quick_reduce = _make_quick_allreduce_for_test(min_size_mb=0)
72
+
73
+ below_builtin_min = torch.empty(8, dtype=torch.float16)
74
+
75
+ assert quick_reduce.should_quick_allreduce(below_builtin_min)
76
+
77
+
78
+def test_quick_allreduce_min_size_env_unset(
79
+ monkeypatch: pytest.MonkeyPatch,
80
+ envs_cache_disabled,
81
+):
82
+ monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB", raising=False)
83
+
84
+ assert QuickAllReduce._get_qr_min_size(qr_max_size=16 * MB) is None
85
+
86
+
87
+def test_quick_allreduce_min_size_env_converts_mb_to_bytes(
88
+ monkeypatch: pytest.MonkeyPatch,
89
+ envs_cache_disabled,
90
+):
91
+ monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB", "4")
92
+
93
+ assert QuickAllReduce._get_qr_min_size(qr_max_size=16 * MB) == 4 * MB
94
+
95
+
96
+def test_quick_allreduce_min_size_env_rejects_negative(
97
+ monkeypatch: pytest.MonkeyPatch,
98
+ envs_cache_disabled,
99
+):
100
+ monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB", "-1")
101
+
102
+ with pytest.raises(ValueError, match="must be non-negative"):
103
+ QuickAllReduce._get_qr_min_size(qr_max_size=16 * MB)
104
+
105
+
106
+def test_quick_allreduce_min_size_env_allows_equal_to_max(
107
+ monkeypatch: pytest.MonkeyPatch,
108
+ envs_cache_disabled,
109
+):
110
+ monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB", "16")
111
+
112
+ assert QuickAllReduce._get_qr_min_size(qr_max_size=16 * MB) == 16 * MB
113
+
114
+
115
+def test_quick_allreduce_min_size_env_rejects_larger_than_max(
116
+ monkeypatch: pytest.MonkeyPatch,
117
+ envs_cache_disabled,
118
+):
119
+ monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB", "17")
120
+
121
+ with pytest.raises(ValueError, match="effective QuickReduce max size"):
122
+ QuickAllReduce._get_qr_min_size(qr_max_size=16 * MB)
123
+
124
+
125
+def test_quick_allreduce_quantization_min_size_env_unset(
126
+ monkeypatch: pytest.MonkeyPatch,
127
+ envs_cache_disabled,
128
+):
129
+ monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB", raising=False)
130
+
131
+ assert QuickAllReduce._get_qr_quantization_min_size() is None
132
+
133
+
134
+def test_quick_allreduce_quantization_min_size_env_converts_kb_to_bytes(
135
+ monkeypatch: pytest.MonkeyPatch,
136
+ envs_cache_disabled,
137
+):
138
+ monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB", "2048")
139
+
140
+ assert QuickAllReduce._get_qr_quantization_min_size() == 2048 * KB
141
+
142
+
143
+def test_quick_allreduce_quantization_min_size_env_rejects_negative(
144
+ monkeypatch: pytest.MonkeyPatch,
145
+ envs_cache_disabled,
146
+):
147
+ monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB", "-1")
148
+
149
+ with pytest.raises(ValueError, match="must be non-negative"):
150
+ QuickAllReduce._get_qr_quantization_min_size()
151
+
152
+
153
+def test_quick_allreduce_quantization_min_size_unset_uses_configured_codec():
154
+ quick_reduce = _make_quick_allreduce_for_test(quantization_min_size=None)
155
+ inp = torch.empty(8, dtype=torch.float16)
156
+
157
+ assert quick_reduce._get_qr_quant_level(inp) == QuickReduceRegime.INT4.value
158
+
159
+
160
+def test_quick_allreduce_quantization_min_size_uses_fp_below_threshold():
161
+ quick_reduce = _make_quick_allreduce_for_test(quantization_min_size=2048)
162
+ inp = torch.empty(1024 // 2, dtype=torch.float16)
163
+
164
+ assert quick_reduce._get_qr_quant_level(inp) == QuickReduceRegime.FP.value
165
+
166
+
167
+def test_quick_allreduce_quantization_min_size_uses_configured_codec_at_threshold():
168
+ quick_reduce = _make_quick_allreduce_for_test(quantization_min_size=2048)
169
+ inp = torch.empty(2048 // 2, dtype=torch.float16)
170
+
171
+ assert quick_reduce._get_qr_quant_level(inp) == QuickReduceRegime.INT4.value
172
+
173
+
174
+def test_quick_allreduce_quantization_min_size_does_not_change_eligibility():
175
+ quick_reduce = _make_quick_allreduce_for_test(quantization_min_size=2 * MB)
176
+
177
+ below_builtin_min = torch.empty(MB // 4, dtype=torch.float16)
178
+ at_builtin_min = torch.empty(MB // 2, dtype=torch.float16)
179
+
180
+ assert not quick_reduce.should_quick_allreduce(below_builtin_min)
181
+ assert quick_reduce.should_quick_allreduce(at_builtin_min)
182
+
183
+
184
+def test_quick_allreduce_passes_dynamic_quant_level(
185
+ monkeypatch: pytest.MonkeyPatch,
186
+):
187
+ quick_reduce = _make_quick_allreduce_for_test(quantization_min_size=2 * KB)
188
+ quick_reduce._ptr = object()
189
+ inp = torch.empty(KB // 2, dtype=torch.float16)
190
+ called_quant_level = None
191
+
192
+ def fake_qr_all_reduce(
193
+ fa,
194
+ inp,
195
+ out,
196
+ quant_level,
197
+ cast_bf2half,
198
+ ):
199
+ nonlocal called_quant_level
200
+ called_quant_level = quant_level
201
+
202
+ monkeypatch.setattr(ops, "qr_all_reduce", fake_qr_all_reduce)
203
+
204
+ quick_reduce.quick_all_reduce(inp)
205
+
206
+ assert called_quant_level == QuickReduceRegime.FP.value
207
+
208
+
209
@ray.remote(num_gpus=1, max_calls=1)
210
def graph_quickreduce(
211
monkeypatch: pytest.MonkeyPatch,
vllm/distributed/device_communicators/quick_all_reduce.py
+65
-8
@@ -39,7 +39,8 @@ class QuickReduceRegime(Enum):
39
NONE = 4
40
41
42
-MB = 1024 * 1024
42
+KB = 1024
43
+MB = 1024 * KB
44
45
46
class QuickAllReduce:
@@ -183,6 +184,7 @@ class QuickAllReduce:
184
)
185
return
186
self.qr_quant_level = QuickReduceRegime[regime_str]
187
+ self.qr_quantization_min_size = self._get_qr_quantization_min_size()
188
vllm_config = get_current_vllm_config_or_none()
189
if (
190
vllm_config is not None
@@ -215,11 +217,56 @@ class QuickAllReduce:
217
"lead to error or degradation to custom allreduce or rccl."
218
)
219
qr_max_size = qr_max_size * MB
220
+ effective_qr_max_size = (
221
+ qr_max_size if qr_max_size is not None else ops.qr_max_size()
222
+ )
223
+ qr_min_size = self._get_qr_min_size(effective_qr_max_size)
224
self._ptr = ops.init_custom_qr(self.rank, self.world_size, qr_max_size)
219
- self.qr_max_size = qr_max_size if qr_max_size is not None else ops.qr_max_size()
225
+ self.qr_max_size = effective_qr_max_size
226
+ self.qr_min_size = qr_min_size
227
+ if qr_min_size is not None:
228
+ logger.info(
229
+ "Custom quick allreduce: min size override = %d MB",
230
+ qr_min_size // MB,
231
+ )
232
+ if self.qr_quantization_min_size is not None:
233
+ logger.info(
234
+ "Custom quick allreduce: quantization codec threshold = %d KB",
235
+ self.qr_quantization_min_size // KB,
236
+ )
237
self.create_shared_buffer()
238
self.disabled = False
239
240
+ @staticmethod
241
+ def _get_qr_min_size(qr_max_size: int | None) -> int | None:
242
+ qr_min_size = envs.VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB
243
+ if qr_min_size is None:
244
+ return None
245
+ if qr_min_size < 0:
246
+ raise ValueError(
247
+ "VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB must be non-negative, "
248
+ f"got {qr_min_size}"
249
+ )
250
+ qr_min_size *= MB
251
+ if qr_max_size is not None and qr_min_size > qr_max_size:
252
+ raise ValueError(
253
+ "VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB must be less than or "
254
+ "equal to the effective QuickReduce max size"
255
+ )
256
+ return qr_min_size
257
+
258
+ @staticmethod
259
+ def _get_qr_quantization_min_size() -> int | None:
260
+ quantization_min_size = envs.VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB
261
+ if quantization_min_size is None:
262
+ return None
263
+ if quantization_min_size < 0:
264
+ raise ValueError(
265
+ "VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB must be "
266
+ f"non-negative, got {quantization_min_size}"
267
+ )
268
+ return quantization_min_size * KB
269
+
270
def _rocm_arch_available(self):
271
if not current_platform.is_rocm():
272
return False
@@ -261,11 +308,12 @@ class QuickAllReduce:
308
dtype = inp.dtype
309
if self.use_fp16_kernels:
310
dtype = torch.float16
264
- return (
265
- inp_size <= self.qr_max_size
266
- and inp_size
267
- >= self._QR_MIN_SIZE[(dtype, self.world_size)][self.qr_quant_level.value]
268
- )
311
+ min_size = self.qr_min_size
312
+ if min_size is None:
313
+ min_size = self._QR_MIN_SIZE[(dtype, self.world_size)][
314
+ self.qr_quant_level.value
315
+ ]
316
+ return inp_size <= self.qr_max_size and inp_size >= min_size
317
318
def quick_all_reduce(self, inp: torch.Tensor, *, out: torch.Tensor = None):
319
"""Performs an out-of-place custom quick all reduce."""
@@ -274,10 +322,19 @@ class QuickAllReduce:
322
if out is None:
323
out = torch.empty_like(inp)
324
ops.qr_all_reduce(
277
- self._ptr, inp, out, self.qr_quant_level.value, self.use_fp16_kernels
325
+ self._ptr, inp, out, self._get_qr_quant_level(inp), self.use_fp16_kernels
326
)
327
return out
328
329
+ def _get_qr_quant_level(self, inp: torch.Tensor) -> int:
330
+ quantization_min_size = self.qr_quantization_min_size
331
+ if (
332
+ quantization_min_size is not None
333
+ and inp.numel() * inp.element_size() < quantization_min_size
334
+ ):
335
+ return QuickReduceRegime.FP.value
336
+ return self.qr_quant_level.value
337
+
338
def close(self):
339
if not self.disabled and getattr(self, "_ptr", None):
340
if ops is not None:
vllm/envs.py
+15
@@ -209,6 +209,8 @@ if TYPE_CHECKING:
209
] = "NONE"
210
VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16: bool = True
211
VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB: int | None = None
212
+ VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB: int | None = None
213
+ VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB: int | None = None
214
VLLM_MORIIO_CONNECTOR_READ_MODE: bool = False
215
VLLM_MORIIO_QP_PER_TRANSFER: int = 1
216
VLLM_MORIIO_POST_BATCH_SIZE: int = -1
@@ -1122,6 +1124,19 @@ environment_variables: dict[str, Callable[[], Any]] = {
1124
"VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB": lambda: maybe_convert_int(
1125
os.environ.get("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", None)
1126
),
1127
+ # Custom quick allreduce kernel for MI3* cards.
1128
+ # Controls the minimum allowed number of data bytes(MB) required to use
1129
+ # custom quick allreduce communication.
1130
+ # If unset, use the built-in threshold table.
1131
+ "VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB": lambda: maybe_convert_int(
1132
+ os.environ.get("VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB", None)
1133
+ ),
1134
+ # Controls the minimum tensor size (KB, where 1 KB = 1024 bytes) required
1135
+ # to use the configured QuickReduce codec. Smaller tensors use FP
1136
+ # QuickReduce. This does not affect QuickReduce eligibility.
1137
+ "VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB": lambda: maybe_convert_int(
1138
+ os.environ.get("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB", None)
1139
+ ),
1140
# Divisor for dynamic query scale factor calculation for FP8 KV Cache
1141
"Q_SCALE_CONSTANT": lambda: int(os.getenv("Q_SCALE_CONSTANT", "200")),
1142
# Divisor for dynamic key scale factor calculation for FP8 KV Cache