diff --git a/benchmarks/kernels/cpu/benchmark_cpu_attn.py b/benchmarks/kernels/cpu/benchmark_cpu_attn.py index 63d034278c..08afd693c3 100644 --- a/benchmarks/kernels/cpu/benchmark_cpu_attn.py +++ b/benchmarks/kernels/cpu/benchmark_cpu_attn.py @@ -12,7 +12,6 @@ from vllm._custom_ops import ( cpu_attn_get_scheduler_metadata, cpu_attn_reshape_and_cache, ) -from vllm.platforms import CpuArchEnum, current_platform from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed from vllm.v1.attention.backends.cpu_attn import CPUAttentionBackend, _get_attn_isa @@ -22,15 +21,14 @@ def get_attn_isa( block_size: int | None = None, dtype: torch.dtype | None = None, ): - if block_size and dtype: - return _get_attn_isa(dtype, block_size) - else: - if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: - return "neon" - elif torch.cpu._is_amx_tile_supported(): - return "amx" - else: - return "vec" + # Delegate to _get_attn_isa so the fallback path applies the same arch + # gating (e.g. RISC-V RVV is only chosen when the build's hardcoded + # VLEN=128 kernel is actually present; on VLEN=256 / scalar hosts it + # correctly falls through to vec/vec16). + return _get_attn_isa( + dtype if dtype is not None else torch.bfloat16, + block_size if block_size else 32, + ) # rand number generation takes too much time, cache rand tensors @@ -235,7 +233,7 @@ if __name__ == "__main__": ) parser.add_argument("--use-sink", action="store_true") parser.add_argument( - "--isa", type=str, choices=["vec", "neon", "amx", "vec16"], default=None + "--isa", type=str, choices=["vec", "neon", "amx", "vec16", "rvv"], default=None ) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--iters", type=int, default=20) diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 361f08b510..ffab4015f4 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -195,10 +195,12 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") endif() if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0) message(STATUS "RISC-V RVV VLEN=${VLLM_RVV_VLEN}") + # Sources gate FP16/BF16 paths on the compiler-provided + # __riscv_zvfh / __riscv_zvfbfmin macros, which GCC and clang + # define automatically when those extensions appear in -march. if(RVV_BF16_FOUND) message(STATUS "BF16 extension detected") set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d) - add_compile_definitions(RISCV_BF16_SUPPORT) elseif(RVV_FP16_FOUND) message(WARNING "BF16 functionality is not available") set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d) diff --git a/csrc/cpu/cpu_arch_macros.h b/csrc/cpu/cpu_arch_macros.h index 9be45a3efc..53ae70497c 100644 --- a/csrc/cpu/cpu_arch_macros.h +++ b/csrc/cpu/cpu_arch_macros.h @@ -156,4 +156,20 @@ #endif // __aarch64__ +// RISC-V RVV +#ifdef __riscv_v + #include + + #ifdef __riscv_zihintpause + #define FAST_SPINNING __riscv_pause(); + #endif + + // FP32Vec16::exp() in cpu_types_riscv.hpp already implements the full + // polynomial approximation for RVV, so we simply delegate to it. + #define DEFINE_FAST_EXP \ + auto fast_exp = [&](const vec_op::FP32Vec16& vec) \ + __attribute__((always_inline)) { return vec.exp(); }; + +#endif // __riscv_v + #endif diff --git a/csrc/cpu/cpu_attn.cpp b/csrc/cpu/cpu_attn.cpp index 4750dd7883..26b881f4f1 100644 --- a/csrc/cpu/cpu_attn.cpp +++ b/csrc/cpu/cpu_attn.cpp @@ -29,6 +29,8 @@ torch::Tensor get_scheduler_metadata( isa = cpu_attention::ISA::NEON; } else if (isa_hint == "vxe") { isa = cpu_attention::ISA::VXE; + } else if (isa_hint == "rvv") { + isa = cpu_attention::ISA::RVV; } else if (isa_hint == "vsx") { isa = cpu_attention::ISA::VSX; } else { @@ -131,6 +133,8 @@ void cpu_attn_reshape_and_cache( return cpu_attention::ISA::NEON; } else if (isa == "vxe") { return cpu_attention::ISA::VXE; + } else if (isa == "rvv") { + return cpu_attention::ISA::RVV; } else if (isa == "vsx") { return cpu_attention::ISA::VSX; } else { diff --git a/csrc/cpu/cpu_attn_impl.hpp b/csrc/cpu/cpu_attn_impl.hpp index b9987fb26c..72c96252b6 100644 --- a/csrc/cpu/cpu_attn_impl.hpp +++ b/csrc/cpu/cpu_attn_impl.hpp @@ -12,7 +12,7 @@ #include "cpu/utils.hpp" namespace cpu_attention { -enum class ISA { AMX, VEC, VEC16, NEON, VXE, VSX }; +enum class ISA { AMX, VEC, VEC16, NEON, VXE, RVV, VSX }; // Mirrors csrc/attention/dtype_fp8.cuh Fp8KVCacheDataType exactly. enum class Fp8KVCacheDataType { @@ -164,6 +164,9 @@ struct AttentionMetadata { case ISA::VXE: ss << "VXE, "; break; + case ISA::RVV: + ss << "RVV, "; + break; case ISA::VSX: ss << "VSX, "; break; diff --git a/csrc/cpu/cpu_attn_rvv.hpp b/csrc/cpu/cpu_attn_rvv.hpp new file mode 100644 index 0000000000..ec1230547e --- /dev/null +++ b/csrc/cpu/cpu_attn_rvv.hpp @@ -0,0 +1,445 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#ifndef CPU_ATTN_RVV_HPP +#define CPU_ATTN_RVV_HPP + +// This kernel is currently hardcoded to VLEN=128 (m1/m2 intrinsics, vl=8). +// The fixed-width typedefs below use `riscv_rvv_vector_bits(128)`, which +// only matches `vfloat16m1_t`/`vuint16m1_t` register layout when VLEN==128; +// at VLEN>=256 those typedefs fail to compile. Scalar RISC-V builds +// (-march=rv64gc) additionally don't have . For both +// cases we omit the file entirely and let the dispatcher fall back to the +// scalar VEC / VEC16 implementations. TODO: migrate to RVVI() macros + +// semantic names in cpu_types_riscv_defs.hpp to support VLEN>=256 natively. +#if defined(__riscv_v_min_vlen) && __riscv_v_min_vlen == 128 + + #include "cpu_attn_impl.hpp" + #include + #include + +namespace cpu_attention { + +namespace { + +// File-local concrete-LMUL typedefs. The shared _defs.hpp exposes +// VLEN-independent semantic names (fixed_fp32x8_t, fixed_fp16x8_t, ...), +// but this kernel is currently hardcoded to VLEN=128 (m1/m2 intrinsics), +// so keep the legacy concrete aliases scoped to this file. +typedef vfloat16m1_t fixed_vfloat16m1_t + __attribute__((riscv_rvv_vector_bits(128))); +typedef vfloat32m2_t fixed_vfloat32m2_t + __attribute__((riscv_rvv_vector_bits(256))); +typedef vuint16m1_t fixed_vuint16m1_t + __attribute__((riscv_rvv_vector_bits(128))); +typedef vuint32m2_t fixed_vuint32m2_t + __attribute__((riscv_rvv_vector_bits(256))); + #ifdef __riscv_zvfbfmin +typedef vbfloat16m1_t fixed_vbfloat16m1_t + __attribute__((riscv_rvv_vector_bits(128))); + #endif + + #define BLOCK_SIZE_ALIGNMENT 32 + #define HEAD_SIZE_ALIGNMENT 32 + #define MAX_Q_HEAD_NUM_PER_ITER 16 + +// ============================================================================ +// B-matrix row loading: load 8 elements as FP32 (using m2 LMUL at VLEN=128) +// ============================================================================ + +template +FORCE_INLINE fixed_vfloat32m2_t load_row8_B_as_f32(const kv_cache_t* p); + +template <> +FORCE_INLINE fixed_vfloat32m2_t load_row8_B_as_f32(const float* p) { + return __riscv_vle32_v_f32m2(p, 8); +} + +template <> +FORCE_INLINE fixed_vfloat32m2_t +load_row8_B_as_f32(const c10::Half* p) { + #ifdef __riscv_zvfh + fixed_vfloat16m1_t h = + __riscv_vle16_v_f16m1(reinterpret_cast(p), 8); + return __riscv_vfwcvt_f_f_v_f32m2(h, 8); + #else + // Fallback for hardware without Zvfh: scalar half->float conversion. + // c10::Half provides operator float() so this is correct on any RVV CPU + // that has only the base V extension. Slower than the Zvfh path, but + // keeps the kernel buildable on Zvfhmin-only / no-fp16 hardware. + alignas(16) float tmp[8]; + for (int i = 0; i < 8; ++i) { + tmp[i] = static_cast(p[i]); + } + return __riscv_vle32_v_f32m2(tmp, 8); + #endif +} + +template <> +FORCE_INLINE fixed_vfloat32m2_t +load_row8_B_as_f32(const c10::BFloat16* p) { + #ifdef __riscv_zvfbfmin + fixed_vbfloat16m1_t bf = + __riscv_vle16_v_bf16m1(reinterpret_cast(p), 8); + return __riscv_vfwcvtbf16_f_f_v_f32m2(bf, 8); + #else + // Fallback: load as uint16, zero-extend to uint32, shift left by 16 + fixed_vuint16m1_t raw = + __riscv_vle16_v_u16m1(reinterpret_cast(p), 8); + fixed_vuint32m2_t wide = __riscv_vzext_vf2_u32m2(raw, 8); + fixed_vuint32m2_t shifted = __riscv_vsll_vx_u32m2(wide, 16, 8); + return __riscv_vreinterpret_v_u32m2_f32m2(shifted); + #endif +} + +// ============================================================================ +// Micro kernel: Mx8 tile, K unrolled by 4, RVV scalar-broadcast FMA +// ============================================================================ +// +// NEON uses vfmaq_laneq_f32 (lane-indexed FMA from a preloaded A vector). +// RVV has no lane-indexed FMA; instead we load A elements as scalars and +// use __riscv_vfmacc_vf (scalar * vector + accumulator), which is equally +// efficient and avoids the need for vrgather/vslidedown. +// +// At VLEN=128, m2 holds 8 x FP32, matching the 8-column tile width. +// Register budget: M accumulators (m2 each) + 1 B temp = 2M+2 regs. +// M=8 => 18 regs out of 32 available — no spills. + +template +FORCE_INLINE void gemm_micro_rvv_fma_Mx8_Ku4( + const float* __restrict A, // [M x K] + const kv_cache_t* __restrict B, // [K x 8] + float* __restrict C, // [M x 8] + int64_t lda, int64_t ldb, int64_t ldc, int32_t K, bool accumulate) { + static_assert(1 <= M && M <= 8, "M must be in [1,8]"); + + constexpr size_t vl = 8; + + // helpers for per-M codegen + #define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7) + #define IF_M(i) if constexpr (M > (i)) + + // A row base pointers + #define DECL_A(i) const float* a##i = A + (i) * lda; + ROWS_APPLY(DECL_A) + #undef DECL_A + + // declare one m2 accumulator per row + #define DECL_ACC(i) fixed_vfloat32m2_t acc##i; + ROWS_APPLY(DECL_ACC) + #undef DECL_ACC + + // initialize accumulators + #define INIT_ACC(i) \ + IF_M(i) { \ + if (accumulate) { \ + acc##i = __riscv_vle32_v_f32m2(C + (i) * ldc, vl); \ + } else { \ + acc##i = __riscv_vfmv_v_f_f32m2(0.f, vl); \ + } \ + } + ROWS_APPLY(INIT_ACC) + #undef INIT_ACC + + int32_t k = 0; + + // K unrolled by 4 + for (; k + 3 < K; k += 4) { + // k + 0 + { + fixed_vfloat32m2_t b = + load_row8_B_as_f32(B + (int64_t)(k + 0) * ldb); + #define STEP_K0(i) \ + IF_M(i) { \ + acc##i = __riscv_vfmacc_vf_f32m2(acc##i, *(a##i + k + 0), b, vl); \ + } + ROWS_APPLY(STEP_K0) + #undef STEP_K0 + } + // k + 1 + { + fixed_vfloat32m2_t b = + load_row8_B_as_f32(B + (int64_t)(k + 1) * ldb); + #define STEP_K1(i) \ + IF_M(i) { \ + acc##i = __riscv_vfmacc_vf_f32m2(acc##i, *(a##i + k + 1), b, vl); \ + } + ROWS_APPLY(STEP_K1) + #undef STEP_K1 + } + // k + 2 + { + fixed_vfloat32m2_t b = + load_row8_B_as_f32(B + (int64_t)(k + 2) * ldb); + #define STEP_K2(i) \ + IF_M(i) { \ + acc##i = __riscv_vfmacc_vf_f32m2(acc##i, *(a##i + k + 2), b, vl); \ + } + ROWS_APPLY(STEP_K2) + #undef STEP_K2 + } + // k + 3 + { + fixed_vfloat32m2_t b = + load_row8_B_as_f32(B + (int64_t)(k + 3) * ldb); + #define STEP_K3(i) \ + IF_M(i) { \ + acc##i = __riscv_vfmacc_vf_f32m2(acc##i, *(a##i + k + 3), b, vl); \ + } + ROWS_APPLY(STEP_K3) + #undef STEP_K3 + } + } + + // K tail + for (; k < K; ++k) { + fixed_vfloat32m2_t b = load_row8_B_as_f32(B + (int64_t)k * ldb); + #define TAIL_ROW(i) \ + IF_M(i) { acc##i = __riscv_vfmacc_vf_f32m2(acc##i, *(a##i + k), b, vl); } + ROWS_APPLY(TAIL_ROW) + #undef TAIL_ROW + } + + // store accumulators to C + #define STORE_ROW(i) \ + IF_M(i) { __riscv_vse32_v_f32m2(C + (i) * ldc, acc##i, vl); } + ROWS_APPLY(STORE_ROW) + #undef STORE_ROW + + #undef ROWS_APPLY + #undef IF_M +} + +// ============================================================================ +// Macro kernel: dispatch M tiles of {8,4,2,1}, step N by 8 +// ============================================================================ + +template +FORCE_INLINE void gemm_macro_rvv_fma_Mx8_Ku4(const float* __restrict A, + const kv_cache_t* __restrict B, + float* __restrict C, int32_t M, + int32_t K, int64_t lda, + int64_t ldb, int64_t ldc, + bool accumulate) { + static_assert(N % 8 == 0, "N must be a multiple of 8"); + for (int32_t m = 0; m < M;) { + int32_t mb = (M - m >= 8) ? 8 : (M - m >= 4) ? 4 : (M - m >= 2) ? 2 : 1; + const float* Ab = A + m * lda; + float* Cb = C + m * ldc; + + for (int32_t n = 0; n < N; n += 8) { + const kv_cache_t* Bn = B + n; + float* Cn = Cb + n; + switch (mb) { + case 8: + gemm_micro_rvv_fma_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, + K, accumulate); + break; + case 4: + gemm_micro_rvv_fma_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, + K, accumulate); + break; + case 2: + gemm_micro_rvv_fma_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, + K, accumulate); + break; + default: + gemm_micro_rvv_fma_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, + K, accumulate); + break; + } + } + m += mb; + } +} + +// ============================================================================ +// TileGemm wrapper — plugs into AttentionMainLoop +// ============================================================================ + +template +class TileGemmRVV { + public: + template + FORCE_INLINE static void gemm(const int32_t m_size, + float* __restrict__ a_tile, + kv_cache_t* __restrict__ b_tile, + float* __restrict__ c_tile, const int64_t lda, + const int64_t ldb, const int64_t ldc, + const int32_t block_size, + const int32_t dynamic_k_size, + const bool accum_c) { + if constexpr (phase == AttentionGemmPhase::QK) { + gemm_macro_rvv_fma_Mx8_Ku4( + a_tile, b_tile, c_tile, m_size, k_size, lda, ldb, ldc, accum_c); + } else { + gemm_macro_rvv_fma_Mx8_Ku4( + a_tile, b_tile, c_tile, m_size, dynamic_k_size, lda, ldb, ldc, + accum_c); + } + } +}; + +} // namespace + +// ============================================================================ +// AttentionImpl — mirrors ISA::NEON specialization +// ============================================================================ + +template +class AttentionImpl { + public: + using query_t = scalar_t; + using q_buffer_t = float; + using kv_cache_t = scalar_t; + using logits_buffer_t = float; + using partial_output_buffer_t = float; + using prob_buffer_t = float; + + constexpr static int64_t BlockSizeAlignment = BLOCK_SIZE_ALIGNMENT; + constexpr static int64_t HeadDimAlignment = HEAD_SIZE_ALIGNMENT; + constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER; + constexpr static int64_t HeadDim = head_dim; + constexpr static ISA ISAType = ISA::RVV; + constexpr static bool scale_on_logits = false; + + static_assert(HeadDim % HeadDimAlignment == 0); + static_assert(HeadDimAlignment % 8 == 0); + static_assert(BlockSizeAlignment % 8 == 0); + + public: + template