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

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

Browse Source 
Co-authored-by: liqichao <liqichao@baidu.com>
Co-authored-by: yhyang201 <yhyang201@gmail.com>
This commit is contained in:
Qichao Li
2026-06-18 14:41:56 -07:00
committed by GitHub
co-authored by liqichao yhyang201
parent 27a374eaef
commit bea282cede
6 changed files with 70 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit_v2.cuh
+17 -9
View File
@@ -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);
scale_output = output_s + offset_num_groups;
static_assert(!SCALE_UE8M0 || std::is_same_v<scale_packed_t, float>);
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
+23 -11
View File
@@ -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
python/sglang/srt/models/deepseek_v4.py
+1 -1
View File
@@ -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
+27
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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