From 990c2a0187228b6596e355a7045f7acf3f2f0ca1 Mon Sep 17 00:00:00 2001 From: velonica0 <47554626+velonica0@users.noreply.github.com> Date: Mon, 6 Jul 2026 14:05:16 +0800 Subject: [PATCH] [RISC-V] Enable BF16 on VLEN=256 hardware (#45243) Signed-off-by: velonica0 --- cmake/cpu_extension.cmake | 16 ++++++++++++---- csrc/cpu/cpu_types_riscv_impl.hpp | 17 ++++++++++++----- 2 files changed, 24 insertions(+), 9 deletions(-) diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index ce8f5da2af..3aca9bcea9 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -15,6 +15,7 @@ endif() # set(ENABLE_X86_ISA $ENV{VLLM_CPU_X86}) set(ENABLE_ARM_BF16 $ENV{VLLM_CPU_ARM_BF16}) +set(ENABLE_RVV_BF16 $ENV{VLLM_CPU_RVV_BF16}) include_directories("${CMAKE_SOURCE_DIR}/csrc") @@ -110,6 +111,13 @@ else() set(ARM_BF16_FOUND ON) message(STATUS "ARM BF16 support enabled via VLLM_CPU_ARM_BF16 environment variable") endif() + # Some kernels (e.g. Bianbu on Spacemit X100) do not report zvfbfmin + # in /proc/cpuinfo despite hardware support. VLLM_CPU_RVV_BF16=1 + # overrides the detection result. + if (ENABLE_RVV_BF16) + set(RVV_BF16_FOUND ON) + message(STATUS "RVV BF16 support enabled via VLLM_CPU_RVV_BF16 environment variable") + endif() endif() if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64" OR ENABLE_X86_ISA) @@ -205,9 +213,9 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") if(NOT DEFINED VLLM_RVV_VLEN AND (RVV_FP16_FOUND OR RVV_BF16_FOUND)) message(FATAL_ERROR "RISC-V RVV is available but VLEN could not be auto-detected. " - "Please specify VLEN explicitly:\n" - " -DVLLM_RVV_VLEN=128 (for VLEN=128 hardware)\n" - " -DVLLM_RVV_VLEN=256 (for VLEN=256 hardware, e.g. Spacemit X100)") + "Please specify VLEN explicitly via CMAKE_ARGS:\n" + " CMAKE_ARGS='-DVLLM_RVV_VLEN=128' (for VLEN=128 hardware)\n" + " CMAKE_ARGS='-DVLLM_RVV_VLEN=256' (for VLEN=256 hardware, e.g. Spacemit X100)") endif() endif() if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0) @@ -219,7 +227,7 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") message(STATUS "BF16 extension detected") set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d) elseif(RVV_FP16_FOUND) - message(WARNING "BF16 functionality is not available") + message(WARNING "BF16 functionality is not available.") set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d) else() message(STATUS "compile riscv with scalar (no FP16/BF16)") diff --git a/csrc/cpu/cpu_types_riscv_impl.hpp b/csrc/cpu/cpu_types_riscv_impl.hpp index dd29606c20..ae17272968 100644 --- a/csrc/cpu/cpu_types_riscv_impl.hpp +++ b/csrc/cpu/cpu_types_riscv_impl.hpp @@ -214,11 +214,18 @@ struct BF16Vec32 : public Vec { explicit BF16Vec32(const BF16Vec8& v) { fixed_u16x8_t u16_val = bf16_to_u16(v.reg); - fixed_u16x32_t u16_combined = - RVVI4(__riscv_vcreate_v_u16, LMUL_128, _u16, LMUL_512)( - u16_val, u16_val, u16_val, u16_val); - reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16, - LMUL_512)(u16_combined); + // Widen LMUL_128 → LMUL_256 so vslideup operands share a type. + // At VLEN=256 this is mf2→m1 (both integer); at VLEN=128 it is m1→m2. + fixed_u16x16_t ext = + RVVI4(__riscv_vlmul_ext_v_u16, LMUL_128, _u16, LMUL_256)(u16_val); + // Build 16-element half: place the 8 elements at offsets 0 and 8. + fixed_u16x16_t half = RVVI(__riscv_vmv_v_x_u16, LMUL_256)(0, 16); + half = RVVI(__riscv_vslideup_vx_u16, LMUL_256)(half, ext, 0, 8); + half = RVVI(__riscv_vslideup_vx_u16, LMUL_256)(half, ext, 8, 16); + // Double to LMUL_512 (m1→m2 at VLEN=256, m2→m4 at VLEN=128). + fixed_u16x32_t dst = + RVVI4(__riscv_vcreate_v_u16, LMUL_256, _u16, LMUL_512)(half, half); + reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16, LMUL_512)(dst); }; void save(void* ptr) const {