diff --git a/python/sglang/kernels/aot/csrc/allreduce/quick_all_reduce.cuh b/python/sglang/kernels/aot/csrc/allreduce/quick_all_reduce.cuh index 4032e6666..53bdb8f4a 100644 --- a/python/sglang/kernels/aot/csrc/allreduce/quick_all_reduce.cuh +++ b/python/sglang/kernels/aot/csrc/allreduce/quick_all_reduce.cuh @@ -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 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 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 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(&packed); +} + // Twoshot All Reduce template 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(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) { - float2 f = __bfloat1622float2(bf_buf[j]); - half_buf[j] = __float22half2_rn(f); + 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(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(bf16_buf), dst_buffer.descriptor, dst_offset, 0, 0); diff --git a/python/sglang/kernels/aot/csrc/allreduce/quick_all_reduce_base.h b/python/sglang/kernels/aot/csrc/allreduce/quick_all_reduce_base.h index c22a7fb3e..c6a134380 100644 --- a/python/sglang/kernels/aot/csrc/allreduce/quick_all_reduce_base.h +++ b/python/sglang/kernels/aot/csrc/allreduce/quick_all_reduce_base.h @@ -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 } diff --git a/test/registered/kernels/test_quick_allreduce_bf16_range.py b/test/registered/kernels/test_quick_allreduce_bf16_range.py new file mode 100644 index 000000000..347a40693 --- /dev/null +++ b/test/registered/kernels/test_quick_allreduce_bf16_range.py @@ -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()