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

[CPU] Add GPT-OSS model optimization for CPU (#16775)

Browse Source 
Co-authored-by: mingfeima <mingfei.ma@intel.com>
Co-authored-by: jianan-gu <jianan.gu@intel.com>
This commit is contained in:
blzheng
2026-05-29 16:05:26 +08:00
committed by GitHub
co-authored by mingfeima jianan-gu
parent 5601b7139d
commit 3ecf2c76ad
35 changed files with 2000 additions and 530 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/cpu/gemm_fp8.cpp
+83
View File
@@ -62,6 +62,23 @@ inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__
}
}
template <>
inline void
copy_add_stub(float* __restrict__ out, const float* __restrict__ input, const float* __restrict__ bias, int64_t size) {
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = fVec::size();
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
fVec data = fVec::loadu(input + d) + fVec::loadu(bias + d);
data.store(out + d);
}
for (; d < size; ++d) {
out[d] = input[d] + bias[d];
}
}
inline void unpack_B(
at::BFloat16* __restrict__ Btmp,
const at::Float8_e4m3fn* __restrict__ packed_B,
@@ -913,6 +930,7 @@ void tinygemm_kernel(
scalar_t* __restrict__ C,
scalar_t* __restrict__ Btmp,
float* __restrict__ Ctmp,
const float* __restrict__ Bbias,
const float* __restrict__ scale,
int64_t M,
int64_t N,
@@ -923,6 +941,11 @@ void tinygemm_kernel(
bool brg,
int64_t block_size_K,
bool do_unpack) {
if (Bbias != nullptr) {
tinygemm_kernel<scalar_t, at::Float8_e4m3fn, float, true>(
A, B, C, Btmp, Ctmp, scale, Bbias, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
return;
}
tinygemm_kernel<scalar_t, at::Float8_e4m3fn, float, false>(
A, B, C, Btmp, Ctmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
}
@@ -952,6 +975,7 @@ void tinygemm_kernel(
scalar_t* __restrict__ C,
scalar_t* __restrict__ Btmp,
float* __restrict__ Ctmp,
const float* __restrict__ Bbias,
const uint8_t* __restrict__ scale,
int64_t M,
int64_t N,
@@ -962,10 +986,68 @@ void tinygemm_kernel(
bool brg,
int64_t block_size_K,
bool do_unpack) {
if (Bbias != nullptr) {
tinygemm_kernel<scalar_t, uint8_t, uint8_t, true>(
A, B, C, Btmp, Ctmp, scale, Bbias, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
return;
}
tinygemm_kernel<scalar_t, uint8_t, uint8_t, false>(
A, B, C, Btmp, Ctmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
}
// tinygemm interface
template <typename scalar_t>
void tinygemm_kernel(
const scalar_t* __restrict__ A,
const at::Float8_e4m3fn* __restrict__ B,
float* __restrict__ C,
scalar_t* __restrict__ Btmp,
const float* __restrict__ Bbias,
const float* __restrict__ scale,
int64_t M,
int64_t N,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc,
bool brg,
int64_t block_size_K,
bool do_unpack) {
if (Bbias != nullptr) {
tinygemm_kernel<scalar_t, at::Float8_e4m3fn, float, true>(
A, B, C, Btmp, scale, Bbias, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
return;
}
tinygemm_kernel<scalar_t, at::Float8_e4m3fn, float, false>(
A, B, C, Btmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
}
template <typename scalar_t>
void tinygemm_kernel(
const scalar_t* __restrict__ A,
const uint8_t* __restrict__ B,
float* __restrict__ C,
scalar_t* __restrict__ Btmp,
const float* __restrict__ Bbias,
const uint8_t* __restrict__ scale,
int64_t M,
int64_t N,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc,
bool brg,
int64_t block_size_K,
bool do_unpack) {
if (Bbias != nullptr) {
tinygemm_kernel<scalar_t, uint8_t, uint8_t, true>(
A, B, C, Btmp, scale, Bbias, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
return;
}
tinygemm_kernel<scalar_t, uint8_t, uint8_t, false>(
A, B, C, Btmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
}
#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE_A, TYPE_B, TYPE_S) \
template void tinygemm_kernel<TYPE_A>( \
const TYPE_A* __restrict__ A, \
@@ -973,6 +1055,7 @@ void tinygemm_kernel(
TYPE_A* __restrict__ C, \
TYPE_A* __restrict__ Btmp, \
float* __restrict__ Ctmp, \
const float* __restrict__ Bbias, \
const TYPE_S* __restrict__ scale, \
int64_t M, \
int64_t N, \