A high-throughput and memory-efficient inference and serving engine for LLMs
[CPU][RISC-V] Enable oneDNN W8A8 INT8 to run on RISC-V (#44478)
Signed-off-by: velonica0 <like@mail.nankai.edu.cn>
velonica0 committed
Jun 11, 2026 at 12:09 UTC
f31bc2ea60f685a65885fc3c4c7753e7f1ca5a61
4 files changed
+42
-5
cmake/cpu_extension.cmake
+6
-3
@@ -166,6 +166,10 @@ elseif (S390_FOUND)
166
"-mtune=native")
167
elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
168
message(STATUS "RISC-V detected")
169
+ if(DEFINED VLLM_RVV_VLEN AND NOT VLLM_RVV_VLEN GREATER 0)
170
+ message(FATAL_ERROR
171
+ "VLLM_RVV_VLEN must be a positive integer; got '${VLLM_RVV_VLEN}'")
172
+ endif()
173
# VLLM_RVV_VLEN selects the target VLEN. Auto-detected from /proc/cpuinfo
174
# by default; override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256.
175
if(NOT DEFINED VLLM_RVV_VLEN)
@@ -189,8 +193,7 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
193
"RISC-V RVV is available but VLEN could not be auto-detected. "
194
"Please specify VLEN explicitly:\n"
195
" -DVLLM_RVV_VLEN=128 (for VLEN=128 hardware)\n"
192
- " -DVLLM_RVV_VLEN=256 (for VLEN=256 hardware, e.g. Spacemit X100)\n"
193
- " -DVLLM_RVV_VLEN=0 (force scalar, no RVV)")
196
+ " -DVLLM_RVV_VLEN=256 (for VLEN=256 hardware, e.g. Spacemit X100)")
197
endif()
198
endif()
199
if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0)
@@ -219,7 +222,7 @@ endif()
222
223
224
# Build oneDNN for GEMM kernels
222
-if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
225
+if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND OR RVV_FP16_FOUND OR RVV_BF16_FOUND)
226
# Fetch and build Arm Compute Library (ACL) as oneDNN's backend for AArch64
227
# TODO [fadara01]: remove this once ACL can be fetched and built automatically as a dependency of oneDNN
228
set(ONEDNN_AARCH64_USE_ACL OFF CACHE BOOL "")
csrc/cpu/cpu_types_riscv_defs.hpp
+4
@@ -57,6 +57,10 @@ typedef RVVTYPE(vfloat32, LMUL_512, _t) fixed_fp32x16_t
57
typedef RVVTYPE(vfloat32, LMUL_1024, _t) fixed_fp32x32_t
58
__attribute__((riscv_rvv_vector_bits(1024)));
59
60
+// int8
61
+typedef RVVTYPE(vint8, LMUL_128, _t) fixed_i8x16_t
62
+ __attribute__((riscv_rvv_vector_bits(128)));
63
+
64
// int32
65
typedef RVVTYPE(vint32, LMUL_256, _t) fixed_i32x8_t
66
__attribute__((riscv_rvv_vector_bits(256)));
csrc/cpu/cpu_types_riscv_impl.hpp
+29
@@ -734,10 +734,18 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
734
return FP32Vec16(
735
RVVI(__riscv_vfmax_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
736
}
737
+ FP32Vec16 max(const FP32Vec16& b, const int elem_num) const {
738
+ return FP32Vec16(
739
+ RVVI(__riscv_vfmax_vv_f32, LMUL_512)(reg, b.reg, elem_num));
740
+ }
741
FP32Vec16 min(const FP32Vec16& b) const {
742
return FP32Vec16(
743
RVVI(__riscv_vfmin_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
744
}
745
+ FP32Vec16 min(const FP32Vec16& b, const int elem_num) const {
746
+ return FP32Vec16(
747
+ RVVI(__riscv_vfmin_vv_f32, LMUL_512)(reg, b.reg, elem_num));
748
+ }
749
FP32Vec16 abs() const {
750
return FP32Vec16(RVVI(__riscv_vfabs_v_f32, LMUL_512)(reg, VEC_ELEM_NUM));
751
}
@@ -867,6 +875,27 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
875
}
876
};
877
878
+struct INT8Vec16 : public Vec<INT8Vec16> {
879
+ constexpr static int VEC_ELEM_NUM = 16;
880
+ fixed_i8x16_t reg;
881
+
882
+ explicit INT8Vec16(const FP32Vec16& vec) {
883
+ auto i32_vec =
884
+ RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_512)(vec.reg, VEC_ELEM_NUM);
885
+ auto i16_vec = RVVI(__riscv_vnclip_wx_i16, LMUL_256)(
886
+ i32_vec, 0, __RISCV_VXRM_RNU, VEC_ELEM_NUM);
887
+ reg = RVVI(__riscv_vnclip_wx_i8, LMUL_128)(i16_vec, 0, __RISCV_VXRM_RNU,
888
+ VEC_ELEM_NUM);
889
+ }
890
+
891
+ void save(int8_t* ptr) const {
892
+ RVVI(__riscv_vse8_v_i8, LMUL_128)(ptr, reg, VEC_ELEM_NUM);
893
+ }
894
+ void save(int8_t* ptr, int elem_num) const {
895
+ RVVI(__riscv_vse8_v_i8, LMUL_128)(ptr, reg, elem_num);
896
+ }
897
+};
898
+
899
// ============================================================================
900
// Type Traits & Global Helpers
901
// ============================================================================
csrc/cpu/torch_bindings.cpp
+3
-2
@@ -329,8 +329,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
329
ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);
330
331
// Quantization
332
-#if defined(__AVX512F__) || defined(__AVX2__) || \
333
- (defined(__aarch64__) && !defined(__APPLE__)) || defined(__powerpc64__)
332
+#if defined(__AVX512F__) || defined(__AVX2__) || \
333
+ (defined(__aarch64__) && !defined(__APPLE__)) || defined(__powerpc64__) || \
334
+ defined(__riscv_v)
335
// Helper function to release oneDNN handlers
336
ops.def("release_dnnl_matmul_handler(int handler) -> ()",
337
&release_dnnl_matmul_handler);