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

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 <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

@@ -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

@@ -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()