minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[ROCm] Fix QuickReduce fp16 saturation corrupting bf16 all-reduces (106M non-finite -> 0, +0.3%) (#34484)

Browse Source
This commit is contained in:
Alex Nails
2026-08-29 22:08:17 -07:00
committed by GitHub
parent 67bd163a48
commit 78fa921189
3 changed files with 184 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/aot/csrc/allreduce/quick_all_reduce.cuh
+41
View File
@@ -22,6 +22,9 @@ struct CodecFP : public CodecBase {
static constexpr int kWorldSize = world_size;
static constexpr int kRankAtoms = kAtoms / kWorldSize;
// No block scale to protect, so the bf16 -> fp16 cast scale is the only range guard.
static constexpr int kCastScaleLog2 = kQRFp16CastScaleLog2Fp;
// Codec tile size process by this workgroup.
// Each thread processes atoms of f16x8_t (16B).
static constexpr int kRankTransmittedTileSize = kBlockSize * kRankAtoms * sizeof(int32x4_t);
@@ -54,6 +57,9 @@ template <typename T, int world_size>
struct CodecQ4 : public CodecBase {
static constexpr int kWorldSize = world_size;
// Block-scaled: the cast scale would only cost the low end. See quick_all_reduce_base.h.
static constexpr int kCastScaleLog2 = kQRFp16CastScaleLog2Quant;
// Codec tile size process by this workgroup.
// Each threads processes a fragment of fp16x8_t (16B),
// into a int4x8_t (4B) and a fp16 scale shared among 32 values.
@@ -192,6 +198,9 @@ template <typename T, int world_size>
struct CodecQ6 : public CodecBase {
static constexpr int kWorldSize = world_size;
// Block-scaled: the cast scale would only cost the low end. See quick_all_reduce_base.h.
static constexpr int kCastScaleLog2 = kQRFp16CastScaleLog2Quant;
// Codec tile size process by this workgroup.
// Each threads processes a fragment of fp16x8_t (16B),
// into a int6x8_t (4B + 2B) and a fp16 scale shared among 32 values.
@@ -350,6 +359,9 @@ template <typename T, int world_size>
struct CodecQ8 : public CodecBase {
static constexpr int kWorldSize = world_size;
// Block-scaled: the cast scale would only cost the low end. See quick_all_reduce_base.h.
static constexpr int kCastScaleLog2 = kQRFp16CastScaleLog2Quant;
// Codec tile size process by this workgroup.
// Each threads processes a fragment of f16x8_t (16B),
// into a int8x8_t (8B) and a f16 scale shared among 32 values.
@@ -486,6 +498,21 @@ struct CodecQ8 : public CodecBase {
}
};
// Keep the scale on the f32 side of the narrowing conversion. With the HIP
// intrinsic, LLVM can reassociate (bf16_as_f32 * scale) -> fp16 into
// fp16(bf16_as_f32) * scale, which clips values above 65504 before the range
// guard is applied. The opaque ISA conversion makes the scaled f32 values
// explicit inputs and prevents that transform.
__quickreduce_device_inline__ half2 scaled_bfloat162_to_half2(nv_bfloat162 value, float scale) {
float2 scaled = __bfloat1622float2(value);
scaled.x *= scale;
scaled.y *= scale;
int packed;
asm volatile("v_cvt_pk_f16_f32 %0, %1, %2" : "=v"(packed) : "v"(scaled.x), "v"(scaled.y));
return *reinterpret_cast<half2*>(&packed);
}
// Twoshot All Reduce
template <typename T, class Codec, bool cast_bf2half>
struct AllReduceTwoshot {
@@ -493,6 +520,10 @@ struct AllReduceTwoshot {
static constexpr int kWorldSize = Codec::kWorldSize;
// Power of two, so both multiplies are exact.
static constexpr float kCastScale = static_cast<float>(1 << Codec::kCastScaleLog2);
static constexpr float kCastInvScale = 1.0f / kCastScale;
__device__ static void
run(T const* __restrict__ input,
T* __restrict__ output,
@@ -524,8 +555,15 @@ struct AllReduceTwoshot {
half2 half_buf[4];
#pragma unroll
for (int j = 0; j < 4; ++j) {
if constexpr (Codec::kCastScaleLog2 == 0) {
float2 f = __bfloat1622float2(bf_buf[j]);
// S=1 for quantized codecs; preserve their existing conversion path.
f.x *= kCastInvScale;
f.y *= kCastInvScale;
half_buf[j] = __float22half2_rn(f);
} else {
half_buf[j] = scaled_bfloat162_to_half2(bf_buf[j], kCastInvScale);
}
}
tA[i] = *reinterpret_cast<const int32x4_t*>(half_buf);
}
@@ -620,6 +658,9 @@ struct AllReduceTwoshot {
#pragma unroll
for (int j = 0; j < 4; ++j) {
float2 f = __half22float2(half_buf[j]);
// Undo the load-side scale; the fp32 intermediate cannot overflow.
f.x *= kCastScale;
f.y *= kCastScale;
bf16_buf[j] = __float22bfloat162_rn(f);
}
buffer_store_dwordx4(*reinterpret_cast<const int32x4_t*>(bf16_buf), dst_buffer.descriptor, dst_offset, 0, 0);
python/sglang/kernels/aot/csrc/allreduce/quick_all_reduce_base.h
+19 -3
View File
@@ -32,6 +32,19 @@ using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int;
static constexpr int kNegOne = 0xBC00BC00; // {-1, -1}, fp16x2_t
// Range guard for the bf16 -> fp16 fast path (AllReduceTwoshot<..., true>): fp16 saturates at
// 65504. Divide by a power of two on load and multiply it back on store; the shift is exact, so
// sum(x_i / S) * S == sum(x_i). Per codec, because only CodecFP needs it: it carries no block
// scale, so S is its only range guard, and is free there. The quantized codecs already normalize
// each 32 values by their own block scale, and MODE.FP16_OVFL (armed in CodecBase) keeps an
// over-ceiling element from becoming an inf that poisons its block through the block max. S buys
// them nothing and costs the low end: encoding_scale = rcp(decoding_scale) saturates at 65504,
// past which encode and decode stop being reciprocals and the block is attenuated wholesale.
// That cliff sits at blockmax = S * L / 65504 (L = 8 / 32 / 128 for Q4 / Q6 / Q8), so raising S
// walks it into real data -- 1 -> 2 measures 62% perplexity on GLM-5.2. Leave it at 1.
static constexpr int kQRFp16CastScaleLog2Fp = 4; // S = 16, CodecFP
static constexpr int kQRFp16CastScaleLog2Quant = 0; // S = 1, CodecQ4 / CodecQ6 / CodecQ8
// Number of atoms (4xf16x2_t) processed by a single thread
static constexpr int kAtoms = 8;
@@ -95,12 +108,15 @@ __quickreduce_device_inline__ static void
buffer_store_dwordx4(int32x4_t data, int32x4_t srsrc, int32_t voffset, int32_t soffset, int32_t aux) {}
#endif
// MODE.FP16_OVFL clamps overflowing fp16 results to +/-MAX_FP16 instead of inf, the f32 -> f16
// conversion (v_cvt_pk_f16_f32) included. The memory clobber is load-bearing: without it the
// compiler may hoist a conversion above the s_setreg, and that conversion still produces inf.
__quickreduce_device_inline__ static void set_fp16_ovfl(bool const value) {
#if defined(__gfx942__)
#if defined(__gfx942__) || defined(__gfx950__)
if (value) {
asm volatile("s_setreg_imm32_b32 0xdc1, 1;" ::);
asm volatile("s_setreg_imm32_b32 0xdc1, 1;" ::: "memory");
} else {
asm volatile("s_setreg_imm32_b32 0xdc1, 0;" ::);
asm volatile("s_setreg_imm32_b32 0xdc1, 0;" ::: "memory");
}
#endif
}
test/registered/kernels/test_quick_allreduce_bf16_range.py
+122
View File
@@ -0,0 +1,122 @@
import multiprocessing
import os
import socket
import time
import unittest
import torch
import torch.distributed as dist
from sglang.srt.distributed.device_communicators.quick_all_reduce import (
QuickAllReduce,
qr_rocm_arch_available,
)
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.test_utils import CustomTestCase
register_amd_ci(est_time=30, suite="stage-c-test-4-gpu-amd")
register_amd_ci(est_time=30, suite="stage-c-test-large-8-gpu-amd-mi35x")
def _get_open_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.bind(("", 0))
return sock.getsockname()[1]
def _run_bf16_range_test(rank: int, world_size: int, port: int) -> None:
os.environ["ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16"] = "1"
device = torch.device(f"cuda:{rank}")
torch.cuda.set_device(device)
dist.init_process_group(
backend="gloo",
init_method=f"tcp://127.0.0.1:{port}",
rank=rank,
world_size=world_size,
)
try:
numel = 1 << 20
cases = [
("low", 2**-8, False),
("ordinary", 100.0, False),
("sum_above_fp16", 20_000.0, False),
("input_above_fp16", 80_000.0, False),
("negative_sum_above_fp16", -20_000.0, False),
("mixed_input_above_fp16", 1.0, True),
]
for quant_mode in ("FP", "INT8", "INT6", "INT4"):
os.environ["ROCM_QUICK_REDUCE_QUANTIZATION"] = quant_mode
quick_all_reduce = QuickAllReduce(group=dist.group.WORLD, device=device)
assert not quick_all_reduce.disabled
assert quick_all_reduce.use_fp16_kernels
try:
for case_name, value, mixed in cases:
inp = torch.full(
(numel,), value, dtype=torch.bfloat16, device=device
)
if mixed:
inp[::32] = 80_000.0
expected = (inp.float() * world_size).to(torch.bfloat16)
dist.barrier()
out = quick_all_reduce.quick_all_reduce(inp)
torch.cuda.synchronize()
assert (
torch.isfinite(out).all().item()
), f"{quant_mode=} {case_name=} produced non-finite output"
if quant_mode == "FP" or case_name in ("low", "ordinary"):
torch.testing.assert_close(
out,
expected,
rtol=0,
atol=0,
msg=lambda msg: f"{quant_mode=} {case_name=}\n{msg}",
)
finally:
quick_all_reduce.close()
finally:
dist.destroy_process_group()
class TestQuickAllReduceBf16Range(CustomTestCase):
@unittest.skipUnless(
qr_rocm_arch_available() and torch.cuda.device_count() >= 4,
"QuickReduce range test requires at least four supported ROCm GPUs",
)
def test_bf16_range(self):
world_size = 4
port = _get_open_port()
context = multiprocessing.get_context("spawn")
processes = [
context.Process(
target=_run_bf16_range_test,
args=(rank, world_size, port),
)
for rank in range(world_size)
]
for process in processes:
process.start()
deadline = time.monotonic() + 120
for process in processes:
process.join(max(0, deadline - time.monotonic()))
if any(process.is_alive() for process in processes):
for process in processes:
if process.is_alive():
process.terminate()
process.join()
self.fail("QuickReduce bf16 range test timed out")
for rank, process in enumerate(processes):
self.assertEqual(
process.exitcode,
0,
f"QuickReduce bf16 range test failed on rank {rank}",
)
if __name__ == "__main__":
unittest.main()