[CPU][RISC-V] Add RVV-optimized attention kernels for RISC-V Vector Extension (#40119)
Signed-off-by: liuyudong <liuyudong@iscas.ac.cn> Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
lyd1992
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -156,4 +156,20 @@
|
||||
|
||||
#endif // __aarch64__
|
||||
|
||||
// RISC-V RVV
|
||||
#ifdef __riscv_v
|
||||
#include <riscv_vector.h>
|
||||
|
||||
#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
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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 <riscv_vector.h>. 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 <riscv_vector.h>
|
||||
#include <type_traits>
|
||||
|
||||
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 <typename kv_cache_t>
|
||||
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<float>(const float* p) {
|
||||
return __riscv_vle32_v_f32m2(p, 8);
|
||||
}
|
||||
|
||||
template <>
|
||||
FORCE_INLINE fixed_vfloat32m2_t
|
||||
load_row8_B_as_f32<c10::Half>(const c10::Half* p) {
|
||||
#ifdef __riscv_zvfh
|
||||
fixed_vfloat16m1_t h =
|
||||
__riscv_vle16_v_f16m1(reinterpret_cast<const _Float16*>(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<float>(p[i]);
|
||||
}
|
||||
return __riscv_vle32_v_f32m2(tmp, 8);
|
||||
#endif
|
||||
}
|
||||
|
||||
template <>
|
||||
FORCE_INLINE fixed_vfloat32m2_t
|
||||
load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p) {
|
||||
#ifdef __riscv_zvfbfmin
|
||||
fixed_vbfloat16m1_t bf =
|
||||
__riscv_vle16_v_bf16m1(reinterpret_cast<const __bf16*>(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<const uint16_t*>(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 <int32_t M, typename kv_cache_t>
|
||||
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<kv_cache_t>(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<kv_cache_t>(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<kv_cache_t>(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<kv_cache_t>(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<kv_cache_t>(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 <int32_t N, typename kv_cache_t>
|
||||
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 <typename kv_cache_t>
|
||||
class TileGemmRVV {
|
||||
public:
|
||||
template <AttentionGemmPhase phase, int32_t k_size>
|
||||
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<BLOCK_SIZE_ALIGNMENT, kv_cache_t>(
|
||||
a_tile, b_tile, c_tile, m_size, k_size, lda, ldb, ldc, accum_c);
|
||||
} else {
|
||||
gemm_macro_rvv_fma_Mx8_Ku4<HEAD_SIZE_ALIGNMENT, kv_cache_t>(
|
||||
a_tile, b_tile, c_tile, m_size, dynamic_k_size, lda, ldb, ldc,
|
||||
accum_c);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
// ============================================================================
|
||||
// AttentionImpl<ISA::RVV> — mirrors ISA::NEON specialization
|
||||
// ============================================================================
|
||||
|
||||
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
|
||||
class AttentionImpl<ISA::RVV, scalar_t, head_dim, kv_cache_scalar_t> {
|
||||
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 <template <typename tile_gemm_t> typename attention>
|
||||
FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
|
||||
attention<TileGemmRVV<kv_cache_t>> attention_iteration;
|
||||
attention_iteration(CPU_ATTENTION_PARAMS);
|
||||
}
|
||||
|
||||
constexpr static int64_t k_cache_token_group_stride(
|
||||
const int32_t block_size) {
|
||||
return BlockSizeAlignment;
|
||||
}
|
||||
|
||||
constexpr static int64_t v_cache_token_group_stride(
|
||||
const int32_t block_size) {
|
||||
return head_dim * BlockSizeAlignment;
|
||||
}
|
||||
|
||||
constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
|
||||
return HeadDimAlignment;
|
||||
}
|
||||
|
||||
static void copy_q_heads_tile(scalar_t* __restrict__ src,
|
||||
float* __restrict__ q_buffer,
|
||||
const int32_t q_num,
|
||||
const int32_t q_heads_per_kv,
|
||||
const int64_t q_num_stride,
|
||||
const int64_t q_head_stride, float scale) {
|
||||
static_assert(head_dim % 16 == 0);
|
||||
constexpr int32_t unroll_size = head_dim / 16;
|
||||
using load_vec_t = typename VecTypeTrait<scalar_t>::vec_t;
|
||||
|
||||
vec_op::FP32Vec16 scale_vec(scale);
|
||||
for (int32_t q_num_idx = 0; q_num_idx < q_num; ++q_num_idx) {
|
||||
for (int32_t q_head_idx = 0; q_head_idx < q_heads_per_kv; ++q_head_idx) {
|
||||
scalar_t* __restrict__ curr_q =
|
||||
src + q_num_idx * q_num_stride + q_head_idx * q_head_stride;
|
||||
float* __restrict__ curr_q_buffer =
|
||||
q_buffer + q_num_idx * q_heads_per_kv * head_dim +
|
||||
q_head_idx * head_dim;
|
||||
|
||||
vec_op::unroll_loop<int32_t, unroll_size>([&](int32_t i) {
|
||||
load_vec_t vec(curr_q);
|
||||
vec_op::FP32Vec16 fp32_vec(vec);
|
||||
fp32_vec = fp32_vec * scale_vec;
|
||||
fp32_vec.save(curr_q_buffer);
|
||||
|
||||
curr_q += 16;
|
||||
curr_q_buffer += 16;
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static void reshape_and_cache(
|
||||
const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
|
||||
scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
|
||||
const int64_t* __restrict__ slot_mapping, const int64_t token_num,
|
||||
const int64_t key_token_num_stride, const int64_t value_token_num_stride,
|
||||
const int64_t head_num, const int64_t key_head_num_stride,
|
||||
const int64_t value_head_num_stride, const int64_t num_blocks,
|
||||
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
|
||||
const int64_t block_size, const int64_t block_size_stride,
|
||||
const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
|
||||
#pragma omp parallel for collapse(2)
|
||||
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
|
||||
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
|
||||
const int64_t pos = slot_mapping[token_idx];
|
||||
if (pos < 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const int64_t block_idx = pos / block_size;
|
||||
const int64_t block_offset = pos % block_size;
|
||||
{
|
||||
// Write Key (transpose to column-major: [head_dim, block_size])
|
||||
const scalar_t* key_start_ptr = key +
|
||||
token_idx * key_token_num_stride +
|
||||
head_idx * key_head_num_stride;
|
||||
scalar_t* key_cache_start_ptr =
|
||||
key_cache + block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride + block_offset;
|
||||
|
||||
// Strided vector store for efficient transpose.
|
||||
// Load contiguous key elements, store with stride = block_size.
|
||||
{
|
||||
const ptrdiff_t byte_stride = block_size * sizeof(scalar_t);
|
||||
int64_t i = 0;
|
||||
for (; i < head_dim;) {
|
||||
size_t vl;
|
||||
if constexpr (std::is_same_v<scalar_t, float>) {
|
||||
vl = __riscv_vsetvl_e32m2(head_dim - i);
|
||||
vfloat32m2_t v = __riscv_vle32_v_f32m2(
|
||||
reinterpret_cast<const float*>(key_start_ptr + i), vl);
|
||||
__riscv_vsse32_v_f32m2(
|
||||
reinterpret_cast<float*>(key_cache_start_ptr +
|
||||
i * block_size),
|
||||
byte_stride, v, vl);
|
||||
} else {
|
||||
// Half and BFloat16 are both 16-bit types
|
||||
vl = __riscv_vsetvl_e16m1(head_dim - i);
|
||||
vuint16m1_t v = __riscv_vle16_v_u16m1(
|
||||
reinterpret_cast<const uint16_t*>(key_start_ptr + i), vl);
|
||||
__riscv_vsse16_v_u16m1(
|
||||
reinterpret_cast<uint16_t*>(key_cache_start_ptr +
|
||||
i * block_size),
|
||||
byte_stride, v, vl);
|
||||
}
|
||||
i += vl;
|
||||
}
|
||||
}
|
||||
}
|
||||
{
|
||||
// Write Value (row-major: [block_size, head_dim])
|
||||
const scalar_t* value_start_ptr = value +
|
||||
token_idx * value_token_num_stride +
|
||||
head_idx * value_head_num_stride;
|
||||
scalar_t* value_cache_start_ptr =
|
||||
value_cache + block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride + block_offset * head_dim;
|
||||
std::memcpy(value_cache_start_ptr, value_start_ptr,
|
||||
sizeof(scalar_t) * head_dim);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace cpu_attention
|
||||
|
||||
#undef BLOCK_SIZE_ALIGNMENT
|
||||
#undef HEAD_SIZE_ALIGNMENT
|
||||
#undef MAX_Q_HEAD_NUM_PER_ITER
|
||||
|
||||
#endif // __riscv_v_min_vlen == 128
|
||||
|
||||
#endif // CPU_ATTN_RVV_HPP
|
||||
@@ -72,7 +72,7 @@ typedef RVVTYPE(vuint16, LMUL_512, _t) fixed_u16x32_t
|
||||
__attribute__((riscv_rvv_vector_bits(512)));
|
||||
|
||||
// bfloat16
|
||||
#ifdef RISCV_BF16_SUPPORT
|
||||
#ifdef __riscv_zvfbfmin
|
||||
typedef RVVTYPE(vbfloat16, LMUL_128, _t) fixed_bf16x8_t
|
||||
__attribute__((riscv_rvv_vector_bits(128)));
|
||||
typedef RVVTYPE(vbfloat16, LMUL_256, _t) fixed_bf16x16_t
|
||||
|
||||
@@ -15,8 +15,17 @@
|
||||
#include <torch/all.h>
|
||||
namespace vec_op {
|
||||
|
||||
// BFloat16 is always supported on RISC-V: natively when RISCV_BF16_SUPPORT
|
||||
// is defined, otherwise via the FP32-simulation fallback path.
|
||||
// FP8 KV cache is not supported on RISC-V. These tag types and the
|
||||
// corresponding BF16Vec32 stub constructors below exist solely so that
|
||||
// templates referencing vec_op::fp8_*_tag in their bodies (e.g. in
|
||||
// cpu_attn_vec.hpp) compile under GCC's -Wtemplate-body lookup. The
|
||||
// stubs are never instantiated by CPU_ATTN_DISPATCH on __riscv.
|
||||
struct fp8_e4m3_tag {};
|
||||
struct fp8_e5m2_tag {};
|
||||
|
||||
// BFloat16 is always supported on RISC-V: natively when __riscv_zvfbfmin
|
||||
// is defined (compiler-provided when -march includes zvfbfmin), otherwise
|
||||
// via the FP32-simulation fallback path.
|
||||
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
|
||||
@@ -106,7 +115,7 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
|
||||
// BF16 Implementation
|
||||
// ============================================================================
|
||||
|
||||
#ifdef RISCV_BF16_SUPPORT
|
||||
#ifdef __riscv_zvfbfmin
|
||||
|
||||
FORCE_INLINE fixed_u16x8_t bf16_to_u16(fixed_bf16x8_t v) {
|
||||
return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_128, _u16, LMUL_128)(v);
|
||||
@@ -185,6 +194,13 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
|
||||
|
||||
explicit BF16Vec32(fixed_bf16x32_t data) : reg(data) {};
|
||||
|
||||
// FP8 KV cache stubs: never instantiated on RISC-V (CPU_ATTN_DISPATCH
|
||||
// omits FP8 cases on __riscv); exist only so name lookup succeeds.
|
||||
explicit BF16Vec32(const uint8_t* ptr, fp8_e4m3_tag)
|
||||
: BF16Vec32(static_cast<const void*>(ptr)) {}
|
||||
explicit BF16Vec32(const uint8_t* ptr, fp8_e5m2_tag)
|
||||
: BF16Vec32(static_cast<const void*>(ptr)) {}
|
||||
|
||||
explicit BF16Vec32(const BF16Vec8& v) {
|
||||
fixed_u16x8_t u16_val = bf16_to_u16(v.reg);
|
||||
fixed_u16x32_t u16_combined =
|
||||
@@ -323,6 +339,13 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
|
||||
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp, 32);
|
||||
}
|
||||
|
||||
// FP8 KV cache stubs: never instantiated on RISC-V (CPU_ATTN_DISPATCH
|
||||
// omits FP8 cases on __riscv); exist only so name lookup succeeds.
|
||||
explicit BF16Vec32(const uint8_t* ptr, fp8_e4m3_tag)
|
||||
: BF16Vec32(static_cast<const void*>(ptr)) {}
|
||||
explicit BF16Vec32(const uint8_t* ptr, fp8_e5m2_tag)
|
||||
: BF16Vec32(static_cast<const void*>(ptr)) {}
|
||||
|
||||
explicit BF16Vec32(const BF16Vec8& v) {
|
||||
float tmp_small[8];
|
||||
RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp_small, v.reg_fp32, 8);
|
||||
@@ -410,7 +433,7 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
|
||||
explicit FP32Vec8(fixed_fp16x8_t v)
|
||||
: reg(RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(v, VEC_ELEM_NUM)) {};
|
||||
|
||||
#ifdef RISCV_BF16_SUPPORT
|
||||
#ifdef __riscv_zvfbfmin
|
||||
explicit FP32Vec8(fixed_bf16x8_t v)
|
||||
: reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(v, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec8(const BF16Vec8& v)
|
||||
@@ -608,7 +631,7 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
|
||||
explicit FP32Vec16(const FP16Vec16& v);
|
||||
|
||||
#ifdef RISCV_BF16_SUPPORT
|
||||
#ifdef __riscv_zvfbfmin
|
||||
explicit FP32Vec16(fixed_bf16x16_t v)
|
||||
: reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_512)(v, VEC_ELEM_NUM)) {};
|
||||
explicit FP32Vec16(const BF16Vec16& v)
|
||||
@@ -868,7 +891,7 @@ inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
|
||||
acc = acc.fma(a, b);
|
||||
}
|
||||
|
||||
#ifdef RISCV_BF16_SUPPORT
|
||||
#ifdef __riscv_zvfbfmin
|
||||
template <>
|
||||
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
|
||||
*ptr = static_cast<__bf16>(v);
|
||||
|
||||
@@ -20,7 +20,8 @@ ISA_TYPES = {
|
||||
"VEC16": 2,
|
||||
"NEON": 3,
|
||||
"VXE": 4,
|
||||
"VSX": 5,
|
||||
"RVV": 5,
|
||||
"VSX": 6,
|
||||
}
|
||||
|
||||
# KV cache index: 0 = auto (same as scalar_t), 1 = fp8_e4m3, 2 = fp8_e5m2
|
||||
@@ -38,7 +39,7 @@ KV_CACHE_CPP_TYPES = {
|
||||
}
|
||||
|
||||
# ISAs supported for head_dims divisible by 32
|
||||
ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE", "VSX"]
|
||||
ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE", "RVV", "VSX"]
|
||||
|
||||
# ISAs supported for head_dims divisible by 16 only
|
||||
ISA_FOR_16 = ["VEC16"]
|
||||
@@ -149,6 +150,15 @@ def generate_header_file() -> str:
|
||||
#include "cpu_attn_vxe.hpp"
|
||||
#endif
|
||||
|
||||
// cpu_attn_rvv.hpp is hardcoded to VLEN==128 (m1/m2 intrinsics, vl=8) and
|
||||
// itself includes <riscv_vector.h>, which is unavailable on scalar
|
||||
// (-march=rv64gc) builds. Gate the include the same way as the dispatch
|
||||
// macro below, so non-128 / scalar RISC-V builds skip it entirely.
|
||||
#if defined(__riscv) && defined(__riscv_v_min_vlen) && \
|
||||
__riscv_v_min_vlen == 128
|
||||
#include "cpu_attn_rvv.hpp"
|
||||
#endif
|
||||
|
||||
#ifdef __powerpc__
|
||||
#include "cpu_attn_vsx.hpp"
|
||||
#endif
|
||||
@@ -212,6 +222,23 @@ def generate_header_file() -> str:
|
||||
["VXE", "VEC", "VEC16"],
|
||||
fp8=False,
|
||||
)
|
||||
# RISC-V with RVV. cpu_attn_rvv.hpp is hardcoded to VLEN==128
|
||||
# (riscv_rvv_vector_bits(128) typedefs + vl=8 m1/m2 intrinsics), so
|
||||
# we split the dispatch into two top-level branches: VLEN==128 builds
|
||||
# get the full RVV+VEC+VEC16 case set, other VLEN builds get a
|
||||
# VEC/VEC16-only fallback. Preprocessor directives cannot appear
|
||||
# inside a #define body, so this duplication is necessary.
|
||||
header += _macro_block(
|
||||
"#elif defined(__riscv) && defined(__riscv_v_min_vlen) "
|
||||
"&& __riscv_v_min_vlen == 128",
|
||||
["RVV", "VEC", "VEC16"],
|
||||
fp8=False,
|
||||
)
|
||||
header += _macro_block(
|
||||
"#elif defined(__riscv)",
|
||||
["VEC", "VEC16"],
|
||||
fp8=False,
|
||||
)
|
||||
header += _macro_block(
|
||||
"#elif defined(__powerpc__)",
|
||||
["VSX", "VEC", "VEC16"],
|
||||
@@ -233,8 +260,8 @@ def generate_header_file() -> str:
|
||||
fp8=False,
|
||||
)
|
||||
header += (
|
||||
"#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / "
|
||||
"__s390x__ / __powerpc__ */\n\n"
|
||||
"#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ /"
|
||||
" __riscv / __powerpc__ */\n\n"
|
||||
"#endif // CPU_ATTN_DISPATCH_GENERATED_H\n"
|
||||
)
|
||||
|
||||
|
||||
@@ -49,15 +49,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
|
||||
@@ -583,6 +582,52 @@ def test_varlen_with_paged_kv_normal_neon(
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3"])
|
||||
@pytest.mark.parametrize("seq_lens", SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("block_size", [96, 128])
|
||||
@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
|
||||
@pytest.mark.parametrize("dtype", QTYPES)
|
||||
@pytest.mark.parametrize("soft_cap", [None])
|
||||
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
|
||||
@pytest.mark.parametrize("use_alibi", [False])
|
||||
@pytest.mark.parametrize("use_sink", [False])
|
||||
@pytest.mark.parametrize("isa", ["rvv"])
|
||||
@pytest.mark.skipif(
|
||||
current_platform.get_cpu_architecture() != CpuArchEnum.RISCV,
|
||||
reason="Not a RISC-V CPU.",
|
||||
)
|
||||
def test_varlen_with_paged_kv_normal_rvv(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
sliding_window: int | None,
|
||||
dtype: torch.dtype,
|
||||
block_size: int,
|
||||
soft_cap: float | None,
|
||||
num_blocks: int,
|
||||
use_alibi: bool,
|
||||
use_sink: bool,
|
||||
isa: str,
|
||||
kv_cache_dtype: str,
|
||||
) -> None:
|
||||
varlen_with_paged_kv(
|
||||
seq_lens=seq_lens,
|
||||
num_heads=num_heads,
|
||||
head_size=head_size,
|
||||
sliding_window=sliding_window,
|
||||
dtype=dtype,
|
||||
block_size=block_size,
|
||||
soft_cap=soft_cap,
|
||||
num_blocks=num_blocks,
|
||||
use_alibi=use_alibi,
|
||||
use_sink=use_sink,
|
||||
isa=isa,
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3"])
|
||||
@pytest.mark.parametrize("seq_lens", SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import functools
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, ClassVar
|
||||
|
||||
@@ -34,6 +35,7 @@ _CPU_ARCH_PREFER_MIXED_BATCH = (
|
||||
CpuArchEnum.X86,
|
||||
CpuArchEnum.ARM,
|
||||
CpuArchEnum.S390X,
|
||||
CpuArchEnum.RISCV,
|
||||
CpuArchEnum.POWERPC,
|
||||
)
|
||||
|
||||
@@ -506,6 +508,29 @@ def _make_sliding_window_bias(
|
||||
return attn_biases
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def _riscv_supports_rvv_vlen128() -> bool:
|
||||
"""Whether the C++ RVV attention path (hardcoded to VLEN==128) is usable.
|
||||
|
||||
The kernel in csrc/cpu/cpu_attn_rvv.hpp uses riscv_rvv_vector_bits(128)
|
||||
typedefs and m1/m2 intrinsics with vl=8; CMake's auto-detect picks the
|
||||
largest zvl<N>b advertised by /proc/cpuinfo, so the binary contains the
|
||||
RVV path only when the build host advertised exactly zvl128b. Mirror
|
||||
that here so the Python dispatch doesn't request ISA::RVV on builds
|
||||
where it wasn't compiled in (would TORCH_CHECK at first attention call).
|
||||
"""
|
||||
try:
|
||||
with open("/proc/cpuinfo") as f:
|
||||
cpuinfo = f.read()
|
||||
except OSError:
|
||||
return False
|
||||
if "zvl128b" not in cpuinfo:
|
||||
return False
|
||||
# CMake auto-detect picks the largest advertised VLEN; if the host
|
||||
# advertises zvl256b or higher, the build skipped the RVV-128 path.
|
||||
return all(f"zvl{n}b" not in cpuinfo for n in (256, 512, 1024))
|
||||
|
||||
|
||||
def _get_attn_isa(
|
||||
dtype: torch.dtype,
|
||||
block_size: int,
|
||||
@@ -523,6 +548,7 @@ def _get_attn_isa(
|
||||
arch = current_platform.get_cpu_architecture()
|
||||
supports_arm = arch == CpuArchEnum.ARM
|
||||
supports_vxe = arch == CpuArchEnum.S390X
|
||||
supports_riscv = arch == CpuArchEnum.RISCV
|
||||
supports_vsx = arch == CpuArchEnum.POWERPC
|
||||
supports_avx512 = torch.cpu._is_avx512_supported()
|
||||
if fp8_kv and not supports_amx and not supports_avx512:
|
||||
@@ -535,6 +561,8 @@ def _get_attn_isa(
|
||||
if supports_arm:
|
||||
# support ARM NEON FMLA and BFMMLA (bf16) for block size 32
|
||||
return "neon"
|
||||
elif supports_riscv and _riscv_supports_rvv_vlen128():
|
||||
return "rvv"
|
||||
elif supports_vxe:
|
||||
return "vxe"
|
||||
elif supports_vsx:
|
||||
|
||||
Reference in New Issue
Block a user