A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 08:30 UTC
[CPU][RISC-V] Add missing RVV cpu_types helpers for WNA16 (#42730)
Signed-off-by: wcy <233313160abc@gmail.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
wcy committed
Jun 1, 2026 at 14:56 UTC
98f1279815db9201f3ce437843c002f0b5ec5733
3 files changed
+59
-1
cmake/cpu_extension.cmake
+7
@@ -396,6 +396,13 @@ set(VLLM_EXT_SRC
396
"csrc/cpu/cpu_attn.cpp"
397
"csrc/cpu/torch_bindings.cpp")
398
399
+if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64" AND VLLM_RVV_VLEN AND
400
+ VLLM_RVV_VLEN GREATER 0 AND (RVV_FP16_FOUND OR RVV_BF16_FOUND))
401
+ set(VLLM_EXT_SRC
402
+ "csrc/cpu/cpu_wna16.cpp"
403
+ ${VLLM_EXT_SRC})
404
+endif()
405
+
406
if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
407
set(VLLM_EXT_SRC
408
"csrc/cpu/shm.cpp"
csrc/cpu/cpu_types_riscv_impl.hpp
+51
@@ -94,6 +94,10 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
94
: reg(RVVI(__riscv_vle16_v_f16, LMUL_256)(
95
static_cast<const _Float16*>(ptr), VEC_ELEM_NUM)) {};
96
97
+ explicit FP16Vec16(const c10::Half v)
98
+ : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _f16, LMUL_256)(
99
+ RVVI(__riscv_vmv_v_x_u16, LMUL_256)(v.x, VEC_ELEM_NUM))) {};
100
+
101
explicit FP16Vec16(const FP32Vec16& vec);
102
103
void save(void* ptr) const {
@@ -165,6 +169,9 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
169
reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
170
171
explicit BF16Vec16(fixed_bf16x16_t data) : reg(data) {};
172
+ explicit BF16Vec16(const c10::BFloat16 v)
173
+ : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _bf16, LMUL_256)(
174
+ RVVI(__riscv_vmv_v_x_u16, LMUL_256)(v.x, VEC_ELEM_NUM))) {};
175
explicit BF16Vec16(const FP32Vec16&);
176
177
void save(void* ptr) const {
@@ -290,6 +297,9 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
297
}
298
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16);
299
}
300
+ explicit BF16Vec16(const c10::BFloat16 v)
301
+ : reg_fp32(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(static_cast<float>(v),
302
+ VEC_ELEM_NUM)) {}
303
explicit BF16Vec16(const FP32Vec16&);
304
void save(void* ptr) const {
305
float tmp[16];
@@ -629,6 +639,19 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
639
: reg(RVVI4(__riscv_vcreate_v_f32, LMUL_256, _f32, LMUL_512)(
640
data.reg, data.reg)) {};
641
explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
642
+ explicit FP32Vec16(int64_t value, const FP32Vec16& lut) {
643
+ const uint64_t q_values = static_cast<uint64_t>(value);
644
+ auto packed = RVVI(__riscv_vmv_v_x_u64, LMUL_1024)(q_values, VEC_ELEM_NUM);
645
+ auto lane_ids = RVVI(__riscv_vid_v_u64, LMUL_1024)(VEC_ELEM_NUM);
646
+ auto shifts =
647
+ RVVI(__riscv_vsll_vx_u64, LMUL_1024)(lane_ids, 2, VEC_ELEM_NUM);
648
+ auto shifted =
649
+ RVVI(__riscv_vsrl_vv_u64, LMUL_1024)(packed, shifts, VEC_ELEM_NUM);
650
+ auto idx64 =
651
+ RVVI(__riscv_vand_vx_u64, LMUL_1024)(shifted, 0xF, VEC_ELEM_NUM);
652
+ auto idx32 = RVVI(__riscv_vnsrl_wx_u32, LMUL_512)(idx64, 0, VEC_ELEM_NUM);
653
+ reg = RVVI(__riscv_vrgather_vv_f32, LMUL_512)(lut.reg, idx32, VEC_ELEM_NUM);
654
+ }
655
explicit FP32Vec16(const FP16Vec16& v);
656
657
#ifdef __riscv_zvfbfmin
@@ -641,6 +664,10 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
664
explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {};
665
#endif
666
667
+ // FP8 stub: dead code on RISC-V (fp8 KV cache is x86-only), needed for
668
+ // load_b_pair_vec template to compile on all platforms.
669
+ explicit FP32Vec16(const BF16Vec32&, int) : FP32Vec16() {}
670
+
671
FP32Vec16 operator+(const FP32Vec16& b) const {
672
return FP32Vec16(
673
RVVI(__riscv_vfadd_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
@@ -891,6 +918,30 @@ inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
918
acc = acc.fma(a, b);
919
}
920
921
+template <typename VecT>
922
+static void interleave_save_16b(const VecT& vec0, const VecT& vec1, void* ptr) {
923
+ alignas(64) uint16_t values0[VecT::VEC_ELEM_NUM];
924
+ alignas(64) uint16_t values1[VecT::VEC_ELEM_NUM];
925
+ vec0.save(values0);
926
+ vec1.save(values1);
927
+
928
+ auto* packed = reinterpret_cast<uint32_t*>(ptr);
929
+ for (int32_t i = 0; i < VecT::VEC_ELEM_NUM; ++i) {
930
+ packed[i] = static_cast<uint32_t>(values0[i]) |
931
+ (static_cast<uint32_t>(values1[i]) << 16);
932
+ }
933
+}
934
+
935
+static void interleave_save(const FP16Vec16& vec0, const FP16Vec16& vec1,
936
+ void* ptr) {
937
+ interleave_save_16b(vec0, vec1, ptr);
938
+}
939
+
940
+static void interleave_save(const BF16Vec16& vec0, const BF16Vec16& vec1,
941
+ void* ptr) {
942
+ interleave_save_16b(vec0, vec1, ptr);
943
+}
944
+
945
#ifdef __riscv_zvfbfmin
946
template <>
947
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
csrc/cpu/torch_bindings.cpp
+1
-1
@@ -518,7 +518,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
518
ops.def("dynamic_per_token_scaled_fp8_quant() -> ()", placeholder_op);
519
520
// WNA16
521
-#if defined(__AVX512F__)
521
+#if defined(__AVX512F__) || defined(__riscv_v)
522
ops.def(
523
"cpu_gemm_wna16(Tensor input, Tensor q_weight, Tensor(a2!) output, "
524
"Tensor scales, Tensor? zeros, Tensor? g_idx, Tensor? bias, SymInt "