@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 1744adc256

[ROCM] [Communication] Add INT3 quantization method for quickreduce (#45666)

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

haoyangli0109 committed Jun 25, 2026 at 23:14 UTC 1744adc256b81b8a7d2371a5e21f88174d60cc93
7 files changed +264 -20
csrc/custom_quickreduce.cu
+10
@@ -97,18 +97,28 @@ int64_t qr_max_size() {
97 cast_bf2half>; \
98 template struct quickreduce::AllReduceTwoshot<T, Codec<T, 8>, cast_bf2half>;
99
100 + // INT3 (CodecQ3) is restricted to TP2 only, so we only instantiate the
101 + // world_size == 2 kernel for it.
102 + #define INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(T, Codec, cast_bf2half) \
103 + template struct quickreduce::AllReduceTwoshot<T, Codec<T, 2>, cast_bf2half>;
104 +
105 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecFP, false)
106 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ4, false)
107 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ6, false)
108 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ8, false)
109 +INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(quickreduce::nv_bfloat16,
110 + quickreduce::CodecQ3, false)
111 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecFP, true)
112 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ4, true)
113 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ6, true)
114 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ8, true)
115 +INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(quickreduce::nv_bfloat16,
116 + quickreduce::CodecQ3, true)
117
118 INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecFP, false)
119 INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ4, false)
120 INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ6, false)
121 INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ8, false)
122 +INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(half, quickreduce::CodecQ3, false)
123
124 #endif // USE_ROCM
\ No newline at end of file
csrc/quickreduce/base.h
+23
@@ -283,6 +283,29 @@ __quickreduce_device_inline__ int packed_rcp<nv_bfloat16>(int a) {
283 return R.i;
284 }
285
286 +template <typename T>
287 +__quickreduce_device_inline__ int packed_from_int16_pair(int16_t low,
288 + int16_t high);
289 +
290 +template <>
291 +__quickreduce_device_inline__ int packed_from_int16_pair<half>(int16_t low,
292 + int16_t high) {
293 + // Convert two signed integers to one fp16x2 packed 32-bit lane.
294 + half2 h = __halves2half2(__int2half_rn(static_cast<int>(low)),
295 + __int2half_rn(static_cast<int>(high)));
296 + return __builtin_bit_cast(int, h);
297 +}
298 +
299 +template <>
300 +__quickreduce_device_inline__ int packed_from_int16_pair<nv_bfloat16>(
301 + int16_t low, int16_t high) {
302 + // Convert two signed integers to one bf16x2 packed 32-bit lane.
303 + nv_bfloat16 bf_low = __float2bfloat16(static_cast<float>(low));
304 + nv_bfloat16 bf_high = __float2bfloat16(static_cast<float>(high));
305 + nv_bfloat162 bf2 = __halves2bfloat162(bf_low, bf_high);
306 + return *reinterpret_cast<int*>(&bf2);
307 +}
308 +
309 // changes dtype
310 __quickreduce_device_inline__ float T2float_cast(half a) {
311 return __half2float(a);
csrc/quickreduce/quick_reduce.h
+26 -4
@@ -59,11 +59,30 @@ allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks,
59 flag_color, this->kMaxProblemSize); \
60 }
61
62 +// INT3 only retains good performance on TP2 (world_size == 2). On TP4/TP8
63 +// the 3-bit codec's pack/unpack overhead outweighs the reduced communication
64 +// volume, so INT3 is restricted to a TP2-only dispatch here.
65 +#define TWOSHOT_DISPATCH_TP2_ONLY(__codec) \
66 + if (world_size == 2) { \
67 + using LineCodec = __codec<T, 2>; \
68 + using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
69 + hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
70 + dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
71 + num_blocks, rank, dbuffer_list, data_offset, \
72 + flag_color, this->kMaxProblemSize); \
73 + } else { \
74 + throw std::runtime_error( \
75 + "INT3 quick all-reduce is only supported for world_size == 2 " \
76 + "(TP2); use INT4/NONE for larger world sizes."); \
77 + }
78 +
79 enum QuickReduceQuantLevel {
63 - F16 = 0,
64 - INT8 = 1,
65 - INT6 = 2,
66 - INT4 = 3,
80 + // Keep these ids in sync with Python QuickReduceRegime enum.
81 + F16 = 0, // full-precision fp16/bf16 communication
82 + INT8 = 1, // symmetric int8 + per-block scale
83 + INT6 = 2, // symmetric int6 + per-block scale
84 + INT4 = 3, // symmetric int4 + per-block scale
85 + INT3 = 4, // symmetric int3 + per-block scale (TP2 only)
86 };
87
88 struct DeviceComms {
@@ -184,6 +203,9 @@ struct DeviceComms {
203 case QuickReduceQuantLevel::INT4:
204 TWOSHOT_DISPATCH(CodecQ4)
205 break;
206 + case QuickReduceQuantLevel::INT3:
207 + TWOSHOT_DISPATCH_TP2_ONLY(CodecQ3)
208 + break;
209 default:
210 TWOSHOT_DISPATCH(CodecFP)
211 break;
csrc/quickreduce/quick_reduce_impl.cuh
+163 -1
@@ -2,6 +2,7 @@
2
3 #include <hip/hip_runtime.h>
4 #include "base.h"
5 +#include <hip/hip_fp16.h>
6
7 namespace quickreduce {
8
@@ -206,6 +207,168 @@ struct CodecQ4 : public CodecBase {
207 }
208 };
209
210 +// Int3 symmetric quantization codec.
211 +// We quantize the FP16 data to block-scaled Int3 in blocks of 4 *
212 +// kThreadGroupSize. Uniform symmetric quantization (round-to-int + clip),
213 +// matching the structure of CodecQ4. Signed range is [-4, +3].
214 +template <typename T, int world_size>
215 +struct CodecQ3 : public CodecBase {
216 + static constexpr int kWorldSize = world_size;
217 +
218 + // Layout per quantization block (32 values = 8 threads * 4 fp16x2 lanes):
219 + // - each thread owns 8 values and writes:
220 + // * q2 payload : 8 * 2 bits -> uint16 (2 bytes)
221 + // * q1 payload : 8 * 1 bit -> uint8 (1 byte)
222 + // - one scale is shared per 32 values and written by group leader.
223 + //
224 + // kRankTileStride is split as:
225 + // [0 .. 511] : q2 payload region (256 threads * 2 bytes)
226 + // [512 .. 767] : q1 payload region (256 threads * 1 byte)
227 + // [768 .. 895] : scale region (32 groups * 4 bytes)
228 + static constexpr int kRankAtoms = kAtoms / kWorldSize;
229 + static constexpr int kRankTileStride = 896;
230 + static constexpr int kRankTileQ1Offset = 512;
231 + static constexpr int kRankTileScaleOffset = 768;
232 + static constexpr int kRankTransmittedTileSize = kRankTileStride * kRankAtoms;
233 + static_assert(kRankTransmittedTileSize % 16 == 0,
234 + "kRankTransmittedTileSize must be 16B aligned.");
235 +
236 + static constexpr int kRankBufferTileStride =
237 + kRankTileStride / sizeof(int32x4_t);
238 +
239 + static constexpr int kTransmittedTileSize =
240 + kRankTransmittedTileSize * kWorldSize;
241 +
242 + // {-1/4.0h, -1/4.0h}, f16x2_t / bf16x2_t. Sign-flipped so absmax maps
243 + // to -4; the sign cancels with decoding_scale on the recv side.
244 + static constexpr int kScaleFactor =
245 + std::is_same<T, half>::value ? 0xB400B400 : 0xBE80BE80;
246 +
247 + // {1e-7, 1e-7}, f16x2_t
248 + static constexpr int kScaleEpsilon =
249 + std::is_same<T, half>::value ? 0x00010001 : 0x33D733D7;
250 +
251 + // {-4, -4}, f16x2_t / bf16x2_t
252 + static constexpr int kRangeMin =
253 + std::is_same<T, half>::value ? 0xC400C400 : 0xC080C080;
254 +
255 + // {+3, +3}, f16x2_t / bf16x2_t
256 + static constexpr int kRangeMax =
257 + std::is_same<T, half>::value ? 0x42004200 : 0x40404040;
258 +
259 + // {+4, +4}, int16x2_t -- shifts signed [-4, +3] to unsigned [0, 7].
260 + static constexpr int kRangeBias = 0x00040004;
261 +
262 + __quickreduce_device_inline__ CodecQ3(int thread, int rank)
263 + : CodecBase(thread, rank) {}
264 +
265 + __quickreduce_device_inline__ void send(int32x4_t* __restrict__ send_buffer,
266 + const int32x4_t* __restrict__ data) {
267 + for (int k = 0; k < kRankAtoms; k++) {
268 + int32x4_t const atom = data[k];
269 +
270 + // 1) Per-group dynamic scale (shared across 32 values).
271 + int wblockmax = group_abs_max<T>(atom);
272 + int decoding_scale = packed_mul<T>(wblockmax, kScaleFactor);
273 + int encoding_scale = packed_add<T>(decoding_scale, kScaleEpsilon);
274 + encoding_scale = packed_rcp<T>(encoding_scale);
275 +
276 + // 2) Scale + clip to signed int3 range [-4, +3].
277 + int32x4_t w;
278 + for (int i = 0; i < 4; i++) {
279 + w[i] = packed_mul<T>(atom[i], encoding_scale);
280 + w[i] = packed_max<T>(w[i], kRangeMin);
281 + w[i] = packed_min<T>(w[i], kRangeMax);
282 + }
283 +
284 + // 3) Round to integer and bias to unsigned domain [0, 7].
285 + int32x4_t q;
286 + {
287 + int16_t* qi = reinterpret_cast<int16_t*>(&q);
288 + T* wh = reinterpret_cast<T*>(&w);
289 + for (int i = 0; i < 8; i++) qi[i] = (int16_t)rintf(T2float_cast(wh[i]));
290 +
291 + for (int i = 0; i < 4; i++) {
292 + q[i] = packed_add<int16_t>(q[i], kRangeBias);
293 + }
294 + }
295 +
296 + // 4) Split each 3-bit unsigned value into low-2-bit and high-1-bit
297 + // halves, packed into one uint16 (low 2 bits per value) plus one
298 + // uint8 (high 1 bit per value).
299 + uint16_t q2w = 0;
300 + uint8_t q1w = 0;
301 + {
302 + int16_t* tw = reinterpret_cast<int16_t*>(&q);
303 +#pragma unroll
304 + for (int i = 0; i < 8; i++) {
305 + uint32_t v = static_cast<uint32_t>(tw[i]) & 0x7u;
306 + q2w |= static_cast<uint16_t>((v & 0x3u) << (i * 2));
307 + q1w |= static_cast<uint8_t>(((v >> 2) & 0x1u) << i);
308 + }
309 + }
310 +
311 + uint8_t* atom_ptr =
312 + reinterpret_cast<uint8_t*>(send_buffer + k * kRankBufferTileStride);
313 + uint16_t* q2w_ptr = reinterpret_cast<uint16_t*>(atom_ptr) + thread;
314 + uint8_t* q1w_ptr =
315 + reinterpret_cast<uint8_t*>(atom_ptr + kRankTileQ1Offset) + thread;
316 + int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
317 + (thread / 8);
318 +
319 + __builtin_nontemporal_store(q2w, q2w_ptr);
320 + *q1w_ptr = q1w;
321 + if (threadIdx.x == group_leader) {
322 + __builtin_nontemporal_store(decoding_scale, qs_ptr);
323 + }
324 + }
325 + }
326 +
327 + __quickreduce_device_inline__ void recv(int32x4_t** __restrict__ recv_buffer,
328 + int32x4_t* __restrict__ data) {
329 + for (int k = 0; k < kRankAtoms; k++) {
330 + uint8_t* atom_ptr = reinterpret_cast<uint8_t*>(*recv_buffer);
331 + uint16_t* q2w_ptr = reinterpret_cast<uint16_t*>(atom_ptr) + thread;
332 + uint8_t* q1w_ptr =
333 + reinterpret_cast<uint8_t*>(atom_ptr + kRankTileQ1Offset) + thread;
334 + int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
335 + (thread / 8);
336 +
337 + uint16_t q2w = __builtin_nontemporal_load(q2w_ptr);
338 + uint8_t q1w = *q1w_ptr;
339 + int qs = __builtin_nontemporal_load(qs_ptr);
340 +
341 + *recv_buffer += kRankBufferTileStride;
342 +
343 + // Unpack unsigned values [0, 7] then shift back to signed domain
344 + // [-4, +3] by adding kRangeMin.
345 + int32x4_t w;
346 + {
347 + int16_t qv[8];
348 +#pragma unroll
349 + for (int i = 0; i < 8; i++) {
350 + uint32_t low2 = (q2w >> (2 * i)) & 0x3u;
351 + uint32_t high1 = (q1w >> i) & 0x1u;
352 + qv[i] = static_cast<int16_t>(low2 | (high1 << 2));
353 + }
354 +
355 +#pragma unroll
356 + for (int i = 0; i < 4; i++) {
357 + int qpack = packed_from_int16_pair<T>(qv[2 * i], qv[2 * i + 1]);
358 + w[i] = packed_add<T>(qpack, kRangeMin);
359 + }
360 + }
361 +
362 + // Apply decode scale to reconstruct fp16/bf16 lanes.
363 + for (int i = 0; i < 4; i++) {
364 + w[i] = packed_mul<T>(w[i], qs);
365 + }
366 +
367 + data[k] = w;
368 + }
369 + }
370 +};
371 +
372 // Int6 symmetric quantization codec.
373 // We quantize the FP16 data to block-scaled Int6 in blocks of 4 *
374 // kThreadGroupSize.
@@ -377,7 +540,6 @@ struct CodecQ6 : public CodecBase {
540 w[i] = packed_mul<T>(w[i], qs);
541 }
542
380 - // That's pretty much it...
543 data[k] = w;
544 }
545 }
tests/distributed/test_quick_all_reduce.py
+2 -2
@@ -350,7 +350,7 @@ def bf16_cast_quickreduce(
350 @pytest.mark.skipif(
351 not current_platform.is_rocm(), reason="only test quick allreduce for rocm"
352 )
353 -@pytest.mark.parametrize("quant_mode", ["FP", "INT8", "INT6", "INT4"])
353 +@pytest.mark.parametrize("quant_mode", ["FP", "INT8", "INT6", "INT4", "INT3"])
354 @pytest.mark.parametrize("tp_size", [2])
355 @pytest.mark.parametrize("pipeline_parallel_size", [1, 2])
356 @pytest.mark.parametrize("test_target", [graph_quickreduce, eager_quickreduce])
@@ -438,7 +438,7 @@ def qr_variable_input(rank, world_size):
438 s2 = 2048
439 inp1 = torch.ones((s1, s2), dtype=dtype, device=device_idx)
440 result = torch.empty_like(inp1)
441 - # FP = 0 INT8 = 1 INT6 = 2 INT4 = 3 NONE = 4
441 + # FP = 0 INT8 = 1 INT6 = 2 INT4 = 3 INT3 = 4
442 ops.qr_all_reduce(_ptr, inp1, result, 3, cast_bf2half=True)
443 try:
444 if inp1[0, 0] == 0:
vllm/distributed/device_communicators/quick_all_reduce.py
+37 -10
@@ -28,11 +28,13 @@ from vllm.distributed.utils import is_weak_contiguous # noqa: E402, F401
28
29
30 class QuickReduceRegime(Enum):
31 + # Keep integer ids aligned with csrc/quickreduce/quick_reduce.h
32 FP = 0
33 INT8 = 1
34 INT6 = 2
35 INT4 = 3
35 - NONE = 4
36 + INT3 = 4
37 + NONE = 5
38
39
40 KB = 1024
@@ -43,14 +45,20 @@ class QuickAllReduce:
45 _SUPPORTED_WORLD_SIZES = [2, 4, 8]
46 _SUPPORTED_DTYPES = [torch.float16, torch.bfloat16]
47 # The following data is based on kernel tests.
46 - # In this order [FP, INT8, INT6, INT4].
48 + # In this order [FP, INT8, INT6, INT4, INT3].
49 _QR_MIN_SIZE = {
48 - (torch.float16, 2): [1 * MB, 2 * MB, 2 * MB, 1 * MB],
49 - (torch.float16, 4): [1 * MB, 16 * MB, 4 * MB, 2 * MB],
50 - (torch.float16, 8): [16 * MB, 4 * MB, 4 * MB, 2 * MB],
51 - (torch.bfloat16, 2): [2 * MB, 8 * MB, 8 * MB, 8 * MB],
52 - (torch.bfloat16, 4): [8 * MB, 64 * MB, 64 * MB, 16 * MB],
53 - (torch.bfloat16, 8): [16 * MB, 2048 * MB, 2048 * MB, 2048 * MB],
50 + (torch.float16, 2): [1 * MB, 2 * MB, 2 * MB, 1 * MB, 1 * MB],
51 + (torch.float16, 4): [1 * MB, 16 * MB, 4 * MB, 2 * MB, 2 * MB],
52 + (torch.float16, 8): [16 * MB, 4 * MB, 4 * MB, 2 * MB, 2 * MB],
53 + (torch.bfloat16, 2): [2 * MB, 8 * MB, 8 * MB, 8 * MB, 8 * MB],
54 + (torch.bfloat16, 4): [8 * MB, 64 * MB, 64 * MB, 16 * MB, 16 * MB],
55 + (torch.bfloat16, 8): [
56 + 16 * MB,
57 + 2048 * MB,
58 + 2048 * MB,
59 + 2048 * MB,
60 + 2048 * MB,
61 + ],
62 }
63
64 def __init__(self, group: ProcessGroup, device: int | str | torch.device) -> None:
@@ -59,8 +67,10 @@ class QuickAllReduce:
67 available for CUDA and ROCm MI300 series.
68
69 Custom quick allreduce leverages quantization for further
62 - acceleration on ROCm. It currently supports Q8, Q6, and Q4
63 - quantization formats and FP(float16, bfloat16).
70 + acceleration on ROCm. It currently supports Q8, Q6, Q4, and Q3
71 + quantization formats and FP(float16, bfloat16). Q3 (INT3) is
72 + restricted to TP2 (world_size == 2) due to poor performance on
73 + larger world sizes.
74
75 Quick allreduce is designed as a complement to custom allreduce.
76 Its initialization requires even stricter conditions.
@@ -178,6 +188,23 @@ class QuickAllReduce:
188 )
189 return
190 self.qr_quant_level = QuickReduceRegime[regime_str]
191 +
192 + # INT3 is only enabled for TP2 (world_size == 2).
193 + # Kernel benchmarks show INT3 all-reduce on TP4/TP8 has poor
194 + # performance (the extra ranks make the 3-bit codec's pack/unpack
195 + # overhead outweigh the reduced communication volume), so INT3 is
196 + # restricted to 2-GPU tensor parallelism. For TP4/TP8 use a wider
197 + # codec (e.g. INT4) or NONE instead.
198 + if self.qr_quant_level == QuickReduceRegime.INT3 and self.world_size != 2:
199 + logger.warning(
200 + "Custom quick allreduce is disabled: INT3 quantization is "
201 + "only supported for TP2 (world_size == 2), but world_size "
202 + "is %d. INT3 on TP4/TP8 is disabled due to poor kernel "
203 + "performance. Use INT4/NONE for this world size.",
204 + self.world_size,
205 + )
206 + return
207 +
208 self.qr_quantization_min_size = self._get_qr_quantization_min_size()
209 vllm_config = get_current_vllm_config_or_none()
210 if (
vllm/envs.py
+3 -3
@@ -214,7 +214,7 @@ if TYPE_CHECKING:
214 VLLM_SSM_CONV_STATE_LAYOUT: Literal["SD", "DS"] | None = None
215 VLLM_COMPUTE_NANS_IN_LOGITS: bool = False
216 VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: Literal[
217 - "FP", "INT8", "INT6", "INT4", "NONE"
217 + "FP", "INT8", "INT6", "INT4", "INT3", "NONE"
218 ] = "NONE"
219 VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16: bool = True
220 VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB: int | None = None
@@ -1224,12 +1224,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
1224 os.getenv("VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT", "False").lower() in ("true", "1")
1225 ),
1226 # Custom quick allreduce kernel for MI3* cards
1227 - # Choice of quantization level: FP, INT8, INT6, INT4 or NONE
1227 + # Choice of quantization level: FP, INT8, INT6, INT4, INT3 or NONE
1228 # Recommended for large models to get allreduce
1229 "VLLM_ROCM_QUICK_REDUCE_QUANTIZATION": env_with_choices(
1230 "VLLM_ROCM_QUICK_REDUCE_QUANTIZATION",
1231 "NONE",
1232 - ["FP", "INT8", "INT6", "INT4", "NONE"],
1232 + ["FP", "INT8", "INT6", "INT4", "INT3", "NONE"],
1233 ),
1234 # Custom quick allreduce kernel for MI3* cards
1235 # Due to the lack of the bfloat16 asm instruction, bfloat16