[DeepSeek-V4] Fuse UE8M0 scale rounding into FP8 group quantization (#26766)

Co-authored-by: liqichao <liqichao@baidu.com>
Co-authored-by: yhyang201 <yhyang201@gmail.com>

python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit_v2.cuh

@@ -214,7 +214,7 @@ template <
     bool SCALE_UE8M0,
     bool FUSE_SILU_AND_MUL,
     bool kUsePDL,
-    typename scale_packed_t = std::conditional_t<SCALE_UE8M0, uint32_t, float>>
+    typename scale_packed_t = std::conditional_t<SCALE_UE8M0 && IS_COLUMN_MAJOR, uint32_t, float>>
 __global__ void per_token_group_quant_8bit_v2_kernel(
     const T* __restrict__ input,
     DST_DTYPE* __restrict__ output_q,
@@ -226,7 +226,7 @@ __global__ void per_token_group_quant_8bit_v2_kernel(
     const int scale_hidden_stride,
     const int num_tokens_per_expert) {
   using dst_dtype_info = DtypeInfo<DST_DTYPE>;
-  using scale_element_t = std::conditional_t<SCALE_UE8M0, uint8_t, float>;
+  using scale_element_t = std::conditional_t<SCALE_UE8M0 && IS_COLUMN_MAJOR, uint8_t, float>;
   static_assert(sizeof(scale_packed_t) % sizeof(scale_element_t) == 0);
 
   device::PDLWaitPrimary<kUsePDL>();
@@ -279,8 +279,8 @@ __global__ void per_token_group_quant_8bit_v2_kernel(
                           hidden_idx_packed * scale_hidden_stride * num_elems_per_pack +
                           token_idx * scale_token_stride * num_elems_per_pack + pack_idx);
         } else {
-          static_assert(!SCALE_UE8M0);
+          static_assert(!SCALE_UE8M0 || std::is_same_v<scale_packed_t, float>);
-          scale_output = output_s + offset_num_groups;
+          scale_output = reinterpret_cast<scale_element_t*>(output_s) + offset_num_groups;
         }
 
         if constexpr (IS_COLUMN_MAJOR and SCALE_UE8M0) {
@@ -309,12 +309,16 @@ __global__ void per_token_group_quant_8bit_v2_kernel(
         local_absmax = GroupReduceMax<THREADS_PER_SUBWARP>(local_absmax);
 
         float y_scale, y_scale_inv;
+        // When SCALE_UE8M0, always quantize with the rounded (power-of-2) scale
+        // — not with the exact scale followed by post-hoc rounding.
+        // This matches the official DeepSeek-V4 kernel.py act_quant(scale_fmt="ue8m0")
+        // and avoids a scale mismatch between quantization and downstream GEMM dequant,
+        // which otherwise amplifies error ~14x and degrades EAGLE accept rate on Blackwell.
         calculate_fp8_scales<SCALE_UE8M0, dst_dtype_info>(local_absmax, y_scale, y_scale_inv);
-        float2 y_scale_repeated = {y_scale, y_scale};
-
         if (lane_id == 0) {
-          *scale_output = extract_required_scale_format<SCALE_UE8M0>(y_scale_inv);
+          *scale_output = extract_required_scale_format < SCALE_UE8M0 && IS_COLUMN_MAJOR > (y_scale_inv);
         }
+        float2 y_scale_repeated = {y_scale, y_scale};
 
         int4 output_buf;
         if constexpr (std::is_same_v<DST_DTYPE, fp8_e4m3_t>) {
@@ -378,7 +382,7 @@ struct PerTokenGroupQuant8bitV2Kernel {
       void* output_q,
       void* output_s,
       const int32_t* masked_m) {
-    using scale_packed_t = std::conditional_t<SCALE_UE8M0, uint32_t, float>;
+    using scale_packed_t = std::conditional_t<SCALE_UE8M0 && IS_COLUMN_MAJOR, uint32_t, float>;
     auto kernel = per_token_group_quant_8bit_v2_kernel<
         SCHEDULER,
         GROUP_SIZE,
@@ -463,7 +467,11 @@ struct PerTokenGroupQuant8bitV2Kernel {
         launch_with_config(TypeTag<NaiveScheduler>{}, std::true_type{}, std::false_type{}, std::false_type{});
       }
     } else {
-      launch_with_config(TypeTag<NaiveScheduler>{}, std::false_type{}, std::false_type{}, std::false_type{});
+      if (scale_ue8m0) {
+        launch_with_config(TypeTag<NaiveScheduler>{}, std::false_type{}, std::true_type{}, std::false_type{});
+      } else {
+        launch_with_config(TypeTag<NaiveScheduler>{}, std::false_type{}, std::false_type{}, std::false_type{});
+      }
     }
   }
 

python/sglang/srt/layers/quantization/fp8_kernel.py

@@ -464,17 +464,29 @@ def create_per_token_group_quant_fp8_output_scale(
     scale_ue8m0: bool,
 ):
     if scale_ue8m0:
-        assert column_major_scales and scale_tma_aligned
+        if column_major_scales and scale_tma_aligned:
-        *x_batch, x_q_mn, x_q_k = x_shape
+            *x_batch, x_q_mn, x_q_k = x_shape
-        x_s_mn, x_s_k = x_q_mn, x_q_k // 128
+            x_s_mn, x_s_k = x_q_mn, x_q_k // 128
-        aligned_mn = ceil_align(x_s_mn, 4)
+            aligned_mn = ceil_align(x_s_mn, 4)
-        aligned_k = ceil_align(x_s_k, 4)
+            aligned_k = ceil_align(x_s_k, 4)
-        # TODO(FIXME): Fix cuda kernel and recover here to empty.
+            # TODO(FIXME): Fix cuda kernel and recover here to empty.
-        return torch.empty(
+            return torch.empty(
-            (*x_batch, aligned_k // 4, aligned_mn),
+                (*x_batch, aligned_k // 4, aligned_mn),
-            device=device,
+                device=device,
-            dtype=torch.int,
+                dtype=torch.int,
-        ).transpose(-1, -2)[..., :x_s_mn, :]
+            ).transpose(-1, -2)[..., :x_s_mn, :]
+        else:
+            assert not column_major_scales, (
+                "column_major_scales requires scale_tma_aligned=True "
+                "when scale_ue8m0 is enabled"
+            )
+            # Row-major UE8M0 keeps the scale as float32 power-of-two values,
+            # matching deep_gemm.ceil_to_ue8m0 and deep_gemm.fp8_einsum.
+            return torch.empty(
+                x_shape[:-1] + (x_shape[-1] // group_size,),
+                device=device,
+                dtype=torch.float32,
+            )
     elif column_major_scales:
         if scale_tma_aligned:
             # TODO extract "align" function

python/sglang/srt/models/deepseek_v4.py

@@ -1066,8 +1066,8 @@ class MQALayer(nn.Module):
             o_fp8, o_s = sglang_per_token_group_quant_fp8(
                 o.reshape(T * G, D).contiguous(),
                 group_size=128,
+                scale_ue8m0=True,
             )
-            o_s = deep_gemm.ceil_to_ue8m0(o_s)
             output = torch.empty(T, G, R, device=o.device, dtype=torch.bfloat16)
             deep_gemm.fp8_einsum(
                 "bhr,hdr->bhd",

test/registered/jit/test_per_token_group_quant_8bit_v2.py

@@ -10,6 +10,7 @@ from sglang.srt.layers.quantization.fp8_kernel import (
     fp8_dtype,
     fp8_max,
     fp8_min,
+    sglang_per_token_group_quant_fp8,
 )
 from sglang.test.ci.ci_register import register_cuda_ci
 
@@ -82,6 +83,32 @@ def test_v2_jit_matches_aot(dtype, num_tokens, hidden, fuse_silu_and_mul, scale_
     assert torch.equal(x_s, s_ref), "scales differ"
 
 
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("num_tokens", [1, 33, 128])
+@pytest.mark.parametrize("hidden", [128, 512, 4096, 7168])
+def test_sglang_per_token_group_quant_fp8_row_major_ue8m0(dtype, num_tokens, hidden):
+    """Row-major scale_ue8m0=True quantizes WITH the rounded (power-of-2) scale.
+    Verify: (1) scales are exact powers of 2, (2) dequant ≈ original within FP8 tolerance.
+    """
+    torch.manual_seed(num_tokens * 1000 + hidden)
+    x = torch.randn(num_tokens, hidden, device="cuda", dtype=dtype)
+
+    x_q, x_s = sglang_per_token_group_quant_fp8(x, G, scale_ue8m0=True)
+    torch.cuda.synchronize()
+
+    # Scales must be exact powers of 2
+    log2_s = torch.log2(x_s.abs())
+    assert torch.equal(log2_s, log2_s.round()), "scales are not power-of-2"
+
+    # Dequant should approximate original within FP8 precision
+    x_deq = x_q.float().view(num_tokens, -1, G) * x_s.unsqueeze(-1)
+    x_deq = x_deq.view(num_tokens, hidden)
+    rel_err = (x.float() - x_deq).abs() / (x.float().abs() + 1e-6)
+    assert (
+        rel_err.mean() < 0.05
+    ), f"mean relative dequant error too large: {rel_err.mean():.4f}"
+
+
 # Masked (EP-MoE) path: the v2 op only has a masked scheduler for the
 # column-major + ue8m0 + fused-silu+mul + masked combination. Input is 3D
 # [num_experts, tokens_padded, hidden*2]; only tokens < masked_m[e] are processed

test/registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py

@@ -41,7 +41,7 @@ class TestDSV4FlashFP4B200(
     """LowLatency recipe: TP=4, FP4 (mxfp4), EAGLE spec decoding."""
 
     gsm8k_accuracy_thres = 0.93
-    accept_length_thres = 2.6
+    accept_length_thres = 2.8
     bs_1_speed_thres = 220
 
     @classmethod

test/registered/models_e2e/test_deepseek_v4_flash_fp4_megamoe_b200.py

@@ -95,7 +95,7 @@ class TestDSV4FlashFP4B200W4A4MegaMoE(
     """Balanced recipe: TP=4, DP=4, MegaMoE."""
 
     gsm8k_accuracy_thres = 0.93
-    accept_length_thres = 2.6
+    accept_length_thres = 2.8
     bs_1_speed_thres = 100
 
     @classmethod