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
+43 -2
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) {
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<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
}