[RISC-V] Enable BF16 on VLEN=256 hardware (#45243)
Signed-off-by: velonica0 <like@mail.nankai.edu.cn>
velonica0 committed
Jul 6, 2026 at 14:05 UTC
990c2a0187228b6596e355a7045f7acf3f2f0ca1
2 files changed
+24
-9
cmake/cpu_extension.cmake
+12
-4
@@ -15,6 +15,7 @@ endif()
15
#
16
set(ENABLE_X86_ISA $ENV{VLLM_CPU_X86})
17
set(ENABLE_ARM_BF16 $ENV{VLLM_CPU_ARM_BF16})
18
+set(ENABLE_RVV_BF16 $ENV{VLLM_CPU_RVV_BF16})
19
20
include_directories("${CMAKE_SOURCE_DIR}/csrc")
21
@@ -110,6 +111,13 @@ else()
111
set(ARM_BF16_FOUND ON)
112
message(STATUS "ARM BF16 support enabled via VLLM_CPU_ARM_BF16 environment variable")
113
endif()
114
+ # Some kernels (e.g. Bianbu on Spacemit X100) do not report zvfbfmin
115
+ # in /proc/cpuinfo despite hardware support. VLLM_CPU_RVV_BF16=1
116
+ # overrides the detection result.
117
+ if (ENABLE_RVV_BF16)
118
+ set(RVV_BF16_FOUND ON)
119
+ message(STATUS "RVV BF16 support enabled via VLLM_CPU_RVV_BF16 environment variable")
120
+ endif()
121
endif()
122
123
if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64" OR ENABLE_X86_ISA)
@@ -205,9 +213,9 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
213
if(NOT DEFINED VLLM_RVV_VLEN AND (RVV_FP16_FOUND OR RVV_BF16_FOUND))
214
message(FATAL_ERROR
215
"RISC-V RVV is available but VLEN could not be auto-detected. "
208
- "Please specify VLEN explicitly:\n"
209
- " -DVLLM_RVV_VLEN=128 (for VLEN=128 hardware)\n"
210
- " -DVLLM_RVV_VLEN=256 (for VLEN=256 hardware, e.g. Spacemit X100)")
216
+ "Please specify VLEN explicitly via CMAKE_ARGS:\n"
217
+ " CMAKE_ARGS='-DVLLM_RVV_VLEN=128' (for VLEN=128 hardware)\n"
218
+ " CMAKE_ARGS='-DVLLM_RVV_VLEN=256' (for VLEN=256 hardware, e.g. Spacemit X100)")
219
endif()
220
endif()
221
if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0)
@@ -219,7 +227,7 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
227
message(STATUS "BF16 extension detected")
228
set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
229
elseif(RVV_FP16_FOUND)
222
- message(WARNING "BF16 functionality is not available")
230
+ message(WARNING "BF16 functionality is not available.")
231
set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
232
else()
233
message(STATUS "compile riscv with scalar (no FP16/BF16)")
csrc/cpu/cpu_types_riscv_impl.hpp
+12
-5
@@ -214,11 +214,18 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
214
215
explicit BF16Vec32(const BF16Vec8& v) {
216
fixed_u16x8_t u16_val = bf16_to_u16(v.reg);
217
- fixed_u16x32_t u16_combined =
218
- RVVI4(__riscv_vcreate_v_u16, LMUL_128, _u16, LMUL_512)(
219
- u16_val, u16_val, u16_val, u16_val);
220
- reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16,
221
- LMUL_512)(u16_combined);
217
+ // Widen LMUL_128 → LMUL_256 so vslideup operands share a type.
218
+ // At VLEN=256 this is mf2→m1 (both integer); at VLEN=128 it is m1→m2.
219
+ fixed_u16x16_t ext =
220
+ RVVI4(__riscv_vlmul_ext_v_u16, LMUL_128, _u16, LMUL_256)(u16_val);
221
+ // Build 16-element half: place the 8 elements at offsets 0 and 8.
222
+ fixed_u16x16_t half = RVVI(__riscv_vmv_v_x_u16, LMUL_256)(0, 16);
223
+ half = RVVI(__riscv_vslideup_vx_u16, LMUL_256)(half, ext, 0, 8);
224
+ half = RVVI(__riscv_vslideup_vx_u16, LMUL_256)(half, ext, 8, 16);
225
+ // Double to LMUL_512 (m1→m2 at VLEN=256, m2→m4 at VLEN=128).
226
+ fixed_u16x32_t dst =
227
+ RVVI4(__riscv_vcreate_v_u16, LMUL_256, _u16, LMUL_512)(half, half);
228
+ reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16, LMUL_512)(dst);
229
};
230
231
void save(void* ptr) const {