diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index e3e9b750303..b39112d24c6 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -438,6 +438,12 @@ if(USE_ONEDNN) ${VLLM_EXT_SRC}) endif() +if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") + set(VLLM_EXT_SRC + "csrc/cpu/sgl-kernels/gemm_int4.cpp" + ${VLLM_EXT_SRC}) +endif() + if (ENABLE_X86_ISA) set(VLLM_EXT_SRC_SGL "csrc/cpu/sgl-kernels/conv.cpp" diff --git a/csrc/cpu/sgl-kernels/gemm_int4.cpp b/csrc/cpu/sgl-kernels/gemm_int4.cpp index 5b66b2a5aee..1fec14c956f 100644 --- a/csrc/cpu/sgl-kernels/gemm_int4.cpp +++ b/csrc/cpu/sgl-kernels/gemm_int4.cpp @@ -268,6 +268,23 @@ void _dequant_gemm_accum_small_M( _dequant_gemm_accum_small_M(C, A, scales_a, qzeros_a, B, scales_b, qzeros_b, K, lda, ldc); #endif +template +inline int32_t load_uint4_vnni(const uint8_t* __restrict__ B, int64_t k, int64_t n) { + // B is packed as [_block_k / 4, N / 2, 4] for VNNI4. Each byte stores two + // columns from adjacent 8-column groups for one K lane. + constexpr int64_t n_group_size = 8; + constexpr int64_t vnni_size = 4; + static_assert(N % (2 * n_group_size) == 0); + + int64_t n_group = n / n_group_size; + int64_t ni = n % n_group_size; + int64_t ki = k % vnni_size; + int64_t k_base = k - ki; + int64_t packed_n = (n_group / 2) * n_group_size + ni; + uint8_t packed = B[k_base * ldb + packed_n * vnni_size + ki]; + return (n_group % 2 == 0) ? (packed & 0x0f) : ((packed >> 4) & 0x0f); +} + template void _dequant_gemm_accum( float* C, @@ -321,7 +338,24 @@ void _dequant_gemm_accum( } else #endif { - TORCH_CHECK(false, "tinygemm_kernel: scalar path not implemented!"); + for (int64_t m = 0; m < M; ++m) { + for (int64_t n = 0; n < N; ++n) { + int32_t acc = 0; + for (int64_t k = 0; k < K; ++k) { + int32_t b = load_uint4_vnni(B, k, n) - qzeros_b[n]; + if constexpr (sym_quant_act) { + const int8_t* A_s8 = reinterpret_cast(A); + acc += static_cast(A_s8[m * lda + k]) * b; + } else { + acc += static_cast(A[m * lda + k]) * b; + } + } + if constexpr (!sym_quant_act) { + acc -= qzeros_a[m] * compensation[n]; + } + C[m * ldc + n] += static_cast(acc) * scales_a[m] * scales_b[n]; + } + } } } @@ -496,9 +530,11 @@ void _da8w4_linear_impl( store_out(C_tmp, output + mci * block_m * N + nc * BLOCK_N, m_size, N /*lda*/); } } +#if defined(CPU_CAPABILITY_AVX512) if (use_brgemm) { at::native::cpublas::brgemm_release(); } +#endif }); } diff --git a/csrc/cpu/sgl-kernels/vec.h b/csrc/cpu/sgl-kernels/vec.h index 77ffeec9fe7..72143fedc69 100644 --- a/csrc/cpu/sgl-kernels/vec.h +++ b/csrc/cpu/sgl-kernels/vec.h @@ -245,7 +245,7 @@ quantize_row_int8(uint8_t* __restrict__ Aq, float& As, const scalar_t* __restric for (int64_t k = 0; k < K; ++k) { const float val = static_cast(A[k]) * inv_scale; - Aq[k] = (uint8_t)(std::round(val)) + 128; + Aq[k] = static_cast(static_cast(std::round(val)) + 128); } As = scale; } diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index c5ce7c46bb9..495185769ba 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -429,19 +429,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.impl("int8_scaled_mm_with_quant", torch::kCPU, &int8_scaled_mm_with_quant); - // Adapted from sglang: INT4 W4A8 kernels - ops.def( - "convert_weight_packed_scale_zp(Tensor weight, Tensor qzeros, Tensor " - "scales, int quant_method_4bit) -> (Tensor, " - "Tensor, Tensor)"); - ops.impl("convert_weight_packed_scale_zp", torch::kCPU, - &convert_weight_packed_scale_zp); - - ops.def( - "int4_scaled_mm_cpu(Tensor(a0!) x, Tensor(a1!) w, Tensor(a2!) w_zeros, " - "Tensor(a3!) w_scales, Tensor? bias) -> Tensor"); - ops.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu); - // Adapted from sglang: FP8 W8A16 kernel ops.def( "fp8_scaled_mm_cpu(Tensor(a0!) mat1, Tensor(a1!) mat2, Tensor(a2!) " @@ -468,6 +455,23 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu); #endif +#if (defined(__AVX512BF16__) && defined(__AVX512F__) && \ + defined(__AVX512VNNI__)) || \ + defined(__riscv) + // Adapted from sglang: INT4 W4A8 kernels + ops.def( + "convert_weight_packed_scale_zp(Tensor weight, Tensor qzeros, Tensor " + "scales, int quant_method_4bit) -> (Tensor, " + "Tensor, Tensor)"); + ops.impl("convert_weight_packed_scale_zp", torch::kCPU, + &convert_weight_packed_scale_zp); + + ops.def( + "int4_scaled_mm_cpu(Tensor(a0!) x, Tensor(a1!) w, Tensor(a2!) w_zeros, " + "Tensor(a3!) w_scales, Tensor? bias) -> Tensor"); + ops.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu); +#endif + // Adapted from sglang: GDN kernels ops.def( "chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, "