Skip to content

Support fp8 block gemm with fp8_e8m0 scale #398

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
baodii merged 1 commit into from
Jun 11, 2026
Merged

Support fp8 block gemm with fp8_e8m0 scale #398

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
107 changes: 44 additions & 63 deletions csrc/xpu/onednn/fp8_gemm_w8a8.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,20 +25,19 @@ static inline void dnnl_matmul_w8a8_fp8(
const int n = o_sz.back(); // presume channel last format
const int k = *(src_sz.end() - 1);

bool is_block_quant =
(m1_sc.scalar_type() != at::ScalarType::Float8_e8m0fnu) &&
(m2_sc.scalar_type() != at::ScalarType::Float8_e8m0fnu) &&
(m1_sc.dim() == 2) && (m2_sc.dim() == 2) && (m1_sc.size(1) != 1) &&
(m2_sc.size(0) != 1);
int64_t group_size = -1;
bool is_block_quant = (m1_sc.dim() == 2) && (m1_sc.size(1) > 1);

int64_t wei_group_k = -1;
int64_t wei_group_n = -1;
if (is_block_quant) {
TORCH_CHECK(
m1_sc.size(1) == m2_sc.size(0),
"Mismatch group size in input and weight.",
m1_sc.size(1),
" vs ",
m2_sc.size(0));
group_size = k / m1_sc.size(1);
wei_group_k = k / m2_sc.size(0);
wei_group_n = n / m2_sc.size(1);
}

// get joint dtypes
Expand Down Expand Up @@ -84,74 +83,53 @@ static inline void dnnl_matmul_w8a8_fp8(
: mat2.strides()[mat2.dim() - 1];
int64_t ldc = result.strides()[leading_dim];

auto m1_sc_dtype = m1_sc.scalar_type();
auto m2_sc_dtype = m2_sc.scalar_type();
auto f_attr = [&](dnnl::primitive_attr& pattr) {
pattr.set_scratchpad_mode(dnnl::scratchpad_mode::user);

if (m1_sc_dtype == at::ScalarType::Float8_e8m0fnu) {
TORCH_CHECK(
m2_sc_dtype == at::ScalarType::Float8_e8m0fnu,
"Mismatched scale data types in mxfp8 matmul: ",
m1_sc_dtype,
" vs ",
m2_sc_dtype);
if (is_block_quant) {
pattr.set_scales(
DNNL_ARG_SRC,
/* mask */ (1 << 0) + (1 << 1),
{1, wei_group_k},
get_onednn_dtype(m1_sc));
/* per block quant. MXFP8 (float32 or e8m0 scales) */
} else if (m1_sc.numel() == 1) {
pattr.set_scales(
DNNL_ARG_SRC,
/* mask */ 0,
{},
get_onednn_dtype(m1_sc));
/* per tensor quant */
} else {
pattr.set_scales(
DNNL_ARG_SRC,
/* mask */ (1 << 0) + (1 << 1),
{1, 32},
{1, k},
get_onednn_dtype(m1_sc));
/* per token quant */
}

if (is_block_quant) {
pattr.set_scales(
DNNL_ARG_WEIGHTS,
/* mask */ (1 << 0) + (1 << 1),
{32, 1},
{wei_group_k, wei_group_n},
get_onednn_dtype(m2_sc));
/* per block quant. MXFP8 (float32 or e8m0 scales) */
} else if (m2_sc.numel() == 1) {
pattr.set_scales(
DNNL_ARG_WEIGHTS,
/* mask */ 0,
{},
get_onednn_dtype(m2_sc));
/* per tensor quant */
} else {
if (m1_sc.numel() == 1) {
pattr.set_scales(
DNNL_ARG_SRC,
/* mask */ 0,
{},
get_onednn_dtype(m1_sc));
/* per tensor quant */
} else if (is_block_quant) {
pattr.set_scales(
DNNL_ARG_SRC,
/* mask */ (1 << 0) + (1 << 1),
{1, group_size},
get_onednn_dtype(m1_sc));
/* per block quant */
} else {
pattr.set_scales(
DNNL_ARG_SRC,
/* mask */ (1 << 0) + (1 << 1),
{1, k},
get_onednn_dtype(m1_sc));
/* per token quant */
}

if (m2_sc.numel() == 1) {
pattr.set_scales(
DNNL_ARG_WEIGHTS,
/* mask */ 0,
{},
get_onednn_dtype(m2_sc));
/* per tensor quant */
} else if (is_block_quant) {
pattr.set_scales(
DNNL_ARG_WEIGHTS,
/* mask */ (1 << 0) + (1 << 1),
{group_size, group_size},
get_onednn_dtype(m2_sc));
/* per block quant */
} else {
pattr.set_scales(
DNNL_ARG_WEIGHTS,
/* mask */ (1 << 1),
{},
get_onednn_dtype(m2_sc));
/* per channel quant */
}
pattr.set_scales(
DNNL_ARG_WEIGHTS,
/* mask */ (1 << 1),
{},
get_onednn_dtype(m2_sc));
/* per channel quant */
}
};

Expand All @@ -166,6 +144,9 @@ static inline void dnnl_matmul_w8a8_fp8(
int m1_sc_group_size = m1_sc.numel();
int m2_sc_group_size = m2_sc.numel();
int sc_group_size = (m1_sc_group_size << 8) | m2_sc_group_size;
if (m1_sc.scalar_type() == at::ScalarType::Float8_e8m0fnu) {
sc_group_size |= (1 << 30);
}
Comment on lines 144 to +149
auto& matmul_ext = matmul_primitive_create_and_cache(
jd, tt, b_type, m, n, k, lda, ldb, ldc, dev_id, f_attr, sc_group_size);

Expand Down
12 changes: 11 additions & 1 deletion tests/ops/fp8_quant_op.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ def fp8_block_quant_2d(
block_n: int,
fp8_dtype=torch.float8_e4m3fn,
eps: float = 1e-6,
use_ue8m0: bool = False,
):
"""
Reference FP8 2D block quantization
Expand All @@ -28,9 +29,11 @@ def fp8_block_quant_2d(
block_m: block rows
block_n: block cols
fp8_dtype: torch.float8_e4m3fn
use_ue8m0: return scales as torch.float8_e8m0fnu
Returns:
q: FP8 tensor [M, N]
scales: FP32 tensor [ceil(M/BM), ceil(N/BN)]
scales: FP32 tensor [ceil(M/BM), ceil(N/BN)], or float8_e8m0fnu when
use_ue8m0 is True
"""
assert x.dim() == 2
M, N = x.shape
Expand Down Expand Up @@ -60,13 +63,18 @@ def fp8_block_quant_2d(
amax = block.abs().max()
scale = amax / FP8_MAX
scale = torch.clamp(scale, min=eps)
if use_ue8m0:
scale = torch.exp2(torch.ceil(torch.log2(scale)))

scales[gm, gn] = scale

# quantize
q_block = (block / scale).to(fp8_dtype)
q[m0:m1, n0:n1] = q_block

if use_ue8m0:
scales = scales.to(torch.float8_e8m0fnu)

return q, scales


Expand All @@ -93,6 +101,7 @@ def fp8_block_dequant_2d(
M, N = q.shape
grid_m, grid_n = scales.shape

scales = scales.to(torch.float32)
return (q.to(torch.float32).reshape(grid_m, block_m, grid_n, block_n) *
scales.reshape(grid_m, 1, grid_n, 1)).reshape(M, N).to(dtype)

Expand All @@ -118,6 +127,7 @@ def per_token_group_dequant_fp8(
M, K = q.shape
num_groups = K // group_size

scales = scales.to(torch.float32)
return (q.to(torch.float32).reshape(M, num_groups, group_size) *
scales.unsqueeze(-1)).reshape(M, K).to(dtype)

Expand Down
11 changes: 7 additions & 4 deletions tests/test_fp8_gemm_onednn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -304,9 +304,10 @@ def test_mxfp8_gemm(mnk_factors, out_dtype):
@pytest.mark.parametrize("is_nt", [True, False])
@pytest.mark.parametrize("batch", BATCHES)
@pytest.mark.parametrize("group_size", GROUP_SIZE)
@pytest.mark.parametrize("use_ue8m0", [False, True])
@pytest.mark.parametrize("mnk_factors", MNK_BLOCK_FACTORS)
def test_fp8_gemm_per_block(fp8_dtype, dtype, is_nt, batch, group_size,
mnk_factors):
use_ue8m0, mnk_factors):
seed = 1234
torch.manual_seed(seed)

Expand All @@ -319,9 +320,11 @@ def test_fp8_gemm_per_block(fp8_dtype, dtype, is_nt, batch, group_size,
input_fp8, scale_src_fp8 = per_token_group_quant_fp8(input.reshape(-1, k),
group_size,
dtype=fp8_dtype,
use_ue8m0=False)
weight_fp8, scale_wei_fp8 = fp8_block_quant_2d(weight, group_size,
group_size)
use_ue8m0=use_ue8m0)
weight_fp8, scale_wei_fp8 = fp8_block_quant_2d(weight,
group_size,
group_size,
use_ue8m0=use_ue8m0)

# reference: dequantize FP8 data, then multiply
input_deq = per_token_group_dequant_fp8(input_fp8, scale_src_fp8,
Expand Down
Loading