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

[sgl-kernel/cpu] support w8a8 int8 model for arm cpu (#16045)

Browse Source 
skip gpu test as this one is not related to gpu backend.
This commit is contained in:
Yibo Cai
2026-05-08 14:47:06 +08:00
committed by GitHub
parent 47e9ec11ad
commit 55d8223c2b
7 changed files with 832 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/w8a8_int8.py
+12 -10
View File
@@ -28,6 +28,7 @@ from sglang.srt.utils import (
cpu_has_amx_support,
is_cpu,
is_cuda,
is_host_cpu_arm64,
set_weight_attrs,
use_intel_amx_backend,
)
@@ -39,6 +40,7 @@ if TYPE_CHECKING:
_is_cuda = is_cuda()
_is_cpu_amx_available = cpu_has_amx_support()
_is_cpu = is_cpu()
_is_cpu_arm64 = is_host_cpu_arm64()
if _is_cuda:
from sgl_kernel import int8_scaled_mm
@@ -159,10 +161,12 @@ class W8A8Int8LinearMethod(LinearMethodBase):
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
if _is_cpu:
assert (
_is_cpu_amx_available
), "W8A8Int8LinearMethod on CPU requires that CPU has AMX support"
_amx_process_weight_after_loading(layer, ["weight"])
if _is_cpu_amx_available:
_amx_process_weight_after_loading(layer, ["weight"])
elif _is_cpu_arm64:
layer.weight = Parameter(layer.weight.data, requires_grad=False)
else:
assert False, "W8A8Int8LinearMethod on CPU only works on AMX or Arm64"
else:
layer.weight = Parameter(layer.weight.t(), requires_grad=False)
layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False)
@@ -204,7 +208,7 @@ class W8A8Int8LinearMethod(LinearMethodBase):
x: torch.Tensor,
bias: Optional[torch.Tensor] = None,
):
if use_intel_amx_backend(layer):
if use_intel_amx_backend(layer) or _is_cpu_arm64:
return torch.ops.sgl_kernel.int8_scaled_mm_with_quant(
x,
layer.weight,
@@ -310,10 +314,7 @@ class W8A8Int8MoEMethod(FusedMoEMethodBase):
layer.register_parameter("w2_input_scale", w2_input_scale)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
if _is_cpu:
assert (
_is_cpu_amx_available
), "W8A8Int8MoEMethod on CPU requires that CPU has AMX support"
if _is_cpu_amx_available:
_amx_process_weight_after_loading(layer, ["w13_weight", "w2_weight"])
else:
layer.w13_weight = Parameter(layer.w13_weight, requires_grad=False)
@@ -353,10 +354,11 @@ class W8A8Int8MoEMethod(FusedMoEMethodBase):
x = dispatch_output.hidden_states
topk_output = dispatch_output.topk_output
if use_intel_amx_backend(layer):
if use_intel_amx_backend(layer) or _is_cpu_arm64:
from sglang.srt.layers.moe.topk import apply_topk_weights_cpu
topk_weights, topk_ids, _ = topk_output
topk_ids = topk_ids.int()
x, topk_weights = apply_topk_weights_cpu(
self.moe_runner_config.apply_router_weight_on_input, topk_weights, x
)
sgl-kernel/csrc/cpu/CMakeLists.txt
+18 -4
View File
@@ -27,17 +27,32 @@ include_directories(
${CMAKE_CURRENT_SOURCE_DIR}
)
# Platform-specific library directory
# Build all cpp files in current and sub dirs
file(GLOB_RECURSE SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp")
# Exclude all arch dependent files, then add back only files for this arch
set(ARCH_DIRS "x86_64|aarch64|ppc64")
list(FILTER SOURCES EXCLUDE REGEX "^${CMAKE_CURRENT_SOURCE_DIR}/(${ARCH_DIRS})/")
# Platform-specific source and library directory
set(MY_ARCH_DIR "")
if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|AMD64")
set(PLAT_LIB_DIR "/usr/lib/x86_64-linux-gnu")
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64")
set(MY_ARCH_DIR "x86_64")
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64")
set(PLAT_LIB_DIR "/usr/lib/aarch64-linux-gnu")
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "ppc64le|ppc64")
set(MY_ARCH_DIR "aarch64")
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "ppc64le|ppc64")
set(PLAT_LIB_DIR "/usr/lib/powerpc64le-linux-gnu")
set(MY_ARCH_DIR "ppc64")
else()
set(PLAT_LIB_DIR "/usr/lib/${CMAKE_SYSTEM_PROCESSOR}-linux-gnu")
endif()
link_directories(${PLAT_LIB_DIR})
if(MY_ARCH_DIR)
file(GLOB_RECURSE ARCH_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/${MY_ARCH_DIR}/*.cpp")
list(APPEND SOURCES ${ARCH_SOURCES})
endif()
# Conda library path support
if(DEFINED ENV{CONDA_PREFIX})
@@ -73,7 +88,6 @@ else()
endif()
endif()
file(GLOB_RECURSE SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp")
# These kernels still rely on x86-specific AMX/AVX512 implementations.
# Keep them out of Arm64 bootstrap builds until native Arm paths land.
sgl-kernel/csrc/cpu/aarch64/gemm_int8.cpp
+132
View File
@@ -0,0 +1,132 @@
#include "../common.h"
#include "op.h"
namespace {
// out = mat1 @ mat2 + bias
template <typename scalar_t>
void int8_scaled_mm_impl(
scalar_t* __restrict__ out, // [M, N], row major
const int8_t* __restrict__ mat1, // [M, K], row major
const int8_t* __restrict__ mat2, // [K, N], column major
const float* __restrict__ scales1, // [M, 1], mat1 scales
const float* __restrict__ scales2, // [1, N], mat2 scales
const float* __restrict__ bias, // [1, N]
int64_t M,
int64_t N,
int64_t K) {
TORCH_CHECK(false, "not supported yet");
}
template <>
void int8_scaled_mm_impl<at::BFloat16>(
at::BFloat16* __restrict__ out,
const int8_t* __restrict__ mat1,
const int8_t* __restrict__ mat2,
const float* __restrict__ scales1,
const float* __restrict__ scales2,
const float* __restrict__ bias,
int64_t M,
int64_t N,
int64_t K) {
const int slice_size = (M * K * sizeof(int8_t)) > kL2Size ? 64 : 8;
const int num_slices = (N + slice_size - 1) / slice_size;
auto mm = [mat1, mat2, out, M, N, K, scales1, scales2, bias, slice_size](int64_t begin, int64_t end) {
for (int64_t slice_idx = begin; slice_idx < end; ++slice_idx) {
const int64_t n_start = slice_idx * slice_size;
const int64_t n_end = std::min(n_start + slice_size, N);
const int slice_width = static_cast<int>(n_end - n_start);
const int8_t* a_ptr = mat1;
const int8_t* b_ptr = mat2 + n_start * K;
bfloat16_t* c_ptr = reinterpret_cast<bfloat16_t*>(out) + n_start;
op::i8mm_matmul(a_ptr, b_ptr, c_ptr, M, K, N, slice_width, scales1, scales2 + n_start);
// NOTE: matmul reduces matrix values to BF16, may influence precision
if (bias) {
op::add_bias(c_ptr, bias + n_start, M, N, slice_width);
}
}
};
at::parallel_for(0, num_slices, 0, mm);
}
} // anonymous namespace
// weight : static, per-channel, symmetric
// activation : dynamic, per-token, symmetric
//
// mat1 : [M, K]
// mat2 : [N, K]
// scales1 : [M]
// scales2 : [N]
// bias : [N]
// out : [M, N]
//
// fused activation quantization and matmul
at::Tensor int8_scaled_mm_with_quant(
at::Tensor& mat1,
at::Tensor& mat2,
at::Tensor& scales2,
const std::optional<at::Tensor>& bias,
at::ScalarType out_dtype,
bool /*is_vnni*/) {
CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
CHECK_INPUT(mat2);
CHECK_INPUT(scales2);
CHECK_DIM(2, mat1);
CHECK_DIM(2, mat2);
int64_t M = mat1.size(0);
int64_t N = mat2.size(0);
int64_t K = mat1.size(1);
int64_t lda = mat1.stride(0);
CHECK_EQ(mat2.size(1), K);
CHECK_EQ(scales2.numel(), N);
const auto st = mat1.scalar_type();
TORCH_CHECK(st == at::kBFloat16, "int8_scaled_mm_with_quant: expect A to be bfloat16.");
TORCH_CHECK(st == out_dtype, "int8_scaled_mm_with_quant: expect A has same dtype with out_dtype.");
TORCH_CHECK(mat2.scalar_type() == at::kChar, "int8_scaled_mm_with_quant: expect mat2 to be int8.");
TORCH_CHECK(scales2.scalar_type() == at::kFloat, "int8_scaled_mm_with_quant: expect scales to be float32.");
const int64_t buffer_size = M * K + M * sizeof(float);
auto buffer = at::empty({buffer_size}, mat1.options().dtype(at::kChar));
auto out = at::empty({M, N}, mat1.options().dtype(out_dtype));
const bool has_bias = bias.has_value();
const float* bias_data = nullptr;
if (has_bias) {
CHECK_EQ(bias.value().size(0), N);
bias_data = bias.value().data_ptr<float>();
}
AT_DISPATCH_REDUCED_FLOATING_TYPES(out_dtype, "int8_scaled_mm_with_quant_kernel_impl", [&] {
int8_t* __restrict__ Aq_data = buffer.data_ptr<int8_t>();
float* __restrict__ As_data = (float*)((void*)(Aq_data + M * K));
const scalar_t* __restrict__ A_data = mat1.data_ptr<scalar_t>();
const int64_t grain = kL1Size / (K * sizeof(scalar_t));
at::parallel_for(0, M, grain, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
op::quantize_row_int8(Aq_data + m * K, As_data + m, A_data + m * lda, K);
}
});
int8_scaled_mm_impl<scalar_t>(
out.data_ptr<scalar_t>(),
Aq_data,
mat2.data_ptr<int8_t>(),
As_data,
scales2.data_ptr<float>(),
bias_data,
M,
N,
K);
});
return out;
}
sgl-kernel/csrc/cpu/aarch64/moe.cpp
+323
View File
@@ -0,0 +1,323 @@
#include "../common.h"
#include "op.h"
namespace {
// key: expert id, value: input rows and weights for this expert
using expert_to_rows_t = std::map<int, std::vector<std::tuple<int, float>>>;
// for expert_id, row_weight_list in x_per_expert.items():
// rows, weights = zip(*row_weight_list)
// x_rows = x[rows]
// w1, w3 = torch.chunk(w13[expert_id], chunks=2)
// gate = x_rows @ w1
// up = x_rows @ w3
// up *= silu(gate)
// down = up @ w2[expert_id]
// down *= weights
// y.index_add_(0, rows, down)
template <typename scalar_t>
void fused_experts_int8_kernel_impl(
scalar_t* __restrict__ y, // [M, K], row major
const int8_t* __restrict__ x, // [M, K], row major
const int8_t* __restrict__ w13, // [E, K, 2N], per expert [K, N], column major, w1 before w3
const int8_t* __restrict__ w2, // [E, N, K], per expert [N, K], column major
const float* __restrict__ x_scale, // [M, 1]
const float* __restrict__ w13_scale, // [E, 1, 2N], per expert [1, N], w1 before w3
const float* __restrict__ w2_scale, // [E, 1, K], per expert [1, K]
const expert_to_rows_t& x_per_expert, // expert id -> related x rows and weights
int64_t M,
int64_t N,
int64_t K,
int64_t E,
int64_t topk) {
TORCH_CHECK(false, "not implemented yet");
}
template <>
void fused_experts_int8_kernel_impl<at::BFloat16>(
at::BFloat16* __restrict__ y,
const int8_t* __restrict__ x,
const int8_t* __restrict__ w13,
const int8_t* __restrict__ w2,
const float* __restrict__ x_scale,
const float* __restrict__ w13_scale,
const float* __restrict__ w2_scale,
const expert_to_rows_t& x_per_expert,
int64_t M,
int64_t N,
int64_t K,
int64_t E,
int64_t topk) {
// x dispatch buffer to aggregate all rows per expert
int64_t max_agg_rows = 0;
for (const auto& [eid, rows] : x_per_expert) {
max_agg_rows = std::max<int64_t>(max_agg_rows, rows.size());
}
// x_scale_agg[max_agg_rows] + up_scale[max_agg_rows] +
// gate[max_agg_rows,N] + up[max_agg_rows,N] + down[max_agg_rows,K]
auto f32_buffer = at::empty({max_agg_rows, 1 + 1 + N + N + K}, at::kFloat);
float* x_scale_agg = f32_buffer.data_ptr<float>();
float* up_scale = x_scale_agg + max_agg_rows;
float* gate = up_scale + max_agg_rows;
float* up = gate + max_agg_rows * N;
float* down = up + max_agg_rows * N;
// x_agg[max_agg_rows,K] + up_q8[max_agg_rows,N]
auto int8_buffer = at::empty({max_agg_rows, K + N}, at::kChar);
int8_t* x_agg = int8_buffer.data_ptr<int8_t>();
int8_t* up_q8 = x_agg + max_agg_rows * K;
// out[M,K]: accumulated output
auto out_buffer = at::zeros({M, K}, at::kFloat);
float* out = out_buffer.data_ptr<float>();
// iterate used experts
for (const auto& [eid, rows] : x_per_expert) {
const int64_t n_agg = rows.size();
// copy input rows using this expert to contiguous buffer
{
int8_t* x_agg_ptr = x_agg;
float* x_scale_agg_ptr = x_scale_agg;
for (const auto [row, weight] : rows) {
// int row; float weight;
std::memcpy(x_agg_ptr, x + row * K, K * sizeof(int8_t));
*x_scale_agg_ptr = x_scale[row];
x_agg_ptr += K;
++x_scale_agg_ptr;
}
}
// gate = x_agg @ w1
// up = x_agg @ w3
// up *= silu(gate)
{
// expert specific tensors
const int8_t* w1e = w13 + eid * 2 * N * K;
const int8_t* w3e = w1e + N * K;
const float* w1e_scale = w13_scale + eid * 2 * N;
const float* w3e_scale = w1e_scale + N;
// tensor shapes
// - x_agg: [n_agg, K], int8, row major
// - x_scale_agg: [n_agg, 1], float
// - w{1,3}e: [K, N], int8, col major
// - w{1,3}e_scale: [1, N], float
// - gate: [n_agg, N], float, row major
// - up: [n_agg, N], float, row major
// - up_q8: [n_agg, N], int8, row major
// - up_scale: [n_agg, 1], float
const int slice_size = (n_agg * K * sizeof(int8_t)) > kL2Size ? 64 : 8;
const int num_slices = (N + slice_size - 1) / slice_size;
auto mm = [&](int64_t begin, int64_t end) {
for (int64_t slice_idx = begin; slice_idx < end; ++slice_idx) {
const int64_t n_start = slice_idx * slice_size;
const int64_t n_end = std::min(n_start + slice_size, N);
const int slice_width = static_cast<int>(n_end - n_start);
const int8_t* w1e_ptr = w1e + n_start * K;
const int8_t* w3e_ptr = w3e + n_start * K;
const float* w1e_scale_ptr = w1e_scale + n_start;
const float* w3e_scale_ptr = w3e_scale + n_start;
float* gate_ptr = gate + n_start;
float* up_ptr = up + n_start;
op::i8mm_matmul(x_agg, w1e_ptr, gate_ptr, n_agg, K, N, slice_width, x_scale_agg, w1e_scale_ptr);
op::i8mm_matmul(x_agg, w3e_ptr, up_ptr, n_agg, K, N, slice_width, x_scale_agg, w3e_scale_ptr);
for (int i = 0; i < n_agg; ++i) {
const float* __restrict__ gate_ptr = gate + n_start + i * N;
float* __restrict__ up_ptr = up + n_start + i * N;
// TODO: vectorize
for (int j = 0; j < slice_width; ++j) {
up_ptr[j] *= gate_ptr[j] / (1 + std::exp(-gate_ptr[j]));
}
}
}
};
at::parallel_for(0, num_slices, 0, mm);
}
// quantize
{
const int64_t grain = kL1Size / (K * sizeof(float));
at::parallel_for(0, n_agg, grain, [&](int64_t begin, int64_t end) {
for (int64_t i = begin; i < end; ++i) {
op::quantize_row_int8(up_q8 + i * N, up_scale + i, up + i * N, N);
}
});
}
// down = up @ w2
{
// expert specific tensors
const int8_t* w2e = w2 + eid * K * N;
const float* w2e_scale = w2_scale + eid * K;
// tensor shapes
// - up_q8: [n_agg, N], int8, row major
// - up_scale: [n_agg, 1], float
// - w2e: [N, K], int8, col major
// - w2e_scale: [1, K], float
// - down: [n_agg, K], float, row major
// - out: [M, K], float, row major
const int slice_size = (n_agg * N * sizeof(int8_t)) > kL2Size ? 64 : 8;
const int num_slices = (K + slice_size - 1) / slice_size;
auto mm = [&](int64_t begin, int64_t end) {
for (int64_t slice_idx = begin; slice_idx < end; ++slice_idx) {
const int64_t n_start = slice_idx * slice_size;
const int64_t n_end = std::min(n_start + slice_size, K);
const int slice_width = static_cast<int>(n_end - n_start);
{
const int8_t* w2e_ptr = w2e + n_start * N;
const float* w2e_scale_ptr = w2e_scale + n_start;
float* down_ptr = down + n_start;
op::i8mm_matmul(up_q8, w2e_ptr, down_ptr, n_agg, N, K, slice_width, up_scale, w2e_scale_ptr);
}
// accumulate to out buffer
{
const float* __restrict__ down_ptr = down + n_start;
for (const auto [row, weight] : rows) {
// int row; float weight;
float* __restrict__ out_ptr = out + n_start + row * K;
// auto vectorizable
for (int i = 0; i < slice_width; ++i) {
out_ptr[i] += down_ptr[i] * weight;
}
down_ptr += K;
}
}
}
};
at::parallel_for(0, num_slices, 0, mm);
}
}
// copy output: float -> bf16
{
// tensor shapes
// - out: [M, K], float, row major
// - y: [M, K], bf16, row major
const int64_t grain = kL1Size / (K * sizeof(float));
at::parallel_for(0, M, grain, [&](int64_t begin, int64_t end) {
const float* out_ptr = out + begin * K;
bfloat16_t* y_ptr = reinterpret_cast<bfloat16_t*>(y) + begin * K;
op::f32_to_bf16(out_ptr, y_ptr, (end - begin) * K);
});
}
}
} // anonymous namespace
// hidden_states: [M, K]
// w13: [E, 2N, K]
// w2: [E, K, N]
// topk_weights: [M, topk]
// topk_ids: [M, topk] (int32_t)
// w13_scale: [E, 2N]
// w2_scale: [E, K]
at::Tensor fused_experts_cpu(
at::Tensor& hidden_states,
at::Tensor& w13,
at::Tensor& w2,
at::Tensor& topk_weights,
at::Tensor& topk_ids,
bool inplace,
int64_t moe_comp_method,
const std::optional<at::Tensor>& w13_scale,
const std::optional<at::Tensor>& w2_scale,
const std::optional<at::Tensor>& /*w13_zero*/,
const std::optional<at::Tensor>& /*w2_zero*/,
const std::optional<std::vector<int64_t>> block_size,
bool /*is_vnni*/) {
const auto st = hidden_states.scalar_type();
CHECK_INPUT(hidden_states);
CHECK_INPUT(w13);
CHECK_INPUT(w2);
CHECK_EQ(topk_weights.sizes(), topk_ids.sizes());
CHECK_DIM(2, hidden_states);
CHECK_DIM(3, w13);
CHECK_DIM(3, w2);
CHECK_DIM(2, topk_weights);
CHECK_DIM(2, topk_ids);
CHECK_EQ(topk_ids.scalar_type(), at::kInt);
// TODO: support topk_weights to be bf16 or fp16 in the kernel
auto topk_weights_ = topk_weights.to(at::kFloat);
int64_t M = hidden_states.size(0);
int64_t K = hidden_states.size(1);
int64_t N = w13.size(1) / 2;
int64_t E = w13.size(0);
int64_t topk = topk_weights_.size(1);
// check weight shapes
CHECK_EQ(w2.size(0), E);
CHECK_EQ(w2.size(1), K);
CHECK_EQ(w13.size(2), K);
CHECK_EQ(w2.size(2), N);
CHECK_EQ(inplace, false);
at::Tensor out = at::empty_like(hidden_states);
// expert id -> related input rows and weights
expert_to_rows_t x_per_expert; // std::map<int, std::vector<std::tuple<int, float>>>
{
const int* ids = topk_ids.data_ptr<int>();
const float* weights = topk_weights_.data_ptr<float>();
for (int i = 0; i < M; ++i) {
for (int j = 0; j < topk; ++j) {
x_per_expert[*ids].emplace_back(i, *weights);
++ids;
++weights;
}
}
}
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "fused_experts_kernel_impl", [&] {
auto& w13s = w13_scale.value();
auto& w2s = w2_scale.value();
TORCH_CHECK(w13s.numel() == E * 2 * N);
TORCH_CHECK(w2s.numel() == E * K);
// quantize hidden_states
auto x_buffer = at::empty({M * K}, hidden_states.options().dtype(at::kChar));
auto x_scale_buffer = at::empty({M}, at::kFloat);
int8_t* x = x_buffer.data_ptr<int8_t>();
float* x_scale = x_scale_buffer.data_ptr<float>();
scalar_t* in = hidden_states.data_ptr<scalar_t>();
const int64_t grain = kL1Size / (K * sizeof(scalar_t));
at::parallel_for(0, M, grain, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
op::quantize_row_int8(x + m * K, x_scale + m, in + m * K, K);
}
});
fused_experts_int8_kernel_impl<scalar_t>(
out.data_ptr<scalar_t>(),
x,
w13.data_ptr<int8_t>(),
w2.data_ptr<int8_t>(),
x_scale,
w13s.data_ptr<float>(),
w2s.data_ptr<float>(),
x_per_expert,
M,
N,
K,
E,
topk);
});
return out;
}
sgl-kernel/csrc/cpu/aarch64/op.h
+343
View File
@@ -0,0 +1,343 @@
#pragma once
#include <arm_neon.h>
constexpr int64_t kL1Size = 64 * 1024;
constexpr int64_t kL2Size = 1 * 1024 * 1024;
// simd optimized operators
namespace op {
// do matmul in "R rows x C cols" tile with sdot
// - a is the full [M, K] matrix, row major
// - b is one slice of a [K, N] matrix, column major ([N, K] row major)
// - c is one slice of a [M, N] matrix, row major
//
// slice_width slice_width
// a b / c /
// | |----| |----|
// v | | v v
// / ------ \ v v / ------ \
// | ------ | / |||||| \ | ------ |
// M | ------ | @ | |||||| | = | ------ |
// | ------ | | |||||| | | ------ |
// | ------ | \ |||||| / | ------ |
// \ ------ / \ ------ /
// K N
//
template <int R = 4, int C = 8, typename T>
__attribute__((target("+dotprod+bf16"))) void sdot_matmul(
const int8_t* __restrict__ a,
const int8_t* __restrict__ b,
T* c,
int64_t M,
int64_t K,
int64_t N,
int slice_width,
const float* __restrict__ scales1,
const float* __restrict__ scales2) {
static_assert(std::is_same_v<T, float> || std::is_same_v<T, bfloat16_t>);
int row = 0;
for (; row + R <= M; row += R) {
const int8_t* a_rows[R];
T* c_rows[R];
for (int i = 0; i < R; ++i) {
a_rows[i] = a + (row + i) * K;
c_rows[i] = c + (row + i) * N;
}
int col = 0;
for (; col + C <= slice_width; col += C) {
const int8_t* b_cols[C];
for (int i = 0; i < C; ++i) {
b_cols[i] = b + (col + i) * K;
}
int32x4_t vsums[R][C]{};
// TODO: accumulated integer sum may overflow when K >= 65536
int k = 0;
for (; k + 16 <= K; k += 16) {
int8x16_t va[R];
int8x16_t vb[C];
for (int i = 0; i < R; ++i) {
va[i] = vld1q_s8(a_rows[i] + k);
}
for (int i = 0; i < C; ++i) {
vb[i] = vld1q_s8(b_cols[i] + k);
}
for (int i = 0; i < R; ++i) {
for (int j = 0; j < C; ++j) {
vsums[i][j] = vdotq_s32(vsums[i][j], va[i], vb[j]);
}
}
}
if (k < K) {
int8_t abuf[16]{};
int8_t bbuf[16]{};
for (int i = 0; i < R; ++i) {
memcpy(abuf, a_rows[i] + k, K - k);
const int8x16_t va = vld1q_s8(abuf);
for (int j = 0; j < C; ++j) {
memcpy(bbuf, b_cols[j] + k, K - k);
const int8x16_t vb = vld1q_s8(bbuf);
vsums[i][j] = vdotq_s32(vsums[i][j], va, vb);
}
}
}
for (int i = 0; i < R; ++i) {
for (int j = 0; j < C; ++j) {
const float sum = vaddvq_s32(vsums[i][j]);
const float sum_scaled = sum * scales1[row + i] * scales2[col + j];
if constexpr (std::is_same_v<T, bfloat16_t>) {
c_rows[i][col + j] = vcvth_bf16_f32(sum_scaled);
} else {
c_rows[i][col + j] = sum_scaled;
}
}
}
}
if (col < slice_width) {
sdot_matmul<R, 1>(a, b + col * K, c + col, M, K, N, slice_width - col, scales1, scales2 + col);
}
}
if (row < M) {
sdot_matmul<1, C>(a + row * K, b, c + row * N, M - row, K, N, slice_width, scales1 + row, scales2);
}
}
// do matmul in "R rows x C cols" tile with i8mm
template <int R = 4, int C = 8, typename T>
__attribute__((target("+i8mm+bf16"))) void i8mm_matmul(
const int8_t* __restrict__ a,
const int8_t* __restrict__ b,
T* c,
int64_t M,
int64_t K,
int64_t N,
int slice_width,
const float* __restrict__ scales1,
const float* __restrict__ scales2) {
static_assert(std::is_same_v<T, float> || std::is_same_v<T, bfloat16_t>);
static_assert(R % 2 == 0 && C % 2 == 0);
int row = 0;
for (; row + R <= M; row += R) {
const int8_t* a_rows[R];
T* c_rows[R];
for (int i = 0; i < R; ++i) {
a_rows[i] = a + (row + i) * K;
c_rows[i] = c + (row + i) * N;
}
int col = 0;
for (; col + C <= slice_width; col += C) {
const int8_t* b_cols[C];
for (int i = 0; i < C; ++i) {
b_cols[i] = b + (col + i) * K;
}
int8x16_t va[R], vb[C];
int32x4_t vsums[R / 2][C / 2]{};
// TODO: accumulated integer sum may overflow when K >= 65536
int k = 0;
for (; k + 16 <= K; k += 16) {
for (int i = 0; i < R; i += 2) {
const int64x2_t va0_s64 = vreinterpretq_s64_s8(vld1q_s8(a_rows[i + 0] + k));
const int64x2_t va1_s64 = vreinterpretq_s64_s8(vld1q_s8(a_rows[i + 1] + k));
va[i + 0] = vreinterpretq_s8_s64(vzip1q_s64(va0_s64, va1_s64));
va[i + 1] = vreinterpretq_s8_s64(vzip2q_s64(va0_s64, va1_s64));
}
for (int i = 0; i < C; i += 2) {
const int64x2_t vb0_s64 = vreinterpretq_s64_s8(vld1q_s8(b_cols[i + 0] + k));
const int64x2_t vb1_s64 = vreinterpretq_s64_s8(vld1q_s8(b_cols[i + 1] + k));
vb[i + 0] = vreinterpretq_s8_s64(vzip1q_s64(vb0_s64, vb1_s64));
vb[i + 1] = vreinterpretq_s8_s64(vzip2q_s64(vb0_s64, vb1_s64));
}
for (int i = 0; i < R / 2; ++i) {
for (int j = 0; j < C / 2; ++j) {
vsums[i][j] = vmmlaq_s32(vsums[i][j], va[i * 2 + 0], vb[j * 2 + 0]);
vsums[i][j] = vmmlaq_s32(vsums[i][j], va[i * 2 + 1], vb[j * 2 + 1]);
}
}
}
if (k < K) {
int8_t buf0[16]{}, buf1[16]{};
for (int i = 0; i < R; i += 2) {
memcpy(buf0, a_rows[i + 0] + k, (K - k) * sizeof(int8_t));
memcpy(buf1, a_rows[i + 1] + k, (K - k) * sizeof(int8_t));
const int64x2_t va0_s64 = vreinterpretq_s64_s8(vld1q_s8(buf0));
const int64x2_t va1_s64 = vreinterpretq_s64_s8(vld1q_s8(buf1));
va[i + 0] = vreinterpretq_s8_s64(vzip1q_s64(va0_s64, va1_s64));
va[i + 1] = vreinterpretq_s8_s64(vzip2q_s64(va0_s64, va1_s64));
}
for (int i = 0; i < C; i += 2) {
memcpy(buf0, b_cols[i + 0] + k, (K - k) * sizeof(int8_t));
memcpy(buf1, b_cols[i + 1] + k, (K - k) * sizeof(int8_t));
const int64x2_t vb0_s64 = vreinterpretq_s64_s8(vld1q_s8(buf0));
const int64x2_t vb1_s64 = vreinterpretq_s64_s8(vld1q_s8(buf1));
vb[i + 0] = vreinterpretq_s8_s64(vzip1q_s64(vb0_s64, vb1_s64));
vb[i + 1] = vreinterpretq_s8_s64(vzip2q_s64(vb0_s64, vb1_s64));
}
for (int i = 0; i < R / 2; ++i) {
for (int j = 0; j < C / 2; ++j) {
vsums[i][j] = vmmlaq_s32(vsums[i][j], va[i * 2 + 0], vb[j * 2 + 0]);
vsums[i][j] = vmmlaq_s32(vsums[i][j], va[i * 2 + 1], vb[j * 2 + 1]);
}
}
}
for (int i = 0; i < R; i += 2) {
for (int j = 0; j < C; j += 2) {
float32x4_t vsum_f32 = vcvtq_f32_s32(vsums[i / 2][j / 2]);
const float32x4_t scales = {
scales1[row + i + 0] * scales2[col + j + 0],
scales1[row + i + 0] * scales2[col + j + 1],
scales1[row + i + 1] * scales2[col + j + 0],
scales1[row + i + 1] * scales2[col + j + 1],
};
vsum_f32 = vmulq_f32(vsum_f32, scales);
if constexpr (std::is_same_v<T, bfloat16_t>) {
const bfloat16x4_t vsum_bf16 = vcvt_bf16_f32(vsum_f32);
c_rows[i + 0][col + j + 0] = vget_lane_bf16(vsum_bf16, 0);
c_rows[i + 0][col + j + 1] = vget_lane_bf16(vsum_bf16, 1);
c_rows[i + 1][col + j + 0] = vget_lane_bf16(vsum_bf16, 2);
c_rows[i + 1][col + j + 1] = vget_lane_bf16(vsum_bf16, 3);
} else {
c_rows[i + 0][col + j + 0] = vgetq_lane_f32(vsum_f32, 0);
c_rows[i + 0][col + j + 1] = vgetq_lane_f32(vsum_f32, 1);
c_rows[i + 1][col + j + 0] = vgetq_lane_f32(vsum_f32, 2);
c_rows[i + 1][col + j + 1] = vgetq_lane_f32(vsum_f32, 3);
}
}
}
}
if (col < slice_width) {
sdot_matmul<R, 1>(a, b + col * K, c + col, M, K, N, slice_width - col, scales1, scales2 + col);
}
}
if (row < M) {
sdot_matmul<1, C>(a + row * K, b, c + row * N, M - row, K, N, slice_width, scales1 + row, scales2);
}
}
__attribute__((target("+bf16"))) inline void
add_bias(bfloat16_t* __restrict__ out, const float* __restrict__ bias, int64_t M, int64_t N, int width) {
int col = 0;
for (; col + 4 <= width; col += 4) {
const float32x4_t vbias32 = vld1q_f32(bias + col);
bfloat16_t* out_ptr = out + col;
for (int64_t i = 0; i < M; ++i) {
bfloat16x4_t vout16 = vld1_bf16(out_ptr);
float32x4_t vout32 = vcvt_f32_bf16(vout16);
vout32 = vaddq_f32(vout32, vbias32);
vout16 = vcvt_bf16_f32(vout32);
vst1_bf16(out_ptr, vout16);
out_ptr += N;
}
}
for (; col < width; ++col) {
const float vbias32 = bias[col];
bfloat16_t* out_ptr = out + col;
for (int64_t i = 0; i < M; ++i) {
bfloat16_t vout16 = *out_ptr;
float vout32 = vcvtah_f32_bf16(vout16);
vout32 += vbias32;
vout16 = vcvth_bf16_f32(vout32);
*out_ptr = vout16;
out_ptr += N;
}
}
}
constexpr float eps = 1e-7;
template <typename scalar_t>
inline void quantize_row_int8(int8_t* __restrict__ q, float* scale, const scalar_t* __restrict__ x, int64_t n) {
float max_abs_val = eps;
for (int64_t i = 0; i < n; ++i) {
max_abs_val = std::max(std::abs(static_cast<float>(x[i])), max_abs_val);
}
*scale = max_abs_val / 127.0f;
const float inv_scale = 127.0f / max_abs_val;
for (int64_t i = 0; i < n; ++i) {
q[i] = static_cast<int8_t>(std::round(static_cast<float>(x[i]) * inv_scale));
}
}
// manually optimize for bf16
template <>
__attribute__((target("+bf16"))) inline void
quantize_row_int8<bfloat16_t>(int8_t* __restrict__ q, float* scale, const bfloat16_t* __restrict__ x, int64_t n) {
float max_abs_val = eps;
for (int64_t i = 0; i < n; ++i) {
max_abs_val = std::max(std::abs(vcvtah_f32_bf16(x[i])), max_abs_val);
}
*scale = max_abs_val / 127.0f;
const float inv_scale = 127.0f / max_abs_val;
int64_t i = 0;
for (; i + 16 <= n; i += 16) {
int32x4_t qv_s32[4];
{
const bfloat16x8x2_t xv_bf16 = vld1q_bf16_x2(x + i);
const float32x4_t xv_f32[4] = {
vcvtq_low_f32_bf16(xv_bf16.val[0]),
vcvtq_high_f32_bf16(xv_bf16.val[0]),
vcvtq_low_f32_bf16(xv_bf16.val[1]),
vcvtq_high_f32_bf16(xv_bf16.val[1]),
};
for (int j = 0; j < 4; ++j) {
float32x4_t qv_f32 = vmulq_n_f32(xv_f32[j], inv_scale);
qv_f32 = vrndaq_f32(qv_f32);
qv_s32[j] = vcvtq_s32_f32(qv_f32);
}
}
const int16x8_t qv_s16[2] = {
vuzp1q_s16(vreinterpretq_s16_s32(qv_s32[0]), vreinterpretq_s16_s32(qv_s32[1])),
vuzp1q_s16(vreinterpretq_s16_s32(qv_s32[2]), vreinterpretq_s16_s32(qv_s32[3])),
};
const int8x16_t qv_s8 = vuzp1q_s8(vreinterpretq_s8_s16(qv_s16[0]), vreinterpretq_s8_s16(qv_s16[1]));
vst1q_s8(q + i, qv_s8);
}
for (; i < n; ++i) {
q[i] = static_cast<int8_t>(std::round(vcvtah_f32_bf16(x[i]) * inv_scale));
}
}
template <>
inline void
quantize_row_int8<at::BFloat16>(int8_t* __restrict__ q, float* scale, const at::BFloat16* __restrict__ x, int64_t n) {
quantize_row_int8(q, scale, reinterpret_cast<const bfloat16_t*>(x), n);
}
__attribute__((target("+bf16"))) inline void f32_to_bf16(const float* f32, bfloat16_t* bf16, int64_t n) {
int64_t i = 0;
for (; i + 4 <= n; i += 4) {
const float32x4_t vf32 = vld1q_f32(f32 + i);
const bfloat16x4_t vbf16 = vcvt_bf16_f32(vf32);
vst1_bf16(bf16 + i, vbf16);
}
for (; i < n; ++i) {
bf16[i] = vcvth_bf16_f32(f32[i]);
}
}
} // namespace op
sgl-kernel/csrc/cpu/gemm_int8.cpp
+2
View File
@@ -473,6 +473,7 @@ at::Tensor int8_scaled_mm_cpu(
return out;
}
#ifndef __aarch64__
// fused `per_token_quant_int8_cpu` and `int8_scaled_mm_cpu`
at::Tensor int8_scaled_mm_with_quant(
at::Tensor& mat1,
@@ -539,3 +540,4 @@ at::Tensor int8_scaled_mm_with_quant(
});
return out;
}
#endif // #ifndef __aarch64__
sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
+2 -2
View File
@@ -211,7 +211,6 @@ at::Tensor fused_linear_sigmoid_mul(
// bmm
void bmm_cpu(at::Tensor& out, at::Tensor& mat1, at::Tensor& mat2, bool is_vnni, const std::optional<at::Tensor>& scale);
#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS)
// fused moe
at::Tensor fused_experts_cpu(
at::Tensor& hidden_states,
@@ -228,6 +227,7 @@ at::Tensor fused_experts_cpu(
const std::optional<std::vector<int64_t>> block_size,
bool is_vnni);
#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS)
at::Tensor shared_expert_cpu(
at::Tensor& hidden_states,
at::Tensor& w1,
@@ -546,7 +546,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
m.def("bmm_cpu(Tensor(a!) out, Tensor mat1, Tensor mat2, bool is_vnni, Tensor? scale) -> ()");
m.impl("bmm_cpu", torch::kCPU, &bmm_cpu);
#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS)
// moe
m.def(
"fused_experts_cpu(Tensor hidden_states, Tensor w1, Tensor w2, Tensor topk_weights, Tensor topk_ids, bool "
@@ -554,6 +553,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
"Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, bool is_vnni) -> Tensor");
m.impl("fused_experts_cpu", torch::kCPU, &fused_experts_cpu);
#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS)
// weight absorption
m.def(
"qkv_proj_with_rope(Tensor hidden_states, Tensor q_a_proj_weight, Tensor q_b_proj_weight, Tensor "