From c35aa0238c730e68810b422bcc91d2a0790538ea Mon Sep 17 00:00:00 2001 From: jianan-gu Date: Fri, 30 Jan 2026 14:30:13 +0800 Subject: [PATCH] [CPU][INT4] Add INT4 kernels for CPU (#8226) Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- python/sglang/srt/layers/amx_utils.py | 9 + python/sglang/srt/layers/quantization/fp8.py | 12 +- .../sglang/srt/layers/quantization/unquant.py | 12 +- .../srt/layers/quantization/w8a8_int8.py | 12 +- python/sglang/srt/models/deepseek_v2.py | 2 - sgl-kernel/csrc/cpu/gemm.h | 69 ++ sgl-kernel/csrc/cpu/gemm_int4.cpp | 811 ++++++++++++++++++ sgl-kernel/csrc/cpu/moe.cpp | 113 ++- sgl-kernel/csrc/cpu/moe_int4.cpp | 484 +++++++++++ sgl-kernel/csrc/cpu/torch_extension_cpu.cpp | 34 +- test/srt/cpu/test_gemm.py | 50 ++ test/srt/cpu/test_moe.py | 125 ++- test/srt/cpu/test_shared_expert.py | 6 - test/srt/cpu/utils.py | 102 +++ 14 files changed, 1769 insertions(+), 72 deletions(-) create mode 100644 sgl-kernel/csrc/cpu/gemm_int4.cpp create mode 100644 sgl-kernel/csrc/cpu/moe_int4.cpp diff --git a/python/sglang/srt/layers/amx_utils.py b/python/sglang/srt/layers/amx_utils.py index 1fb704f95..81a455976 100644 --- a/python/sglang/srt/layers/amx_utils.py +++ b/python/sglang/srt/layers/amx_utils.py @@ -6,6 +6,15 @@ from sglang.srt.utils import cpu_has_amx_support logger = logging.getLogger(__name__) +from enum import IntEnum + + +class CPUQuantMethod(IntEnum): + UNQUANT = 0 + INT8_W8A8 = 1 + FP8_W8A16 = 2 + INT4_W4A8 = 3 + def amx_process_weight_after_loading(weight, is_conv=False): if weight.device != torch.device("cpu"): diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 679e62fdb..8b5d2ea0a 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -15,7 +15,10 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, ) from sglang.srt.environ import envs -from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading +from sglang.srt.layers.amx_utils import ( + CPUQuantMethod, + _amx_process_weight_after_loading, +) from sglang.srt.layers.dp_attention import is_allocation_symmetric from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig from sglang.srt.layers.moe.moe_runner.deep_gemm import DeepGemmMoeQuantInfo @@ -1355,13 +1358,12 @@ class Fp8MoEMethod(FusedMoEMethodBase): topk_weights, topk_ids, False, # inplace See [Note] inplace should be False in fused_experts. - False, # use_int8_w8a8 - True, # use_fp8_w8a16 + CPUQuantMethod.FP8_W8A16, layer.w13_weight_scale_inv, # w1_scale layer.w2_weight_scale_inv, # w2_scale + None, # w1_zp + None, # w2_zp self.quant_config.weight_block_size, # block_size - None, # a1_scale - None, # a2_scale True, # is_vnni ) return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 628fadbd1..4be793fad 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -6,7 +6,10 @@ import torch import torch.nn.functional as F from torch.nn.parameter import Parameter -from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading +from sglang.srt.layers.amx_utils import ( + CPUQuantMethod, + _amx_process_weight_after_loading, +) from sglang.srt.layers.moe import ( MoeRunner, MoeRunnerBackend, @@ -447,13 +450,12 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): topk_weights, topk_ids, False, # inplace # See [Note] inplace should be False in fused_experts. - False, # use_int8_w8a8 - False, # use_fp8_w8a16 + CPUQuantMethod.UNQUANT, None, # w1_scale None, # w2_scale + None, # w1_zp + None, # w2_zp None, # block_size - None, # a1_scale - None, # a2_scale True, # is_vnni ) return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/layers/quantization/w8a8_int8.py b/python/sglang/srt/layers/quantization/w8a8_int8.py index 7c9ff6385..fb12bc4fe 100644 --- a/python/sglang/srt/layers/quantization/w8a8_int8.py +++ b/python/sglang/srt/layers/quantization/w8a8_int8.py @@ -8,7 +8,10 @@ import torch from torch.nn.parameter import Parameter from sglang.srt.distributed import get_tensor_model_parallel_world_size -from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading +from sglang.srt.layers.amx_utils import ( + CPUQuantMethod, + _amx_process_weight_after_loading, +) from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo from sglang.srt.layers.parameter import ChannelQuantScaleParameter, ModelWeightParameter @@ -352,13 +355,12 @@ class W8A8Int8MoEMethod(FusedMoEMethodBase): topk_weights, topk_ids, False, # inplace See [Note] inplace should be False in fused_experts. - True, # use_int8_w8a8 - False, # use_fp8_w8a16 + CPUQuantMethod.INT8_W8A8, layer.w13_weight_scale, # w1_scale layer.w2_weight_scale, # w2_scale + None, # w1_zp + None, # w2_zp None, # block_size - layer.w13_input_scale, # a1_scale - layer.w2_input_scale, # a2_scale True, # is_vnni ) return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 329ee7fdc..624b7112f 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -742,8 +742,6 @@ class DeepseekV2MoE(nn.Module): if self.shared_experts_is_fp8 else None ), # block_size - None, # a1_scale - None, # a2_scale True, # is_vnni ) if self.tp_size > 1 and not should_allreduce_fusion: diff --git a/sgl-kernel/csrc/cpu/gemm.h b/sgl-kernel/csrc/cpu/gemm.h index 6a16a2985..cbed7edea 100644 --- a/sgl-kernel/csrc/cpu/gemm.h +++ b/sgl-kernel/csrc/cpu/gemm.h @@ -56,9 +56,27 @@ inline int64_t get_row_size(int64_t K, bool use_int8_w8a8) { return use_int8_w8a8 ? K + sizeof(int32_t) : K; } +enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3 }; + +constexpr bool operator==(CPUQuantMethod a, int64_t b) { + return static_cast(a) == b; +} + +constexpr bool operator==(int64_t a, CPUQuantMethod b) { + return a == static_cast(b); +} + +inline int64_t get_4bit_block_k_size(int64_t group_size) { + return group_size > 128 ? 128 : group_size; +} + // pack weight to vnni format at::Tensor convert_weight_packed(at::Tensor& weight); +// pack weight to vnni format for int4 +std::tuple +convert_weight_packed_scale_zp(at::Tensor qweight, at::Tensor qzeros, at::Tensor scales); + // moe implementations for int8 w8a8 template void fused_experts_int8_kernel_impl( @@ -132,6 +150,37 @@ void shared_expert_int8_kernel_impl( int64_t N, int64_t K); +template +void fused_experts_int4_w4a8_kernel_impl( + scalar_t* __restrict__ output, + scalar_t* __restrict__ ic0, + scalar_t* __restrict__ ic1, + scalar_t* __restrict__ ic2, + uint8_t* __restrict__ A_tmp, + uint8_t* __restrict__ Aq_tmp, + float* __restrict__ As_tmp, + int32_t* __restrict__ Azp_tmp, + float* __restrict__ C_tmp, + int8_t* __restrict__ dqB_tmp, + const scalar_t* __restrict__ input, + const uint8_t* __restrict__ packed_w1, + const uint8_t* __restrict__ packed_w2, + const int8_t* __restrict__ w1z, + const int8_t* __restrict__ w2z, + const float* __restrict__ w1s, + const float* __restrict__ w2s, + int group_size, + const float* __restrict__ topk_weights, + const int32_t* __restrict__ sorted_ids, + const int32_t* __restrict__ expert_ids, + const int32_t* __restrict__ offsets, + int64_t M, + int64_t N, + int64_t K, + int64_t E, + int64_t topk, + int64_t num_tokens_post_pad); + template void shared_expert_fp8_kernel_impl( scalar_t* __restrict__ output, @@ -200,3 +249,23 @@ void tinygemm_kernel( bool brg, int64_t block_size_K, bool do_unpack = true); + +template +void tinygemm_kernel( + scalar_t* C, + float* C_temp, + const uint8_t* A, + const float* scales_a, + const int32_t* qzeros_a, + const uint8_t* B, + const float* scales_b, + const int8_t* qzeros_b, + const int32_t* compensation, + int8_t* dqB_tmp, + int64_t M, + int64_t K, + int64_t lda, + int64_t ldc_f, + int64_t ldc_s, + bool store_out, + bool use_brgemm); diff --git a/sgl-kernel/csrc/cpu/gemm_int4.cpp b/sgl-kernel/csrc/cpu/gemm_int4.cpp new file mode 100644 index 000000000..e541f1bfc --- /dev/null +++ b/sgl-kernel/csrc/cpu/gemm_int4.cpp @@ -0,0 +1,811 @@ +#include + +#include "gemm.h" +#include "vec.h" + +namespace { + +#define BLOCK_N block_size_n() +#define BLOCK_M 128 + +template +struct ActDtype; +template <> +struct ActDtype { + using type = int8_t; +}; +template <> +struct ActDtype { + using type = uint8_t; +}; + +struct alignas(32) m256i_wrapper { + __m256i data; +}; + +#if defined(CPU_CAPABILITY_AVX512) +inline std::array load_zps_4vnni(const int8_t* __restrict__ zps) { + // broadcast 01234567 to + // 01234567012345670123456701234567 + __m256i vzps_low = _mm256_set1_epi64x(*reinterpret_cast(zps)); + __m256i vzps_high = _mm256_set1_epi64x(*reinterpret_cast(zps + 8)); + // shuffle from + // 01234567012345670123456701234567 + // to + // 00001111222233334444555566667777 + __m256i shuffle_mask = + _mm256_set_epi8(7, 7, 7, 7, 6, 6, 6, 6, 5, 5, 5, 5, 4, 4, 4, 4, 3, 3, 3, 3, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0); + vzps_low = _mm256_shuffle_epi8(vzps_low, shuffle_mask); + vzps_high = _mm256_shuffle_epi8(vzps_high, shuffle_mask); + m256i_wrapper vzps_low_wp, vzps_high_wp; + vzps_low_wp.data = vzps_low; + vzps_high_wp.data = vzps_high; + return {vzps_low_wp, vzps_high_wp}; +} + +inline std::array load_uint4_as_int8(const uint8_t* __restrict__ qB) { + __m256i packed = _mm256_loadu_si256(reinterpret_cast(qB)); + const __m256i low_mask = _mm256_set1_epi8(0x0f); + __m256i high = _mm256_srli_epi16(packed, 4); + high = _mm256_and_si256(high, low_mask); + __m256i low = _mm256_and_si256(packed, low_mask); + m256i_wrapper low_wp, high_wp; + low_wp.data = low; + high_wp.data = high; + return {low_wp, high_wp}; +} + +template +void _dequant_weight_zp_only(const uint8_t* __restrict__ B, int8_t* dqB, const int8_t* __restrict__ qzeros, int64_t K) { + // unpack weight int8 -> two int4 + // subtract zero point + // B shape = [K, ldb] = [K, N / 2], actual shape = [K / 4, N / 2, 4] + // dqB shape = [K, N], actual shape = [K / 4, N, 4] +#pragma GCC unroll 2 + for (int n = 0; n < N; n += 16) { + auto [zps_low_wp, zps_high_wp] = load_zps_4vnni(&qzeros[n]); + auto zps_low = zps_low_wp.data; + auto zps_high = zps_high_wp.data; + for (int k = 0; k < K; k += 4) { + auto [vb_low_wp, vb_high_wp] = load_uint4_as_int8(B + ldb * k + n / 2 * 4); + auto vb_low = vb_low_wp.data; + auto vb_high = vb_high_wp.data; + vb_high = _mm256_sub_epi8(vb_high, zps_high); + vb_low = _mm256_sub_epi8(vb_low, zps_low); + // store vb to B + _mm256_storeu_si256(reinterpret_cast<__m256i_u*>(dqB + N * k + n * 4), vb_low); + _mm256_storeu_si256(reinterpret_cast<__m256i_u*>(dqB + N * k + (n + 8) * 4), vb_high); + } + } +} + +template +void _dequant_and_store( + float* __restrict__ output, + const int32_t* __restrict__ input, + const float* __restrict__ scale_a, + const int32_t* __restrict__ zp_a, + const float* __restrict__ scale_b, + const int32_t* __restrict__ comp_b, + int M, + int ldi, + int ldo, + int ldsa = 1) { + for (int m = 0; m < M; ++m) { + float a_scale = *(scale_a + m * ldsa); + __m512 va_scale = _mm512_set1_ps(a_scale); + int32_t a_zp; + __m512i va_zp; + if constexpr (!sym_quant_act) { + a_zp = *(zp_a + m * ldsa); + va_zp = _mm512_set1_epi32(a_zp); + } + int n = 0; +#pragma GCC unroll 2 + for (; n < N; n += 16) { + __m512i vc = _mm512_loadu_si512(input + m * ldi + n); + if constexpr (!sym_quant_act) { + __m512i vb_comp = _mm512_loadu_si512(comp_b + n); + vc = _mm512_sub_epi32(vc, _mm512_mullo_epi32(vb_comp, va_zp)); + } + __m512 vc_f = _mm512_cvtepi32_ps(vc); + __m512 vc_f_mul = _mm512_mul_ps(vc_f, va_scale); + __m512 vb_s = _mm512_loadu_ps(scale_b + n); + vc_f_mul = _mm512_mul_ps(vc_f_mul, vb_s); + if constexpr (accum) { + __m512 vo = _mm512_loadu_ps(output + m * ldo + n); + _mm512_storeu_ps(output + m * ldo + n, _mm512_add_ps(vo, vc_f_mul)); + } else { + _mm512_storeu_ps(output + m * ldo + n, vc_f_mul); + } + } + for (; n < N; ++n) { + float dq_val; + if constexpr (sym_quant_act) { + dq_val = (float)input[m * ldi + n] * a_scale * scale_b[n]; + } else { + dq_val = (float)(input[m * ldi + n] - a_zp * comp_b[n]) * a_scale * scale_b[n]; + } + if constexpr (accum) { + output[m * ldo + n] += dq_val; + } else { + output[m * ldo + n] = dq_val; + } + } + } +} + +#else +template +void _dequant_weight_zp_only(const uint8_t* B, int8_t* dqB, const int8_t* qzeros, int64_t K) { + // B shape = [K, N / 2] + // dqB shape = [K, N] + for (int k = 0; k < K; ++k) { + for (int n = 0; n < N / 2; ++n) { + int32_t b = (int32_t)B[k * ldb + n]; + dqB[k * N + n * 2] = (b & 0xf) - qzeros[n]; + dqB[k * N + n * 2 + 1] = (b >> 4) - qzeros[n]; + } + } +} +#endif + +#if defined(CPU_CAPABILITY_AVX512) +inline __m512i combine_m256i(__m256i a, __m256i b) { + __m512i c = _mm512_castsi256_si512(a); + return _mm512_inserti64x4(c, b, 1); +} + +inline __m512i combine_m256i(std::array two_256) { + return combine_m256i(two_256[0].data, two_256[1].data); +} + +// negate elements in a according to b's sign +static inline __m512i _mm512_sign_epi8(__m512i a, __m512i b) { + __m512i zero = _mm512_setzero_si512(); + __mmask64 blt0 = _mm512_movepi8_mask(b); + return _mm512_mask_sub_epi8(a, blt0, zero, a); +} + +template +void _dequant_gemm_accum_small_M( + float* __restrict__ C, + const uint8_t* A, + const float* scales_a, + const int32_t* qzeros_a, + const uint8_t* B, + const float* scales_b, + const int8_t* qzeros_b, + int64_t K, + int64_t lda, + int64_t ldc) { + // if sym_quant_act is true, A pointer type is passed in as uint8_t* but actually int8_t*. + + constexpr int COLS = N / 16; + // Computing compensation is faster than loading it for small M + // because it's memory bound. + __m512i ones = _mm512_set1_epi8(1); // used for computing compensation + __m512i va; + __m512i vb[COLS]; + __m512i vc[M * COLS]; + __m512 vscales[COLS]; + __m512i vzps[COLS]; + __m512i vcompensate[COLS]; + + // Load scales and zps + Unroll{}([&](auto i) { + vscales[i] = _mm512_loadu_ps(scales_b + i * 16); + vzps[i] = combine_m256i(load_zps_4vnni(qzeros_b + i * 16)); + if constexpr (!sym_quant_act) { + vcompensate[i] = _mm512_setzero_epi32(); + } + }); + Unroll{}([&](auto i) { vc[i] = _mm512_setzero_epi32(); }); + + auto compute = [&](auto i, int k) { + constexpr const int row = i / COLS; + constexpr const int col = i % COLS; + + if constexpr (col == 0) { + va = _mm512_set1_epi32(*(int32_t*)(A + row * lda + k)); + } + + if constexpr (row == 0) { + int B_offset = k * ldb + col * 16 * 2; + vb[col] = combine_m256i(load_uint4_as_int8(B + B_offset)); + vb[col] = _mm512_sub_epi8(vb[col], vzps[col]); + if constexpr (!sym_quant_act) { + vcompensate[col] = _mm512_dpbusd_epi32(vcompensate[col], ones, vb[col]); + } + _mm_prefetch(B + B_offset + 128 * ldb, _MM_HINT_T0); + } + if constexpr (sym_quant_act) { + auto vsb = _mm512_sign_epi8(vb[col], va); + auto vabsa = _mm512_sign_epi8(va, va); + vc[i] = _mm512_dpbusds_epi32(vc[i], vabsa, vsb); + } else { + vc[i] = _mm512_dpbusd_epi32(vc[i], va, vb[col]); + } + }; + + // Accumulate along k + constexpr const int unroll = 4; + int k = 0; + for (; k < K / 4 / unroll; k++) { + Unroll{}([&](auto i) { Unroll{}(compute, 4 * (k * unroll + i)); }); + } + k *= 4 * unroll; + for (; k < K; k += 4) { + Unroll{}(compute, k); + } + + // Store to C + auto store = [&](auto i) { + constexpr const int row = i / COLS; + constexpr const int col = i % COLS; + // compute (qC - compensate * zp_a) * scale_a * scale_b + __m512 vc_float; + if constexpr (!sym_quant_act) { + vc[i] = _mm512_sub_epi32(vc[i], _mm512_mullo_epi32(vcompensate[col], _mm512_set1_epi32(*(qzeros_a + row)))); + } + vc_float = _mm512_cvtepi32_ps(vc[i]); + vc_float = _mm512_mul_ps(vc_float, _mm512_set1_ps(*(scales_a + row))); + + vc_float = _mm512_mul_ps(vc_float, vscales[col]); + auto vc_old = _mm512_loadu_ps(C + row * ldc + col * 16); + vc_float = _mm512_add_ps(vc_float, vc_old); + _mm512_storeu_ps(C + row * ldc + col * 16, vc_float); + }; + Unroll{}(store); +} + +#define CALL_DEQUANT_GEMM_ACCUM_SMALL_M(M) \ + _dequant_gemm_accum_small_M(C, A, scales_a, qzeros_a, B, scales_b, qzeros_b, K, lda, ldc); +#endif + +template +void _dequant_gemm_accum( + float* C, + const uint8_t* A, + const float* scales_a, + const int32_t* qzeros_a, + const uint8_t* B, + const float* scales_b, + const int8_t* qzeros_b, + const int32_t* compensation, + int8_t* dqB, + int64_t M, + int64_t K, + int64_t lda, + int64_t ldc, + bool use_brgemm) { + // Compute GEMM int8 * int8 -> int32 + // dequant result to float by applying scales/qzeros +#if defined(CPU_CAPABILITY_AVX512) + if (!use_brgemm) { + switch (M) { + case 1: + CALL_DEQUANT_GEMM_ACCUM_SMALL_M(1); + break; + case 2: + CALL_DEQUANT_GEMM_ACCUM_SMALL_M(2); + break; + case 3: + CALL_DEQUANT_GEMM_ACCUM_SMALL_M(3); + break; + case 4: + CALL_DEQUANT_GEMM_ACCUM_SMALL_M(4); + break; + default: + TORCH_CHECK(false, "tinygemm_kernel: unexpected M for AVX path!"); + } + return; + } + + _dequant_weight_zp_only(B, dqB, qzeros_b, K); + using Tin = typename ActDtype::type; + Tin* A_ptr = (Tin*)A; + if (use_brgemm) { + int32_t C_i32[M * N]; + at::native::cpublas::brgemm( + M, N, K, lda, N /*ldb*/, N /*ldc*/, false /* add_C */, A_ptr, dqB, C_i32, true /* is_vnni */); + _mm_prefetch(B + N * K / 2, _MM_HINT_T0); + _mm_prefetch(A + K, _MM_HINT_T0); + _dequant_and_store( + C, C_i32, scales_a, qzeros_a, scales_b, compensation, M, N /*ldi*/, ldc, 1 /*ldsa*/); + } else +#endif + { + TORCH_CHECK(false, "tinygemm_kernel: scalar path not implemented!"); + } +} + +template +inline void copy_bias(const float* bias_ptr, float* y_buf, int64_t m) { + if (bias_ptr) { + for (int i = 0; i < m; ++i) { + int j = 0; +#if defined(CPU_CAPABILITY_AVX512) +#pragma GCC unroll 2 + for (; j < N; j += 16) { + __m512 bias_vec = _mm512_loadu_ps(bias_ptr + j); + _mm512_storeu_ps(y_buf + i * N + j, bias_vec); + } +#endif + for (; j < N; ++j) { + y_buf[i * N + j] = bias_ptr[j]; + } + } + } else { // initialize to zero + for (int i = 0; i < m; ++i) { + int j = 0; +#if defined(CPU_CAPABILITY_AVX512) +#pragma GCC unroll 2 + for (; j < N; j += 16) { + __m512 zero_vec = _mm512_setzero_ps(); + _mm512_storeu_ps(y_buf + i * N + j, zero_vec); + } +#endif + for (; j < N; ++j) { + y_buf[i * N + j] = 0; + } + } + } +} + +template +inline void store_out(const float* y_buf, out_dtype* c_ptr, int64_t m, /* int64_t n, */ int64_t lda) { + for (int i = 0; i < m; ++i) { + int j = 0; + if constexpr (std::is_same::value) { +#if defined(CPU_CAPABILITY_AVX512) +#pragma GCC unroll 2 + for (; j < N; j += 16) { + __m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j); + _mm512_storeu_ps(c_ptr + i * lda + j, y_vec); + } +#endif + for (; j < N; ++j) { + c_ptr[i * lda + j] = y_buf[i * N + j]; + } + } else if constexpr (std::is_same::value) { +#if defined(CPU_CAPABILITY_AVX512) +#pragma GCC unroll 2 + for (; j < N; j += 16) { + __m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j); + __m256i y_bf16_vec = at::vec::cvtfp32_bf16(y_vec); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_bf16_vec); + } +#endif + for (; j < N; ++j) { + c_ptr[i * lda + j] = at::BFloat16(y_buf[i * N + j]); + } + } else if constexpr (std::is_same::value) { +#if defined(CPU_CAPABILITY_AVX512) +#pragma GCC unroll 2 + for (; j < N; j += 16) { + __m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j); + __m256i y_fp16_vec = at::vec::cvtfp32_fp16(y_vec); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_fp16_vec); + } +#endif + for (; j < N; ++j) { + c_ptr[i * lda + j] = at::Half(y_buf[i * N + j]); + } + } else { + TORCH_CHECK(false, "Unsupported output dtype"); + } + } +} + +void fill_val_stub(int32_t* __restrict__ output, int32_t value, int64_t size) { + using iVec = at::vec::Vectorized; + constexpr int VecSize = iVec::size(); + const iVec fill_val_vec = iVec(value); + int64_t d; +#pragma GCC unroll 4 + for (d = 0; d <= size - VecSize; d += VecSize) { + fill_val_vec.store(output + d); + } + for (; d < size; ++d) { + output[d] = value; + } +} + +template +void _da8w4_linear_impl( + act_dtype* __restrict__ input, + const float* __restrict__ input_scales, + const int32_t* __restrict__ input_qzeros, + const uint8_t* __restrict__ weight, + const float* __restrict__ weight_scales, + const int8_t* __restrict__ weight_qzeros, + const float* __restrict__ bias, + out_dtype* __restrict__ output, + float* __restrict__ output_temp, + int8_t* __restrict__ dequant_weight_temp, + int64_t M, + int64_t N, + int64_t K, + int64_t num_groups) { + // weight + compensation shape = [Nc, Kc, BLOCK_N * _block_k / 2 + BLOCK_N*sizeof(int32_t)] + // scales/qzeros shape = [Nc, G, BLOCK_N] + const bool use_brgemm = can_use_brgemm(M); + int64_t block_m = [&]() -> long { + if (M <= 48) { + return M; + } else if (M < 64) { + return 32; + } else if (M < 96) { + return 64; + } else { + return 128; + } + }(); + int64_t Mc = div_up(M, block_m); + bool parallel_on_M = M > 128; + int64_t Nc = N / BLOCK_N; + int64_t num_blocks = parallel_on_M ? Mc * Nc : Nc; + int64_t group_size = div_up(K, num_groups); + int64_t _block_k = get_4bit_block_k_size(group_size); + int64_t Kc = K / _block_k; + int64_t block_per_group = group_size / _block_k; + + at::parallel_for(0, num_blocks, 1, [&](int64_t begin, int64_t end) { + int tid = get_thread_num(); + float* C_tmp = output_temp + tid * block_m * BLOCK_N; + int8_t* dqB_tmp = dequant_weight_temp + tid * _block_k * BLOCK_N; + for (const auto i : c10::irange(begin, end)) { + int64_t mc = parallel_on_M ? i / Nc : 0; + int64_t nc = parallel_on_M ? i % Nc : i; + int64_t mc_end = parallel_on_M ? mc + 1 : Mc; + + for (int mci = mc; mci < mc_end; ++mci) { + int64_t m_size = mci * block_m + block_m > M ? M - mci * block_m : block_m; + // copy bias to y_buf if bias is not None + auto bias_data = bias ? bias + nc * BLOCK_N : nullptr; + copy_bias(bias_data, C_tmp, m_size); + for (int kci = 0; kci < Kc; ++kci) { + int32_t* compensation_ptr = + sym_quant_act + ? nullptr + : (int32_t*)(void*)(weight + (nc * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) + + _block_k * BLOCK_N / 2) /*Bcomp*/; + _dequant_gemm_accum( + /*C*/ C_tmp, + /*A*/ (uint8_t*)input + mci * block_m * K + kci * _block_k, + /*scales_a*/ input_scales + mci * block_m, + /*qzeros_a*/ input_qzeros + mci * block_m, + /*B*/ weight + (nc * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))), + /*scales_b*/ weight_scales + nc * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N, + /*qzeros_b*/ weight_qzeros + nc * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N, + /*Bcomp*/ compensation_ptr, + /*dqB_tmp*/ dqB_tmp, + /*M*/ m_size, + /*K*/ _block_k, + /*lda*/ K, + /*ldc*/ BLOCK_N, + /*use_brgemm*/ use_brgemm); + } + // store y_buf to output with dtype conversion + store_out(C_tmp, output + mci * block_m * N + nc * BLOCK_N, m_size, N /*lda*/); + } + } + if (use_brgemm) { + at::native::cpublas::brgemm_release(); + } + }); +} + +} // anonymous namespace + +/* +return: packed_weight, packed_scales, packed_qzeros +*/ +std::tuple convert_int4_weight_packed_with_compensation( + const at::Tensor& weight, const at::Tensor& scales, const at::Tensor& qzeros) { + // weight shape = [N, K] + // scales shape = [N, G] + // qzeros shape = [N, G] + TORCH_CHECK(weight.dim() == 2, "DA8W4 CPU: Weight should be a 2D tensor for packing"); + TORCH_CHECK(weight.size(1) % 2 == 0, "DA8W4 CPU: Weight should have even number of columns for packing"); + + auto new_scales = scales; + auto new_qzeros = qzeros; + if (new_scales.dim() == 1) { + new_scales.unsqueeze_(1); + } + new_scales = new_scales.to(at::kFloat); + if (new_qzeros.dim() == 1) { + new_qzeros.unsqueeze_(1); + } + new_qzeros = new_qzeros.to(at::kChar); + int64_t N = weight.size(0); + int64_t K = weight.size(1); + int64_t G = scales.size(1); + int64_t group_size = K / G; + int64_t _block_k = get_4bit_block_k_size(group_size); + constexpr int block_n = block_size_n(); + int64_t Nc = N / block_n; + int64_t Kc = K / _block_k; + + // Reorder weight to [N/block_n, K/_block_k, _block_k, block_n] + // Reorder scales/qzeros to [N/block_n, G, block_n] + // weight + compensation shape = [Nc, Kc, block_n * _block_k / 2 + block_n*sizeof(int32_t)] + // scales/qzeros shape = [Nc, G, block_n] + auto weight_view = weight.view({Nc, block_n, Kc, _block_k}); + at::Tensor weight_reordered = weight_view.permute({0, 2, 3, 1}).contiguous(); + at::Tensor blocked_weight; + at::Tensor blocked_scales = new_scales.view({Nc, block_n, G}).permute({0, 2, 1}).contiguous(); + at::Tensor blocked_qzeros = new_qzeros.view({Nc, block_n, G}).permute({0, 2, 1}).contiguous(); + // Compensation = Σ(k)(W[k][n] - ZP[n]) for each block. + auto weight_sub_qzero = weight.view({Nc, block_n, G, -1}).to(at::kInt) - new_qzeros.view({Nc, block_n, G, -1}); + weight_sub_qzero = weight_sub_qzero.view({Nc, block_n, Kc, _block_k}); + at::Tensor compensation = weight_sub_qzero.sum(-1); + compensation = compensation.permute({0, 2, 1}).contiguous().to(at::kInt); + int64_t buffer_size_nbytes = _block_k * block_n / 2 + block_n * sizeof(int32_t); + blocked_weight = at::empty({Nc, Kc, buffer_size_nbytes}, weight.options()); + + auto weight_ptr = weight_reordered.data_ptr(); + auto compensation_ptr = compensation.data_ptr(); + auto blocked_weight_ptr = blocked_weight.data_ptr(); + int64_t num_blocks = Nc * Kc; + at::parallel_for(0, num_blocks, 1, [&](int64_t begin, int64_t end) { + for (const auto i : c10::irange(begin, end)) { + auto in_ptr = weight_ptr + i * _block_k * block_n; + auto out_ptr = blocked_weight_ptr + i * block_n * (_block_k / 2 + sizeof(int32_t)); + int32_t* comp_in_prt = compensation_ptr + i * block_n; + int32_t* comp_out_prt = (int32_t*)(void*)(blocked_weight_ptr + i * block_n * (_block_k / 2 + sizeof(int32_t)) + + _block_k * block_n / 2); + // Reorder weight block to VNNI4 and pack two lanes along N + // N=16 viewed as two lanes: a0, ...a7, b0, ...b7 + // pack two lanes: [a0, b0], ..., [a7, b7] + // plain shape = [_block_k, block_n] + // packed shape = [_block_k / 4, block_n / 2, 4] viewed as [_block_k, block_n / 2] + constexpr int n_group_size = 8; + constexpr int vnni_size = 4; + constexpr int n_group = block_n / n_group_size; // 4 + for (int nb = 0; nb < n_group; nb += 2) { + for (int k = 0; k < _block_k; k += vnni_size) { + for (int ni = 0; ni < n_group_size; ++ni) { + for (int ki = 0; ki < vnni_size; ++ki) { + int src_idx_1 = nb * n_group_size + ni + (k + ki) * block_n; + int src_idx_2 = (nb + 1) * n_group_size + ni + (k + ki) * block_n; + int dst_idx = (nb / 2 * n_group_size + ni) * vnni_size + k * block_n / 2 + ki; + uint8_t src_1 = *(in_ptr + src_idx_1); + uint8_t src_2 = *(in_ptr + src_idx_2); + uint8_t dst = (src_1 & 0x0f) | ((src_2 & 0x0f) << 4); + *(out_ptr + dst_idx) = dst; + } + } + } + } + // compensation [block_n] + for (int nb = 0; nb < block_n; nb++) { + *(comp_out_prt + nb) = *(comp_in_prt + nb); + } + } + }); + + return std::make_tuple(std::move(blocked_weight), std::move(blocked_scales), std::move(blocked_qzeros)); +} + +std::tuple autoawq_to_int4pack( + at::Tensor qweight, // (*, K, N / 8), int32 + at::Tensor qzeros) // (*, K / group_size, N / 8), int32 +{ + // bitshifts: [0, 4, 1, 5, 2, 6, 3, 7] * 4 + auto bitshifts = at::tensor({0, 4, 1, 5, 2, 6, 3, 7}, at::kInt) * 4; + // qweight: assumed shape [..., K, N/8] (int32) + auto qweight_unsq = qweight.unsqueeze(-1); // [..., K, N/8, 1] + auto shape = qweight_unsq.sizes().vec(); // shape: [A, B, C, 1] + shape[3] = 8; + auto unpacked = at::bitwise_right_shift(qweight_unsq, bitshifts) & 0xF; + auto qweight_final = unpacked.flatten(-2).transpose(-1, -2).to(at::kByte); + + auto qzeros_unsq = qzeros.unsqueeze(-1); + auto qzeros_unpacked = at::bitwise_right_shift(qzeros_unsq, bitshifts) & 0xF; + auto qzeros_final = qzeros_unpacked.flatten(-2).to(at::kByte); + + return std::make_tuple(qweight_final, qzeros_final); +} + +std::tuple convert_weight_packed_scale_zp( + at::Tensor qweight, // (*, K, N / 8), int32 + at::Tensor qzeros, // (*, K / group_size, N / 8), int32 + at::Tensor scales // (*, K / group_size, N), bfloat16 +) { + auto res = autoawq_to_int4pack(qweight, qzeros); + auto _qweight = std::get<0>(res); + auto _qzeros = std::get<1>(res); + auto _scales = scales; + _qzeros = _qzeros.transpose(-2, -1).contiguous(); // .T + _scales = _scales.transpose(-2, -1).contiguous(); + if (_qweight.dim() == 3) { // Dim=3 for MOE packing, TODO: refine a unified loop + int64_t E = _qweight.size(0); + int64_t K = _qweight.size(2); + int64_t G = _scales.size(2); + int64_t group_size = K / G; + int64_t _block_k = get_4bit_block_k_size(group_size); + int64_t block_n = block_size_n(); + int64_t Nc = _qweight.size(1) / block_n; + int64_t Kc = K / _block_k; + int64_t buffer_size_nbytes = _block_k * block_n / 2 + block_n * sizeof(int32_t); + auto blocked_weight = at::empty({E, Nc, Kc, buffer_size_nbytes}, _qweight.options()); + auto blocked_scales = at::empty({E, Nc, G, block_n}, _scales.options()).to(at::kFloat); + auto blocked_qzeros = at::empty({E, Nc, G, block_n}, _qzeros.options()).to(at::kChar); + for (int i = 0; i < _qweight.size(0); i++) { + auto res_ = convert_int4_weight_packed_with_compensation(_qweight[i], _scales[i], _qzeros[i]); + blocked_weight[i] = std::get<0>(res_); + blocked_scales[i] = std::get<1>(res_); + blocked_qzeros[i] = std::get<2>(res_); + } + _qweight = blocked_weight; + _scales = blocked_scales; + _qzeros = blocked_qzeros; + } else { + auto res_ = convert_int4_weight_packed_with_compensation(_qweight, _scales, _qzeros); + _qweight = std::get<0>(res_); + _scales = std::get<1>(res_); + _qzeros = std::get<2>(res_); + } + + return std::make_tuple(_qweight, _qzeros, _scales); +} + +at::Tensor int4_scaled_mm_cpu_with_quant( + const at::Tensor& input, + const at::Tensor& weight, + const at::Tensor& weight_scales, + const at::Tensor& weight_qzeros, + const std::optional& bias, + at::ScalarType output_dtype) { + RECORD_FUNCTION("sgl-kernel::int4_scaled_mm_cpu_with_quant", std::vector({input, weight})); + + int64_t M_a = input.size(0); + int64_t K_a = input.size(1); + int64_t lda = input.stride(0); + + const auto st = input.scalar_type(); + TORCH_CHECK( + st == at::kBFloat16 || st == at::kHalf, "int4_scaled_mm_cpu_with_quant: expect A to be bfloat16 or half."); + + constexpr bool sym_quant_act = false; // TODO: add sym quant path + using Tin = typename ActDtype::type; + int64_t act_buffer_size = /* act quant */ M_a * K_a + + /* act scale */ M_a * sizeof(float) + + /* act zp */ M_a * sizeof(int32_t); + auto act_buffer = at::empty({act_buffer_size}, input.options().dtype(at::kByte)); + // asym path, activation quants into uint8_t + auto Aq_data = act_buffer.data_ptr(); + auto As_data = reinterpret_cast(Aq_data + M_a * K_a); + auto Azp_data = reinterpret_cast(As_data + M_a); + fill_val_stub(Azp_data, 128, M_a); // sym_a s8s8 is unified to u8s8 with compensation (128) + + auto out_sizes = input.sizes().vec(); + int64_t N = weight_scales.size(0) * weight_scales.size(-1); + out_sizes.back() = N; + auto output = at::empty(out_sizes, input.options()); + // weight + compensation shape = [Nc, Kc, BLOCK_N * _block_k / 2 + BLOCK_N*sizeof(int32_t)] + // scales/qzeros shape = [Nc, G, BLOCK_N] + int64_t Nc = weight.size(0); + int64_t Kc = weight.size(1); + int64_t _block_k = K_a / Kc; + TORCH_CHECK(N == Nc * BLOCK_N, "DA8W4: weight and input shapes mismatch"); + // scales/qzeros shape = [Nc, G, BLOCK_N] + int64_t num_groups = weight_scales.size(1); + + const uint8_t* b_ptr = weight.data_ptr(); + const float* b_scales_ptr = weight_scales.data_ptr(); + const int8_t* b_qzeros_ptr = weight_qzeros.data_ptr(); + const float* bias_ptr = bias.has_value() ? bias.value().data_ptr() : nullptr; + int num_threads = at::get_num_threads(); + int64_t temp_buffer_size = /* output temp */ num_threads * BLOCK_M * BLOCK_N * sizeof(float) + + /* weight dequant temp */ num_threads * _block_k * BLOCK_N; + auto c_temp_buffer = at::empty({temp_buffer_size}, input.options().dtype(at::kChar)); + float* c_temp_ptr = (float*)((void*)(c_temp_buffer.data_ptr())); + int8_t* dqB_temp_ptr = (int8_t*)((void*)(c_temp_ptr + num_threads * BLOCK_M * BLOCK_N)); + +#define LAUNCH_DA8W4_LINEAR_WITH_QUANT_IMPL(sym_quant_act) \ + AT_DISPATCH_FLOATING_TYPES_AND2( \ + at::ScalarType::BFloat16, at::ScalarType::Half, output_dtype, "int4_scaled_mm_cpu_with_quant", [&] { \ + const scalar_t* __restrict__ A_data = input.data_ptr(); \ + scalar_t* __restrict__ c_ptr = output.data_ptr(); \ + at::parallel_for(0, M_a, 0, [&](int64_t begin, int64_t end) { \ + for (int64_t m = begin; m < end; ++m) { \ + quantize_row_int8(Aq_data + m * K_a, As_data[m], A_data + m * lda, K_a); \ + } \ + }); \ + _da8w4_linear_impl( \ + Aq_data, \ + As_data, \ + Azp_data, \ + b_ptr, \ + b_scales_ptr, \ + b_qzeros_ptr, \ + bias_ptr, \ + c_ptr, \ + c_temp_ptr, \ + dqB_temp_ptr, \ + M_a, \ + N, \ + K_a, \ + num_groups); \ + }); + + LAUNCH_DA8W4_LINEAR_WITH_QUANT_IMPL(sym_quant_act); + + return output; +} +template +inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int64_t size) { + using Vec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; +// no remainder +#pragma GCC unroll 4 + for (int64_t d = 0; d < size; d += Vec::size()) { + fVec x0 = fVec::loadu(input + d); + fVec x1 = fVec::loadu(input + d + fVec::size()); + Vec res = convert_from_float_ext(x0, x1); + res.store(out + d); + } +} + +template +void tinygemm_kernel( + scalar_t* C, + float* C_temp, + const uint8_t* A, + const float* scales_a, + const int32_t* qzeros_a, + const uint8_t* B, + const float* scales_b, + const int8_t* qzeros_b, + const int32_t* compensation, + int8_t* dqB_tmp, + int64_t M, + int64_t K, + int64_t lda, + int64_t ldc_f, + int64_t ldc_s, + bool store_out, + bool use_brgemm) { + // TODO: add sym quant act, now only asym + _dequant_gemm_accum( + C_temp, A, scales_a, qzeros_a, B, scales_b, qzeros_b, compensation, dqB_tmp, M, K, lda, ldc_f, use_brgemm); + if (store_out) { + // copy from Ctmp to C + for (int64_t m = 0; m < M; ++m) { + copy_stub(C + m * ldc_s, C_temp + m * ldc_f, BLOCK_N); + } + } +} + +#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE) \ + template void tinygemm_kernel( \ + TYPE * C, \ + float* C_temp, \ + const uint8_t* A, \ + const float* scales_a, \ + const int32_t* qzeros_a, \ + const uint8_t* B, \ + const float* scales_b, \ + const int8_t* qzeros_b, \ + const int32_t* compensation, \ + int8_t* dqB_tmp, \ + int64_t M, \ + int64_t K, \ + int64_t lda, \ + int64_t ldc_f, \ + int64_t ldc_s, \ + bool store_out, \ + bool use_brgemm) + +INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16); +INSTANTIATE_TINYGEMM_TEMPLATE(at::Half); + +// int4 gemm dispatch api register +at::Tensor int4_scaled_mm_cpu( + at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros, at::Tensor& w_scales, std::optional bias) { + return int4_scaled_mm_cpu_with_quant(x, w, w_scales, w_zeros, bias, x.scalar_type()); +} diff --git a/sgl-kernel/csrc/cpu/moe.cpp b/sgl-kernel/csrc/cpu/moe.cpp index c3d66cec7..eaf700bf9 100644 --- a/sgl-kernel/csrc/cpu/moe.cpp +++ b/sgl-kernel/csrc/cpu/moe.cpp @@ -908,14 +908,10 @@ static inline void check_moe_scales( bool use_fp8_w8a16, const std::optional& w1_scale, const std::optional& w2_scale, - const std::optional> block_size, - const std::optional& a1_scale, - const std::optional& a2_scale) { + const std::optional> block_size) { if (use_int8_w8a8) { TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for int8 w8a8."); TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for int8 w8a8."); - TORCH_CHECK(!a1_scale.has_value(), "static quantization for activation not supported."); - TORCH_CHECK(!a2_scale.has_value(), "static quantization for activation not supported."); } if (use_fp8_w8a16) { TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for fp8 w8a16."); @@ -942,6 +938,7 @@ static inline void check_moe_scales( // topk_weights: [M, topk] // topk_ids: [M, topk] (int32_t) // + at::Tensor fused_experts_cpu( at::Tensor& hidden_states, at::Tensor& w1, @@ -949,13 +946,12 @@ at::Tensor fused_experts_cpu( at::Tensor& topk_weights, at::Tensor& topk_ids, bool inplace, - bool use_int8_w8a8, - bool use_fp8_w8a16, + int64_t moe_comp_method, const std::optional& w1_scale, const std::optional& w2_scale, + const std::optional& w1_zero, + const std::optional& w2_zero, const std::optional> block_size, - const std::optional& a1_scale, - const std::optional& a2_scale, bool is_vnni) { RECORD_FUNCTION( "sgl-kernel::fused_experts_cpu", std::vector({hidden_states, w1, w2, topk_weights, topk_ids})); @@ -972,8 +968,13 @@ at::Tensor fused_experts_cpu( CHECK_INPUT(w2); CHECK_EQ(topk_weights.sizes(), topk_ids.sizes()); CHECK_DIM(2, hidden_states); - CHECK_DIM(3, w1); - CHECK_DIM(3, w2); + if (moe_comp_method == CPUQuantMethod::INT4_W4A8 && is_vnni) { + CHECK_DIM(4, w1); + CHECK_DIM(4, w2); + } else { + CHECK_DIM(3, w1); + CHECK_DIM(3, w2); + } CHECK_DIM(2, topk_weights); CHECK_DIM(2, topk_ids); @@ -987,22 +988,29 @@ at::Tensor fused_experts_cpu( int64_t M = hidden_states.size(0); int64_t K = hidden_states.size(1); - int64_t N = w1.size(1) / 2; + int64_t N = moe_comp_method == CPUQuantMethod::INT4_W4A8 ? w1_scale.value().size(1) * w1_scale.value().size(3) / 2 + : w1.size(1) / 2; int64_t E = w1.size(0); int64_t topk = topk_weights_.size(1); // we use int32_t compensation for int8 w8a8 - int64_t packed_K = get_row_size(K, use_int8_w8a8); - int64_t packed_N = get_row_size(N, use_int8_w8a8); + int64_t packed_K = get_row_size(K, moe_comp_method == CPUQuantMethod::INT8_W8A8); + int64_t packed_N = get_row_size(N, moe_comp_method == CPUQuantMethod::INT8_W8A8); // check weight shapes CHECK_EQ(w2.size(0), E); - CHECK_EQ(w2.size(1), K); - CHECK_EQ(packed_w1.size(2), packed_K); - CHECK_EQ(packed_w2.size(2), packed_N); - + if (!(moe_comp_method == CPUQuantMethod::INT4_W4A8)) { + CHECK_EQ(w2.size(1), K); + CHECK_EQ(packed_w1.size(2), packed_K / (moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 2 : 1)); + CHECK_EQ(packed_w2.size(2), packed_N / (moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 2 : 1)); + } // check scales - check_moe_scales(use_int8_w8a8, use_fp8_w8a16, w1_scale, w2_scale, block_size, a1_scale, a2_scale); + check_moe_scales( + moe_comp_method == CPUQuantMethod::INT8_W8A8, + moe_comp_method == CPUQuantMethod::FP8_W8A16, + w1_scale, + w2_scale, + block_size); at::Tensor out_hidden_states = inplace ? hidden_states : at::empty_like(hidden_states); @@ -1058,24 +1066,29 @@ at::Tensor fused_experts_cpu( // 7. intermediate_cache0 : [M * topk, 2N] // 8. B_tmp : [T, MAX_CACHE_BLOCK_SIZE, BLOCK_N, std::max(K, N)] // - int64_t buffer_size_nbytes = M * topk * N * 2 + M * topk * K * 2 + - num_threads * BLOCK_M * K * (use_int8_w8a8 ? 1 : 2) + - num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float); + int64_t buffer_size_nbytes = + M * topk * N * 2 + M * topk * K * 2 + + num_threads * BLOCK_M * K * + (moe_comp_method == CPUQuantMethod::INT8_W8A8 | moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 1 : 2) + + num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float); - if (use_int8_w8a8) { + if (moe_comp_method == CPUQuantMethod::INT8_W8A8) { buffer_size_nbytes += std::max(M * K, M * topk * N) + M * topk * sizeof(float); } - if (use_fp8_w8a16) { + if (moe_comp_method == CPUQuantMethod::FP8_W8A16) { buffer_size_nbytes += M * topk * 2 * N * 2 + num_threads * MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N) * 2; } - + if (moe_comp_method == CPUQuantMethod::INT4_W4A8) { + buffer_size_nbytes += M * topk * 2 * N * 2 + std::max(M * K, M * topk * N) + M * topk * sizeof(float) + + num_threads * 2 * get_4bit_block_k_size(K / w1_scale.value().size(2)) * BLOCK_N; + } auto buffer2 = at::empty({buffer_size_nbytes}, hidden_states.options().dtype(at::kChar)); AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "fused_experts_kernel_impl", [&] { scalar_t* __restrict__ intermediate_cache1 = (scalar_t*)((void*)(buffer2.data_ptr())); scalar_t* __restrict__ intermediate_cache2 = intermediate_cache1 + M * topk * N; - if (use_int8_w8a8) { + if (moe_comp_method == CPUQuantMethod::INT8_W8A8) { uint8_t* __restrict__ A_tmp = (uint8_t*)((void*)(intermediate_cache2 + M * topk * K)); float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K)); uint8_t* __restrict__ Aq_tmp = (uint8_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N)); @@ -1109,7 +1122,7 @@ at::Tensor fused_experts_cpu( E, topk, num_tokens_post_pad); - } else if (use_fp8_w8a16) { + } else if (moe_comp_method == CPUQuantMethod::FP8_W8A16) { // here we just ignore C_tmp as it is not used scalar_t* __restrict__ A_tmp = (scalar_t*)((void*)(intermediate_cache2 + M * topk * K)); float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K)); @@ -1142,6 +1155,48 @@ at::Tensor fused_experts_cpu( E, topk, num_tokens_post_pad); + } else if (moe_comp_method == CPUQuantMethod::INT4_W4A8) { + uint8_t* __restrict__ A_tmp = (uint8_t*)((void*)(intermediate_cache2 + M * topk * K)); + float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K)); + scalar_t* __restrict__ intermediate_cache0 = (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N)); + uint8_t* __restrict__ Aq_tmp = (uint8_t*)((void*)(intermediate_cache0 + M * topk * 2 * N)); + float* __restrict__ As_tmp = (float*)((void*)(Aq_tmp + std::max(M * K, M * topk * N))); + int8_t* __restrict__ dqB_tmp = (int8_t*)((void*)(As_tmp + M * topk)); + + // weight + compensation shape = [Nc, Kc, block_n * block_k / 2 + block_n*sizeof(int32_t)] + // scales/qzeros shape = [E, Nc, G, block_n] + int64_t num_groups = w1_scale.value().size(2); + const int group_size = K / num_groups; + // TODO: check scales and zeros + fused_experts_int4_w4a8_kernel_impl( + out_hidden_states.data_ptr(), + intermediate_cache0, + intermediate_cache1, + intermediate_cache2, + A_tmp, + Aq_tmp, + As_tmp, + nullptr, + C_tmp, + dqB_tmp, + hidden_states.data_ptr(), + packed_w1.data_ptr(), + packed_w2.data_ptr(), + w1_zero.value().data_ptr(), + w2_zero.value().data_ptr(), + w1_scale.value().data_ptr(), + w2_scale.value().data_ptr(), + group_size, + topk_weights.data_ptr(), + sorted_ids, + expert_ids, + offsets, + M, + N, + K, + E, + topk, + num_tokens_post_pad); } else { scalar_t* __restrict__ A_tmp = intermediate_cache2 + M * topk * K; float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K)); @@ -1188,8 +1243,6 @@ at::Tensor shared_expert_cpu( const std::optional& w1_scale, const std::optional& w2_scale, const std::optional> block_size, - const std::optional& a1_scale, - const std::optional& a2_scale, bool is_vnni) { RECORD_FUNCTION("sgl-kernel::shared_expert_cpu", std::vector({hidden_states, w1, w2})); @@ -1224,7 +1277,7 @@ at::Tensor shared_expert_cpu( CHECK_EQ(packed_w2.size(1), packed_N); // check scales - check_moe_scales(use_int8_w8a8, use_fp8_w8a16, w1_scale, w2_scale, block_size, a1_scale, a2_scale); + check_moe_scales(use_int8_w8a8, use_fp8_w8a16, w1_scale, w2_scale, block_size); at::Tensor out_hidden_states = inplace ? hidden_states : at::empty_like(hidden_states); diff --git a/sgl-kernel/csrc/cpu/moe_int4.cpp b/sgl-kernel/csrc/cpu/moe_int4.cpp new file mode 100644 index 000000000..10956769b --- /dev/null +++ b/sgl-kernel/csrc/cpu/moe_int4.cpp @@ -0,0 +1,484 @@ +#include "common.h" +#include "gemm.h" +#include "vec.h" +namespace { + +template +inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) { + using Vec = at::vec::Vectorized; +// no remainder +#pragma GCC unroll 4 + for (int64_t d = 0; d < size; d += Vec::size()) { + Vec data = Vec::loadu(input + d); + data.store(out + d); + } +} + +template +inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int64_t size) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int kVecSize = bVec::size(); + int64_t d; +#pragma GCC unroll 4 + for (d = 0; d <= size - kVecSize; d += kVecSize) { + bVec x = bVec::loadu(input + d); + fVec x0, x1; + std::tie(x0, x1) = at::vec::convert_to_float(x); + bVec out_vec = convert_from_float_ext(x0, x1); + out_vec.store(out + d); + } + for (; d < size; ++d) { + out[d] = static_cast(input[d]); + } +} + +template +inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input, float weight, int64_t size) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int kVecSize = bVec::size(); + const fVec weight_vec = fVec(weight); + int64_t d; +#pragma GCC unroll 4 + for (d = 0; d <= size - kVecSize; d += kVecSize) { + fVec data0 = fVec::loadu(input + d) * weight_vec; + fVec data1 = fVec::loadu(input + d + fVec::size()) * weight_vec; + bVec out_vec = convert_from_float_ext(data0, data1); + out_vec.store(out + d); + } + for (; d < size; ++d) { + out[d] = static_cast(input[d] * weight); + } +} + +// acc from [topk, K] to [K] +template +inline void sum_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t topk, int64_t K) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int kVecSize = bVec::size(); + if (topk == 1) { + // do copy for topk = 1 + copy_stub(out, input, K); + } else { + // do sum for topk != 1 + int64_t d; +#pragma GCC unroll 4 + for (d = 0; d <= K - kVecSize; d += kVecSize) { + fVec sum_fvec0 = fVec(0.f); + fVec sum_fvec1 = fVec(0.f); + for (int t = 0; t < topk; ++t) { + bVec x_bvec = bVec::loadu(input + t * K + d); + fVec x_fvec0, x_fvec1; + std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec); + + sum_fvec0 += x_fvec0; + sum_fvec1 += x_fvec1; + } + bVec out_bvec = convert_from_float_ext(sum_fvec0, sum_fvec1); + out_bvec.store(out + d); + } + for (; d < K; ++d) { + float sum_val = 0.f; + for (int t = 0; t < topk; ++t) { + sum_val += static_cast(input[t * K + d]); + } + out[d] = static_cast(sum_val); + } + } +} + +// out = input + input2 * scale +template +inline void add_mul_stub( + scalar_t* __restrict__ out, + const scalar_t* __restrict__ input, + const scalar_t* __restrict__ input2, + float scale, + int64_t size) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int kVecSize = bVec::size(); + const fVec s_vec = fVec(scale); + + int64_t d; +#pragma GCC unroll 4 + for (d = 0; d <= size - kVecSize; d += kVecSize) { + bVec x_bvec = bVec::loadu(input + d); + fVec x0, x1; + std::tie(x0, x1) = at::vec::convert_to_float(x_bvec); + + bVec y_bvec = bVec::loadu(input2 + d); + fVec y0, y1; + std::tie(y0, y1) = at::vec::convert_to_float(y_bvec); + + x0 = x0 + y0 * s_vec; + x1 = x1 + y1 * s_vec; + bVec out_vec = convert_from_float_ext(x0, x1); + out_vec.store(out + d); + } + for (; d < size; ++d) { + out[d] = static_cast(input[d] + float(input2[d]) * scale); + } +} + +template +inline void silu_and_mul_stub( + scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const scalar_t* __restrict__ input2, int64_t size) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + const fVec one = fVec(1.f); + + // no remainder +#pragma GCC unroll 4 + for (int64_t d = 0; d < size; d += bVec::size()) { + bVec x = bVec::loadu(input + d); + fVec x0, x1; + std::tie(x0, x1) = at::vec::convert_to_float(x); + bVec y = bVec::loadu(input2 + d); + fVec y0, y1; + std::tie(y0, y1) = at::vec::convert_to_float(y); + x0 = x0 / (one + x0.neg().exp_u20()); + x1 = x1 / (one + x1.neg().exp_u20()); + x0 = x0 * y0; + x1 = x1 * y1; + bVec out_vec = convert_from_float_ext(x0, x1); + out_vec.store(out + d); + } +} + +} // anonymous namespace + +// TODO: stride access +template +inline void copy_bias(const float* bias_ptr, float* y_buf, int64_t m, int64_t ldn) { + if (bias_ptr) { + for (int i = 0; i < m; ++i) { + int j = 0; +#if defined(CPU_CAPABILITY_AVX512) +#pragma GCC unroll 2 + for (; j < N; j += 16) { + __m512 bias_vec = _mm512_loadu_ps(bias_ptr + j); + _mm512_storeu_ps(y_buf + i * ldn + j, bias_vec); + } +#endif + for (; j < N; ++j) { + y_buf[i * ldn + j] = bias_ptr[j]; + } + } + } else { // initialize to zero + for (int i = 0; i < m; ++i) { + int j = 0; +#if defined(CPU_CAPABILITY_AVX512) +#pragma GCC unroll 2 + for (; j < N; j += 16) { + __m512 zero_vec = _mm512_setzero_ps(); + _mm512_storeu_ps(y_buf + i * ldn + j, zero_vec); + } +#endif + for (; j < N; ++j) { + y_buf[i * ldn + j] = 0; + } + } + } +} + +template +void fused_experts_int4_w4a8_kernel_impl( + scalar_t* __restrict__ output, + scalar_t* __restrict__ ic0, + scalar_t* __restrict__ ic1, + scalar_t* __restrict__ ic2, + uint8_t* __restrict__ A_tmp, + uint8_t* __restrict__ Aq_tmp, + float* __restrict__ As_tmp, + int32_t* __restrict__ Azp_tmp, + float* __restrict__ C_tmp, + int8_t* __restrict__ dqB_tmp, + const scalar_t* __restrict__ input, + const uint8_t* __restrict__ packed_w1, + const uint8_t* __restrict__ packed_w2, + const int8_t* __restrict__ w1z, + const int8_t* __restrict__ w2z, + const float* __restrict__ w1s, + const float* __restrict__ w2s, + int group_size, + const float* __restrict__ topk_weights, + const int32_t* __restrict__ sorted_ids, + const int32_t* __restrict__ expert_ids, + const int32_t* __restrict__ offsets, + int64_t M, + int64_t N, + int64_t K, + int64_t E, + int64_t topk, + int64_t num_tokens_post_pad) { + constexpr int64_t BLOCK_M = block_size_m(); + constexpr int64_t BLOCK_N = block_size_n(); + int num_threads = at::get_num_threads(); + // int64_t buffer_size_nbytes = M * topk * N * 2 + // M * topk * K * 2 + + // num_threads * BLOCK_M * K + + // num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float) + + // M * topk * 2 * N * 2 + + // max(M * K, M * topk * N) + + // M * topk * sizeof(float); + + // intermediate_cache1 (scalar_t): START + M * topk * N + // intermediate_cache2 (scalar_t): + M * topk * K + // A_tmp (uint8_t): + num_threads * BLOCK_M * K + // C_tmp (float): + num_threads * 2 * BLOCK_M * BLOCK_N + // intermediate_cache0 (scalar_t): + M * topk * 2 * N + // Aq_tmp (uint8_t): + max(M * K, M * topk * N) + // As_tmp (float): + M * topk + // dqB_tmp (int8_t) + num_threads * _block_k * BlOCK_N + + // stage 0: quantize input to uint8, [M, K] + at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) { + for (int64_t m = begin; m < end; ++m) { + quantize_row_int8(Aq_tmp + m * K, As_tmp[m], input + m * K, K); + } + }); + int64_t _block_k = get_4bit_block_k_size(group_size); + auto Azp = at::ones({M * topk}).to(at::kInt).mul(128); + auto Azp_ptr = Azp.data_ptr(); + // stage 1: intermediate_cache0 = hidden_states @ w1 + const int64_t MB = div_up(num_tokens_post_pad, BLOCK_M); + const int64_t NB = div_up(N, BLOCK_N); + + int64_t block_per_group = group_size / _block_k; + int64_t Kc = K / _block_k; + int64_t num_groups = K / group_size; + + const int64_t stride_e = 2 * NB * Kc * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))); + const bool sym_quant_act = false; + // weight + compensation shape = [E, Nc, Kc, block_n * _block_k / 2 + block_n*sizeof(int32_t)] + // scales/qzeros shape = [E, Nc, G, block_n] + + // here we only parallel on half of 2N to fuse silu_and_mul with gemm + at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) { + // get local pointers + int tid = at::get_thread_num(); + int8_t* dqB_tmp1 = dqB_tmp + tid * 2 * _block_k * BLOCK_N; + int8_t* dqB_tmp2 = dqB_tmp1 + _block_k * BLOCK_N; + alignas(64) float As[BLOCK_M]; + uint8_t* __restrict__ A = A_tmp + tid * BLOCK_M * K; + float* __restrict__ C0 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N; + float* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N; + bool is_brgemm_used = false; + for (int64_t i = begin; i < end; ++i) { + int64_t mb = i / NB; + int64_t nb = i % NB; + int64_t nb1 = nb + NB; + int64_t n_size = std::min(N - nb * BLOCK_N, BLOCK_N); + // B shape [K, n_size] in vnni format + int32_t expert_id = expert_ids[mb]; + const uint8_t* __restrict__ B = packed_w1 + expert_id * stride_e; + // Bz and Bs: [E, K/gs, 2N] + const int8_t* __restrict__ Bz = w1z + expert_id * (num_groups) * (2 * N); + const float* __restrict__ Bs = w1s + expert_id * (num_groups) * (2 * N); + + // 1.a load A + const int32_t* A_ids = sorted_ids + mb * BLOCK_M; + int64_t m_size = offsets[mb + 1] - offsets[mb]; + const bool use_brgemm = can_use_brgemm(m_size); + is_brgemm_used = is_brgemm_used || use_brgemm; + // copy to A [BLOCK_M, K] + for (int64_t m = 0; m < m_size; ++m) { + int32_t index = A_ids[m] / topk; + copy_stub(A + m * K, Aq_tmp + index * K, K); + As[m] = As_tmp[index]; + } + const int64_t offset = offsets[mb]; + copy_bias(nullptr, C0, m_size, BLOCK_N); + copy_bias(nullptr, C1, m_size, BLOCK_N); + for (int kci = 0; kci < Kc; ++kci) { + int32_t* compensation_ptr = + sym_quant_act ? nullptr + : (int32_t*)(void*)(B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) + + _block_k * BLOCK_N / 2) /*Bcomp*/; + tinygemm_kernel( + ic0 + offset * 2 * N + nb * BLOCK_N, + C0, + A + kci * _block_k, + As, + Azp_ptr, + B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) /*B*/, + Bs + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*scales_b*/, + Bz + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*qzeros_b*/, + compensation_ptr, + dqB_tmp1, + m_size, + _block_k, + K, + BLOCK_N, + 2 * N, + kci == Kc - 1, + use_brgemm); + } + + for (int kci = 0; kci < Kc; ++kci) { + int32_t* compensation_ptr = + sym_quant_act ? nullptr + : (int32_t*)(void*)(B + (nb1 * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) + + _block_k * BLOCK_N / 2) /*Bcomp*/; + tinygemm_kernel( + ic0 + offset * 2 * N + nb1 * BLOCK_N, + C1, + A + kci * _block_k, + As, + Azp_ptr, + B + (nb1 * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) /*B*/, + Bs + nb1 * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*scales_b*/, + Bz + nb1 * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*qzeros_b*/, + compensation_ptr, + dqB_tmp2, + m_size, + _block_k, + K, + BLOCK_N, + 2 * N, + kci == Kc - 1, + use_brgemm); + } + } + + if (is_brgemm_used) { + at::native::cpublas::brgemm_release(); + } + }); + + // stage 1.5: intermediate_cache1 = silu(intermediate_cache0) + at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) { + for (int64_t m = begin; m < end; ++m) { + silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N); + } + }); + + // stage 1.5: quantize ic1 to uint8, [M * topk, N] + at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) { + for (int64_t m = begin; m < end; ++m) { + quantize_row_int8(Aq_tmp + m * N, As_tmp[m], ic1 + m * N, N); + } + }); + // stage 2: intermediate_cache2 = intermediate_cache1 @ w2 + // w2 : [E, K, N] as [E, OC, IC] + const int64_t OC = K; // rename K as OC + const int64_t IC = N; // rename N as IC + const int64_t MB2 = MB; + const int64_t NB2 = div_up(OC, BLOCK_N); + const int64_t stride_oc = IC; + num_groups = IC / group_size; + Kc = IC / _block_k; + const int64_t stride_e2 = NB2 * Kc * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))); + // parallel on [MB2, NB2] + at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) { + int tid = at::get_thread_num(); + int8_t* dqB_tmp1 = dqB_tmp + tid * 2 * _block_k * BLOCK_N; + float* __restrict__ C2 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N; + bool is_brgemm_used = false; + for (int64_t i = begin; i < end; ++i) { + int64_t mb = i / NB2; + int64_t nb = i % NB2; + + int64_t m_size = offsets[mb + 1] - offsets[mb]; + int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N); + const bool use_brgemm = can_use_brgemm(m_size); + is_brgemm_used = is_brgemm_used || use_brgemm; + const int32_t* A_ids = sorted_ids + mb * BLOCK_M; + + // B shape [IC, n_size] in vnni format + int32_t expert_id = expert_ids[mb]; + const uint8_t* __restrict__ B = packed_w2 + expert_id * stride_e2; + + // Bz and Bs: [E, IC/gs, OC] + const int8_t* __restrict__ Bz = w2z + expert_id * (num_groups)*OC; + const float* __restrict__ Bs = w2s + expert_id * (num_groups)*OC; + + // A ptr from ic1 of [M * topk, N] in sorted order + // so as to avoid copy A to tmp buffer again + const uint8_t* __restrict__ A = Aq_tmp + offsets[mb] * IC; + const float* __restrict__ As = As_tmp + offsets[mb]; + copy_bias(nullptr, C2, m_size, BLOCK_N); + for (int kci = 0; kci < Kc; ++kci) { + int32_t* compensation_ptr = + sym_quant_act ? nullptr + : (int32_t*)(void*)(B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) + + _block_k * BLOCK_N / 2) /*Bcomp*/; + tinygemm_kernel( + nullptr, /*store_out is false*/ + C2, + A + kci * _block_k, + As, + Azp_ptr, + B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))), + Bs + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*scales_b*/, + Bz + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*zeros_b*/, + compensation_ptr, + dqB_tmp1, + m_size, + _block_k, + IC, + BLOCK_N, + BLOCK_N, + false, + use_brgemm); + } + + // 2.b copy from C to ic2 in original order + // and also mul topk_weights in float32 + for (int64_t m = 0; m < m_size; ++m) { + int32_t index = A_ids[m]; + float weight = topk_weights[index]; + copy_mul_stub(ic2 + index * K + nb * BLOCK_N, C2 + m * BLOCK_N, weight, n_size); + } + } + + if (is_brgemm_used) { + at::native::cpublas::brgemm_release(); + } + }); + + // stage 3: out = intermediate_cache2.sum(dim=1) + // from [M, topk, K] to [M, K] + at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) { + for (int64_t m = begin; m < end; ++m) { + sum_stub(output + m * K, ic2 + m * topk * K, topk, K); + } + }); +} + +#define INSTANTIATE_MOE_INT4_W4A8_TEMPLATE(TYPE) \ + template void fused_experts_int4_w4a8_kernel_impl( \ + TYPE* __restrict__ output, \ + TYPE* __restrict__ ic0, \ + TYPE* __restrict__ ic1, \ + TYPE* __restrict__ ic2, \ + uint8_t* __restrict__ A_tmp, \ + uint8_t* __restrict__ Aq_tmp, \ + float* __restrict__ As_tmp, \ + int32_t* __restrict__ Azp_tmp, \ + float* __restrict__ C_tmp, \ + int8_t* __restrict__ dqB_tmp, \ + const TYPE* __restrict__ input, \ + const uint8_t* __restrict__ packed_w1, \ + const uint8_t* __restrict__ packed_w2, \ + const int8_t* __restrict__ w1z, \ + const int8_t* __restrict__ w2z, \ + const float* __restrict__ w1s, \ + const float* __restrict__ w2s, \ + int group_size, \ + const float* __restrict__ topk_weights, \ + const int32_t* __restrict__ sorted_ids, \ + const int32_t* __restrict__ expert_ids, \ + const int32_t* __restrict__ offsets, \ + int64_t M, \ + int64_t N, \ + int64_t K, \ + int64_t E, \ + int64_t topk, \ + int64_t num_tokens_post_pad) + +INSTANTIATE_MOE_INT4_W4A8_TEMPLATE(at::BFloat16); +INSTANTIATE_MOE_INT4_W4A8_TEMPLATE(at::Half); diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp index 8c66e1d2b..da7d09cc3 100644 --- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp +++ b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp @@ -181,6 +181,14 @@ at::Tensor int8_scaled_mm_with_quant( at::ScalarType out_dtype, bool is_vnni); +// int4 gemm +at::Tensor int4_scaled_mm_cpu( + at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros, at::Tensor& w_scales, std::optional bias); + +// weight prepack for int4 weights +std::tuple +convert_weight_packed_scale_zp(at::Tensor qweight, at::Tensor qzeros, at::Tensor scales); + // bmm void bmm_cpu(at::Tensor& out, at::Tensor& mat1, at::Tensor& mat2, bool is_vnni, const std::optional& scale); @@ -192,13 +200,12 @@ at::Tensor fused_experts_cpu( at::Tensor& topk_weights, at::Tensor& topk_ids, bool inplace, - bool use_int8_w8a8, - bool use_fp8_w8a16, + int64_t moe_comp_method, const std::optional& w1_scale, const std::optional& w2_scale, + const std::optional& w1_zero, + const std::optional& w2_zero, const std::optional> block_size, - const std::optional& a1_scale, - const std::optional& a2_scale, bool is_vnni); at::Tensor shared_expert_cpu( @@ -213,8 +220,6 @@ at::Tensor shared_expert_cpu( const std::optional& w1_scale, const std::optional& w2_scale, const std::optional> block_size, - const std::optional& a1_scale, - const std::optional& a2_scale, bool is_vnni); // weight absorption @@ -441,6 +446,16 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "is_vnni) -> Tensor"); m.impl("int8_scaled_mm_with_quant", torch::kCPU, &int8_scaled_mm_with_quant); + // int4 gemm + m.def("int4_scaled_mm_cpu(Tensor x, Tensor w, Tensor w_zeros, Tensor w_scales, Tensor? bias) -> Tensor"); + m.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu); + + // weight prepack for int4 weights + m.def( + "convert_weight_packed_scale_zp(Tensor weight, Tensor qzeros, Tensor scales) -> (Tensor, Tensor, " + "Tensor)"); + m.impl("convert_weight_packed_scale_zp", torch::kCPU, &convert_weight_packed_scale_zp); + // bmm m.def("bmm_cpu(Tensor(a!) out, Tensor mat1, Tensor mat2, bool is_vnni, Tensor? scale) -> ()"); m.impl("bmm_cpu", torch::kCPU, &bmm_cpu); @@ -448,9 +463,8 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { // moe m.def( "fused_experts_cpu(Tensor hidden_states, Tensor w1, Tensor w2, Tensor topk_weights, Tensor topk_ids, bool " - "inplace, bool use_int8_w8a8, bool use_fp8_w8a16, Tensor? w1_scale, Tensor? w2_scale, int[]? block_size, Tensor? " - "a1_scale, Tensor? a2_scale, bool " - "is_vnni) -> Tensor"); + "inplace, int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, " + "Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, bool is_vnni) -> Tensor"); m.impl("fused_experts_cpu", torch::kCPU, &fused_experts_cpu); // weight absorption @@ -474,7 +488,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { m.def( "shared_expert_cpu(Tensor hidden_states, Tensor w1, Tensor w2, Tensor fused_experts_out, float " "routed_scaling_factor, bool inplace, bool use_int8_w8a8, bool use_fp8_w8a16, Tensor? w1_scale, Tensor? " - "w2_scale, int[]? block_size, Tensor? a1_scale, Tensor? a2_scale, bool is_vnni) -> Tensor"); + "w2_scale, int[]? block_size, bool is_vnni) -> Tensor"); m.impl("shared_expert_cpu", torch::kCPU, &shared_expert_cpu); // causal conv1d diff --git a/test/srt/cpu/test_gemm.py b/test/srt/cpu/test_gemm.py index 0d3113c6e..1019a06f2 100644 --- a/test/srt/cpu/test_gemm.py +++ b/test/srt/cpu/test_gemm.py @@ -9,6 +9,7 @@ from utils import ( native_w8a8_per_token_matmul, per_token_quant_int8, precision, + unpack_and_dequant_awq, ) from sglang.test.test_utils import CustomTestCase @@ -39,6 +40,10 @@ class TestGemm(CustomTestCase): N_fp8 = [128, 224] K_fp8 = [512, 576] + M_awq = [1, 32] + N_awq = [4096] + K_awq = [4096] + def _bf16_gemm(self, M, N, K, has_bias): mat1 = torch.randn(M, K, dtype=torch.bfloat16) @@ -227,6 +232,51 @@ class TestGemm(CustomTestCase): ): self._fp8_gemm(*params) + def _int4_awq_gemm(self, M, N, K, group_size, has_bias): + awq_weight = torch.randint(-128, 128, (K, N // 8)).to(torch.int) + awq_zero = torch.randint(0, 10, (K // group_size, N // 8)).to(torch.int) + awq_scales = torch.rand(int(K // group_size), N).to(torch.bfloat16) + bf16_weight, _ = unpack_and_dequant_awq( + awq_weight, awq_zero, awq_scales, 4, 128 + ) + if has_bias: + bias = torch.rand(bf16_weight.shape[0]).to(torch.float) + else: + bias = None + x = torch.rand(M, bf16_weight.size(-1)).to(torch.bfloat16) + ref_res = torch.nn.functional.linear( + x, bf16_weight, bias=bias.to(torch.bfloat16) if has_bias else None + ) + + packed_weight, packed_zero, packed_scales = ( + torch.ops.sgl_kernel.convert_weight_packed_scale_zp( + awq_weight, awq_zero, awq_scales + ) + ) + target_res = torch.ops.sgl_kernel.int4_scaled_mm_cpu( + x, + packed_weight, + packed_zero, + packed_scales, + bias, + ) + + atol = rtol = precision[ref_res.dtype] + torch.testing.assert_close(ref_res, target_res, atol=atol, rtol=rtol) + + def test_int4_awq_gemm(self): + for params in itertools.product( + self.M_awq, self.N_awq, self.K_awq, [128], self.has_bias + ): + with self.subTest( + M=params[0], + N=params[1], + K=params[2], + group_size=params[3], + has_bias=params[4], + ): + self._int4_awq_gemm(*params) + if __name__ == "__main__": unittest.main() diff --git a/test/srt/cpu/test_moe.py b/test/srt/cpu/test_moe.py index 7babd5167..a1f6bf288 100644 --- a/test/srt/cpu/test_moe.py +++ b/test/srt/cpu/test_moe.py @@ -5,9 +5,11 @@ import unittest # TODO: use interface in cpu.py import torch +from sglang.srt.layers.amx_utils import CPUQuantMethod + kernel = torch.ops.sgl_kernel -torch.manual_seed(1234) +torch.manual_seed(128) from utils import ( BLOCK_K, @@ -20,6 +22,7 @@ from utils import ( scaled_weight, torch_naive_fused_moe, torch_w8a8_per_column_fused_moe, + unpack_and_dequant_awq, ) from sglang.test.test_utils import CustomTestCase @@ -48,8 +51,7 @@ def fused_moe(a, w1, w2, score, topk, renormalize, prepack): topk_weights, topk_ids, inplace, - False, - False, + CPUQuantMethod.UNQUANT, None, None, None, @@ -79,6 +81,12 @@ class TestFusedExperts(CustomTestCase): E_fp8 = [8] topk_fp8 = [4] + M_int4 = [1, 6] + N_int4 = [512] + K_int4 = [256] + E_int4 = [8] + topk_int4 = [4] + def _bf16_moe(self, m, n, k, e, topk, renormalize): dtype = torch.bfloat16 prepack = True @@ -156,8 +164,7 @@ class TestFusedExperts(CustomTestCase): topk_weight, topk_ids.to(torch.int32), inplace, - True, - False, + CPUQuantMethod.INT8_W8A8, w1_s, w2_s, None, @@ -229,13 +236,12 @@ class TestFusedExperts(CustomTestCase): topk_weight, topk_ids.to(torch.int32), False, - False, - True, + CPUQuantMethod.FP8_W8A16, w1s, w2s, + None, + None, [BLOCK_N, BLOCK_K], - None, - None, True, ) @@ -259,6 +265,107 @@ class TestFusedExperts(CustomTestCase): ): self._fp8_moe(*params) + def _int4_moe(self, M, N, K, E, topk, group_size=128): + dtype = torch.bfloat16 + + a = torch.rand(M, K, dtype=dtype) / math.sqrt(K) + + awq_w13_weight = torch.randint(-127, 128, (E, K, 2 * N // 8)).to(torch.int) + awq_w13_zero = torch.randint(0, 10, (E, K // group_size, 2 * N // 8)).to( + torch.int + ) + awq_w13_scales = torch.rand(E, int(K // group_size), 2 * N).to(torch.bfloat16) + + awq_w2_weight = torch.randint(-127, 128, (E, N, K // 8)).to(torch.int) + awq_w2_zero = torch.randint(0, 10, (E, N // group_size, K // 8)).to(torch.int) + awq_w2_scales = torch.rand(E, int(N // group_size), K).to(torch.bfloat16) + bf16_w13_weight = [] + bf16_w2_weight = [] + for i in range(E): + bf16_w13_weight_i, _ = unpack_and_dequant_awq( + awq_w13_weight[i], awq_w13_zero[i], awq_w13_scales[i], 4, 128 + ) + bf16_w2_weight_i, _ = unpack_and_dequant_awq( + awq_w2_weight[i], awq_w2_zero[i], awq_w2_scales[i], 4, 128 + ) + bf16_w13_weight.append(bf16_w13_weight_i) + bf16_w2_weight.append(bf16_w2_weight_i) + bf16_w13_weight = torch.stack(bf16_w13_weight).detach() + bf16_w2_weight = torch.stack(bf16_w2_weight).detach() + + score = torch.rand((M, E), dtype=dtype) + + ref_out = torch_naive_fused_moe( + a, bf16_w13_weight, bf16_w2_weight, score, topk, False + ) + score = torch.softmax(score, dim=-1, dtype=torch.float32) + topk_weight, topk_ids = torch.topk(score, topk) + awq_w13_weight_pack = [] + awq_w13_zero_pack = [] + awq_w13_scales_pack = [] + awq_w2_weight_pack = [] + awq_w2_zero_pack = [] + awq_w2_scales_pack = [] + for i in range(E): + packed_weight_13_i, packed_zero_13_i, packed_scales_13_i = ( + torch.ops.sgl_kernel.convert_weight_packed_scale_zp( + awq_w13_weight[i], awq_w13_zero[i], awq_w13_scales[i] + ) + ) + awq_w13_weight_pack.append(packed_weight_13_i) + awq_w13_zero_pack.append(packed_zero_13_i) + awq_w13_scales_pack.append(packed_scales_13_i) + packed_weight_2_i, packed_zero_2_i, packed_scales_2_i = ( + torch.ops.sgl_kernel.convert_weight_packed_scale_zp( + awq_w2_weight[i], awq_w2_zero[i], awq_w2_scales[i] + ) + ) + awq_w2_weight_pack.append(packed_weight_2_i) + awq_w2_zero_pack.append(packed_zero_2_i) + awq_w2_scales_pack.append(packed_scales_2_i) + awq_w13_weight_pack = torch.stack(awq_w13_weight_pack).detach() + awq_w13_zero_pack = torch.stack(awq_w13_zero_pack).detach() + awq_w13_scales_pack = torch.stack(awq_w13_scales_pack).detach() + awq_w2_weight_pack = torch.stack(awq_w2_weight_pack).detach() + awq_w2_zero_pack = torch.stack(awq_w2_zero_pack).detach() + awq_w2_scales_pack = torch.stack(awq_w2_scales_pack).detach() + + out = kernel.fused_experts_cpu( + a, + awq_w13_weight_pack, + awq_w2_weight_pack, + topk_weight, + topk_ids.to(torch.int32), + False, + CPUQuantMethod.INT4_W4A8, + awq_w13_scales_pack, + awq_w2_scales_pack, + awq_w13_zero_pack, + awq_w2_zero_pack, + None, + True, + ) + + atol = rtol = precision[dtype] + torch.testing.assert_close(ref_out.bfloat16(), out, atol=atol, rtol=rtol) + + def test_int4_moe(self): + for params in itertools.product( + self.M_int4, + self.N_int4, + self.K_int4, + self.E_int4, + self.topk_int4, + ): + with self.subTest( + M=params[0], + N=params[1], + K=params[2], + E=params[3], + topk=params[4], + ): + self._int4_moe(*params) + if __name__ == "__main__": unittest.main() diff --git a/test/srt/cpu/test_shared_expert.py b/test/srt/cpu/test_shared_expert.py index 358709a6a..f5a8f3632 100644 --- a/test/srt/cpu/test_shared_expert.py +++ b/test/srt/cpu/test_shared_expert.py @@ -66,8 +66,6 @@ class TestSharedExpert(CustomTestCase): None, None, None, - None, - None, False, ) @@ -124,8 +122,6 @@ class TestSharedExpert(CustomTestCase): w1_s, w2_s, None, - None, - None, False, ) @@ -196,8 +192,6 @@ class TestSharedExpert(CustomTestCase): w1s, w2s, [BLOCK_N, BLOCK_K], - None, - None, True, ) diff --git a/test/srt/cpu/utils.py b/test/srt/cpu/utils.py index 8f03c1bc9..f90967ed7 100644 --- a/test/srt/cpu/utils.py +++ b/test/srt/cpu/utils.py @@ -286,3 +286,105 @@ def make_non_contiguous(x: torch.Tensor) -> torch.Tensor: """ last_dim = x.shape[-1] return x[..., : last_dim // 2] if x.is_contiguous() else x + + +def awq_reverse_reorder_int_tensor(int_tensor, bits: int): + assert bits == 4 + + int_tensor = int_tensor.T.contiguous() + compress_ratio = 32 // bits + assert int_tensor.shape[-1] % compress_ratio == 0 + + order_map = [0, 2, 4, 6, 1, 3, 5, 7] + order_tensor = torch.tensor( + order_map, dtype=torch.int32, device=int_tensor.device + ).reshape(1, -1) + order_tensor = order_tensor.repeat(int_tensor.shape[1] // compress_ratio, 1) + order_tensor = order_tensor + torch.arange( + 0, + int_tensor.shape[1], + compress_ratio, + dtype=torch.int32, + device=int_tensor.device, + ).reshape(-1, 1) + order_tensor = order_tensor.reshape(-1) + + reverse_order_tensor = torch.arange(order_tensor.shape[0])[order_tensor] + reverse_order_tensor = reverse_order_tensor[order_tensor] + int_tensor = int_tensor[:, reverse_order_tensor] + return int_tensor + + +def unpack_and_dequant_awq( + awq_qweight: torch.Tensor, + awq_qzeros: torch.Tensor, + awq_scales: torch.Tensor, + bits: int, + group_size: int, +): + """ + Args: + awq_qweight (`torch.LongTensor`): + Expected shape: (in_features, out_features // (32 // bits)) + awq_qzeros (`torch.LongTensor`): + Expected shape: (in_features // group_size, out_features // (32 // bits)) + awq_scales (`torch.LongTensor`): + Expected shape: (in_features // group_size, out_features) + + Returns: + fp16_weight (`torch.LongTensor`): + With shape (in_features, out_features). + zeros (`torch.LongTensor`): + With shape (in_features // group_size, out_features). + """ + assert bits == 4 + + qzeros = awq_qzeros + qweight = awq_qweight + qweight = qweight.T.contiguous() + + scales = awq_scales + scales = scales.reshape(-1, 1, scales.shape[-1]) + + infeatures = awq_qweight.shape[0] + + wf = torch.tensor( + list(range(0, 32, bits)), dtype=torch.int32, device=qzeros.device + ).unsqueeze(0) + zeros = torch.bitwise_right_shift(torch.unsqueeze(qzeros, 2), wf.unsqueeze(0)).to( + torch.int16 if bits == 8 else torch.int8 + ) + + torch.bitwise_and(zeros, (2**bits) - 1, out=zeros) + + zeros = zeros.reshape(-1, 1, zeros.shape[1] * zeros.shape[2]) + + weight = torch.bitwise_right_shift( + torch.unsqueeze(qweight, 1), wf.unsqueeze(-1) + ).to(torch.int16 if bits == 8 else torch.int8) + torch.bitwise_and(weight, (2**bits) - 1, out=weight) + weight = weight.reshape(-1, group_size, weight.shape[2]) + + weight = weight.view(-1, weight.shape[-1]) + zeros = zeros.view(-1, zeros.shape[-1]) + + zeros = zeros.T.contiguous() + zeros = awq_reverse_reorder_int_tensor(zeros, bits) + weight = awq_reverse_reorder_int_tensor(weight, bits) + + # Dequantize weights. + scales = awq_scales + zeros = zeros.contiguous() + scale_zeros = zeros * scales + + g_idx = torch.tensor( + [i // group_size for i in range(infeatures)], dtype=torch.int32 + ) + scale_mat = scales[g_idx] + scale_zeros_mat = scale_zeros[g_idx].to(torch.bfloat16) + + qdq_weight_T = weight * scale_mat - scale_zeros_mat.to(torch.bfloat16) + + fp16_weight = qdq_weight_T.T + + return fp16_weight, zeros