chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[CPU] Add FP8 Bmm support (#9744)

Browse Source 
Co-authored-by: Fan Yin <1106310035@qq.com>
This commit is contained in:
blzheng
2026-03-19 13:19:48 +08:00
committed by GitHub
parent c2b01bd2fc
commit cd22aa27a9
14 changed files with 584 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/sglang/srt/models/bailing_moe_linear.py
View File

@@ -1208,14 +1208,6 @@ class BailingMoELinearForCausalLM(nn.Module):
)
if _is_hip:
self_attn.w_scale *= 2.0
# TODO: remove this after adding FP8 support in bmm cpu kernel
if _is_cpu and _is_cpu_amx_available and w.dtype == torch.float8_e4m3fn:
self_attn.w_kc = (
self_attn.w_kc.to(torch.bfloat16) * self_attn.w_scale
)
self_attn.w_vc = (
self_attn.w_vc.to(torch.bfloat16) * self_attn.w_scale
)
else:
num_tiles_k = self_attn.qk_nope_head_dim // weight_block_size[1]
num_tiles_n = self_attn.v_head_dim // weight_block_size[0]

72
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py
View File

@@ -16,6 +16,7 @@ from sglang.srt.layers.quantization.fp8_kernel import (
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.models.deepseek_common.utils import (
FORWARD_ABSORB_CORE_ATTENTION_BACKENDS,
_is_cpu,
_is_cublas_ge_129,
_is_cuda,
_is_gfx95_supported,
@@ -268,18 +269,24 @@ class DeepseekMLAForwardMixin:
)
elif self.w_kc.dtype == torch.float8_e4m3fn:
# fix bmm_fp8 error under cublas12.9 caused by bumpallocator, detail in pr#11612
q_nope_val, q_nope_scale = per_tensor_quant_mla_fp8(
q_nope.transpose(0, 1),
(
torch.zeros((1,), dtype=torch.float32, device=q_nope.device)
if _is_cublas_ge_129
else zero_allocator.allocate(1)
),
)
q_nope_out = bmm_fp8(
q_nope_val, self.w_kc, q_nope_scale, self.w_scale, torch.bfloat16
)
if _is_cpu:
q_nope_out = torch.bmm(
q_nope.to(torch.bfloat16).transpose(0, 1),
self.w_kc.to(torch.bfloat16) * self.w_scale,
)
else:
# fix bmm_fp8 error under cublas12.9 caused by bumpallocator, detail in pr#11612
q_nope_val, q_nope_scale = per_tensor_quant_mla_fp8(
q_nope.transpose(0, 1),
(
torch.zeros((1,), dtype=torch.float32, device=q_nope.device)
if _is_cublas_ge_129
else zero_allocator.allocate(1)
),
)
q_nope_out = bmm_fp8(
q_nope_val, self.w_kc, q_nope_scale, self.w_scale, torch.bfloat16
)
else:
q_nope_out = torch.bmm(q_nope.transpose(0, 1), self.w_kc)
@@ -455,22 +462,31 @@ class DeepseekMLAForwardMixin:
attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2)
elif self.w_vc.dtype == torch.float8_e4m3fn:
attn_output_val, attn_output_scale = per_tensor_quant_mla_fp8(
attn_output.transpose(0, 1),
(
torch.zeros((1,), dtype=torch.float32, device=attn_output.device)
if _is_cublas_ge_129
else zero_allocator.allocate(1)
),
)
attn_bmm_output = bmm_fp8(
attn_output_val,
self.w_vc,
attn_output_scale,
self.w_scale,
torch.bfloat16,
)
attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2)
if _is_cpu:
attn_bmm_output = torch.bmm(
attn_output.to(torch.bfloat16).transpose(0, 1),
self.w_vc.to(torch.bfloat16) * self.w_scale,
)
attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2)
else:
attn_output_val, attn_output_scale = per_tensor_quant_mla_fp8(
attn_output.transpose(0, 1),
(
torch.zeros(
(1,), dtype=torch.float32, device=attn_output.device
)
if _is_cublas_ge_129
else zero_allocator.allocate(1)
),
)
attn_bmm_output = bmm_fp8(
attn_output_val,
self.w_vc,
attn_output_scale,
self.w_scale,
torch.bfloat16,
)
attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2)
else:
if is_in_piecewise_cuda_graph():
# torch dynamo requires out= op was called where output tensor was non-contiguous

3
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py
View File

@@ -100,6 +100,7 @@ class DeepseekMLACpuForwardMixin:
else None
)
),
self.w_scale,
True, # is_vnni
self.weight_block_size,
self.q_lora_rank,
@@ -144,7 +145,7 @@ class DeepseekMLACpuForwardMixin:
attn_output.transpose(0, 1),
self.w_vc,
True, # is_vnni
None, # scale
self.w_scale, # scale
)
attn_output = output
output, _ = self.o_proj(attn_output)

10
python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py
View File

@@ -46,8 +46,6 @@ from sglang.srt.model_loader.utils import (
)
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.deepseek_common.utils import (
_is_cpu,
_is_cpu_amx_available,
_is_cuda,
_is_fp8_fnuz,
_is_hip,
@@ -583,14 +581,6 @@ class DeepseekV2WeightLoaderMixin:
)
if _is_hip:
self_attn.w_scale *= 2.0
# TODO: remove this after adding FP8 support in bmm cpu kernel
if _is_cpu and _is_cpu_amx_available and w.dtype == torch.float8_e4m3fn:
self_attn.w_kc = (
self_attn.w_kc.to(torch.bfloat16) * self_attn.w_scale
)
self_attn.w_vc = (
self_attn.w_vc.to(torch.bfloat16) * self_attn.w_scale
)
else:
num_tiles_k = self_attn.qk_nope_head_dim // weight_block_size[1]
num_tiles_n = self_attn.v_head_dim // weight_block_size[0]

12
python/sglang/srt/models/longcat_flash.py
View File

@@ -776,18 +776,6 @@ class LongcatFlashForCausalLM(nn.Module):
)
if _is_hip:
self_attn.w_scale *= 2.0
# TODO: remove this after adding FP8 support in bmm cpu kernel
if (
_is_cpu
and _is_cpu_amx_available
and w.dtype == torch.float8_e4m3fn
):
self_attn.w_kc = (
self_attn.w_kc.to(torch.bfloat16) * self_attn.w_scale
)
self_attn.w_vc = (
self_attn.w_vc.to(torch.bfloat16) * self_attn.w_scale
)
else:
num_tiles_k = self_attn.qk_nope_head_dim // weight_block_size[1]
num_tiles_n = self_attn.v_head_dim // weight_block_size[0]

4
python/sglang/srt/models/longcat_flash_nextn.py
View File

@@ -426,10 +426,6 @@ class LongcatFlashForCausalLMNextN(LongcatFlashForCausalLM):
)
if _is_hip:
self_attn.w_scale *= 2.0
# TODO: remove this after adding FP8 support in bmm cpu kernel
if _is_cpu and _is_cpu_amx_available and w.dtype == torch.float8_e4m3fn:
self_attn.w_kc = self_attn.w_kc.to(torch.bfloat16) * self_attn.w_scale
self_attn.w_vc = self_attn.w_vc.to(torch.bfloat16) * self_attn.w_scale
else:
num_tiles_k = self_attn.qk_nope_head_dim // weight_block_size[1]
num_tiles_n = self_attn.v_head_dim // weight_block_size[0]

102
sgl-kernel/csrc/cpu/bmm.cpp
View File

@@ -4,11 +4,11 @@
namespace {
template <typename scalar_t>
template <typename scalar_t, typename packed_t>
void bmm_kernel_impl(
scalar_t* __restrict__ out,
const scalar_t* __restrict__ mat1,
const scalar_t* __restrict__ mat2,
const packed_t* __restrict__ mat2,
int64_t B,
int64_t M,
int64_t N,
@@ -67,6 +67,67 @@ void bmm_kernel_impl(
});
}
template <>
void bmm_kernel_impl(
at::BFloat16* __restrict__ out,
const at::BFloat16* __restrict__ mat1,
const at::Float8_e4m3fn* __restrict__ mat2,
int64_t B,
int64_t M,
int64_t N,
int64_t K,
int64_t mat1_strideB,
int64_t mat1_strideM,
int64_t out_strideB,
int64_t out_strideM,
float scale) {
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);
// mat2 contiguous in [B, N, K]
int64_t mat2_strideB = N * K;
int64_t mat2_strideN = K;
const bool use_brgemm = can_use_brgemm<at::BFloat16>(M);
// parallel on [B, MB, NB]
parallel_2d(B * MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
// for brgemm, use float32 for accumulate
alignas(64) float Ctmp[BLOCK_M * BLOCK_N];
// for brgemm when mat2 is float8_e4m3
alignas(64) at::BFloat16 Btmp[BLOCK_N * BLOCK_K];
loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
int64_t bs = mb / MB;
int64_t mb_start = (mb % 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);
tinygemm_kernel(
/* A */ mat1 + bs * mat1_strideB + mb_start * mat1_strideM,
/* B */ mat2 + bs * mat2_strideB + nb_start * mat2_strideN /* nb * BLOCK_N * K */,
/* C */ out + bs * out_strideB + mb_start * out_strideM + nb_start,
/* Btmp*/ Btmp,
/* Ctmp*/ Ctmp,
/*scale*/ scale,
/* M */ mb_size,
/* N */ nb_size,
/* K */ K,
/* lda */ mat1_strideM,
/* ldb */ nb_size,
/* ldc */ out_strideM,
/* brg */ use_brgemm);
});
if (use_brgemm) {
at::native::cpublas::brgemm_release();
}
});
}
} // anonymous namespace
// mat1 : [B, M, K]
@@ -94,7 +155,7 @@ void bmm_cpu(
int64_t N = mat2.size(1);
int64_t K = mat1.size(2);
TORCH_CHECK(!scale.has_value(), "bmm: do not support fp8 weight for now.")
const bool use_fp8_w8a16 = scale.has_value();
TORCH_CHECK(N % 32 == 0, "tinygemm requires N to be 32x.");
int64_t mat1_strideB = mat1.stride(0);
@@ -105,12 +166,32 @@ void bmm_cpu(
// check shapes
TORCH_CHECK(mat2.size(0) == B && mat2.size(2) == K, "bmm: mat2 shape mismatch!");
TORCH_CHECK(out.size(0) == B && out.size(1) == M, "bmm: out shape mismatch!");
if (!use_fp8_w8a16) {
AT_DISPATCH_REDUCED_FLOATING_TYPES(mat1.scalar_type(), "bmm_kernel_impl", [&] {
bmm_kernel_impl<scalar_t, scalar_t>(
out.data_ptr<scalar_t>(),
mat1.data_ptr<scalar_t>(),
packed_w.data_ptr<scalar_t>(),
B,
M,
N,
K,
mat1_strideB,
mat1_strideM,
out_strideB,
out_strideM);
});
} else { // fp8 bmm
float scale_val = 0.f;
AT_DISPATCH_REDUCED_FLOATING_TYPES(mat1.scalar_type(), "bmm_kernel_impl", [&] {
bmm_kernel_impl<scalar_t>(
out.data_ptr<scalar_t>(),
mat1.data_ptr<scalar_t>(),
packed_w.data_ptr<scalar_t>(),
auto scale_tensor = scale.value();
TORCH_CHECK(scale_tensor.ndimension() == 0, "bmm: expect scale to be 0-dim tensor.");
scale_val = scale_tensor.item<float>();
bmm_kernel_impl<at::BFloat16, at::Float8_e4m3fn>(
out.data_ptr<at::BFloat16>(),
mat1.data_ptr<at::BFloat16>(),
packed_w.data_ptr<at::Float8_e4m3fn>(),
B,
M,
N,
@@ -118,6 +199,7 @@ void bmm_cpu(
mat1_strideB,
mat1_strideM,
out_strideB,
out_strideM);
});
out_strideM,
scale_val);
}
}

18
sgl-kernel/csrc/cpu/gemm.h
View File

@@ -232,6 +232,7 @@ void tinygemm_kernel(
int64_t ldc,
bool brg);
// block quantization
template <typename scalar_t>
void tinygemm_kernel(
const scalar_t* __restrict__ A,
@@ -250,6 +251,23 @@ void tinygemm_kernel(
int64_t block_size_K,
bool do_unpack = true);
// per tensor quantization
template <typename scalar_t>
void tinygemm_kernel(
const scalar_t* __restrict__ A,
const at::Float8_e4m3fn* __restrict__ B,
scalar_t* __restrict__ C,
scalar_t* __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);
template <typename scalar_t>
void tinygemm_kernel(
scalar_t* C,

311
sgl-kernel/csrc/cpu/gemm_fp8.cpp
View File

@@ -42,6 +42,25 @@ inline void copy_add_stub(
out[d] = static_cast<scalar_t>(input[d] + bias[d]);
}
}
template <typename scalar_t>
inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int 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);
int d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
fVec data0 = fVec::loadu(input + d) * vscale;
fVec data1 = fVec::loadu(input + d + fVec::size()) * vscale;
bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
out_vec.store(out + d);
}
for (; d < size; ++d) {
out[d] = static_cast<scalar_t>(input[d] * scale);
}
}
inline void unpack_B(
at::BFloat16* __restrict__ Btmp,
@@ -100,6 +119,41 @@ inline void unpack_B(
#endif
}
inline void unpack_B(
at::BFloat16* __restrict__ Btmp,
const at::Float8_e4m3fn* __restrict__ packed_B,
int N,
int K,
int ldb,
int ldb_tmp) {
#if defined(CPU_CAPABILITY_AVX512)
// [K/2, N, 2]
const int K2 = K >> 1;
const int ldb2 = ldb; // ldb * 2 >> 1;
const uint16_t* b_ptr = reinterpret_cast<const uint16_t*>(packed_B);
// prefetch distance
constexpr int PREFETCH_SIZE_K = 64;
#pragma GCC unroll 4
for (int k = 0; k < K2; ++k) {
__m512i b8 = _mm512_loadu_si512(b_ptr + k * ldb2);
if constexpr (PREFETCH_SIZE_K > 0) {
_mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb2, _MM_HINT_T0);
}
__m256i b8_0 = _mm512_extracti32x8_epi32(b8, 0);
__m256i b8_1 = _mm512_extracti32x8_epi32(b8, 1);
__m512bh bf16_0 = CVT_FP8_TO_BF16(b8_0);
__m512bh bf16_1 = CVT_FP8_TO_BF16(b8_1);
_mm512_storeu_si512(Btmp + k * ldb_tmp * 2 + 0, (__m512i)bf16_0);
_mm512_storeu_si512(Btmp + k * ldb_tmp * 2 + 32, (__m512i)bf16_1);
}
#else
TORCH_CHECK(false, "unpack_B: scalar path not implemented!");
#endif
}
template <typename scalar_t, typename packed_t, bool has_bias, int BLOCK_M, int BLOCK_N>
struct tinygemm_kernel_nn {
static inline void apply(
@@ -117,6 +171,20 @@ struct tinygemm_kernel_nn {
}
};
template <typename scalar_t, 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,
scalar_t* __restrict__ C,
float scale,
int K,
int lda,
int ldb,
int ldc) {
TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
}
};
#if defined(CPU_CAPABILITY_AVX512)
template <bool has_bias, int BLOCK_M, int BLOCK_N>
struct tinygemm_kernel_nn<at::BFloat16, at::Float8_e4m3fn, has_bias, BLOCK_M, BLOCK_N> {
@@ -221,6 +289,76 @@ struct tinygemm_kernel_nn<at::BFloat16, at::Float8_e4m3fn, has_bias, BLOCK_M, BL
Unroll<ROWS * COLS>{}(storec);
}
};
template <int BLOCK_M, int BLOCK_N>
struct tinygemm_kernel_nn2<at::BFloat16, 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) {
constexpr int ROWS = BLOCK_M;
constexpr int COLS = BLOCK_N / 16;
// prefetch distance
constexpr int PREFETCH_SIZE_K = 64;
__m512bh va;
__m512bh vb[COLS];
__m512 vc[ROWS * COLS];
const __m512 vscale = _mm512_set1_ps(scale);
auto loadc = [&](auto i) { 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 float* a_ptr = reinterpret_cast<const float*>(A);
const uint16_t* b_ptr = reinterpret_cast<const uint16_t*>(B);
auto compute = [&](auto i, int k) {
constexpr int row = i / COLS;
constexpr int col = i % COLS;
if constexpr (col == 0) {
va = (__m512bh)(_mm512_set1_ps(a_ptr[row * lda2 + k]));
}
if constexpr (row == 0) {
if constexpr (col % 2 == 0) {
__m512i b8 = _mm512_loadu_si512(b_ptr + k * ldb2 + col * 16);
if constexpr (PREFETCH_SIZE_K > 0) {
_mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb2 + col * 16, _MM_HINT_T0);
}
vb[col + 0] = CVT_FP8_TO_BF16(_mm512_extracti32x8_epi32(b8, 0));
vb[col + 1] = CVT_FP8_TO_BF16(_mm512_extracti32x8_epi32(b8, 1));
}
}
vc[i] = _mm512_dpbf16_ps(vc[i], va, vb[col]);
};
for (int k = 0; k < K2; ++k) {
Unroll<ROWS * COLS>{}(compute, k);
}
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)));
}
};
Unroll<ROWS * COLS>{}(storec);
}
};
#endif
#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE) \
@@ -236,6 +374,10 @@ struct tinygemm_kernel_nn<at::BFloat16, at::Float8_e4m3fn, has_bias, BLOCK_M, BL
ldc, \
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);
template <typename scalar_t, typename packed_t, bool has_bias>
struct brgemm {
static inline void apply(
@@ -256,6 +398,8 @@ struct brgemm {
TORCH_CHECK(false, "struct brgemm: primary template not implemented!");
}
};
template <typename scalar_t>
struct brgemm2 {};
template <bool has_bias>
struct brgemm<at::BFloat16, at::Float8_e4m3fn, has_bias> {
@@ -301,6 +445,42 @@ struct brgemm<at::BFloat16, at::Float8_e4m3fn, has_bias> {
}
};
template <>
struct brgemm2<at::BFloat16> {
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,
int M,
int N,
int K,
int lda,
int ldb,
int ldc) {
constexpr int BLOCK_N = block_size_n();
// [BLOCK_K, BLOCK_N] -> [BLOCK_K / 2, BLOCK_N * 2]
const int ldb_tmp = block_size_n();
// accumulate across K per BLOCK_K
for (int k = 0; k < K; k += BLOCK_K) {
int kb_size = std::min(BLOCK_K, K - k);
unpack_B(Btmp, B + k * ldb, N, kb_size, ldb, ldb_tmp);
const bool add_C = (k != 0);
at::native::cpublas::brgemm(M, N, kb_size, lda, ldb_tmp, BLOCK_N, add_C, A + k, Btmp, Ctmp);
}
// 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);
}
}
};
template <typename scalar_t, bool has_bias>
void tinygemm_kernel(
const scalar_t* __restrict__ A,
@@ -356,7 +536,103 @@ void tinygemm_kernel(
}
}
}
template <typename scalar_t>
void tinygemm_kernel2(
const scalar_t* __restrict__ A,
const at::Float8_e4m3fn* __restrict__ B,
scalar_t* __restrict__ C,
scalar_t* __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) {
if (brg) {
brgemm2<scalar_t>::apply(A, B, C, Btmp, Ctmp, scale, M, N, K, lda, ldb, ldc);
return;
}
// pattern: 1-8-8
if (M == 1) {
constexpr int64_t BLOCK_N = 128;
const int64_t NB = div_up(N, BLOCK_N);
int64_t mb_start = 0;
for (int64_t nb = 0; nb < NB; ++nb) {
int64_t nb_start = nb * BLOCK_N;
int64_t nb_size = std::min(BLOCK_N, N - nb_start);
switch (nb_size >> 4) {
case 2:
LAUNCH_TINYGEMM_KERNEL_NN2(1, 32);
break;
case 4:
LAUNCH_TINYGEMM_KERNEL_NN2(1, 64);
break;
case 6:
LAUNCH_TINYGEMM_KERNEL_NN2(1, 96);
break;
case 8:
LAUNCH_TINYGEMM_KERNEL_NN2(1, 128);
break;
default:
TORCH_CHECK(false, "Unexpected block size, 1x", "nb_size");
}
}
return;
}
// pattern: 1-4-16
constexpr int64_t BLOCK_M = 4;
constexpr int64_t BLOCK_N = 64;
const int64_t MB = div_up(M, BLOCK_M);
const int64_t NB = div_up(N, BLOCK_N);
for (int64_t mb = 0; mb < MB; ++mb) {
int64_t mb_start = mb * BLOCK_M;
int64_t mb_size = std::min(BLOCK_M, M - mb_start);
for (int64_t nb = 0; nb < NB; ++nb) {
int64_t nb_start = nb * BLOCK_N;
int64_t nb_size = std::min(BLOCK_N, N - nb_start);
switch (mb_size << 4 | nb_size >> 4) {
// mb_size = 1
case 0x12:
LAUNCH_TINYGEMM_KERNEL_NN2(1, 32);
break;
case 0x14:
LAUNCH_TINYGEMM_KERNEL_NN2(1, 64);
break;
// mb_size = 2
case 0x22:
LAUNCH_TINYGEMM_KERNEL_NN2(2, 32);
break;
case 0x24:
LAUNCH_TINYGEMM_KERNEL_NN2(2, 64);
break;
// mb_size = 3
case 0x32:
LAUNCH_TINYGEMM_KERNEL_NN2(3, 32);
break;
case 0x34:
LAUNCH_TINYGEMM_KERNEL_NN2(3, 64);
break;
// mb_size = 4
case 0x42:
LAUNCH_TINYGEMM_KERNEL_NN2(4, 32);
break;
case 0x44:
LAUNCH_TINYGEMM_KERNEL_NN2(4, 64);
break;
default:
TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
}
}
}
}
template <typename scalar_t>
void fp8_scaled_mm_kernel_impl(
scalar_t* __restrict__ out,
@@ -450,7 +726,23 @@ void tinygemm_kernel(
tinygemm_kernel<scalar_t, false>(
A, B, C, Btmp, Ctmp, 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 at::Float8_e4m3fn* __restrict__ B,
scalar_t* __restrict__ C,
scalar_t* __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) {
tinygemm_kernel2<scalar_t>(A, B, C, Btmp, Ctmp, scale, M, N, K, lda, ldb, ldc, brg);
}
#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE) \
template void tinygemm_kernel<TYPE>( \
const TYPE* __restrict__ A, \
@@ -469,8 +761,25 @@ void tinygemm_kernel(
int64_t block_size_K, \
bool do_unpack)
#define INSTANTIATE_TINYGEMM_TEMPLATE2(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, \
float scale, \
int64_t M, \
int64_t N, \
int64_t K, \
int64_t lda, \
int64_t ldb, \
int64_t ldc, \
bool brg)
INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);
INSTANTIATE_TINYGEMM_TEMPLATE2(at::BFloat16);
at::Tensor fp8_scaled_mm_cpu(
at::Tensor& mat1,

6
sgl-kernel/csrc/cpu/qkv_proj.cpp
View File

@@ -434,6 +434,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> qkv_proj_with_rope(
std::optional<at::Tensor> q_a_proj_scale,
std::optional<at::Tensor> q_b_proj_scale,
std::optional<at::Tensor> kv_a_proj_scale,
std::optional<at::Tensor> w_scale,
bool is_vnni,
std::optional<std::vector<int64_t>> block_size) {
RECORD_FUNCTION(
@@ -601,10 +602,9 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> qkv_proj_with_rope(
qb.as_strided_({num_seqs, num_heads, qk_head_dim}, {num_heads * qk_head_dim, qk_head_dim, 1});
// stage 4: bmm
std::optional<at::Tensor> scale;
auto q_nope = qb.narrow(2, 0, qk_nope_head_dim).transpose_(0, 1);
auto q_nope_out = q_input.narrow(2, 0, kv_lora_rank).transpose_(0, 1);
bmm_cpu(q_nope_out, q_nope, w_kc, is_vnni, scale);
bmm_cpu(q_nope_out, q_nope, w_kc, is_vnni, w_scale);
// stage 5: rope
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "rotary_emb_kernel_impl", [&] {
@@ -643,6 +643,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> qkv_proj_with_rope_fused_weight(
bool use_fp8_w8a16,
std::optional<at::Tensor> qkv_a_proj_scale,
std::optional<at::Tensor> q_b_proj_scale,
std::optional<at::Tensor> w_scale,
bool is_vnni,
std::optional<std::vector<int64_t>> block_size,
int64_t q_lora_rank,
@@ -696,6 +697,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> qkv_proj_with_rope_fused_weight(
q_a_proj_s,
q_b_proj_scale,
kv_a_proj_s,
w_scale,
is_vnni,
block_size);
}

8
sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
View File

@@ -245,6 +245,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> qkv_proj_with_rope(
std::optional<at::Tensor> q_a_proj_scale,
std::optional<at::Tensor> q_b_proj_scale,
std::optional<at::Tensor> kv_a_proj_scale,
std::optional<at::Tensor> w_scale,
bool is_vnni,
std::optional<std::vector<int64_t>> block_size);
@@ -262,6 +263,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> qkv_proj_with_rope_fused_weight(
bool use_fp8_w8a16,
std::optional<at::Tensor> qkv_a_proj_scale,
std::optional<at::Tensor> q_b_proj_scale,
std::optional<at::Tensor> w_scale,
bool is_vnni,
std::optional<std::vector<int64_t>> block_size,
int64_t q_lora_rank,
@@ -515,14 +517,14 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
"qkv_proj_with_rope(Tensor hidden_states, Tensor q_a_proj_weight, Tensor q_b_proj_weight, Tensor "
"kv_a_proj_weight, Tensor w_kc, Tensor q_a_layernorm_weight, Tensor kv_a_layernorm_weight, Tensor positions, "
"Tensor cos_sin_cache, float eps, bool use_int8_w8a8, bool use_fp8_w8a16, Tensor? q_a_proj_scale, Tensor? "
"q_b_proj_scale, Tensor? "
"kv_a_proj_scale, bool is_vnni, int[]? block_size) -> (Tensor, Tensor, Tensor)");
"q_b_proj_scale, Tensor? kv_a_proj_scale, Tensor? w_scale, "
"bool is_vnni, int[]? block_size) -> (Tensor, Tensor, Tensor)");
m.impl("qkv_proj_with_rope", torch::kCPU, &qkv_proj_with_rope);
m.def(
"qkv_proj_with_rope_fused_weight(Tensor hidden_states, Tensor qkv_a_proj_weight, Tensor q_b_proj_weight, "
"Tensor w_kc, Tensor q_a_layernorm_weight, Tensor kv_a_layernorm_weight, Tensor positions, "
"Tensor cos_sin_cache, float eps, bool use_int8_w8a8, bool use_fp8_w8a16, Tensor? qkv_a_proj_scale, Tensor? "
"q_b_proj_scale,"
"q_b_proj_scale, Tensor? w_scale,"
"bool is_vnni, int[]? block_size, int q_lora_rank, int kv_lora_rank,"
"int qk_rope_head_dim) -> (Tensor, Tensor, Tensor)");
m.impl("qkv_proj_with_rope_fused_weight", torch::kCPU, &qkv_proj_with_rope_fused_weight);

95
test/srt/cpu/test_bmm.py Normal file
View File

@@ -0,0 +1,95 @@
import itertools
import unittest
# TODO: use interface in cpu.py
import torch
import torch.nn as nn
from utils import precision
from sglang.srt.layers.quantization.fp8_utils import input_to_float8
from sglang.test.test_utils import CustomTestCase
torch.manual_seed(1234)
class Mod(nn.Module):
def __init__(self, input_channel, output_channel, has_bias):
super(Mod, self).__init__()
self.linear = torch.nn.Linear(input_channel, output_channel, has_bias)
def forward(self, x):
return self.linear(x)
class TestBmm(CustomTestCase):
M = [1, 2, 11, 111]
N = [128 + 32, 512]
K = [512 + 32, 128 + 32]
B = [1, 16, 17]
chunk = [True, False]
def _get_bmm_inputs(self, B, M, N, K, chunk, dtype):
if chunk:
mat1 = (
torch.randn(M, B, K + 64, dtype=dtype).narrow(2, 0, K).transpose_(0, 1)
)
mat2 = torch.randn(B, N, K, dtype=dtype).transpose_(1, 2)
mat3 = (
torch.randn(M, B, N + 64, dtype=dtype).narrow(2, 0, N).transpose_(0, 1)
)
else:
mat1 = torch.randn(M, B, K, dtype=dtype).transpose_(0, 1)
mat2 = torch.randn(B, N, K, dtype=dtype).transpose_(1, 2)
mat3 = torch.randn(M, B, N, dtype=dtype).transpose_(0, 1)
return mat1, mat2, mat3
def _bf16_bmm(self, B, M, N, K, chunk, dtype=torch.bfloat16):
mat1, mat2, mat3 = self._get_bmm_inputs(B, M, N, K, chunk, dtype)
ref = torch.bmm(mat1, mat2)
mat2_t = mat2.transpose_(1, 2)
mat3.zero_()
torch.ops.sgl_kernel.bmm_cpu(mat3, mat1, mat2, False, None)
atol = rtol = precision[ref.dtype]
torch.testing.assert_close(ref, mat3, atol=atol, rtol=rtol)
packed_B = torch.ops.sgl_kernel.convert_weight_packed(mat2_t)
mat3.zero_()
torch.ops.sgl_kernel.bmm_cpu(mat3, mat1, packed_B, True, None)
torch.testing.assert_close(ref, mat3, atol=atol, rtol=rtol)
def _fp8_bmm(self, B, M, N, K, chunk, dtype=torch.bfloat16):
mat1, mat2, mat3 = self._get_bmm_inputs(B, M, N, K, chunk, dtype)
mat2_q, mat2_s = input_to_float8(mat2)
ref = torch.bmm(mat1, mat2_q.to(torch.bfloat16)) * mat2_s
mat2_q_t = mat2_q.transpose_(1, 2).contiguous()
mat3.zero_()
atol = rtol = precision[ref.dtype]
torch.ops.sgl_kernel.bmm_cpu(mat3, mat1, mat2_q_t, False, mat2_s)
torch.testing.assert_close(ref, mat3, atol=atol, rtol=rtol)
packed_B_q = torch.ops.sgl_kernel.convert_weight_packed(mat2_q_t)
mat3.zero_()
torch.ops.sgl_kernel.bmm_cpu(mat3, mat1, packed_B_q, True, mat2_s)
torch.testing.assert_close(ref, mat3, atol=atol, rtol=rtol)
def test_bmm(self):
for params in itertools.product(
self.B,
self.M,
self.N,
self.K,
self.chunk,
):
with self.subTest(
B=params[0],
M=params[1],
N=params[2],
K=params[3],
chunk=params[4],
):
self._bf16_bmm(*params)
self._fp8_bmm(*params)
if __name__ == "__main__":
unittest.main()

17
test/srt/cpu/test_qkv_proj_with_rope.py
View File

@@ -8,6 +8,7 @@ from utils import (
precision,
)
from sglang.srt.layers.quantization.fp8_utils import input_to_float8
from sglang.srt.layers.rotary_embedding.utils import apply_rotary_emb
from sglang.test.test_utils import CustomTestCase
@@ -186,6 +187,7 @@ class TestQKVProjWithROPE(CustomTestCase):
None,
None,
None,
None,
True,
None,
)
@@ -203,6 +205,7 @@ class TestQKVProjWithROPE(CustomTestCase):
False,
None,
None,
None,
True,
None,
q_lora_rank,
@@ -274,6 +277,7 @@ class TestQKVProjWithROPE(CustomTestCase):
w1_s,
w2_s,
w3_s,
None,
True,
None,
)
@@ -294,6 +298,7 @@ class TestQKVProjWithROPE(CustomTestCase):
False,
fused_weight_s,
w2_s,
None,
True,
None,
q_lora_rank,
@@ -320,6 +325,7 @@ class TestQKVProjWithROPE(CustomTestCase):
torch.randn(num_heads * qk_head_dim, q_lora_rank, dtype=dtype) * 0.1
)
w_kc = torch.randn(num_heads, kv_lora_rank, qk_nope_head_dim, dtype=dtype) * 0.1
w_kc_q, w_kc_s = input_to_float8(w_kc)
kv_a_proj_weight = (
torch.randn(kv_lora_rank + qk_rope_head_dim, hidden_size, dtype=dtype) * 0.1
)
@@ -350,13 +356,14 @@ class TestQKVProjWithROPE(CustomTestCase):
) = convert_weight(
kv_a_proj_weight, [scale_block_size_N, scale_block_size_K], torch.bfloat16
)
w_kc_dq = w_kc_q.to(torch.bfloat16) * w_kc_s
q_ref, k_ref, v_ref = native_torch(
q_input,
hidden_states,
q_a_proj_weight_dq,
norm_weight1,
q_b_proj_weight_dq,
w_kc.transpose(1, 2),
w_kc_dq.transpose(1, 2),
kv_a_proj_with_mqa_weight_dq,
norm_weight2,
pos,
@@ -367,13 +374,13 @@ class TestQKVProjWithROPE(CustomTestCase):
fp8_kv_a_proj_with_mqa_weight_packed = convert_weight_packed(
fp8_kv_a_proj_with_mqa_weight
)
w_kc = convert_weight_packed(w_kc)
w_kc_q = convert_weight_packed(w_kc_q)
q_out, k_out, v_out = qkv_proj_with_rope(
hidden_states,
fp8_q_a_proj_weight_packed,
fp8_q_b_proj_weight_packed,
fp8_kv_a_proj_with_mqa_weight_packed,
w_kc,
w_kc_q,
norm_weight1,
norm_weight2,
pos,
@@ -384,6 +391,7 @@ class TestQKVProjWithROPE(CustomTestCase):
q_a_proj_weight_scale_inv.float(),
q_b_proj_weight_scale_inv.float(),
kv_a_proj_with_mqa_weight_scale_inv.float(),
w_kc_s,
True,
[scale_block_size_N, scale_block_size_K],
)
@@ -399,7 +407,7 @@ class TestQKVProjWithROPE(CustomTestCase):
hidden_states,
fused_weight_packed,
fp8_q_b_proj_weight_packed,
w_kc,
w_kc_q,
norm_weight1,
norm_weight2,
pos,
@@ -409,6 +417,7 @@ class TestQKVProjWithROPE(CustomTestCase):
True,
fused_weight_s.float(),
q_b_proj_weight_scale_inv.float(),
w_kc_s,
True,
[scale_block_size_N, scale_block_size_K],
q_lora_rank,

1
test/srt/run_suite.py
View File

@@ -49,6 +49,7 @@ suite_xeon = {
"per-commit-cpu": [
TestFile("cpu/test_activation.py"),
TestFile("cpu/test_binding.py"),
TestFile("cpu/test_bmm.py"),
TestFile("cpu/test_causal_conv1d.py"),
TestFile("cpu/test_cpu_graph.py"),
TestFile("cpu/test_decode.py"),