[JIT] per_token_group_quant: fix the fp16 UE8M0 multiplier overflow, 1.23x on the fused EP-MoE path, unify the scale allocation

python/sglang/kernels/aot/tests/test_per_token_group_quant_8bit.py

@@ -8,6 +8,9 @@
     create_per_token_group_quant_test_data,
 )
 
+from sglang.kernels.ops.quantization.fp8_kernel import (
+    PER_TOKEN_GROUP_QUANT_EPS,
+)
 from sglang.kernels.ops.quantization.fp8_kernel import (
     per_token_group_quant_8bit as triton_per_token_group_quant_8bit,
 )
@@ -136,10 +139,12 @@ def test_per_token_group_quant_with_column_major(
         x=x,
         masked_m=masked_m,
         group_size=group_size,
-        eps=1e-10,
         dst_dtype=dst_dtype,
         **{k: v for k, v in flags.items() if k not in ["masked_layout_mode"]},
     )
+    # The Triton reference still takes the absmax floor per call; the sglang entry
+    # point bakes it in (PER_TOKEN_GROUP_QUANT_EPS), so it has no eps parameter.
+    triton_kwargs = dict(execute_kwargs, eps=PER_TOKEN_GROUP_QUANT_EPS)
 
     def _postprocess(x_q, x_s):
         if masked_m is not None:
@@ -150,7 +155,7 @@ def _postprocess(x_q, x_s):
         return x_q, x_s
 
     x_q_triton, x_s_triton = _postprocess(
-        *triton_per_token_group_quant_8bit(**execute_kwargs)
+        *triton_per_token_group_quant_8bit(**triton_kwargs)
     )
     x_q_sglang, x_s_sglang = _postprocess(
         *sglang_per_token_group_quant_8bit(**execute_kwargs)

python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh

@@ -21,12 +21,18 @@ namespace {
 namespace details {
 
 SGL_DEVICE float silu(const float val) {
-  // silu(x) = x * sigmoid(x)
+  // silu(x) = x * sigmoid(x). Keep the plain divide: under --use_fast_math it is
+  // MUFU.RCP + FMUL, whereas an explicit __frcp_rn(1 + exp) is a Newton
+  // refinement (MUFU.RCP + 2 FFMA + FADD) behind a branch into an out-of-line
+  // special-value fixup. This runs per ELEMENT (not once per group like the
+  // quant multiplier), so on the fused masked path it dominated: 608 -> 384
+  // SASS instructions and 1.23x on the EP-MoE decode shapes when switched back.
+  // The divide also keeps the fused output bit-identical to the AOT v2 op.
 #if SGL_ARCH_BLACKWELL_OR_GREATER
   const float half = 0.5f * val;
   return half * (1.0f + __tanhf(half));
 #else
-  return val * __frcp_rn(1.0f + __expf(-val));
+  return val / (1.0f + __expf(-val));
 #endif
 }
 
@@ -231,6 +237,10 @@ struct QuantTrait {
   static constexpr uint32_t kBlockSize = 256;
   static constexpr uint32_t kVecSize = 32u / sizeof(InputType);
   static constexpr uint32_t kNumLanes = kGroupSize / kVecSize;
+  // Group-absmax floor, baked in here (the host rejects a caller eps != this).
+  // It bounds the quant multiplier at kMaxValue / kAmaxFloor, which is what
+  // decides whether the multiplier is representable in InputType below.
+  static constexpr float kAmaxFloor = 1e-10f;
   static_assert(sizeof(InputType) == 2, "only 16-bit inputs (bf16/fp16) are supported");
   static_assert(16 <= kGroupSize && kGroupSize <= 256, "supported group sizes are 16..256");
   static_assert(kGroupSize % kVecSize == 0 && 1 <= kNumLanes && kNumLanes <= device::kWarpThreads);
@@ -280,27 +290,43 @@ struct QuantTrait {
       local_amax2 = math::max(local_amax2, math::abs(in[i]));
     }
     const auto amax2 = cast<float2>(warp::reduce_max<kNumLanes>(local_amax2));
-    const auto amax = math::max(math::max(amax2.x, amax2.y), 1e-10f);
+    const auto amax = math::max(math::max(amax2.x, amax2.y), kAmaxFloor);
     const float raw_scale = amax * kMaxValueInv;  // the dequant scale the GEMM consumes
 
     out_vec_t out;
     details::scale_t<kUe8m0> scale_inv;
     if constexpr (kUe8m0) {
-      // ue8m0 scale: pow-2 quant multiplier is exact in float16/bfloat16 type
       static_assert(std::is_same_v<Q, fp8_e4m3_t>, "ue8m0 scales imply fp8 quantization");
       const auto exp = cast_to_ue8m0(raw_scale);
       scale_inv = static_cast<uint8_t>(exp);
       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
-      // 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};
+      // The pow-2 multiplier is exact in the packed 16-bit domain, so scaling
+      // there costs one __hmul2 per pair and loses nothing -- but only if the
+      // multiplier itself is representable in InputType. It reaches
+      // kMaxValue / kAmaxFloor (4.5e12), which fits bfloat16 (fp32's exponent
+      // range) and NOT float16 (max 65504): an fp16 group whose absmax falls
+      // below kMaxValue / 65504 = 6.8e-3 would narrow the multiplier to inf and
+      // quantize the whole group to +-448 (0 * inf even yields NaN). Scale in
+      // fp32 there, matching the fp32-scale path below.
+      constexpr bool kScaleFitsInInput = kMaxValue / kAmaxFloor <= DTypeTrait<T>::kFloatMax;
+      if constexpr (kScaleFitsInInput) {
+        const auto scale2 = cast<T2>(float2{quant_scale, quant_scale});
+        // Finite scaled values already lie in +-448 (2^exp >= amax/448), so the
+        // 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(__hmul2(in[i], scale2), max_clip2));
+        for (uint32_t i = 0; i < kVecSize / 2; ++i) {
+          out[i] = static_cast<Q2>(__hmin2(__hmul2(in[i], scale2), max_clip2));
+        }
+      } else {
+        const float2 quant_scale2 = {quant_scale, quant_scale};
+#pragma unroll
+        for (uint32_t i = 0; i < kVecSize / 2; ++i) {
+          out[i] = WTrait::quant(details::mul2(cast<float2>(in[i]), quant_scale2));
+        }
       }
     } else {
       // fp32 scale: multiply in fp32 (hmul2 brings too much precision loss)

python/sglang/kernels/jit/csrc/minimax/per_token_quant_ue8m0.cuh

@@ -19,81 +19,17 @@ namespace {
 using deepseek_v4::fp8::cast_to_ue8m0;
 using deepseek_v4::fp8::pack_fp8;
 
-// Per-token group quant to FP8-e4m3 with a fused UE8M0 scale. Each group of
-// kGroupSize columns gets one UE8M0 exponent byte written contiguously in
-// row-major order into ``x_sf`` (int32 [num_tokens, num_groups/4], 4 group
-// bytes per int32). This is the deep_gemm "transform_sf" pack done inline in
-// the quant, reusing the dsv4 ``cast_to_ue8m0``/``pack_fp8`` primitives -- it
-// is byte-identical to ``per_token_group_quant_fp8(scale_ue8m0=True)`` followed
-// by ``transform_sf_into_required_layout`` (both round via ceil(log2(absmax/
-// FP8_MAX))), but emits no separate transpose/pack kernel.
-struct PerTokenQuantUe8m0Params {
-  const bf16_t* __restrict__ x;  // [num_tokens, hidden]
-  fp8_e4m3_t* __restrict__ x_q;  // [num_tokens, hidden]
-  int32_t* __restrict__ x_sf;    // [num_tokens, num_groups/4]; written as bytes
-  uint32_t num_tokens;
-  uint32_t hidden;
-  uint32_t num_groups;  // hidden / kGroupSize
-};
-
-template <uint32_t kGroupSize, bool kUsePDL>
-__global__ __launch_bounds__(1024, 2) void  //
-    per_token_quant_ue8m0_kernel(const PerTokenQuantUe8m0Params __grid_constant__ params) {
-  using namespace device;
-  constexpr uint32_t kVecElems = 8;  // 8 bf16 = 16B load per thread
-  static_assert(kGroupSize % kVecElems == 0, "group_size must be a multiple of 8");
-  constexpr uint32_t kThreadsPerGroup = kGroupSize / kVecElems;
-  using InputVec = AlignedVector<bf16x2_t, kVecElems / 2>;
-  using OutputVec = AlignedVector<fp8x2_e4m3_t, kVecElems / 2>;
-
-  const uint32_t token_id = blockIdx.x;
-  const uint32_t tid = threadIdx.x;
-  PDLWaitPrimary<kUsePDL>();
-
-  const auto token_in = params.x + static_cast<uint64_t>(token_id) * params.hidden;
-  const auto token_out = params.x_q + static_cast<uint64_t>(token_id) * params.hidden;
-
-  InputVec in_vec;
-  in_vec.load(token_in, tid);
-  float local_max = 0.0f;
-  float vals[kVecElems];
-#pragma unroll
-  for (uint32_t i = 0; i < kVecElems / 2; ++i) {
-    const auto [v0, v1] = cast<fp32x2_t>(in_vec[i]);
-    vals[2 * i + 0] = v0;
-    vals[2 * i + 1] = v1;
-    local_max = fmaxf(local_max, fmaxf(fabsf(v0), fabsf(v1)));
-  }
-  // Absmax across the kThreadsPerGroup threads that cover one group.
-  local_max = warp::reduce_max<kThreadsPerGroup>(local_max);
-  const float absmax = fmaxf(local_max, 1e-10f);
-  const float raw_scale = absmax / math::FP8_E4M3_MAX;
-  const uint32_t ue8m0_exp = cast_to_ue8m0(raw_scale);
-  const float inv_scale = __uint_as_float((127u + 127u - ue8m0_exp) << 23);
-
-  OutputVec out_vec;
-#pragma unroll
-  for (uint32_t i = 0; i < kVecElems / 2; ++i) {
-    out_vec[i] = pack_fp8(vals[2 * i + 0] * inv_scale, vals[2 * i + 1] * inv_scale);
-  }
-  out_vec.store(token_out, tid);
-
-  const uint32_t group_id = tid / kThreadsPerGroup;
-  const uint32_t within_group_id = tid % kThreadsPerGroup;
-  if (within_group_id == 0 && group_id < params.num_groups) {
-    const uint32_t byte_off = token_id * params.num_groups + group_id;
-    reinterpret_cast<uint8_t*>(params.x_sf)[byte_off] = static_cast<uint8_t>(ue8m0_exp);
-  }
-  PDLTriggerSecondary<kUsePDL>();
-}
-
-// Fused quant + scatter: like per_token_quant_ue8m0_kernel, but instead of
-// writing the fp8/scale for the single source token, it scatters them straight
-// into the permuted grouped-GEMM input -- replicating each token to its ``topk``
-// destination rows -- so the separate fill_gateup_input_triton_kernel launch (and
-// the intermediate x_q/x_sf buffers) are eliminated. The fp8 value + UE8M0 scale
-// are computed exactly once per token (identical to the non-fused kernel); only
-// the stores differ.
+// Fused quant + scatter. The quant itself -- per-token group absmax, UE8M0 scale
+// via ceil(log2(absmax / FP8_MAX)), fp8-e4m3 codes -- is byte-identical to
+// per_token_group_quant(scale_ue8m0=True) with a row-major int32-packed scale;
+// what this kernel adds is the OUTPUT MAPPING. Instead of writing one row per
+// source token it scatters straight into the permuted grouped-GEMM input,
+// replicating each token to its ``topk`` destination rows, which eliminates the
+// separate fill_gateup_input_triton_kernel launch and the intermediate x_q/x_sf
+// buffers. The fp8 value + UE8M0 scale are still computed exactly once per
+// token; only the stores differ. (The non-scatter variant that used to live here
+// was a pure duplicate of the shared kernel and was removed -- if this file ever
+// needs a plain quant again, call per_token_group_quant instead of re-adding it.)
 struct PerTokenQuantUe8m0ScatterParams {
   const bf16_t* __restrict__ x;              // [num_tokens, hidden]
   fp8_e4m3_t* __restrict__ gateup_input;     // [E, m_max, hidden]
@@ -243,38 +179,4 @@ void per_token_quant_ue8m0_scatter(
       .enable_pdl(kUsePDL)(kernel, params);
 }
 
-template <int64_t kGroupSize, bool kUsePDL>
-void per_token_quant_ue8m0(tvm::ffi::TensorView x, tvm::ffi::TensorView x_q, tvm::ffi::TensorView x_sf) {
-  using namespace host;
-  auto device = SymbolicDevice{};
-  auto M = SymbolicSize{"num_tokens"};
-  auto H = SymbolicSize{"hidden"};
-  auto G4 = SymbolicSize{"num_groups_div_4"};
-  device.set_options<kDLCUDA>();
-  TensorMatcher({M, H}).with_dtype<bf16_t>().with_device(device).verify(x);
-  TensorMatcher({M, H}).with_dtype<fp8_e4m3_t>().with_device(device).verify(x_q);
-  TensorMatcher({M, G4}).with_dtype<int32_t>().with_device(device).verify(x_sf);
-
-  const uint32_t num_tokens = static_cast<uint32_t>(M.unwrap());
-  const uint32_t hidden = static_cast<uint32_t>(H.unwrap());
-  RuntimeCheck(hidden % kGroupSize == 0, "hidden ", hidden, " not divisible by group_size ", kGroupSize);
-  const uint32_t num_groups = hidden / static_cast<uint32_t>(kGroupSize);
-  RuntimeCheck(static_cast<uint32_t>(G4.unwrap()) * 4 == num_groups);
-  const uint32_t threads = hidden / 8;  // kVecElems
-  RuntimeCheck(threads <= 1024, "hidden/8 must be <= 1024, got ", threads);
-
-  const auto params = PerTokenQuantUe8m0Params{
-      .x = static_cast<const bf16_t*>(x.data_ptr()),
-      .x_q = static_cast<fp8_e4m3_t*>(x_q.data_ptr()),
-      .x_sf = static_cast<int32_t*>(x_sf.data_ptr()),
-      .num_tokens = num_tokens,
-      .hidden = hidden,
-      .num_groups = num_groups,
-  };
-  if (num_tokens == 0) return;
-  constexpr auto kernel = per_token_quant_ue8m0_kernel<kGroupSize, kUsePDL>;
-  LaunchKernel(num_tokens, threads, device.unwrap())  //
-      .enable_pdl(kUsePDL)(kernel, params);
-}
-
 }  // namespace

python/sglang/kernels/ops/attention/dsv4/moe.py

@@ -9,6 +9,7 @@
     load_jit,
     make_cpp_args,
 )
+from sglang.kernels.ops.quantization.quant_format import create_group_quant_outputs
 from sglang.srt.utils import is_xpu
 
 from .utils import make_name
@@ -216,15 +217,21 @@ def silu_and_mul_masked_post_quant(
 
 def silu_and_mul_contig_post_quant(
     input: torch.Tensor,
-    output: torch.Tensor,
-    output_scale: torch.Tensor,
     quant_group_size: int,
     scale_ue8m0: bool = False,
     transposed: bool = False,
     swiglu_limit: Optional[float] = None,
     swizzle: bool = False,
-) -> None:
+) -> Tuple[torch.Tensor, torch.Tensor]:
     apply_swiglu_limit = swiglu_limit is not None
+    output, output_scale = create_group_quant_outputs(
+        x_shape=(input.shape[0], input.shape[1] // 2),
+        device=input.device,
+        group_size=quant_group_size,
+        column_major_scales=transposed,
+        scale_tma_aligned=transposed,
+        scale_ue8m0=scale_ue8m0,
+    )
     module = _jit_silu_mul_quant_contig_module(
         quant_group_size, scale_ue8m0, swizzle, apply_swiglu_limit
     )
@@ -235,3 +242,4 @@ def silu_and_mul_contig_post_quant(
         transposed,
         float(swiglu_limit) if apply_swiglu_limit else 0.0,
     )
+    return output, output_scale