diff --git a/python/sglang/srt/layers/quantization/w8a8_int8.py b/python/sglang/srt/layers/quantization/w8a8_int8.py index d7bfd888b..09de341de 100644 --- a/python/sglang/srt/layers/quantization/w8a8_int8.py +++ b/python/sglang/srt/layers/quantization/w8a8_int8.py @@ -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 ) diff --git a/sgl-kernel/csrc/cpu/CMakeLists.txt b/sgl-kernel/csrc/cpu/CMakeLists.txt index 7b4275ecd..6b0e0d86e 100755 --- a/sgl-kernel/csrc/cpu/CMakeLists.txt +++ b/sgl-kernel/csrc/cpu/CMakeLists.txt @@ -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. diff --git a/sgl-kernel/csrc/cpu/aarch64/gemm_int8.cpp b/sgl-kernel/csrc/cpu/aarch64/gemm_int8.cpp new file mode 100644 index 000000000..8dae500e4 --- /dev/null +++ b/sgl-kernel/csrc/cpu/aarch64/gemm_int8.cpp @@ -0,0 +1,132 @@ +#include "../common.h" +#include "op.h" + +namespace { + +// out = mat1 @ mat2 + bias +template +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* __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(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(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& 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(); + } + + AT_DISPATCH_REDUCED_FLOATING_TYPES(out_dtype, "int8_scaled_mm_with_quant_kernel_impl", [&] { + int8_t* __restrict__ Aq_data = buffer.data_ptr(); + float* __restrict__ As_data = (float*)((void*)(Aq_data + M * K)); + const scalar_t* __restrict__ A_data = mat1.data_ptr(); + + 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( + out.data_ptr(), + Aq_data, + mat2.data_ptr(), + As_data, + scales2.data_ptr(), + bias_data, + M, + N, + K); + }); + return out; +} diff --git a/sgl-kernel/csrc/cpu/aarch64/moe.cpp b/sgl-kernel/csrc/cpu/aarch64/moe.cpp new file mode 100644 index 000000000..3e0774376 --- /dev/null +++ b/sgl-kernel/csrc/cpu/aarch64/moe.cpp @@ -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>>; + +// 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 +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* __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(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* 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* 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(); + + // 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(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(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(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& w13_scale, + const std::optional& w2_scale, + const std::optional& /*w13_zero*/, + const std::optional& /*w2_zero*/, + const std::optional> 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>> + { + const int* ids = topk_ids.data_ptr(); + const float* weights = topk_weights_.data_ptr(); + 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(); + float* x_scale = x_scale_buffer.data_ptr(); + scalar_t* in = hidden_states.data_ptr(); + 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( + out.data_ptr(), + x, + w13.data_ptr(), + w2.data_ptr(), + x_scale, + w13s.data_ptr(), + w2s.data_ptr(), + x_per_expert, + M, + N, + K, + E, + topk); + }); + + return out; +} diff --git a/sgl-kernel/csrc/cpu/aarch64/op.h b/sgl-kernel/csrc/cpu/aarch64/op.h new file mode 100644 index 000000000..eec5d4ce1 --- /dev/null +++ b/sgl-kernel/csrc/cpu/aarch64/op.h @@ -0,0 +1,343 @@ +#pragma once + +#include + +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 +__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 || std::is_same_v); + + 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) { + c_rows[i][col + j] = vcvth_bf16_f32(sum_scaled); + } else { + c_rows[i][col + j] = sum_scaled; + } + } + } + } + + if (col < slice_width) { + sdot_matmul(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 +__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 || std::is_same_v); + 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) { + 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(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 +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(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(std::round(static_cast(x[i]) * inv_scale)); + } +} + +// manually optimize for bf16 +template <> +__attribute__((target("+bf16"))) inline void +quantize_row_int8(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(std::round(vcvtah_f32_bf16(x[i]) * inv_scale)); + } +} + +template <> +inline void +quantize_row_int8(int8_t* __restrict__ q, float* scale, const at::BFloat16* __restrict__ x, int64_t n) { + quantize_row_int8(q, scale, reinterpret_cast(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 diff --git a/sgl-kernel/csrc/cpu/gemm_int8.cpp b/sgl-kernel/csrc/cpu/gemm_int8.cpp index f72616697..1da63f79e 100644 --- a/sgl-kernel/csrc/cpu/gemm_int8.cpp +++ b/sgl-kernel/csrc/cpu/gemm_int8.cpp @@ -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__ diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp index 9ec3ea450..16370d146 100644 --- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp +++ b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp @@ -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& 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> 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 "