From 4678536df6a6ea1e4ea5a229f6ae48b82e8bf166 Mon Sep 17 00:00:00 2001 From: Xinguo Zhu Date: Wed, 16 Sep 2026 12:31:05 +0800 Subject: [PATCH] [CPU] Add fp8_per_tensor_scaled_mm_cpu kernel (#32618) Co-authored-by: AKatydid --- python/sglang/kernels/aot/csrc/cpu/bmm.cpp | 1 + python/sglang/kernels/aot/csrc/cpu/gemm.h | 4 +- .../sglang/kernels/aot/csrc/cpu/gemm_fp8.cpp | 324 ++++++++++++++---- .../aot/csrc/cpu/torch_extension_cpu.cpp | 12 + python/sglang/srt/layers/quantization/fp8.py | 21 ++ .../srt/model_executor/cpu_graph_runner.py | 13 + test/registered/cpu/test_gemm.py | 35 ++ 7 files changed, 350 insertions(+), 60 deletions(-) diff --git a/python/sglang/kernels/aot/csrc/cpu/bmm.cpp b/python/sglang/kernels/aot/csrc/cpu/bmm.cpp index ea496be1b..be12900db 100644 --- a/python/sglang/kernels/aot/csrc/cpu/bmm.cpp +++ b/python/sglang/kernels/aot/csrc/cpu/bmm.cpp @@ -112,6 +112,7 @@ void bmm_kernel_impl( /* C */ out + bs * out_strideB + mb_start * out_strideM + nb_start, /* Btmp*/ Btmp, /* Ctmp*/ Ctmp, + /* bias*/ nullptr, /*scale*/ scale, /* M */ mb_size, /* N */ nb_size, diff --git a/python/sglang/kernels/aot/csrc/cpu/gemm.h b/python/sglang/kernels/aot/csrc/cpu/gemm.h index b8a583170..156bb8a16 100644 --- a/python/sglang/kernels/aot/csrc/cpu/gemm.h +++ b/python/sglang/kernels/aot/csrc/cpu/gemm.h @@ -324,6 +324,7 @@ void tinygemm_kernel( scalar_t* __restrict__ C, scalar_t* __restrict__ Btmp, float* __restrict__ Ctmp, + const float* __restrict__ Bbias, float scale, int64_t M, int64_t N, @@ -331,7 +332,8 @@ void tinygemm_kernel( int64_t lda, int64_t ldb, int64_t ldc, - bool brg); + bool brg, + bool do_unpack = true); // mxfp4 template diff --git a/python/sglang/kernels/aot/csrc/cpu/gemm_fp8.cpp b/python/sglang/kernels/aot/csrc/cpu/gemm_fp8.cpp index 06b0b7f13..ce655c4ce 100644 --- a/python/sglang/kernels/aot/csrc/cpu/gemm_fp8.cpp +++ b/python/sglang/kernels/aot/csrc/cpu/gemm_fp8.cpp @@ -22,6 +22,30 @@ inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ inpu } } +template +inline void copy_mul_add_stub( + scalar_t* __restrict__ out, + const float* __restrict__ input, + const float* __restrict__ bias, + int64_t size, + float scale) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int kVecSize = bVec::size(); + const fVec vscale = fVec(scale); + + int64_t d; +#pragma GCC unroll 4 + for (d = 0; d <= size - kVecSize; d += kVecSize) { + auto [data0, data1] = load_float_vec2(input + d); + auto [bias0, bias1] = load_float_vec2(bias + d); + bVec out_vec = convert_from_float_ext(data0 * vscale + bias0, data1 * vscale + bias1); + out_vec.store(out + d); + } + for (; d < size; ++d) { + out[d] = static_cast(input[d] * scale + bias[d]); + } +} template inline void copy_add_stub( scalar_t* __restrict__ out, const float* __restrict__ input, const float* __restrict__ bias, int64_t size) { @@ -235,17 +259,18 @@ struct tinygemm_kernel_nn { } }; -template +template struct tinygemm_kernel_nn2 { static inline void apply( const scalar_t* __restrict__ A, - const at::Float8_e4m3fn* __restrict__ B, + const packed_t* __restrict__ B, scalar_t* __restrict__ C, + const float* __restrict__ bias, float scale, - int K, - int lda, - int ldb, - int ldc) { + int64_t K, + int64_t lda, + int64_t ldb, + int64_t ldc) { TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!"); } }; @@ -354,35 +379,45 @@ struct tinygemm_kernel_nn -struct tinygemm_kernel_nn2 { +template +struct tinygemm_kernel_nn2 { static inline void apply( const at::BFloat16* __restrict__ A, const at::Float8_e4m3fn* __restrict__ B, at::BFloat16* __restrict__ C, - float scale, - int K, - int lda, - int ldb, - int ldc) { + const float* __restrict__ bias, + const float scale, + int64_t K, + int64_t lda, + int64_t ldb, + int64_t ldc) { constexpr int ROWS = BLOCK_M; constexpr int COLS = BLOCK_N / 16; + const int64_t KB = div_up(K, (int64_t)BLOCK_K); + // prefetch distance constexpr int PREFETCH_SIZE_K = 64; __m512bh va; __m512bh vb[COLS]; __m512 vc[ROWS * COLS]; + __m512 vsum[ROWS * COLS]; const __m512 vscale = _mm512_set1_ps(scale); - auto loadc = [&](auto i) { vc[i] = _mm512_setzero_ps(); }; + auto loadc = [&](auto i) { + constexpr int col = i % COLS; + if constexpr (has_bias) { + vc[i] = _mm512_loadu_ps(bias + col * 16); + } else { + vc[i] = _mm512_setzero_ps(); + } + }; Unroll{}(loadc); - const int K2 = K >> 1; - const int lda2 = lda >> 1; - const int ldb2 = ldb; // ldb * 2 >> 1; + const int64_t lda2 = lda >> 1; + const int64_t ldb2 = ldb; // ldb * 2 >> 1; const float* a_ptr = reinterpret_cast(A); const uint16_t* b_ptr = reinterpret_cast(B); @@ -392,6 +427,9 @@ struct tinygemm_kernel_nn2 { if constexpr (col == 0) { va = (__m512bh)(_mm512_set1_ps(a_ptr[row * lda2 + k])); + if constexpr (PREFETCH_SIZE_K > 0) { + _mm_prefetch(a_ptr + row * lda2 + k + PREFETCH_SIZE_K, _MM_HINT_T0); + } } if constexpr (row == 0) { if constexpr (col % 2 == 0) { @@ -403,10 +441,21 @@ struct tinygemm_kernel_nn2 { vb[col + 1] = CVT_FP8_TO_BF16(_mm512_extracti32x8_epi32(b8, 1)); } } - vc[i] = _mm512_dpbf16_ps(vc[i], va, vb[col]); + vsum[i] = _mm512_dpbf16_ps(vsum[i], va, vb[col]); }; - for (int k = 0; k < K2; ++k) { - Unroll{}(compute, k); + + constexpr int64_t BLOCK_K2 = BLOCK_K >> 1; + for (int64_t kb = 0; kb < KB; ++kb) { + int64_t kb_start = kb * BLOCK_K2; + int64_t kb_end = std::min(K >> 1, kb_start + BLOCK_K2); + // 1. zero vsum for each block + Unroll{}([&](auto i) { vsum[i] = _mm512_setzero_ps(); }); + // 2. accumulate across each block + for (int k = kb_start; k < kb_end; ++k) { + Unroll{}(compute, k); + } + // 3. apply scale + Unroll{}([&](auto i) { vc[i] = _mm512_fmadd_ps(vsum[i], vscale, vc[i]); }); } auto storec = [&](auto i) { @@ -414,10 +463,9 @@ struct tinygemm_kernel_nn2 { constexpr int col = i % COLS; // for COLS = 2, 4 use 512bit store if constexpr (col % 2 == 0) { - __m512 vc0 = _mm512_mul_ps(vc[row * COLS + col + 0], vscale); - __m512 vc1 = _mm512_mul_ps(vc[row * COLS + col + 1], vscale); _mm512_storeu_si512( - reinterpret_cast<__m512i*>((C + row * ldc + col * 16)), (__m512i)(_mm512_cvtne2ps_pbh(vc1, vc0))); + reinterpret_cast<__m512i*>((C + row * ldc + col * 16)), + (__m512i)(_mm512_cvtne2ps_pbh(vc[row * COLS + col + 1], vc[row * COLS + col]))); } }; Unroll{}(storec); @@ -538,9 +586,17 @@ struct tinygemm_kernel_nn::apply( \ - A + mb_start * lda, B + nb_start * 2, C + mb_start * ldc + nb_start, scale, K, lda, ldb, ldc); +#define LAUNCH_TINYGEMM_KERNEL_NN2(MB_SIZE, NB_SIZE) \ + tinygemm_kernel_nn2::apply( \ + A + mb_start * lda, \ + B + nb_start * 2, \ + C + mb_start * ldc + nb_start, \ + has_bias ? bias + nb_start : nullptr, \ + scale, \ + K, \ + lda, \ + ldb, \ + ldc); template struct brgemm { @@ -562,8 +618,27 @@ struct brgemm { TORCH_CHECK(false, "struct brgemm: primary template not implemented!"); } }; -template -struct brgemm2 {}; + +template +struct brgemm2 { + static inline void apply( + const scalar_t* __restrict__ A, + const packed_t* __restrict__ B, + scalar_t* __restrict__ C, + scalar_t* __restrict__ Btmp, + float* __restrict__ Ctmp, + const float* __restrict__ bias, + const float scale, + int M, + int N, + int K, + int lda, + int ldb, + int ldc, + bool do_unpack = true) { + TORCH_CHECK(false, "struct brgemm2: primary template not implemented!"); + } +}; template struct brgemm { @@ -609,21 +684,23 @@ struct brgemm { } }; -template <> -struct brgemm2 { +template +struct brgemm2 { static inline void apply( const at::BFloat16* __restrict__ A, const at::Float8_e4m3fn* __restrict__ B, at::BFloat16* __restrict__ C, at::BFloat16* __restrict__ Btmp, float* __restrict__ Ctmp, - float scale, + const float* __restrict__ bias, + const float scale, int M, int N, int K, int lda, int ldb, - int ldc) { + int ldc, + bool do_unpack = true) { constexpr int BLOCK_N = block_size_n(); // [BLOCK_K, BLOCK_N] -> [BLOCK_K / 2, BLOCK_N * 2] @@ -640,7 +717,11 @@ struct brgemm2 { // copy from Ctmp to C and mul scale for (int m = 0; m < M; ++m) { - copy_mul_stub(C + m * ldc, Ctmp + m * BLOCK_N, N, scale); + if constexpr (has_bias) { + copy_mul_add_stub(C + m * ldc, Ctmp + m * BLOCK_N, bias, N, scale); + } else { + copy_mul_stub(C + m * ldc, Ctmp + m * BLOCK_N, N, scale); + } } } }; @@ -743,23 +824,25 @@ void tinygemm_kernel( } } -template +template void tinygemm_kernel2( const scalar_t* __restrict__ A, - const at::Float8_e4m3fn* __restrict__ B, + const packed_t* __restrict__ B, scalar_t* __restrict__ C, scalar_t* __restrict__ Btmp, float* __restrict__ Ctmp, - float scale, + const float scale, + const float* __restrict__ bias, int64_t M, int64_t N, int64_t K, int64_t lda, int64_t ldb, int64_t ldc, - bool brg) { + bool brg, + bool do_unpack = true) { if (brg) { - brgemm2::apply(A, B, C, Btmp, Ctmp, scale, M, N, K, lda, ldb, ldc); + brgemm2::apply(A, B, C, Btmp, Ctmp, bias, scale, M, N, K, lda, ldb, ldc, do_unpack); return; } @@ -787,7 +870,7 @@ void tinygemm_kernel2( LAUNCH_TINYGEMM_KERNEL_NN2(1, 128); break; default: - TORCH_CHECK(false, "Unexpected block size, 1x", "nb_size"); + TORCH_CHECK(false, "Unexpected block size, 1x", nb_size); } } return; @@ -835,7 +918,7 @@ void tinygemm_kernel2( LAUNCH_TINYGEMM_KERNEL_NN2(4, 64); break; default: - TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size"); + TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", nb_size); } } } @@ -918,6 +1001,68 @@ void fp_scaled_mm_kernel_impl( }); } +template +void fp8_per_tensor_scaled_mm_kernel_impl( + scalar_t* __restrict__ out, + const scalar_t* __restrict__ mat1, + const packed_t* __restrict__ mat2, + const float scale2, + const float* __restrict__ bias, + scalar_t* __restrict__ buffer, + int64_t M, + int64_t N, + int64_t K, + int64_t mat1_strideM, + int64_t out_strideM, + int64_t buffer_size_per_thread) { + constexpr int64_t BLOCK_M = block_size_m(); + constexpr int64_t BLOCK_N = block_size_n(); + const int64_t MB = div_up(M, BLOCK_M); + const int64_t NB = div_up(N, BLOCK_N); + + const bool use_brgemm = can_use_brgemm(M); + const int64_t packed_K = get_row_size(K); + + // parallel on [MB, NB] + AT_DISPATCH_BOOL(bias != nullptr, has_bias, [&] { + parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) { + int tid = get_thread_num(); + scalar_t* __restrict__ Btmp = buffer + tid * buffer_size_per_thread; + float* __restrict__ Ctmp = (float*)((void*)(Btmp + MAX_CACHE_BLOCK_SIZE * BLOCK_N * K)); + + loop_2d(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) { + int64_t mb_start = mb * BLOCK_M; + int64_t mb_size = std::min(M - mb_start, BLOCK_M); + int64_t nb_start = nb * BLOCK_N; + int64_t nb_size = std::min(N - nb_start, BLOCK_N); + + // only do unpacking for the first row + bool do_unpack = (mb == mb0); + + tinygemm_kernel2( + /* A */ mat1 + mb_start * mat1_strideM, + /* B */ mat2 + nb_start * packed_K, + /* C */ out + mb_start * out_strideM + nb_start, + /* Btmp */ Btmp + nb_offset * BLOCK_N * K, + /* Ctmp */ Ctmp, + /* scale */ scale2, + /* bias */ has_bias ? bias + nb_start : nullptr, + /* M */ mb_size, + /* N */ nb_size, + /* K */ K, + /* lda */ mat1_strideM, + /* ldb */ nb_size, + /* ldc */ out_strideM, + /* brg */ use_brgemm, + /* do_unpack */ do_unpack); + }); + if (use_brgemm) { + at::native::cpublas::brgemm_release(); + } + }); + }); +} + } // anonymous namespace // tinygemm interface @@ -948,6 +1093,7 @@ void tinygemm_kernel( A, B, C, Btmp, Ctmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack); } +// tinygemm interface: per tensor quantization template void tinygemm_kernel( const scalar_t* __restrict__ A, @@ -955,15 +1101,20 @@ void tinygemm_kernel( scalar_t* __restrict__ C, scalar_t* __restrict__ Btmp, float* __restrict__ Ctmp, - float scale, + const float* __restrict__ bias, + const float scale2, int64_t M, int64_t N, int64_t K, int64_t lda, int64_t ldb, int64_t ldc, - bool brg) { - tinygemm_kernel2(A, B, C, Btmp, Ctmp, scale, M, N, K, lda, ldb, ldc, brg); + bool brg, + bool do_unpack) { + AT_DISPATCH_BOOL(bias != nullptr, has_bias, [&] { + tinygemm_kernel2( + A, B, C, Btmp, Ctmp, scale2, bias, M, N, K, lda, ldb, ldc, brg, do_unpack); + }); } template @@ -1070,23 +1221,26 @@ INSTANTIATE_TINYGEMM_TEMPLATE(at::Half, at::Float8_e4m3fn, float); INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16, uint8_t, uint8_t); INSTANTIATE_TINYGEMM_TEMPLATE(at::Half, uint8_t, uint8_t); -#define INSTANTIATE_TINYGEMM_TEMPLATE2(TYPE) \ - template void tinygemm_kernel( \ - const TYPE* __restrict__ A, \ - const at::Float8_e4m3fn* __restrict__ B, \ - TYPE* __restrict__ C, \ - TYPE* __restrict__ Btmp, \ - float* __restrict__ Ctmp, \ - float scale, \ - int64_t M, \ - int64_t N, \ - int64_t K, \ - int64_t lda, \ - int64_t ldb, \ - int64_t ldc, \ - bool brg) +#define INSTANTIATE_TINYGEMM_TEMPLATE_PER_TENSOR(TYPE) \ + template void tinygemm_kernel( \ + const TYPE* __restrict__ A, \ + const at::Float8_e4m3fn* __restrict__ B, \ + TYPE* __restrict__ C, \ + TYPE* __restrict__ Btmp, \ + float* __restrict__ Ctmp, \ + const float* __restrict__ bias, \ + float scale, \ + int64_t M, \ + int64_t N, \ + int64_t K, \ + int64_t lda, \ + int64_t ldb, \ + int64_t ldc, \ + bool brg, \ + bool do_unpack) -INSTANTIATE_TINYGEMM_TEMPLATE2(at::BFloat16); +INSTANTIATE_TINYGEMM_TEMPLATE_PER_TENSOR(at::BFloat16); +INSTANTIATE_TINYGEMM_TEMPLATE_PER_TENSOR(at::Half); inline const float* get_bias_data(const std::optional& bias, int64_t N) { if (bias.has_value()) { @@ -1178,6 +1332,58 @@ at::Tensor fp8_scaled_mm_cpu( return out; } +at::Tensor fp8_per_tensor_scaled_mm_cpu( + at::Tensor& mat1, + at::Tensor& mat2, + at::Tensor& scales2, + const std::optional& bias, + at::ScalarType out_dtype, + bool is_vnni) { + auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2); + + CHECK_INPUT(mat1); + CHECK_INPUT(mat2); + CHECK_INPUT(scales2); + + const int64_t M = mat1.size(0); + const int64_t N = mat2.size(0); + const int64_t K = mat2.size(1); + + CHECK_EQ(mat1.size(1), K); + CHECK_DIM(2, mat1); + CHECK_DIM(2, mat2); + + const auto st = mat1.scalar_type(); + // only the bf16 micro-kernels are implemented + TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf, "fp8_per_tensor_scaled_mm_cpu: expect A to be bfloat16 or half."); + TORCH_CHECK(st == out_dtype, "fp8_per_tensor_scaled_mm_cpu: expect A has same dtype with out_dtype."); + TORCH_CHECK(mat2.scalar_type() == at::kFloat8_e4m3fn, "fp8_per_tensor_scaled_mm_cpu: expect mat2 to be fp8_e4m3."); + TORCH_CHECK(scales2.scalar_type() == at::kFloat, "fp8_per_tensor_scaled_mm_cpu: expect scales2 to be float32."); + TORCH_CHECK(scales2.numel() == 1, "fp8_per_tensor_scaled_mm_cpu: expect scales2 to have one element."); + + auto out = at::empty({M, N}, mat1.options().dtype(out_dtype)); + auto buffer = alloc_thread_buffer(mat1.options(), K); + + const float scale_val = scales2.item(); + AT_DISPATCH_REDUCED_FLOATING_TYPES(out_dtype, "fp8_per_tensor_scaled_mm_kernel_impl", [&] { + fp8_per_tensor_scaled_mm_kernel_impl( + out.data_ptr(), + mat1.data_ptr(), + packed_w.data_ptr(), + scale_val, + get_bias_data(bias, N), + buffer.data_ptr(), + M, + N, + K, + mat1.stride(0), + out.stride(0), + buffer.size(-1)); + }); + + return out; +} + // mat1 : [M, K] bfloat16 // mat2 : [N, K / 2] uint8, actual layout: [N / BLOCK_N, K / 2, BLOCK_N, 2] // scales2: [N, K / G], actual layout: [N / BLOCK_N, K / G, BLOCK_N] diff --git a/python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp b/python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp index 6e63b20ad..59254ad7d 100644 --- a/python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp +++ b/python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp @@ -347,6 +347,14 @@ at::Tensor fp8_scaled_mm_cpu( at::ScalarType out_dtype, bool is_vnni); +at::Tensor fp8_per_tensor_scaled_mm_cpu( + at::Tensor& mat1, + at::Tensor& mat2, + at::Tensor& scales2, + const std::optional& bias, + at::ScalarType out_dtype, + bool is_vnni); + // mxfp4 gemm at::Tensor mxfp4_scaled_mm_cpu( at::Tensor& mat1, at::Tensor& mat2, at::Tensor& scales2, const std::optional& bias, bool is_vnni); @@ -840,6 +848,10 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "fp8_scaled_mm_cpu(Tensor mat1, Tensor mat2, Tensor scales2, int[] block_size, Tensor? bias, ScalarType " "out_dtype, bool is_vnni) -> Tensor"); m.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu); + m.def( + "fp8_per_tensor_scaled_mm_cpu(Tensor mat1, Tensor mat2, Tensor scales2, Tensor? bias, ScalarType " + "out_dtype, bool is_vnni) -> Tensor"); + m.impl("fp8_per_tensor_scaled_mm_cpu", torch::kCPU, &fp8_per_tensor_scaled_mm_cpu); // mxfp4 gemm m.def("mxfp4_scaled_mm_cpu(Tensor mat1, Tensor mat2, Tensor scales2, Tensor? bias, bool is_vnni) -> Tensor"); diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 8db9ac7f9..92cf17af6 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -1056,6 +1056,16 @@ class Fp8LinearMethod(LinearMethodBase): layer.input_scale.max(), requires_grad=False ) + if _is_cpu: + assert _is_cpu_amx_available, ( + "Fp8LinearMethod on CPU requires that CPU has AMX support" + ) + layer.weight = Parameter( + layer.weight.data.t().contiguous(), requires_grad=False + ) + _amx_process_weight_after_loading(layer, ["weight"]) + return + if self.use_marlin: if self.block_quant: layer.weight_block_size = self.quant_config.weight_block_size @@ -1141,6 +1151,17 @@ class Fp8LinearMethod(LinearMethodBase): bias=bias, ) + if use_intel_amx_backend(layer): + output = torch.ops.sgl_kernel.fp8_per_tensor_scaled_mm_cpu( + x, + layer.weight, + layer.weight_scale, + bias, + x.dtype, + True, # is_vnni + ) + return output.view(*x.shape[:-1], layer.weight.shape[0]) + if isinstance(x, tuple): # Pre-quantized activation from a fused RMSNorm+FP8 quant kernel: # x = (fp8_input, per_tensor_input_scale[, orig_dtype]). diff --git a/python/sglang/srt/model_executor/cpu_graph_runner.py b/python/sglang/srt/model_executor/cpu_graph_runner.py index 8c2a69e71..7abe4da16 100644 --- a/python/sglang/srt/model_executor/cpu_graph_runner.py +++ b/python/sglang/srt/model_executor/cpu_graph_runner.py @@ -468,6 +468,19 @@ def register_fake_ops(tp_size: int): N = mat2.shape[0] return mat1.new_empty(M, N, dtype=out_dtype) + @register_cpu_compile_fake("fp8_per_tensor_scaled_mm_cpu") + def _( + mat1, + mat2, + scale2, + bias, + out_dtype, + is_vnni, + ): + M = mat1.shape[0] + N = mat2.shape[0] + return mat1.new_empty(M, N, dtype=out_dtype) + @register_cpu_compile_fake("mxfp4_scaled_mm_cpu") def _(mat1, mat2, scales2, bias, is_vnni): sizes = list(mat1.shape) diff --git a/test/registered/cpu/test_gemm.py b/test/registered/cpu/test_gemm.py index d3057c0a2..9ee9365b5 100644 --- a/test/registered/cpu/test_gemm.py +++ b/test/registered/cpu/test_gemm.py @@ -184,6 +184,41 @@ class TestGemm(CustomTestCase): atol = rtol = precision[ref.dtype] torch.testing.assert_close(ref, out, atol=atol, rtol=rtol) + @parametrize( + M=[1, 11, 97], + N=[128, 224], + K=[512, 576], + scale_as_vector=[False, True], + has_bias=[False, True], + prepack=[False, True], + ) + def test_fp8_per_tensor_gemm(self, M, N, K, scale_as_vector, has_bias, prepack): + data = torch.randn(M, K, dtype=torch.bfloat16) / 10 + weight = torch.randn(N, K).to(torch.float8_e4m3fn) + scale = torch.tensor(0.01, dtype=torch.float32) + scales = scale.reshape(1) if scale_as_vector else scale + bias = torch.randn(N, dtype=torch.float32) if has_bias else None + + ref = torch.matmul(data.float(), weight.float().T) * scale + if bias is not None: + ref = ref + bias + ref = ref.bfloat16() + + kernel_weight = ( + torch.ops.sgl_kernel.convert_weight_packed(weight) if prepack else weight + ) + out = torch.ops.sgl_kernel.fp8_per_tensor_scaled_mm_cpu( + data, + kernel_weight, + scales, + bias, + data.dtype, + prepack, + ) + + atol = rtol = precision[ref.dtype] + torch.testing.assert_close(ref, out, atol=atol, rtol=rtol) + @parametrize(M=[1, 11], N=[128, 224], K=[512, 576], has_bias=[False, True]) def test_mxfp4_gemm(self, M, N, K, has_bias): prepack = True