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

[Perf] Halve the non-finite sanitization overhead in per_token_group_quant (#32296)

Browse Source 
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
DarkSharpness
2026-07-25 08:17:12 +08:00
committed by GitHub
co-authored by Claude Fable 5
parent 7de8792758
commit 9402012f0f
2 changed files with 17 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh
+17 -14
View File
@@ -48,14 +48,16 @@ struct WeightTrait<fp8_e4m3_t> {
using packed2_t = fp8x2_e4m3_t;
static constexpr float kMaxValue = DTypeTrait<fp8_e4m3_t>::kFloatMax;
// SATFINITE saturates +-inf / out-of-range values, but converts NaN to an
// fp8 NaN code. IEEE fminf/fmaxf return the non-NaN operand, so clamping
// first quantizes non-finite inputs to +-448 -- matching the v1/v2/Triton
// kernels. CUDA-graph capture warmup runs the model on whatever the
// (reused, uninitialized) buffers contain, and relies on this: an fp8 NaN
// code would poison the downstream GEMM and trip the sampler NaN check.
// For finite inputs the clamp is bit-identical to bare SATFINITE.
// fp8 NaN code. A single upper clamp is enough to sanitize: IEEE fminf
// returns the non-NaN operand, so NaN / +inf quantize to +448, and -inf
// passes through for SATFINITE to saturate to -448 -- non-finite inputs
// never reach an fp8 NaN code, matching the v1/v2/Triton kernels.
// CUDA-graph capture warmup runs the model on whatever the (reused,
// uninitialized) buffers contain, and relies on this: an fp8 NaN code would
// poison the downstream GEMM and trip the sampler NaN check. For finite
// inputs the clamp is bit-identical to bare SATFINITE.
SGL_DEVICE static packed2_t quant(const float2 v) {
return packed2_t{float2{fminf(fmaxf(v.x, -kMaxValue), kMaxValue), fminf(fmaxf(v.y, -kMaxValue), kMaxValue)}};
return packed2_t{float2{fminf(v.x, kMaxValue), fminf(v.y, kMaxValue)}};
}
};
@@ -291,13 +293,14 @@ struct QuantTrait {
const float quant_scale = inv_scale_ue8m0(exp);
const auto scale2 = cast<T2>(float2{quant_scale, quant_scale});
// Finite scaled values already lie in +-448 (2^exp >= amax/448), so the
// clamp only sanitizes non-finite inputs (see WeightTrait<fp8_e4m3_t>);
// __hmin2/__hmax2 return the non-NaN operand.
const auto lo2 = cast<T2>(float2{-kMaxValue, -kMaxValue});
const auto hi2 = cast<T2>(float2{kMaxValue, kMaxValue});
// single __hmin2 only sanitizes NaN / +inf (it returns the non-NaN
// operand); -inf saturates to -448 via the SATFINITE fp8 cast (see
// WeightTrait<fp8_e4m3_t>).
const auto max_clip = cast<T>(kMaxValue);
const auto max_clip2 = T2{max_clip, max_clip};
#pragma unroll
for (uint32_t i = 0; i < kVecSize / 2; ++i) {
out[i] = static_cast<Q2>(__hmin2(__hmax2(__hmul2(in[i], scale2), lo2), hi2));
out[i] = static_cast<Q2>(__hmin2(__hmul2(in[i], scale2), max_clip2));
}
} else {
// fp32 scale: multiply in fp32 (hmul2 brings too much precision loss)
@@ -330,12 +333,12 @@ __global__ __launch_bounds__(Trait::kBlockSize) void per_token_group_quant_flat_
const auto global_warp_id = global_tid / kWarpThreads;
const auto total_work = params.num_tokens * num_groups;
if (global_warp_id * kWorkPerWarp >= total_work) return;
PDLWaitPrimary<kUsePDL>();
// the last partial warp duplicates the tail work (identical-byte stores)
const auto work_id = min(global_tid / kNumLanes, total_work - 1);
const auto lane_id = threadIdx.x % kNumLanes;
const auto token_idx = work_id / num_groups;
const auto group_idx = work_id % num_groups;
PDLWaitPrimary<kUsePDL>();
Trait::run(params, 0, token_idx, group_idx, lane_id);
PDLTriggerSecondary<kUsePDL>();
}
@@ -351,8 +354,8 @@ __global__ __launch_bounds__(Trait::kBlockSize) void per_token_group_quant_maske
constexpr uint32_t kNumLanes = Trait::kNumLanes;
constexpr uint32_t kWorkPerWarp = kWarpThreads / kNumLanes;
const auto num_groups = params.base.scale.num_groups;
PDLWaitPrimary<kUsePDL>();
const auto expert_idx = blockIdx.y;
PDLWaitPrimary<kUsePDL>();
const auto num_expert_tokens = params.masked_m[expert_idx * params.masked_m_stride];
for (uint32_t global_tid = blockIdx.x * Trait::kBlockSize + threadIdx.x;; // initial grid; loop
global_tid += gridDim.x * Trait::kBlockSize) {