From bea282cede6c2eeeb6663d1d73a8357ce1bb2f2f Mon Sep 17 00:00:00 2001 From: Qichao Li Date: Fri, 19 Jun 2026 05:41:56 +0800 Subject: [PATCH] [DeepSeek-V4] Fuse UE8M0 scale rounding into FP8 group quantization (#26766) Co-authored-by: liqichao Co-authored-by: yhyang201 --- .../gemm/per_token_group_quant_8bit_v2.cuh | 26 +++++++++----- .../srt/layers/quantization/fp8_kernel.py | 34 +++++++++++++------ python/sglang/srt/models/deepseek_v4.py | 2 +- .../jit/test_per_token_group_quant_8bit_v2.py | 27 +++++++++++++++ .../test_deepseek_v4_flash_fp4_b200.py | 2 +- ...test_deepseek_v4_flash_fp4_megamoe_b200.py | 2 +- 6 files changed, 70 insertions(+), 23 deletions(-) diff --git a/python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit_v2.cuh b/python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit_v2.cuh index cad7a9a92..022567398 100644 --- a/python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit_v2.cuh +++ b/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> + typename scale_packed_t = std::conditional_t> __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; - using scale_element_t = std::conditional_t; + using scale_element_t = std::conditional_t; static_assert(sizeof(scale_packed_t) % sizeof(scale_element_t) == 0); device::PDLWaitPrimary(); @@ -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); - scale_output = output_s + offset_num_groups; + static_assert(!SCALE_UE8M0 || std::is_same_v); + scale_output = reinterpret_cast(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(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(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(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) { @@ -378,7 +382,7 @@ struct PerTokenGroupQuant8bitV2Kernel { void* output_q, void* output_s, const int32_t* masked_m) { - using scale_packed_t = std::conditional_t; + using scale_packed_t = std::conditional_t; auto kernel = per_token_group_quant_8bit_v2_kernel< SCHEDULER, GROUP_SIZE, @@ -463,7 +467,11 @@ struct PerTokenGroupQuant8bitV2Kernel { launch_with_config(TypeTag{}, std::true_type{}, std::false_type{}, std::false_type{}); } } else { - launch_with_config(TypeTag{}, std::false_type{}, std::false_type{}, std::false_type{}); + if (scale_ue8m0) { + launch_with_config(TypeTag{}, std::false_type{}, std::true_type{}, std::false_type{}); + } else { + launch_with_config(TypeTag{}, std::false_type{}, std::false_type{}, std::false_type{}); + } } } diff --git a/python/sglang/srt/layers/quantization/fp8_kernel.py b/python/sglang/srt/layers/quantization/fp8_kernel.py index 2c5f06b62..0bd72ccb0 100644 --- a/python/sglang/srt/layers/quantization/fp8_kernel.py +++ b/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 - *x_batch, x_q_mn, x_q_k = x_shape - x_s_mn, x_s_k = x_q_mn, x_q_k // 128 - aligned_mn = ceil_align(x_s_mn, 4) - aligned_k = ceil_align(x_s_k, 4) - # TODO(FIXME): Fix cuda kernel and recover here to empty. - return torch.empty( - (*x_batch, aligned_k // 4, aligned_mn), - device=device, - dtype=torch.int, - ).transpose(-1, -2)[..., :x_s_mn, :] + if column_major_scales and scale_tma_aligned: + *x_batch, x_q_mn, x_q_k = x_shape + x_s_mn, x_s_k = x_q_mn, x_q_k // 128 + aligned_mn = ceil_align(x_s_mn, 4) + aligned_k = ceil_align(x_s_k, 4) + # TODO(FIXME): Fix cuda kernel and recover here to empty. + return torch.empty( + (*x_batch, aligned_k // 4, aligned_mn), + device=device, + dtype=torch.int, + ).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 diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 6cf606d07..8519d505c 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/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", diff --git a/test/registered/jit/test_per_token_group_quant_8bit_v2.py b/test/registered/jit/test_per_token_group_quant_8bit_v2.py index 7d5110c7f..ce97a7eee 100644 --- a/test/registered/jit/test_per_token_group_quant_8bit_v2.py +++ b/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 diff --git a/test/registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py b/test/registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py index 62ab37d03..6800b1776 100644 --- a/test/registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py +++ b/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 diff --git a/test/registered/models_e2e/test_deepseek_v4_flash_fp4_megamoe_b200.py b/test/registered/models_e2e/test_deepseek_v4_flash_fp4_megamoe_b200.py index 505241dbf..f24611048 100644 --- a/test/registered/models_e2e/test_deepseek_v4_flash_fp4_megamoe_b200.py +++ b/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