diff --git a/csrc/xpu/onednn/fp8_gemm_w8a8.h b/csrc/xpu/onednn/fp8_gemm_w8a8.h index dbdd30621..22fc631a0 100644 --- a/csrc/xpu/onednn/fp8_gemm_w8a8.h +++ b/csrc/xpu/onednn/fp8_gemm_w8a8.h @@ -25,12 +25,10 @@ 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), @@ -38,7 +36,8 @@ static inline void dnnl_matmul_w8a8_fp8( 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 @@ -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 */ } }; @@ -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); + } 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); diff --git a/tests/ops/fp8_quant_op.py b/tests/ops/fp8_quant_op.py index 9bd1f5c09..2ed45f0b8 100644 --- a/tests/ops/fp8_quant_op.py +++ b/tests/ops/fp8_quant_op.py @@ -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 @@ -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 @@ -60,6 +63,8 @@ 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 @@ -67,6 +72,9 @@ def fp8_block_quant_2d( 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 @@ -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) @@ -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) diff --git a/tests/test_fp8_gemm_onednn.py b/tests/test_fp8_gemm_onednn.py index de89ce45e..2cbef1991 100644 --- a/tests/test_fp8_gemm_onednn.py +++ b/tests/test_fp8_gemm_onednn.py @@ -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) @@ -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,