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",