diff --git a/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh b/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh index deb1cc679..f44489c29 100644 --- a/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh +++ b/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh @@ -48,14 +48,16 @@ struct WeightTrait { using packed2_t = fp8x2_e4m3_t; static constexpr float kMaxValue = DTypeTrait::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(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); - // __hmin2/__hmax2 return the non-NaN operand. - const auto lo2 = cast(float2{-kMaxValue, -kMaxValue}); - const auto hi2 = cast(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). + const auto max_clip = cast(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(__hmin2(__hmax2(__hmul2(in[i], scale2), lo2), hi2)); + out[i] = static_cast(__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(); // 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(); Trait::run(params, 0, token_idx, group_idx, lane_id); PDLTriggerSecondary(); } @@ -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(); const auto expert_idx = blockIdx.y; + PDLWaitPrimary(); 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) { diff --git a/test/registered/kernels/benchmark/quantization/bench_per_token_group_quant.py b/test/registered/kernels/benchmark/quantization/bench_per_token_group_quant.py index a22e6bfed..2d190f3a9 100644 --- a/test/registered/kernels/benchmark/quantization/bench_per_token_group_quant.py +++ b/test/registered/kernels/benchmark/quantization/bench_per_token_group_quant.py @@ -68,7 +68,6 @@ def benchmark(group_size: int, layout: str, num_tokens: int, impl: str): return marker.do_bench( FN[impl], input_args=(group_size, x, x_q, x_s, scale_ue8m0), - graph_clone_args=(1,), memory_args=(x,), memory_output=(x_q, x_s), )