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

[CPU] Add fp8_per_tensor_scaled_mm_cpu kernel (#32618)

Browse Source 
Co-authored-by: AKatydid <xinguojoe@gmail.com>
This commit is contained in:
Xinguo Zhu
2026-09-16 12:31:05 +08:00
committed by GitHub
co-authored by AKatydid
parent f60652a43e
commit 4678536df6
7 changed files with 350 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/aot/csrc/cpu/bmm.cpp
+1
View File
@@ -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,
python/sglang/kernels/aot/csrc/cpu/gemm.h
+3 -1
View File
@@ -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 <typename scalar_t>
python/sglang/kernels/aot/csrc/cpu/gemm_fp8.cpp
+249 -43
View File
@@ -22,6 +22,30 @@ inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ inpu
}
}
template <typename scalar_t>
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<scalar_t>;
using fVec = at::vec::Vectorized<float>;
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<scalar_t>(data0 * vscale + bias0, data1 * vscale + bias1);
out_vec.store(out + d);
}
for (; d < size; ++d) {
out[d] = static_cast<scalar_t>(input[d] * scale + bias[d]);
}
}
template <typename scalar_t>
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 <typename scalar_t, int BLOCK_M, int BLOCK_N>
template <typename scalar_t, typename packed_t, bool has_bias, int BLOCK_M, int BLOCK_N>
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<at::BFloat16, at::Float8_e4m3fn, float, has_bias, BLOC
}
};
template <int BLOCK_M, int BLOCK_N>
struct tinygemm_kernel_nn2<at::BFloat16, BLOCK_M, BLOCK_N> {
template <bool has_bias, int BLOCK_M, int BLOCK_N>
struct tinygemm_kernel_nn2<at::BFloat16, at::Float8_e4m3fn, has_bias, BLOCK_M, BLOCK_N> {
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<ROWS * COLS>{}(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<const float*>(A);
const uint16_t* b_ptr = reinterpret_cast<const uint16_t*>(B);
@@ -392,6 +427,9 @@ struct tinygemm_kernel_nn2<at::BFloat16, BLOCK_M, BLOCK_N> {
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,21 +441,31 @@ struct tinygemm_kernel_nn2<at::BFloat16, BLOCK_M, BLOCK_N> {
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) {
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<ROWS * COLS>{}([&](auto i) { vsum[i] = _mm512_setzero_ps(); });
// 2. accumulate across each block
for (int k = kb_start; k < kb_end; ++k) {
Unroll<ROWS * COLS>{}(compute, k);
}
// 3. apply scale
Unroll<ROWS * COLS>{}([&](auto i) { vc[i] = _mm512_fmadd_ps(vsum[i], vscale, vc[i]); });
}
auto storec = [&](auto i) {
constexpr int row = i / COLS;
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<ROWS * COLS>{}(storec);
@@ -539,8 +587,16 @@ struct tinygemm_kernel_nn<at::BFloat16, uint8_t, uint8_t, has_bias, BLOCK_M, BLO
block_size_K);
#define LAUNCH_TINYGEMM_KERNEL_NN2(MB_SIZE, NB_SIZE) \
tinygemm_kernel_nn2<scalar_t, MB_SIZE, NB_SIZE>::apply( \
A + mb_start * lda, B + nb_start * 2, C + mb_start * ldc + nb_start, scale, K, lda, ldb, ldc);
tinygemm_kernel_nn2<scalar_t, packed_t, has_bias, MB_SIZE, NB_SIZE>::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 <typename scalar_t, typename packed_t, typename param_t, bool has_bias>
struct brgemm {
@@ -562,8 +618,27 @@ struct brgemm {
TORCH_CHECK(false, "struct brgemm: primary template not implemented!");
}
};
template <typename scalar_t>
struct brgemm2 {};
template <typename scalar_t, typename packed_t, bool has_bias>
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 <bool has_bias>
struct brgemm<at::BFloat16, at::Float8_e4m3fn, float, has_bias> {
@@ -609,21 +684,23 @@ struct brgemm<at::BFloat16, at::Float8_e4m3fn, float, has_bias> {
}
};
template <>
struct brgemm2<at::BFloat16> {
template <bool has_bias>
struct brgemm2<at::BFloat16, at::Float8_e4m3fn, has_bias> {
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,9 +717,13 @@ struct brgemm2<at::BFloat16> {
// copy from Ctmp to C and mul scale
for (int m = 0; m < M; ++m) {
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);
}
}
}
};
template <bool has_bias>
@@ -743,23 +824,25 @@ void tinygemm_kernel(
}
}
template <typename scalar_t>
template <typename scalar_t, typename packed_t, bool has_bias>
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<scalar_t>::apply(A, B, C, Btmp, Ctmp, scale, M, N, K, lda, ldb, ldc);
brgemm2<scalar_t, packed_t, has_bias>::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 <typename scalar_t, typename packed_t>
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<packed_t>(M);
const int64_t packed_K = get_row_size<packed_t>(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<packed_t>(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<scalar_t, packed_t, has_bias>(
/* 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 <typename scalar_t>
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<scalar_t>(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<scalar_t, at::Float8_e4m3fn, has_bias>(
A, B, C, Btmp, Ctmp, scale2, bias, M, N, K, lda, ldb, ldc, brg, do_unpack);
});
}
template <typename scalar_t>
@@ -1070,13 +1221,14 @@ 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) \
#define INSTANTIATE_TINYGEMM_TEMPLATE_PER_TENSOR(TYPE) \
template void tinygemm_kernel<TYPE>( \
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, \
@@ -1084,9 +1236,11 @@ INSTANTIATE_TINYGEMM_TEMPLATE(at::Half, uint8_t, uint8_t);
int64_t lda, \
int64_t ldb, \
int64_t ldc, \
bool brg)
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<at::Tensor>& 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<at::Tensor>& 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<float>();
AT_DISPATCH_REDUCED_FLOATING_TYPES(out_dtype, "fp8_per_tensor_scaled_mm_kernel_impl", [&] {
fp8_per_tensor_scaled_mm_kernel_impl<scalar_t, at::Float8_e4m3fn>(
out.data_ptr<scalar_t>(),
mat1.data_ptr<scalar_t>(),
packed_w.data_ptr<at::Float8_e4m3fn>(),
scale_val,
get_bias_data(bias, N),
buffer.data_ptr<scalar_t>(),
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]
python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp
+12
View File
@@ -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<at::Tensor>& 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<at::Tensor>& 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");
python/sglang/srt/layers/quantization/fp8.py
+21
View File
@@ -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]).
python/sglang/srt/model_executor/cpu_graph_runner.py
+13
View File
@@ -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)
test/registered/cpu/test_gemm.py
+35
View File
@@ -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