@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4184653775
Fork 0

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 6, 23:00 UTC

feat: add RISC-V support for CPU backend (v2) (#36578)

Signed-off-by: typer-J <2236066784@qq.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>

typer-J committed Mar 11, 2026 at 12:51 UTC 4184653775fd8cd6b6498e6731f7d2014f0fe05b
5 files changed +851 -30
cmake/cpu_extension.cmake
+12 -3
@@ -79,7 +79,8 @@ else()
79 find_isa(${CPUINFO} "asimd" ASIMD_FOUND) # Check for ARM NEON support
80 find_isa(${CPUINFO} "bf16" ARM_BF16_FOUND) # Check for ARM BF16 support
81 find_isa(${CPUINFO} "S390" S390_FOUND)
82 - find_isa(${CPUINFO} "v" RVV_FOUND) # Check for RISC-V RVV support
82 + find_isa(${CPUINFO} "zvfhmin" RVV_FP16_FOUND) # Check for RISC-V Vector FP16 support
83 + find_isa(${CPUINFO} "zvfbfmin" RVV_BF16_FOUND) # Check for RISC-V Vector BF16 support
84
85 # Support cross-compilation by allowing override via environment variables
86 if (ENABLE_ARM_BF16)
@@ -142,11 +143,19 @@ elseif (S390_FOUND)
143 "-march=native"
144 "-mtune=native")
145 elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
145 - if(RVV_FOUND)
146 - message(FAIL_ERROR "Can't support rvv now.")
146 + message(STATUS "RISC-V detected")
147 + if(RVV_BF16_FOUND)
148 + message(STATUS "BF16 extension detected")
149 + set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl128b -mrvv-vector-bits=zvl -mabi=lp64d)
150 + add_compile_definitions(RISCV_BF16_SUPPORT)
151 + elseif (RVV_FP16_FOUND)
152 + message(WARNING "BF16 functionality is not available")
153 + set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl128b -mrvv-vector-bits=zvl -mabi=lp64d)
154 else()
155 + message(STATUS "compile riscv with scalar")
156 list(APPEND CXX_COMPILE_FLAGS "-march=rv64gc")
157 endif()
158 + list(APPEND CXX_COMPILE_FLAGS ${MARCH_FLAGS})
159 else()
160 message(FATAL_ERROR "vLLM CPU backend requires X86, Power9+ ISA, S390X ISA, ARMv8 or RISC-V support.")
161 endif()
csrc/cpu/cpu_types.hpp
+3
@@ -13,6 +13,9 @@
13 #elif defined(__aarch64__)
14 // arm implementation
15 #include "cpu_types_arm.hpp"
16 +#elif defined(__riscv_v)
17 + // riscv implementation
18 + #include "cpu_types_riscv.hpp"
19 #else
20 #warning "unsupported vLLM cpu implementation, vLLM will compile with scalar"
21 #include "cpu_types_scalar.hpp"
csrc/cpu/cpu_types_riscv.hpp new
+832
@@ -0,0 +1,832 @@
1 +#ifndef CPU_TYPES_RISCV_HPP
2 +#define CPU_TYPES_RISCV_HPP
3 +
4 +#include <algorithm>
5 +#include <cmath>
6 +#include <cstring>
7 +#include <iostream>
8 +#include <limits>
9 +#include <riscv_vector.h>
10 +#include <torch/all.h>
11 +
12 +// ============================================================================
13 +// Vector Register Type Definitions (VLEN=128 bits)
14 +// ============================================================================
15 +
16 +typedef vfloat16m1_t fixed_vfloat16m1_t
17 + __attribute__((riscv_rvv_vector_bits(128)));
18 +typedef vfloat16m2_t fixed_vfloat16m2_t
19 + __attribute__((riscv_rvv_vector_bits(256)));
20 +
21 +typedef vfloat32m1_t fixed_vfloat32m1_t
22 + __attribute__((riscv_rvv_vector_bits(128)));
23 +typedef vfloat32m2_t fixed_vfloat32m2_t
24 + __attribute__((riscv_rvv_vector_bits(256)));
25 +typedef vfloat32m4_t fixed_vfloat32m4_t
26 + __attribute__((riscv_rvv_vector_bits(512)));
27 +typedef vfloat32m8_t fixed_vfloat32m8_t
28 + __attribute__((riscv_rvv_vector_bits(1024)));
29 +
30 +typedef vint32m2_t fixed_vint32m2_t __attribute__((riscv_rvv_vector_bits(256)));
31 +typedef vint32m4_t fixed_vint32m4_t __attribute__((riscv_rvv_vector_bits(512)));
32 +
33 +typedef vuint16m1_t fixed_vuint16m1_t
34 + __attribute__((riscv_rvv_vector_bits(128)));
35 +typedef vuint16m2_t fixed_vuint16m2_t
36 + __attribute__((riscv_rvv_vector_bits(256)));
37 +typedef vuint16m4_t fixed_vuint16m4_t
38 + __attribute__((riscv_rvv_vector_bits(512)));
39 +
40 +#ifdef RISCV_BF16_SUPPORT
41 +typedef vbfloat16m1_t fixed_vbfloat16m1_t
42 + __attribute__((riscv_rvv_vector_bits(128)));
43 +typedef vbfloat16m2_t fixed_vbfloat16m2_t
44 + __attribute__((riscv_rvv_vector_bits(256)));
45 +typedef vbfloat16m4_t fixed_vbfloat16m4_t
46 + __attribute__((riscv_rvv_vector_bits(512)));
47 +#endif
48 +
49 +namespace vec_op {
50 +
51 +#ifdef RISCV_BF16_SUPPORT
52 + #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
53 + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
54 + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
55 + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
56 +#else
57 + #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
58 + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
59 + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
60 +#endif
61 +
62 +#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
63 + AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
64 +
65 +#define FORCE_INLINE __attribute__((always_inline)) inline
66 +
67 +namespace {
68 +template <typename T, T... indexes, typename F>
69 +constexpr void unroll_loop_item(std::integer_sequence<T, indexes...>, F&& f) {
70 + (f(std::integral_constant<T, indexes>{}), ...);
71 +};
72 +} // namespace
73 +
74 +template <typename T, T count, typename F,
75 + typename = std::enable_if_t<std::is_invocable_v<F, T>>>
76 +constexpr void unroll_loop(F&& f) {
77 + unroll_loop_item(std::make_integer_sequence<T, count>{}, std::forward<F>(f));
78 +}
79 +
80 +template <typename T>
81 +struct Vec {
82 + constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; };
83 +};
84 +
85 +struct FP32Vec8;
86 +struct FP32Vec16;
87 +
88 +// ============================================================================
89 +// FP16 Implementation
90 +// ============================================================================
91 +
92 +struct FP16Vec8 : public Vec<FP16Vec8> {
93 + constexpr static int VEC_ELEM_NUM = 8;
94 + fixed_vfloat16m1_t reg;
95 +
96 + explicit FP16Vec8(const void* ptr)
97 + : reg(__riscv_vle16_v_f16m1(static_cast<const _Float16*>(ptr),
98 + VEC_ELEM_NUM)) {};
99 +
100 + explicit FP16Vec8(const FP32Vec8&);
101 +
102 + void save(void* ptr) const {
103 + __riscv_vse16_v_f16m1(static_cast<_Float16*>(ptr), reg, VEC_ELEM_NUM);
104 + }
105 + void save(void* ptr, int elem_num) const {
106 + __riscv_vse16_v_f16m1(static_cast<_Float16*>(ptr), reg, elem_num);
107 + }
108 + void save_strided(void* ptr, ptrdiff_t stride) const {
109 + ptrdiff_t byte_stride = stride * sizeof(_Float16);
110 + __riscv_vsse16_v_f16m1(static_cast<_Float16*>(ptr), byte_stride, reg,
111 + VEC_ELEM_NUM);
112 + }
113 +};
114 +
115 +struct FP16Vec16 : public Vec<FP16Vec16> {
116 + constexpr static int VEC_ELEM_NUM = 16;
117 + fixed_vfloat16m2_t reg;
118 +
119 + explicit FP16Vec16(const void* ptr)
120 + : reg(__riscv_vle16_v_f16m2(static_cast<const _Float16*>(ptr),
121 + VEC_ELEM_NUM)) {};
122 +
123 + explicit FP16Vec16(const FP32Vec16& vec);
124 +
125 + void save(void* ptr) const {
126 + __riscv_vse16_v_f16m2(static_cast<_Float16*>(ptr), reg, VEC_ELEM_NUM);
127 + }
128 + void save(void* ptr, int elem_num) const {
129 + __riscv_vse16_v_f16m2(static_cast<_Float16*>(ptr), reg, elem_num);
130 + }
131 + void save_strided(void* ptr, ptrdiff_t stride) const {
132 + ptrdiff_t byte_stride = stride * sizeof(_Float16);
133 + __riscv_vsse16_v_f16m2(static_cast<_Float16*>(ptr), byte_stride, reg,
134 + VEC_ELEM_NUM);
135 + }
136 +};
137 +
138 +// ============================================================================
139 +// BF16 Implementation
140 +// ============================================================================
141 +
142 +#ifdef RISCV_BF16_SUPPORT
143 +
144 +FORCE_INLINE fixed_vuint16m1_t bf16_to_u16(fixed_vbfloat16m1_t v) {
145 + return __riscv_vreinterpret_v_bf16m1_u16m1(v);
146 +}
147 +FORCE_INLINE fixed_vuint16m2_t bf16_to_u16(fixed_vbfloat16m2_t v) {
148 + return __riscv_vreinterpret_v_bf16m2_u16m2(v);
149 +}
150 +FORCE_INLINE fixed_vuint16m4_t bf16_to_u16(fixed_vbfloat16m4_t v) {
151 + return __riscv_vreinterpret_v_bf16m4_u16m4(v);
152 +}
153 +
154 +struct BF16Vec8 : public Vec<BF16Vec8> {
155 + constexpr static int VEC_ELEM_NUM = 8;
156 + fixed_vbfloat16m1_t reg;
157 +
158 + explicit BF16Vec8(const void* ptr)
159 + : reg(__riscv_vreinterpret_v_u16m1_bf16m1(__riscv_vle16_v_u16m1(
160 + reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
161 +
162 + explicit BF16Vec8(fixed_vbfloat16m1_t data) : reg(data) {};
163 + explicit BF16Vec8(const FP32Vec8&);
164 +
165 + void save(void* ptr) const {
166 + __riscv_vse16_v_u16m1(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
167 + VEC_ELEM_NUM);
168 + }
169 + void save(void* ptr, int elem_num) const {
170 + __riscv_vse16_v_u16m1(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
171 + elem_num);
172 + }
173 + void save_strided(void* ptr, ptrdiff_t stride) const {
174 + ptrdiff_t byte_stride = stride * sizeof(uint16_t);
175 + __riscv_vsse16_v_u16m1(reinterpret_cast<uint16_t*>(ptr), byte_stride,
176 + bf16_to_u16(reg), VEC_ELEM_NUM);
177 + }
178 +};
179 +
180 +struct BF16Vec16 : public Vec<BF16Vec16> {
181 + constexpr static int VEC_ELEM_NUM = 16;
182 + fixed_vbfloat16m2_t reg;
183 +
184 + explicit BF16Vec16(const void* ptr)
185 + : reg(__riscv_vreinterpret_v_u16m2_bf16m2(__riscv_vle16_v_u16m2(
186 + reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
187 +
188 + explicit BF16Vec16(fixed_vbfloat16m2_t data) : reg(data) {};
189 + explicit BF16Vec16(const FP32Vec16&);
190 +
191 + void save(void* ptr) const {
192 + __riscv_vse16_v_u16m2(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
193 + VEC_ELEM_NUM);
194 + }
195 + void save(void* ptr, int elem_num) const {
196 + __riscv_vse16_v_u16m2(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
197 + elem_num);
198 + }
199 + void save_strided(void* ptr, ptrdiff_t stride) const {
200 + ptrdiff_t byte_stride = stride * sizeof(uint16_t);
201 + __riscv_vsse16_v_u16m2(reinterpret_cast<uint16_t*>(ptr), byte_stride,
202 + bf16_to_u16(reg), VEC_ELEM_NUM);
203 + }
204 +};
205 +
206 +struct BF16Vec32 : public Vec<BF16Vec32> {
207 + constexpr static int VEC_ELEM_NUM = 32;
208 + fixed_vbfloat16m4_t reg;
209 +
210 + explicit BF16Vec32(const void* ptr)
211 + : reg(__riscv_vreinterpret_v_u16m4_bf16m4(__riscv_vle16_v_u16m4(
212 + reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
213 +
214 + explicit BF16Vec32(fixed_vbfloat16m4_t data) : reg(data) {};
215 +
216 + explicit BF16Vec32(const BF16Vec8& v) {
217 + fixed_vuint16m1_t u16_val = bf16_to_u16(v.reg);
218 + fixed_vuint16m4_t u16_combined =
219 + __riscv_vcreate_v_u16m1_u16m4(u16_val, u16_val, u16_val, u16_val);
220 + reg = __riscv_vreinterpret_v_u16m4_bf16m4(u16_combined);
221 + };
222 +
223 + void save(void* ptr) const {
224 + __riscv_vse16_v_u16m4(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
225 + VEC_ELEM_NUM);
226 + }
227 + void save(void* ptr, int elem_num) const {
228 + __riscv_vse16_v_u16m4(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
229 + elem_num);
230 + }
231 + void save_strided(void* ptr, ptrdiff_t stride) const {
232 + ptrdiff_t byte_stride = stride * sizeof(uint16_t);
233 + __riscv_vsse16_v_u16m4(reinterpret_cast<uint16_t*>(ptr), byte_stride,
234 + bf16_to_u16(reg), VEC_ELEM_NUM);
235 + }
236 +};
237 +
238 +#else
239 +// ============================================================================
240 +// BF16 Fallback Implementation (FP32 Simulation)
241 +// ============================================================================
242 +
243 +struct BF16Vec8 : public Vec<BF16Vec8> {
244 + constexpr static int VEC_ELEM_NUM = 8;
245 + fixed_vfloat32m2_t reg_fp32;
246 + explicit BF16Vec8(const void* ptr) {
247 + const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
248 + float tmp[8];
249 + for (int i = 0; i < 8; ++i) {
250 + uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
251 + std::memcpy(&tmp[i], &v, 4);
252 + }
253 + reg_fp32 = __riscv_vle32_v_f32m2(tmp, 8);
254 + }
255 + explicit BF16Vec8(const FP32Vec8&);
256 + void save(void* ptr) const {
257 + float tmp[8];
258 + __riscv_vse32_v_f32m2(tmp, reg_fp32, 8);
259 + uint16_t* u16 = static_cast<uint16_t*>(ptr);
260 + for (int i = 0; i < 8; ++i) {
261 + uint32_t v;
262 + std::memcpy(&v, &tmp[i], 4);
263 + u16[i] = static_cast<uint16_t>(v >> 16);
264 + }
265 + }
266 + void save(void* ptr, int elem_num) const {
267 + float tmp[8];
268 + __riscv_vse32_v_f32m2(tmp, reg_fp32, 8);
269 + uint16_t* u16 = static_cast<uint16_t*>(ptr);
270 + for (int i = 0; i < elem_num; ++i) {
271 + uint32_t v;
272 + std::memcpy(&v, &tmp[i], 4);
273 + u16[i] = static_cast<uint16_t>(v >> 16);
274 + }
275 + }
276 + void save_strided(void* ptr, ptrdiff_t stride) const {
277 + float tmp[8];
278 + __riscv_vse32_v_f32m2(tmp, reg_fp32, 8);
279 + uint8_t* u8 = static_cast<uint8_t*>(ptr);
280 + ptrdiff_t byte_stride = stride * sizeof(uint16_t);
281 + for (int i = 0; i < 8; ++i) {
282 + uint32_t v;
283 + std::memcpy(&v, &tmp[i], 4);
284 + uint16_t val = static_cast<uint16_t>(v >> 16);
285 + *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
286 + }
287 + }
288 +};
289 +
290 +struct BF16Vec16 : public Vec<BF16Vec16> {
291 + constexpr static int VEC_ELEM_NUM = 16;
292 + fixed_vfloat32m4_t reg_fp32;
293 + explicit BF16Vec16(const void* ptr) {
294 + const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
295 + float tmp[16];
296 + for (int i = 0; i < 16; ++i) {
297 + uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
298 + std::memcpy(&tmp[i], &v, 4);
299 + }
300 + reg_fp32 = __riscv_vle32_v_f32m4(tmp, 16);
301 + }
302 + explicit BF16Vec16(const FP32Vec16&);
303 + void save(void* ptr) const {
304 + float tmp[16];
305 + __riscv_vse32_v_f32m4(tmp, reg_fp32, 16);
306 + uint16_t* u16 = static_cast<uint16_t*>(ptr);
307 + for (int i = 0; i < 16; ++i) {
308 + uint32_t v;
309 + std::memcpy(&v, &tmp[i], 4);
310 + u16[i] = static_cast<uint16_t>(v >> 16);
311 + }
312 + }
313 + void save(void* ptr, int elem_num) const {
314 + float tmp[16];
315 + __riscv_vse32_v_f32m4(tmp, reg_fp32, 16);
316 + uint16_t* u16 = static_cast<uint16_t*>(ptr);
317 + for (int i = 0; i < elem_num; ++i) {
318 + uint32_t v;
319 + std::memcpy(&v, &tmp[i], 4);
320 + u16[i] = static_cast<uint16_t>(v >> 16);
321 + }
322 + }
323 + void save_strided(void* ptr, ptrdiff_t stride) const {
324 + float tmp[16];
325 + __riscv_vse32_v_f32m4(tmp, reg_fp32, 16);
326 + uint8_t* u8 = static_cast<uint8_t*>(ptr);
327 + ptrdiff_t byte_stride = stride * sizeof(uint16_t);
328 + for (int i = 0; i < 16; ++i) {
329 + uint32_t v;
330 + std::memcpy(&v, &tmp[i], 4);
331 + uint16_t val = static_cast<uint16_t>(v >> 16);
332 + *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
333 + }
334 + }
335 +};
336 +
337 +struct BF16Vec32 : public Vec<BF16Vec32> {
338 + constexpr static int VEC_ELEM_NUM = 32;
339 + fixed_vfloat32m8_t reg_fp32;
340 +
341 + explicit BF16Vec32(const void* ptr) {
342 + const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
343 + float tmp[32];
344 + for (int i = 0; i < 32; ++i) {
345 + uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
346 + std::memcpy(&tmp[i], &v, 4);
347 + }
348 + reg_fp32 = __riscv_vle32_v_f32m8(tmp, 32);
349 + }
350 +
351 + explicit BF16Vec32(const BF16Vec8& v) {
352 + float tmp_small[8];
353 + __riscv_vse32_v_f32m2(tmp_small, v.reg_fp32, 8);
354 + float tmp_large[32];
355 + for (int i = 0; i < 4; ++i) {
356 + std::memcpy(tmp_large + (i * 8), tmp_small, 8 * sizeof(float));
357 + }
358 + reg_fp32 = __riscv_vle32_v_f32m8(tmp_large, 32);
359 + }
360 +
361 + void save(void* ptr) const {
362 + float tmp[32];
363 + __riscv_vse32_v_f32m8(tmp, reg_fp32, 32);
364 + uint16_t* u16 = static_cast<uint16_t*>(ptr);
365 + for (int i = 0; i < 32; ++i) {
366 + uint32_t v;
367 + std::memcpy(&v, &tmp[i], 4);
368 + u16[i] = static_cast<uint16_t>(v >> 16);
369 + }
370 + }
371 +
372 + void save(void* ptr, int elem_num) const {
373 + float tmp[32];
374 + __riscv_vse32_v_f32m8(tmp, reg_fp32, 32);
375 + uint16_t* u16 = static_cast<uint16_t*>(ptr);
376 + for (int i = 0; i < elem_num; ++i) {
377 + uint32_t v;
378 + std::memcpy(&v, &tmp[i], 4);
379 + u16[i] = static_cast<uint16_t>(v >> 16);
380 + }
381 + }
382 +
383 + void save_strided(void* ptr, ptrdiff_t stride) const {
384 + float tmp[32];
385 + __riscv_vse32_v_f32m8(tmp, reg_fp32, 32);
386 + uint8_t* u8 = static_cast<uint8_t*>(ptr);
387 + ptrdiff_t byte_stride = stride * sizeof(uint16_t);
388 + for (int i = 0; i < 32; ++i) {
389 + uint32_t v;
390 + std::memcpy(&v, &tmp[i], 4);
391 + uint16_t val = static_cast<uint16_t>(v >> 16);
392 + *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
393 + }
394 + }
395 +};
396 +#endif
397 +
398 +// ============================================================================
399 +// FP32 Implementation
400 +// ============================================================================
401 +
402 +struct FP32Vec4 : public Vec<FP32Vec4> {
403 + constexpr static int VEC_ELEM_NUM = 4;
404 + fixed_vfloat32m1_t reg;
405 + explicit FP32Vec4(float v) : reg(__riscv_vfmv_v_f_f32m1(v, VEC_ELEM_NUM)) {};
406 + explicit FP32Vec4() : reg(__riscv_vfmv_v_f_f32m1(0.0f, VEC_ELEM_NUM)) {};
407 + explicit FP32Vec4(const float* ptr)
408 + : reg(__riscv_vle32_v_f32m1(ptr, VEC_ELEM_NUM)) {};
409 + explicit FP32Vec4(fixed_vfloat32m1_t data) : reg(data) {};
410 + explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {};
411 + void save(float* ptr) const { __riscv_vse32_v_f32m1(ptr, reg, VEC_ELEM_NUM); }
412 + void save(float* ptr, int elem_num) const {
413 + __riscv_vse32_v_f32m1(ptr, reg, elem_num);
414 + }
415 +};
416 +
417 +struct FP32Vec8 : public Vec<FP32Vec8> {
418 + constexpr static int VEC_ELEM_NUM = 8;
419 + fixed_vfloat32m2_t reg;
420 +
421 + explicit FP32Vec8(float v) : reg(__riscv_vfmv_v_f_f32m2(v, VEC_ELEM_NUM)) {};
422 + explicit FP32Vec8() : reg(__riscv_vfmv_v_f_f32m2(0.0f, VEC_ELEM_NUM)) {};
423 + explicit FP32Vec8(const float* ptr)
424 + : reg(__riscv_vle32_v_f32m2(ptr, VEC_ELEM_NUM)) {};
425 + explicit FP32Vec8(fixed_vfloat32m2_t data) : reg(data) {};
426 + explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {};
427 + explicit FP32Vec8(const FP16Vec8& v)
428 + : reg(__riscv_vfwcvt_f_f_v_f32m2(v.reg, VEC_ELEM_NUM)) {};
429 + explicit FP32Vec8(fixed_vfloat16m1_t v)
430 + : reg(__riscv_vfwcvt_f_f_v_f32m2(v, VEC_ELEM_NUM)) {};
431 +
432 +#ifdef RISCV_BF16_SUPPORT
433 + explicit FP32Vec8(fixed_vbfloat16m1_t v)
434 + : reg(__riscv_vfwcvtbf16_f_f_v_f32m2(v, VEC_ELEM_NUM)) {};
435 + explicit FP32Vec8(const BF16Vec8& v)
436 + : reg(__riscv_vfwcvtbf16_f_f_v_f32m2(v.reg, VEC_ELEM_NUM)) {};
437 +#else
438 + explicit FP32Vec8(const BF16Vec8& v) : reg(v.reg_fp32) {};
439 +#endif
440 +
441 + float reduce_sum() const {
442 + fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
443 + scalar = __riscv_vfredusum_vs_f32m2_f32m1(reg, scalar, VEC_ELEM_NUM);
444 + return __riscv_vfmv_f_s_f32m1_f32(scalar);
445 + }
446 +
447 + FP32Vec8 operator*(const FP32Vec8& b) const {
448 + return FP32Vec8(__riscv_vfmul_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
449 + }
450 + FP32Vec8 operator+(const FP32Vec8& b) const {
451 + return FP32Vec8(__riscv_vfadd_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
452 + }
453 + FP32Vec8 operator-(const FP32Vec8& b) const {
454 + return FP32Vec8(__riscv_vfsub_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
455 + }
456 + FP32Vec8 operator/(const FP32Vec8& b) const {
457 + return FP32Vec8(__riscv_vfdiv_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
458 + }
459 +
460 + FP32Vec8 min(const FP32Vec8& b) const {
461 + return FP32Vec8(__riscv_vfmin_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
462 + }
463 + FP32Vec8 max(const FP32Vec8& b) const {
464 + return FP32Vec8(__riscv_vfmax_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
465 + }
466 + FP32Vec8 abs() const {
467 + return FP32Vec8(__riscv_vfabs_v_f32m2(reg, VEC_ELEM_NUM));
468 + }
469 +
470 + FP32Vec8 min(const FP32Vec8& b, int elem_num) const {
471 + return FP32Vec8(__riscv_vfmin_vv_f32m2(reg, b.reg, elem_num));
472 + }
473 + FP32Vec8 max(const FP32Vec8& b, int elem_num) const {
474 + return FP32Vec8(__riscv_vfmax_vv_f32m2(reg, b.reg, elem_num));
475 + }
476 +
477 + FP32Vec8 clamp(const FP32Vec8& min_v, const FP32Vec8& max_v) const {
478 + fixed_vfloat32m2_t temp =
479 + __riscv_vfmax_vv_f32m2(min_v.reg, reg, VEC_ELEM_NUM);
480 + return FP32Vec8(__riscv_vfmin_vv_f32m2(max_v.reg, temp, VEC_ELEM_NUM));
481 + }
482 +
483 + void save(float* ptr) const { __riscv_vse32_v_f32m2(ptr, reg, VEC_ELEM_NUM); }
484 + void save(float* ptr, int elem_num) const {
485 + __riscv_vse32_v_f32m2(ptr, reg, elem_num);
486 + }
487 + void save_strided(float* ptr, ptrdiff_t stride) const {
488 + ptrdiff_t byte_stride = stride * sizeof(float);
489 + __riscv_vsse32_v_f32m2(ptr, byte_stride, reg, VEC_ELEM_NUM);
490 + }
491 +
492 + FP32Vec8 exp() const {
493 + const float inv_ln2 = 1.44269504088896341f;
494 + fixed_vfloat32m2_t x_scaled =
495 + __riscv_vfmul_vf_f32m2(reg, inv_ln2, VEC_ELEM_NUM);
496 + fixed_vint32m2_t n_int = __riscv_vfcvt_x_f_v_i32m2(x_scaled, VEC_ELEM_NUM);
497 + fixed_vfloat32m2_t n_float = __riscv_vfcvt_f_x_v_f32m2(n_int, VEC_ELEM_NUM);
498 +
499 + fixed_vfloat32m2_t r =
500 + __riscv_vfsub_vv_f32m2(x_scaled, n_float, VEC_ELEM_NUM);
501 +
502 + fixed_vfloat32m2_t poly =
503 + __riscv_vfmv_v_f_f32m2(0.001333355810164f, VEC_ELEM_NUM);
504 + poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
505 + poly = __riscv_vfadd_vf_f32m2(poly, 0.009618129107628f, VEC_ELEM_NUM);
506 + poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
507 + poly = __riscv_vfadd_vf_f32m2(poly, 0.055504108664821f, VEC_ELEM_NUM);
508 + poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
509 + poly = __riscv_vfadd_vf_f32m2(poly, 0.240226506959101f, VEC_ELEM_NUM);
510 + poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
511 + poly = __riscv_vfadd_vf_f32m2(poly, 0.693147180559945f, VEC_ELEM_NUM);
512 + poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
513 + poly = __riscv_vfadd_vf_f32m2(poly, 1.0f, VEC_ELEM_NUM);
514 +
515 + fixed_vint32m2_t biased_exp =
516 + __riscv_vadd_vx_i32m2(n_int, 127, VEC_ELEM_NUM);
517 + biased_exp = __riscv_vmax_vx_i32m2(biased_exp, 0, VEC_ELEM_NUM);
518 + fixed_vint32m2_t exponent_bits =
519 + __riscv_vsll_vx_i32m2(biased_exp, 23, VEC_ELEM_NUM);
520 + fixed_vfloat32m2_t scale =
521 + __riscv_vreinterpret_v_i32m2_f32m2(exponent_bits);
522 +
523 + return FP32Vec8(__riscv_vfmul_vv_f32m2(poly, scale, VEC_ELEM_NUM));
524 + }
525 +
526 + FP32Vec8 tanh() const {
527 + fixed_vfloat32m2_t x_clamped = __riscv_vfmin_vf_f32m2(
528 + __riscv_vfmax_vf_f32m2(reg, -9.0f, VEC_ELEM_NUM), 9.0f, VEC_ELEM_NUM);
529 + fixed_vfloat32m2_t x2 =
530 + __riscv_vfmul_vf_f32m2(x_clamped, 2.0f, VEC_ELEM_NUM);
531 + FP32Vec8 exp_val = FP32Vec8(x2).exp();
532 + fixed_vfloat32m2_t num =
533 + __riscv_vfsub_vf_f32m2(exp_val.reg, 1.0f, VEC_ELEM_NUM);
534 + fixed_vfloat32m2_t den =
535 + __riscv_vfadd_vf_f32m2(exp_val.reg, 1.0f, VEC_ELEM_NUM);
536 + return FP32Vec8(__riscv_vfdiv_vv_f32m2(num, den, VEC_ELEM_NUM));
537 + }
538 +
539 + FP32Vec8 er() const {
540 + const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
541 + a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
542 + fixed_vfloat32m2_t abs_x = __riscv_vfabs_v_f32m2(reg, VEC_ELEM_NUM);
543 +
544 + fixed_vfloat32m2_t t = __riscv_vfadd_vf_f32m2(
545 + __riscv_vfmul_vf_f32m2(abs_x, p, VEC_ELEM_NUM), 1.0f, VEC_ELEM_NUM);
546 + t = __riscv_vfrdiv_vf_f32m2(t, 1.0f, VEC_ELEM_NUM);
547 +
548 + fixed_vfloat32m2_t poly = __riscv_vfmv_v_f_f32m2(a5, VEC_ELEM_NUM);
549 + poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
550 + a4, VEC_ELEM_NUM);
551 + poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
552 + a3, VEC_ELEM_NUM);
553 + poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
554 + a2, VEC_ELEM_NUM);
555 + poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
556 + a1, VEC_ELEM_NUM);
557 + poly = __riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM);
558 +
559 + fixed_vfloat32m2_t exp_val =
560 + FP32Vec8(__riscv_vfneg_v_f32m2(
561 + __riscv_vfmul_vv_f32m2(abs_x, abs_x, VEC_ELEM_NUM),
562 + VEC_ELEM_NUM))
563 + .exp()
564 + .reg;
565 + fixed_vfloat32m2_t res = __riscv_vfrsub_vf_f32m2(
566 + __riscv_vfmul_vv_f32m2(poly, exp_val, VEC_ELEM_NUM), 1.0f,
567 + VEC_ELEM_NUM);
568 +
569 + vbool16_t mask = __riscv_vmflt_vf_f32m2_b16(reg, 0.0f, VEC_ELEM_NUM);
570 + return FP32Vec8(__riscv_vfneg_v_f32m2_m(mask, res, VEC_ELEM_NUM));
571 + }
572 +};
573 +
574 +struct FP32Vec16 : public Vec<FP32Vec16> {
575 + constexpr static int VEC_ELEM_NUM = 16;
576 + fixed_vfloat32m4_t reg;
577 +
578 + explicit FP32Vec16(float v) : reg(__riscv_vfmv_v_f_f32m4(v, VEC_ELEM_NUM)) {};
579 + explicit FP32Vec16() : reg(__riscv_vfmv_v_f_f32m4(0.0f, VEC_ELEM_NUM)) {};
580 + explicit FP32Vec16(const float* ptr)
581 + : reg(__riscv_vle32_v_f32m4(ptr, VEC_ELEM_NUM)) {};
582 + explicit FP32Vec16(fixed_vfloat32m4_t data) : reg(data) {};
583 + explicit FP32Vec16(const FP32Vec8& data)
584 + : reg(__riscv_vcreate_v_f32m2_f32m4(data.reg, data.reg)) {};
585 + explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
586 + explicit FP32Vec16(const FP16Vec16& v);
587 +
588 +#ifdef RISCV_BF16_SUPPORT
589 + explicit FP32Vec16(fixed_vbfloat16m2_t v)
590 + : reg(__riscv_vfwcvtbf16_f_f_v_f32m4(v, VEC_ELEM_NUM)) {};
591 + explicit FP32Vec16(const BF16Vec16& v)
592 + : reg(__riscv_vfwcvtbf16_f_f_v_f32m4(v.reg, VEC_ELEM_NUM)) {};
593 +#else
594 + explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {};
595 +#endif
596 +
597 + FP32Vec16 operator+(const FP32Vec16& b) const {
598 + return FP32Vec16(__riscv_vfadd_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
599 + }
600 + FP32Vec16 operator-(const FP32Vec16& b) const {
601 + return FP32Vec16(__riscv_vfsub_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
602 + }
603 + FP32Vec16 operator*(const FP32Vec16& b) const {
604 + return FP32Vec16(__riscv_vfmul_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
605 + }
606 + FP32Vec16 operator/(const FP32Vec16& b) const {
607 + return FP32Vec16(__riscv_vfdiv_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
608 + }
609 +
610 + FP32Vec16 fma(const FP32Vec16& a, const FP32Vec16& b) const {
611 + return FP32Vec16(__riscv_vfmacc_vv_f32m4(reg, a.reg, b.reg, VEC_ELEM_NUM));
612 + }
613 +
614 + float reduce_sum() const {
615 + fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
616 + scalar = __riscv_vfredusum_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM);
617 + return __riscv_vfmv_f_s_f32m1_f32(scalar);
618 + }
619 +
620 + float reduce_max() const {
621 + fixed_vfloat32m1_t scalar =
622 + __riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::lowest(), 1);
623 + scalar = __riscv_vfredmax_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM);
624 + return __riscv_vfmv_f_s_f32m1_f32(scalar);
625 + }
626 +
627 + float reduce_min() const {
628 + fixed_vfloat32m1_t scalar =
629 + __riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::max(), 1);
630 + scalar = __riscv_vfredmin_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM);
631 + return __riscv_vfmv_f_s_f32m1_f32(scalar);
632 + }
633 +
634 + template <int group_size>
635 + float reduce_sub_sum(int idx) {
636 + static_assert(VEC_ELEM_NUM % group_size == 0);
637 + const int start = idx * group_size;
638 + vuint32m4_t indices = __riscv_vid_v_u32m4(VEC_ELEM_NUM);
639 + vbool8_t mask = __riscv_vmand_mm_b8(
640 + __riscv_vmsgeu_vx_u32m4_b8(indices, start, VEC_ELEM_NUM),
641 + __riscv_vmsltu_vx_u32m4_b8(indices, start + group_size, VEC_ELEM_NUM),
642 + VEC_ELEM_NUM);
643 + fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
644 + scalar =
645 + __riscv_vfredusum_vs_f32m4_f32m1_m(mask, reg, scalar, VEC_ELEM_NUM);
646 + return __riscv_vfmv_f_s_f32m1_f32(scalar);
647 + };
648 +
649 + FP32Vec16 max(const FP32Vec16& b) const {
650 + return FP32Vec16(__riscv_vfmax_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
651 + }
652 + FP32Vec16 min(const FP32Vec16& b) const {
653 + return FP32Vec16(__riscv_vfmin_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
654 + }
655 + FP32Vec16 abs() const {
656 + return FP32Vec16(__riscv_vfabs_v_f32m4(reg, VEC_ELEM_NUM));
657 + }
658 +
659 + FP32Vec16 clamp(const FP32Vec16& min_v, const FP32Vec16& max_v) const {
660 + return FP32Vec16(__riscv_vfmin_vv_f32m4(
661 + max_v.reg, __riscv_vfmax_vv_f32m4(min_v.reg, reg, VEC_ELEM_NUM),
662 + VEC_ELEM_NUM));
663 + }
664 +
665 + void save(float* ptr) const { __riscv_vse32_v_f32m4(ptr, reg, VEC_ELEM_NUM); }
666 + void save(float* ptr, int elem_num) const {
667 + __riscv_vse32_v_f32m4(ptr, reg, elem_num);
668 + }
669 + void save_strided(float* ptr, ptrdiff_t stride) const {
670 + ptrdiff_t byte_stride = stride * sizeof(float);
671 + __riscv_vsse32_v_f32m4(ptr, byte_stride, reg, VEC_ELEM_NUM);
672 + }
673 +
674 + FP32Vec16 exp() const {
675 + const float inv_ln2 = 1.44269504088896341f;
676 + fixed_vfloat32m4_t x_scaled =
677 + __riscv_vfmul_vf_f32m4(reg, inv_ln2, VEC_ELEM_NUM);
678 + fixed_vint32m4_t n_int = __riscv_vfcvt_x_f_v_i32m4(x_scaled, VEC_ELEM_NUM);
679 + fixed_vfloat32m4_t n_float = __riscv_vfcvt_f_x_v_f32m4(n_int, VEC_ELEM_NUM);
680 + fixed_vfloat32m4_t r =
681 + __riscv_vfsub_vv_f32m4(x_scaled, n_float, VEC_ELEM_NUM);
682 +
683 + fixed_vfloat32m4_t poly =
684 + __riscv_vfmv_v_f_f32m4(0.001333355810164f, VEC_ELEM_NUM);
685 + poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
686 + 0.009618129107628f, VEC_ELEM_NUM);
687 + poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
688 + 0.055504108664821f, VEC_ELEM_NUM);
689 + poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
690 + 0.240226506959101f, VEC_ELEM_NUM);
691 + poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
692 + 0.693147180559945f, VEC_ELEM_NUM);
693 + poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
694 + 1.0f, VEC_ELEM_NUM);
695 +
696 + fixed_vint32m4_t biased_exp = __riscv_vmax_vx_i32m4(
697 + __riscv_vadd_vx_i32m4(n_int, 127, VEC_ELEM_NUM), 0, VEC_ELEM_NUM);
698 + fixed_vfloat32m4_t scale = __riscv_vreinterpret_v_i32m4_f32m4(
699 + __riscv_vsll_vx_i32m4(biased_exp, 23, VEC_ELEM_NUM));
700 +
701 + return FP32Vec16(__riscv_vfmul_vv_f32m4(poly, scale, VEC_ELEM_NUM));
702 + }
703 +
704 + FP32Vec16 tanh() const {
705 + fixed_vfloat32m4_t x_clamped = __riscv_vfmin_vf_f32m4(
706 + __riscv_vfmax_vf_f32m4(reg, -9.0f, VEC_ELEM_NUM), 9.0f, VEC_ELEM_NUM);
707 + FP32Vec16 exp_val =
708 + FP32Vec16(__riscv_vfmul_vf_f32m4(x_clamped, 2.0f, VEC_ELEM_NUM)).exp();
709 + return FP32Vec16(__riscv_vfdiv_vv_f32m4(
710 + __riscv_vfsub_vf_f32m4(exp_val.reg, 1.0f, VEC_ELEM_NUM),
711 + __riscv_vfadd_vf_f32m4(exp_val.reg, 1.0f, VEC_ELEM_NUM), VEC_ELEM_NUM));
712 + }
713 +
714 + FP32Vec16 er() const {
715 + const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
716 + a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
717 + fixed_vfloat32m4_t abs_x = __riscv_vfabs_v_f32m4(reg, VEC_ELEM_NUM);
718 + fixed_vfloat32m4_t t = __riscv_vfrdiv_vf_f32m4(
719 + __riscv_vfadd_vf_f32m4(__riscv_vfmul_vf_f32m4(abs_x, p, VEC_ELEM_NUM),
720 + 1.0f, VEC_ELEM_NUM),
721 + 1.0f, VEC_ELEM_NUM);
722 +
723 + fixed_vfloat32m4_t poly = __riscv_vfmv_v_f_f32m4(a5, VEC_ELEM_NUM);
724 + poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
725 + a4, VEC_ELEM_NUM);
726 + poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
727 + a3, VEC_ELEM_NUM);
728 + poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
729 + a2, VEC_ELEM_NUM);
730 + poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
731 + a1, VEC_ELEM_NUM);
732 + poly = __riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM);
733 +
734 + fixed_vfloat32m4_t exp_val =
735 + FP32Vec16(__riscv_vfneg_v_f32m4(
736 + __riscv_vfmul_vv_f32m4(abs_x, abs_x, VEC_ELEM_NUM),
737 + VEC_ELEM_NUM))
738 + .exp()
739 + .reg;
740 + fixed_vfloat32m4_t res = __riscv_vfrsub_vf_f32m4(
741 + __riscv_vfmul_vv_f32m4(poly, exp_val, VEC_ELEM_NUM), 1.0f,
742 + VEC_ELEM_NUM);
743 +
744 + vbool8_t mask = __riscv_vmflt_vf_f32m4_b8(reg, 0.0f, VEC_ELEM_NUM);
745 + return FP32Vec16(__riscv_vfneg_v_f32m4_m(mask, res, VEC_ELEM_NUM));
746 + }
747 +};
748 +
749 +// ============================================================================
750 +// Type Traits & Global Helpers
751 +// ============================================================================
752 +
753 +template <typename T>
754 +struct VecType {
755 + using vec_type = void;
756 + using vec_t = void;
757 +};
758 +
759 +template <typename T>
760 +using vec_t = typename VecType<T>::vec_type;
761 +
762 +template <>
763 +struct VecType<float> {
764 + using vec_type = FP32Vec8;
765 + using vec_t = FP32Vec8;
766 +};
767 +template <>
768 +struct VecType<c10::Half> {
769 + using vec_type = FP16Vec8;
770 + using vec_t = FP16Vec8;
771 +};
772 +template <>
773 +struct VecType<c10::BFloat16> {
774 + using vec_type = BF16Vec8;
775 + using vec_t = BF16Vec8;
776 +};
777 +
778 +template <typename T>
779 +void storeFP32(float v, T* ptr) {
780 + *ptr = v;
781 +}
782 +template <>
783 +inline void storeFP32<c10::Half>(float v, c10::Half* ptr) {
784 + *reinterpret_cast<_Float16*>(ptr) = static_cast<_Float16>(v);
785 +}
786 +
787 +inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
788 + reg = __riscv_vfncvt_f_f_w_f16m2(v.reg, VEC_ELEM_NUM);
789 +}
790 +inline FP16Vec8::FP16Vec8(const FP32Vec8& v) {
791 + reg = __riscv_vfncvt_f_f_w_f16m1(v.reg, VEC_ELEM_NUM);
792 +}
793 +inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
794 + reg = __riscv_vfwcvt_f_f_v_f32m4(v.reg, VEC_ELEM_NUM);
795 +}
796 +inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
797 + acc = acc.fma(a, b);
798 +}
799 +
800 +#ifdef RISCV_BF16_SUPPORT
801 +template <>
802 +inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
803 + *ptr = static_cast<__bf16>(v);
804 +};
805 +inline BF16Vec8::BF16Vec8(const FP32Vec8& v)
806 + : reg(__riscv_vfncvtbf16_f_f_w_bf16m1(v.reg, VEC_ELEM_NUM)) {};
807 +inline BF16Vec16::BF16Vec16(const FP32Vec16& v)
808 + : reg(__riscv_vfncvtbf16_f_f_w_bf16m2(v.reg, VEC_ELEM_NUM)) {};
809 +#else
810 +template <>
811 +inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
812 + uint32_t val;
813 + std::memcpy(&val, &v, 4);
814 + *reinterpret_cast<uint16_t*>(ptr) = static_cast<uint16_t>(val >> 16);
815 +}
816 +inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {}
817 +inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {}
818 +#endif
819 +
820 +inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 1); }
821 +
822 +} // namespace vec_op
823 +
824 +#ifndef CPU_KERNEL_GUARD_IN
825 + #define CPU_KERNEL_GUARD_IN(NAME)
826 +#endif
827 +
828 +#ifndef CPU_KERNEL_GUARD_OUT
829 + #define CPU_KERNEL_GUARD_OUT(NAME)
830 +#endif
831 +
832 +#endif // CPU_TYPES_RISCV_HPP
\ No newline at end of file
requirements/cpu.txt
+3 -3
@@ -7,13 +7,13 @@ numba == 0.61.2; platform_machine != "s390x" # Required for N-gram speculative d
7
8 # Dependencies for CPUs
9 torch==2.10.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x"
10 -torch==2.10.0; platform_machine == "aarch64" or platform_system == "Darwin" or platform_machine == "ppc64le"
10 +torch==2.10.0; platform_machine == "aarch64" or platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64"
11
12 # required for the image processor of minicpm-o-2_6, this must be updated alongside torch
13 -torchaudio; platform_machine != "s390x"
13 +torchaudio; platform_machine != "s390x" and platform_machine != "riscv64"
14
15 # required for the image processor of phi3v, this must be updated alongside torch
16 -torchvision; platform_machine != "s390x"
16 +torchvision; platform_machine != "s390x" and platform_machine != "riscv64"
17
18 # Intel Extension for PyTorch, only for x86_64 CPUs
19 intel-openmp==2024.2.1; platform_machine == "x86_64"
vllm/platforms/cpu.py
+1 -24
@@ -93,30 +93,7 @@ class CpuPlatform(Platform):
93 return [torch.bfloat16, torch.float16, torch.float32]
94 return [torch.float16, torch.float32]
95 elif self.get_cpu_architecture() == CpuArchEnum.RISCV:
96 - # Workaround for Issue #25655: RISC-V scheduler bug with float16
97 - #
98 - # Background:
99 - # - RISC-V currently uses scalar code path
100 - # - There is a latent bug in the vLLM scheduler that provides
101 - # invalid
102 - # physical_block_idx values under certain conditions
103 - # - This bug causes segmentation faults when using float16
104 - # dtype on RISC-V
105 - # - Testing shows that forcing float32 successfully bypasses
106 - # this issue
107 - #
108 - # Technical details:
109 - # - The bug manifests as out-of-bounds physical_block_idx in
110 - # block_tables
111 - # - Only occurs on RISC-V hardware
112 - # tested on Sophgo SG2044
113 - # - Does not reproduce on x86 or other architectures
114 - # - Root cause is in Python-level scheduling logic,
115 - # not C++ kernels
116 - #
117 - # This is a temporary workaround until the scheduler bug is fixed.
118 - # See: https://github.com/vllm-project/vllm/issues/25655
119 - return [torch.float32]
96 + return [torch.bfloat16, torch.float16, torch.float32]
97 # x86/aarch64 CPU has supported both bf16 and fp16 natively.
98 return [torch.bfloat16, torch.float16, torch.float32]
99