[CPU][RISC-V] Add RVV micro GEMM for WNA16 (#44324)
Signed-off-by: wcy <233313160abc@gmail.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
wcy committed
Jun 22, 2026 at 20:53 UTC
d2c671c29be7ffe3c4dc081c90a7ba12fd1a904e
4 files changed
+273
-2
csrc/cpu/cpu_wna16.cpp
+39
@@ -4,6 +4,9 @@
4
#ifdef CPU_CAPABILITY_AMXBF16
5
#include "cpu/micro_gemm/cpu_micro_gemm_amx.hpp"
6
#endif
7
+#if defined(__riscv_v)
8
+ #include "cpu/micro_gemm/cpu_micro_gemm_rvv.hpp"
9
+#endif
10
#include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
11
12
#define VLLM_DISPATCH_CASE_16B_TYPES(...) \
@@ -319,6 +322,8 @@ void cpu_gemm_wna16(
322
return ISA::AMX;
323
} else if (isa_hint == "vec") {
324
return ISA::VEC;
325
+ } else if (isa_hint == "rvv") {
326
+ return ISA::RVV;
327
} else {
328
TORCH_CHECK(false, "unsupported isa hint: " + isa_hint);
329
}
@@ -397,6 +402,40 @@ void cpu_gemm_wna16(
402
pack_factor);
403
return;
404
}
405
+ } else if (isa == ISA::RVV) {
406
+ using gemm_t = cpu_micro_gemm::MicroGemm<ISA::RVV, scalar_t>;
407
+ if (has_zp) {
408
+ using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, true, false>;
409
+ cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
410
+ input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
411
+ output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
412
+ g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
413
+ a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
414
+ scales_group_stride, zeros_group_stride, group_num, group_size,
415
+ pack_factor);
416
+ return;
417
+ }
418
+ if (use_desc_act) {
419
+ using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, false, true>;
420
+ cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
421
+ input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
422
+ output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
423
+ g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
424
+ a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
425
+ scales_group_stride, zeros_group_stride, group_num, group_size,
426
+ pack_factor);
427
+ return;
428
+ } else {
429
+ using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, false, false>;
430
+ cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
431
+ input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
432
+ output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
433
+ g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
434
+ a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
435
+ scales_group_stride, zeros_group_stride, group_num, group_size,
436
+ pack_factor);
437
+ return;
438
+ }
439
}
440
});
441
}
csrc/cpu/micro_gemm/cpu_micro_gemm_rvv.hpp
new
+228
@@ -0,0 +1,228 @@
1
+#ifndef CPU_MICRO_GEMM_RVV_HPP
2
+#define CPU_MICRO_GEMM_RVV_HPP
3
+
4
+#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
5
+
6
+#if defined(__riscv_v)
7
+
8
+namespace cpu_micro_gemm {
9
+namespace {
10
+
11
+constexpr int32_t RVV_MGEMM_N8 = 8;
12
+constexpr int32_t RVV_MGEMM_B_GROUP_STRIDE = 16;
13
+
14
+template <typename scalar_t>
15
+FORCE_INLINE fixed_fp32x8_t load_row8_b_as_f32(const scalar_t* ptr);
16
+
17
+template <>
18
+FORCE_INLINE fixed_fp32x8_t load_row8_b_as_f32<float>(const float* ptr) {
19
+ return RVVI(__riscv_vle32_v_f32, LMUL_256)(ptr, RVV_MGEMM_N8);
20
+}
21
+
22
+template <>
23
+FORCE_INLINE fixed_fp32x8_t
24
+load_row8_b_as_f32<c10::Half>(const c10::Half* ptr) {
25
+ #if defined(__riscv_zvfh)
26
+ fixed_fp16x8_t vec = RVVI(__riscv_vle16_v_f16, LMUL_128)(
27
+ reinterpret_cast<const _Float16*>(ptr), RVV_MGEMM_N8);
28
+ return RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(vec, RVV_MGEMM_N8);
29
+ #else
30
+ alignas(32) float values[RVV_MGEMM_N8];
31
+ for (int32_t i = 0; i < RVV_MGEMM_N8; ++i) {
32
+ values[i] = static_cast<float>(ptr[i]);
33
+ }
34
+ return RVVI(__riscv_vle32_v_f32, LMUL_256)(values, RVV_MGEMM_N8);
35
+ #endif
36
+}
37
+
38
+template <>
39
+FORCE_INLINE fixed_fp32x8_t
40
+load_row8_b_as_f32<c10::BFloat16>(const c10::BFloat16* ptr) {
41
+ #if defined(__riscv_zvfbfmin)
42
+ fixed_u16x8_t raw = RVVI(__riscv_vle16_v_u16, LMUL_128)(
43
+ reinterpret_cast<const uint16_t*>(ptr), RVV_MGEMM_N8);
44
+ fixed_bf16x8_t vec =
45
+ RVVI4(__riscv_vreinterpret_v_u16, LMUL_128, _bf16, LMUL_128)(raw);
46
+ return RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(vec, RVV_MGEMM_N8);
47
+ #else
48
+ fixed_u16x8_t raw = RVVI(__riscv_vle16_v_u16, LMUL_128)(
49
+ reinterpret_cast<const uint16_t*>(ptr), RVV_MGEMM_N8);
50
+ auto wide = RVVI(__riscv_vzext_vf2_u32, LMUL_256)(raw, RVV_MGEMM_N8);
51
+ auto shifted = RVVI(__riscv_vsll_vx_u32, LMUL_256)(wide, 16, RVV_MGEMM_N8);
52
+ return RVVI4(__riscv_vreinterpret_v_u32, LMUL_256, _f32, LMUL_256)(shifted);
53
+ #endif
54
+}
55
+
56
+// Mx8 RVV kernel. B points at one 8-channel half of a 16-channel packed group,
57
+// with rows separated by RVV_MGEMM_B_GROUP_STRIDE scalar elements.
58
+template <int32_t M, typename scalar_t>
59
+FORCE_INLINE void gemm_micro_rvv_fma_mx8_ku4(const scalar_t* __restrict__ a_ptr,
60
+ const scalar_t* __restrict__ b_ptr,
61
+ float* __restrict__ c_ptr,
62
+ const int64_t lda,
63
+ const int64_t ldc, const int32_t k,
64
+ const bool accum_c) {
65
+ static_assert(0 < M && M <= 8);
66
+
67
+ #define RVV_ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
68
+ #define RVV_IF_M(i) if constexpr (M > (i))
69
+
70
+ #define RVV_DECL_A(i) const scalar_t* __restrict__ a##i = a_ptr + (i) * lda;
71
+ RVV_ROWS_APPLY(RVV_DECL_A)
72
+ #undef RVV_DECL_A
73
+
74
+ #define RVV_DECL_ACC(i) fixed_fp32x8_t acc##i;
75
+ RVV_ROWS_APPLY(RVV_DECL_ACC)
76
+ #undef RVV_DECL_ACC
77
+
78
+ #define RVV_INIT_ACC(i) \
79
+ RVV_IF_M(i) { \
80
+ if (accum_c) { \
81
+ acc##i = RVVI(__riscv_vle32_v_f32, LMUL_256)(c_ptr + (i) * ldc, \
82
+ RVV_MGEMM_N8); \
83
+ } else { \
84
+ acc##i = RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.0f, RVV_MGEMM_N8); \
85
+ } \
86
+ }
87
+ RVV_ROWS_APPLY(RVV_INIT_ACC)
88
+ #undef RVV_INIT_ACC
89
+
90
+ int32_t k_idx = 0;
91
+ for (; k_idx + 3 < k; k_idx += 4) {
92
+ #define RVV_FMA_ROW(i, K_OFFSET) \
93
+ RVV_IF_M(i) { \
94
+ acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)( \
95
+ acc##i, static_cast<float>(*(a##i + k_idx + (K_OFFSET))), b, \
96
+ RVV_MGEMM_N8); \
97
+ }
98
+
99
+ #define RVV_STEP_K(K_OFFSET) \
100
+ { \
101
+ fixed_fp32x8_t b = load_row8_b_as_f32<scalar_t>( \
102
+ b_ptr + (k_idx + (K_OFFSET)) * RVV_MGEMM_B_GROUP_STRIDE); \
103
+ RVV_FMA_ROW(0, K_OFFSET) \
104
+ RVV_FMA_ROW(1, K_OFFSET) \
105
+ RVV_FMA_ROW(2, K_OFFSET) \
106
+ RVV_FMA_ROW(3, K_OFFSET) \
107
+ RVV_FMA_ROW(4, K_OFFSET) \
108
+ RVV_FMA_ROW(5, K_OFFSET) \
109
+ RVV_FMA_ROW(6, K_OFFSET) \
110
+ RVV_FMA_ROW(7, K_OFFSET) \
111
+ }
112
+
113
+ RVV_STEP_K(0)
114
+ RVV_STEP_K(1)
115
+ RVV_STEP_K(2)
116
+ RVV_STEP_K(3)
117
+ #undef RVV_STEP_K
118
+ #undef RVV_FMA_ROW
119
+ }
120
+
121
+ for (; k_idx < k; ++k_idx) {
122
+ fixed_fp32x8_t b =
123
+ load_row8_b_as_f32<scalar_t>(b_ptr + k_idx * RVV_MGEMM_B_GROUP_STRIDE);
124
+ #define RVV_TAIL_ROW(i) \
125
+ RVV_IF_M(i) { \
126
+ acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)( \
127
+ acc##i, static_cast<float>(*(a##i + k_idx)), b, RVV_MGEMM_N8); \
128
+ }
129
+ RVV_ROWS_APPLY(RVV_TAIL_ROW)
130
+ #undef RVV_TAIL_ROW
131
+ }
132
+
133
+ #define RVV_STORE_ROW(i) \
134
+ RVV_IF_M(i) { \
135
+ RVVI(__riscv_vse32_v_f32, LMUL_256)(c_ptr + (i) * ldc, acc##i, \
136
+ RVV_MGEMM_N8); \
137
+ }
138
+ RVV_ROWS_APPLY(RVV_STORE_ROW)
139
+ #undef RVV_STORE_ROW
140
+
141
+ #undef RVV_ROWS_APPLY
142
+ #undef RVV_IF_M
143
+}
144
+
145
+template <int32_t M, typename scalar_t>
146
+FORCE_INLINE void gemm_micro_rvv_mx32_ku4(DEFINE_CPU_MICRO_GEMM_PARAMS) {
147
+ static_assert(0 < M && M <= 8);
148
+ scalar_t* __restrict__ curr_b_0 = b_ptr;
149
+ scalar_t* __restrict__ curr_b_1 = b_ptr + b_n_group_stride;
150
+
151
+ gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_0, c_ptr, lda, ldc, k, accum_c);
152
+ gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_0 + RVV_MGEMM_N8,
153
+ c_ptr + RVV_MGEMM_N8, lda, ldc, k, accum_c);
154
+ gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_1, c_ptr + 16, lda, ldc, k,
155
+ accum_c);
156
+ gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_1 + RVV_MGEMM_N8, c_ptr + 24, lda,
157
+ ldc, k, accum_c);
158
+}
159
+
160
+class TileGemmRVV {
161
+ public:
162
+ template <typename scalar_t>
163
+ FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
164
+ switch (m) {
165
+ case 1:
166
+ gemm_micro_rvv_mx32_ku4<1>(CPU_MICRO_GEMM_PARAMS);
167
+ break;
168
+ case 2:
169
+ gemm_micro_rvv_mx32_ku4<2>(CPU_MICRO_GEMM_PARAMS);
170
+ break;
171
+ case 3:
172
+ gemm_micro_rvv_mx32_ku4<3>(CPU_MICRO_GEMM_PARAMS);
173
+ break;
174
+ case 4:
175
+ gemm_micro_rvv_mx32_ku4<4>(CPU_MICRO_GEMM_PARAMS);
176
+ break;
177
+ case 5:
178
+ gemm_micro_rvv_mx32_ku4<5>(CPU_MICRO_GEMM_PARAMS);
179
+ break;
180
+ case 6:
181
+ gemm_micro_rvv_mx32_ku4<6>(CPU_MICRO_GEMM_PARAMS);
182
+ break;
183
+ case 7:
184
+ gemm_micro_rvv_mx32_ku4<7>(CPU_MICRO_GEMM_PARAMS);
185
+ break;
186
+ case 8:
187
+ gemm_micro_rvv_mx32_ku4<8>(CPU_MICRO_GEMM_PARAMS);
188
+ break;
189
+ }
190
+ }
191
+};
192
+
193
+} // namespace
194
+
195
+template <typename scalar_t>
196
+class MicroGemm<cpu_utils::ISA::RVV, scalar_t> {
197
+ public:
198
+ static constexpr int32_t MaxMSize = 8;
199
+ static constexpr int32_t NSize = 32;
200
+
201
+ public:
202
+ void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
203
+ TileGemmRVV::gemm<scalar_t>(CPU_MICRO_GEMM_PARAMS);
204
+ }
205
+
206
+ static void pack_weight(const scalar_t* __restrict__ weight,
207
+ scalar_t* __restrict__ packed_weight,
208
+ const int32_t output_size, const int32_t input_size) {
209
+ TORCH_CHECK_EQ(output_size % 16, 0);
210
+ for (int32_t o_idx = 0; o_idx < output_size; ++o_idx) {
211
+ const scalar_t* __restrict__ curr_weight = weight + o_idx * input_size;
212
+ scalar_t* __restrict__ curr_packed_weight =
213
+ packed_weight + (o_idx / 16) * (16 * input_size) + o_idx % 16;
214
+ for (int32_t i_idx = 0; i_idx < input_size; ++i_idx) {
215
+ *curr_packed_weight = *curr_weight;
216
+
217
+ curr_packed_weight += 16;
218
+ ++curr_weight;
219
+ }
220
+ }
221
+ }
222
+};
223
+
224
+} // namespace cpu_micro_gemm
225
+
226
+#endif // defined(__riscv_v)
227
+
228
+#endif // CPU_MICRO_GEMM_RVV_HPP
csrc/cpu/utils.hpp
+3
-1
@@ -8,13 +8,15 @@
8
#include "cpu/cpu_types.hpp"
9
10
namespace cpu_utils {
11
-enum class ISA { AMX, VEC };
11
+enum class ISA { AMX, VEC, RVV };
12
13
inline ISA get_isa(const std::string& isa) {
14
if (isa == "amx") {
15
return ISA::AMX;
16
} else if (isa == "vec") {
17
return ISA::VEC;
18
+ } else if (isa == "rvv") {
19
+ return ISA::RVV;
20
} else {
21
TORCH_CHECK(false, "Invalid isa type: " + isa);
22
}
vllm/model_executor/kernels/linear/mixed_precision/cpu.py
+3
-1
@@ -9,7 +9,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
9
pack_quantized_values_into_int32,
10
unpack_quantized_values_into_int32,
11
)
12
-from vllm.platforms import current_platform
12
+from vllm.platforms import CpuArchEnum, current_platform
13
from vllm.scalar_type import scalar_types
14
15
from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
@@ -214,5 +214,7 @@ def _get_isa_hint(dtype: torch.dtype) -> str:
214
supports_amx = torch.cpu._is_amx_tile_supported()
215
if supports_amx and dtype in (torch.bfloat16,):
216
return "amx"
217
+ elif current_platform.get_cpu_architecture() == CpuArchEnum.RISCV:
218
+ return "rvv"
219
else:
220
return "vec"