@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f351455f0f
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 3, 15:15 UTC

[CPU][RISC-V] Add RVV-optimized attention kernels for RISC-V Vector Extension (#40119)

[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>

lyd1992 committed May 15, 2026 at 12:08 UTC f351455f0f06c17c723341bc6d78c3d4470b5b09
11 files changed +624 -33
benchmarks/kernels/cpu/benchmark_cpu_attn.py
+9 -11
@@ -12,7 +12,6 @@ from vllm._custom_ops import (
12 cpu_attn_get_scheduler_metadata,
13 cpu_attn_reshape_and_cache,
14 )
15 -from vllm.platforms import CpuArchEnum, current_platform
15 from vllm.utils.argparse_utils import FlexibleArgumentParser
16 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
17 from vllm.v1.attention.backends.cpu_attn import CPUAttentionBackend, _get_attn_isa
@@ -22,15 +21,14 @@ def get_attn_isa(
21 block_size: int | None = None,
22 dtype: torch.dtype | None = None,
23 ):
25 - if block_size and dtype:
26 - return _get_attn_isa(dtype, block_size)
27 - else:
28 - if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
29 - return "neon"
30 - elif torch.cpu._is_amx_tile_supported():
31 - return "amx"
32 - else:
33 - return "vec"
24 + # Delegate to _get_attn_isa so the fallback path applies the same arch
25 + # gating (e.g. RISC-V RVV is only chosen when the build's hardcoded
26 + # VLEN=128 kernel is actually present; on VLEN=256 / scalar hosts it
27 + # correctly falls through to vec/vec16).
28 + return _get_attn_isa(
29 + dtype if dtype is not None else torch.bfloat16,
30 + block_size if block_size else 32,
31 + )
32
33
34 # rand number generation takes too much time, cache rand tensors
@@ -235,7 +233,7 @@ if __name__ == "__main__":
233 )
234 parser.add_argument("--use-sink", action="store_true")
235 parser.add_argument(
238 - "--isa", type=str, choices=["vec", "neon", "amx", "vec16"], default=None
236 + "--isa", type=str, choices=["vec", "neon", "amx", "vec16", "rvv"], default=None
237 )
238 parser.add_argument("--seed", type=int, default=0)
239 parser.add_argument("--iters", type=int, default=20)
cmake/cpu_extension.cmake
+3 -1
@@ -195,10 +195,12 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
195 endif()
196 if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0)
197 message(STATUS "RISC-V RVV VLEN=${VLLM_RVV_VLEN}")
198 + # Sources gate FP16/BF16 paths on the compiler-provided
199 + # __riscv_zvfh / __riscv_zvfbfmin macros, which GCC and clang
200 + # define automatically when those extensions appear in -march.
201 if(RVV_BF16_FOUND)
202 message(STATUS "BF16 extension detected")
203 set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
201 - add_compile_definitions(RISCV_BF16_SUPPORT)
204 elseif(RVV_FP16_FOUND)
205 message(WARNING "BF16 functionality is not available")
206 set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
csrc/cpu/cpu_arch_macros.h
+16
@@ -156,4 +156,20 @@
156
157 #endif // __aarch64__
158
159 +// RISC-V RVV
160 +#ifdef __riscv_v
161 + #include <riscv_vector.h>
162 +
163 + #ifdef __riscv_zihintpause
164 + #define FAST_SPINNING __riscv_pause();
165 + #endif
166 +
167 + // FP32Vec16::exp() in cpu_types_riscv.hpp already implements the full
168 + // polynomial approximation for RVV, so we simply delegate to it.
169 + #define DEFINE_FAST_EXP \
170 + auto fast_exp = [&](const vec_op::FP32Vec16& vec) \
171 + __attribute__((always_inline)) { return vec.exp(); };
172 +
173 +#endif // __riscv_v
174 +
175 #endif
csrc/cpu/cpu_attn.cpp
+4
@@ -29,6 +29,8 @@ torch::Tensor get_scheduler_metadata(
29 isa = cpu_attention::ISA::NEON;
30 } else if (isa_hint == "vxe") {
31 isa = cpu_attention::ISA::VXE;
32 + } else if (isa_hint == "rvv") {
33 + isa = cpu_attention::ISA::RVV;
34 } else if (isa_hint == "vsx") {
35 isa = cpu_attention::ISA::VSX;
36 } else {
@@ -131,6 +133,8 @@ void cpu_attn_reshape_and_cache(
133 return cpu_attention::ISA::NEON;
134 } else if (isa == "vxe") {
135 return cpu_attention::ISA::VXE;
136 + } else if (isa == "rvv") {
137 + return cpu_attention::ISA::RVV;
138 } else if (isa == "vsx") {
139 return cpu_attention::ISA::VSX;
140 } else {
csrc/cpu/cpu_attn_impl.hpp
+4 -1
@@ -12,7 +12,7 @@
12 #include "cpu/utils.hpp"
13
14 namespace cpu_attention {
15 -enum class ISA { AMX, VEC, VEC16, NEON, VXE, VSX };
15 +enum class ISA { AMX, VEC, VEC16, NEON, VXE, RVV, VSX };
16
17 // Mirrors csrc/attention/dtype_fp8.cuh Fp8KVCacheDataType exactly.
18 enum class Fp8KVCacheDataType {
@@ -164,6 +164,9 @@ struct AttentionMetadata {
164 case ISA::VXE:
165 ss << "VXE, ";
166 break;
167 + case ISA::RVV:
168 + ss << "RVV, ";
169 + break;
170 case ISA::VSX:
171 ss << "VSX, ";
172 break;
csrc/cpu/cpu_attn_rvv.hpp new
+445
@@ -0,0 +1,445 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +#ifndef CPU_ATTN_RVV_HPP
5 +#define CPU_ATTN_RVV_HPP
6 +
7 +// This kernel is currently hardcoded to VLEN=128 (m1/m2 intrinsics, vl=8).
8 +// The fixed-width typedefs below use `riscv_rvv_vector_bits(128)`, which
9 +// only matches `vfloat16m1_t`/`vuint16m1_t` register layout when VLEN==128;
10 +// at VLEN>=256 those typedefs fail to compile. Scalar RISC-V builds
11 +// (-march=rv64gc) additionally don't have <riscv_vector.h>. For both
12 +// cases we omit the file entirely and let the dispatcher fall back to the
13 +// scalar VEC / VEC16 implementations. TODO: migrate to RVVI() macros +
14 +// semantic names in cpu_types_riscv_defs.hpp to support VLEN>=256 natively.
15 +#if defined(__riscv_v_min_vlen) && __riscv_v_min_vlen == 128
16 +
17 + #include "cpu_attn_impl.hpp"
18 + #include <riscv_vector.h>
19 + #include <type_traits>
20 +
21 +namespace cpu_attention {
22 +
23 +namespace {
24 +
25 +// File-local concrete-LMUL typedefs. The shared _defs.hpp exposes
26 +// VLEN-independent semantic names (fixed_fp32x8_t, fixed_fp16x8_t, ...),
27 +// but this kernel is currently hardcoded to VLEN=128 (m1/m2 intrinsics),
28 +// so keep the legacy concrete aliases scoped to this file.
29 +typedef vfloat16m1_t fixed_vfloat16m1_t
30 + __attribute__((riscv_rvv_vector_bits(128)));
31 +typedef vfloat32m2_t fixed_vfloat32m2_t
32 + __attribute__((riscv_rvv_vector_bits(256)));
33 +typedef vuint16m1_t fixed_vuint16m1_t
34 + __attribute__((riscv_rvv_vector_bits(128)));
35 +typedef vuint32m2_t fixed_vuint32m2_t
36 + __attribute__((riscv_rvv_vector_bits(256)));
37 + #ifdef __riscv_zvfbfmin
38 +typedef vbfloat16m1_t fixed_vbfloat16m1_t
39 + __attribute__((riscv_rvv_vector_bits(128)));
40 + #endif
41 +
42 + #define BLOCK_SIZE_ALIGNMENT 32
43 + #define HEAD_SIZE_ALIGNMENT 32
44 + #define MAX_Q_HEAD_NUM_PER_ITER 16
45 +
46 +// ============================================================================
47 +// B-matrix row loading: load 8 elements as FP32 (using m2 LMUL at VLEN=128)
48 +// ============================================================================
49 +
50 +template <typename kv_cache_t>
51 +FORCE_INLINE fixed_vfloat32m2_t load_row8_B_as_f32(const kv_cache_t* p);
52 +
53 +template <>
54 +FORCE_INLINE fixed_vfloat32m2_t load_row8_B_as_f32<float>(const float* p) {
55 + return __riscv_vle32_v_f32m2(p, 8);
56 +}
57 +
58 +template <>
59 +FORCE_INLINE fixed_vfloat32m2_t
60 +load_row8_B_as_f32<c10::Half>(const c10::Half* p) {
61 + #ifdef __riscv_zvfh
62 + fixed_vfloat16m1_t h =
63 + __riscv_vle16_v_f16m1(reinterpret_cast<const _Float16*>(p), 8);
64 + return __riscv_vfwcvt_f_f_v_f32m2(h, 8);
65 + #else
66 + // Fallback for hardware without Zvfh: scalar half->float conversion.
67 + // c10::Half provides operator float() so this is correct on any RVV CPU
68 + // that has only the base V extension. Slower than the Zvfh path, but
69 + // keeps the kernel buildable on Zvfhmin-only / no-fp16 hardware.
70 + alignas(16) float tmp[8];
71 + for (int i = 0; i < 8; ++i) {
72 + tmp[i] = static_cast<float>(p[i]);
73 + }
74 + return __riscv_vle32_v_f32m2(tmp, 8);
75 + #endif
76 +}
77 +
78 +template <>
79 +FORCE_INLINE fixed_vfloat32m2_t
80 +load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p) {
81 + #ifdef __riscv_zvfbfmin
82 + fixed_vbfloat16m1_t bf =
83 + __riscv_vle16_v_bf16m1(reinterpret_cast<const __bf16*>(p), 8);
84 + return __riscv_vfwcvtbf16_f_f_v_f32m2(bf, 8);
85 + #else
86 + // Fallback: load as uint16, zero-extend to uint32, shift left by 16
87 + fixed_vuint16m1_t raw =
88 + __riscv_vle16_v_u16m1(reinterpret_cast<const uint16_t*>(p), 8);
89 + fixed_vuint32m2_t wide = __riscv_vzext_vf2_u32m2(raw, 8);
90 + fixed_vuint32m2_t shifted = __riscv_vsll_vx_u32m2(wide, 16, 8);
91 + return __riscv_vreinterpret_v_u32m2_f32m2(shifted);
92 + #endif
93 +}
94 +
95 +// ============================================================================
96 +// Micro kernel: Mx8 tile, K unrolled by 4, RVV scalar-broadcast FMA
97 +// ============================================================================
98 +//
99 +// NEON uses vfmaq_laneq_f32 (lane-indexed FMA from a preloaded A vector).
100 +// RVV has no lane-indexed FMA; instead we load A elements as scalars and
101 +// use __riscv_vfmacc_vf (scalar * vector + accumulator), which is equally
102 +// efficient and avoids the need for vrgather/vslidedown.
103 +//
104 +// At VLEN=128, m2 holds 8 x FP32, matching the 8-column tile width.
105 +// Register budget: M accumulators (m2 each) + 1 B temp = 2M+2 regs.
106 +// M=8 => 18 regs out of 32 available — no spills.
107 +
108 +template <int32_t M, typename kv_cache_t>
109 +FORCE_INLINE void gemm_micro_rvv_fma_Mx8_Ku4(
110 + const float* __restrict A, // [M x K]
111 + const kv_cache_t* __restrict B, // [K x 8]
112 + float* __restrict C, // [M x 8]
113 + int64_t lda, int64_t ldb, int64_t ldc, int32_t K, bool accumulate) {
114 + static_assert(1 <= M && M <= 8, "M must be in [1,8]");
115 +
116 + constexpr size_t vl = 8;
117 +
118 + // helpers for per-M codegen
119 + #define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
120 + #define IF_M(i) if constexpr (M > (i))
121 +
122 + // A row base pointers
123 + #define DECL_A(i) const float* a##i = A + (i) * lda;
124 + ROWS_APPLY(DECL_A)
125 + #undef DECL_A
126 +
127 + // declare one m2 accumulator per row
128 + #define DECL_ACC(i) fixed_vfloat32m2_t acc##i;
129 + ROWS_APPLY(DECL_ACC)
130 + #undef DECL_ACC
131 +
132 + // initialize accumulators
133 + #define INIT_ACC(i) \
134 + IF_M(i) { \
135 + if (accumulate) { \
136 + acc##i = __riscv_vle32_v_f32m2(C + (i) * ldc, vl); \
137 + } else { \
138 + acc##i = __riscv_vfmv_v_f_f32m2(0.f, vl); \
139 + } \
140 + }
141 + ROWS_APPLY(INIT_ACC)
142 + #undef INIT_ACC
143 +
144 + int32_t k = 0;
145 +
146 + // K unrolled by 4
147 + for (; k + 3 < K; k += 4) {
148 + // k + 0
149 + {
150 + fixed_vfloat32m2_t b =
151 + load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 0) * ldb);
152 + #define STEP_K0(i) \
153 + IF_M(i) { \
154 + acc##i = __riscv_vfmacc_vf_f32m2(acc##i, *(a##i + k + 0), b, vl); \
155 + }
156 + ROWS_APPLY(STEP_K0)
157 + #undef STEP_K0
158 + }
159 + // k + 1
160 + {
161 + fixed_vfloat32m2_t b =
162 + load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 1) * ldb);
163 + #define STEP_K1(i) \
164 + IF_M(i) { \
165 + acc##i = __riscv_vfmacc_vf_f32m2(acc##i, *(a##i + k + 1), b, vl); \
166 + }
167 + ROWS_APPLY(STEP_K1)
168 + #undef STEP_K1
169 + }
170 + // k + 2
171 + {
172 + fixed_vfloat32m2_t b =
173 + load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 2) * ldb);
174 + #define STEP_K2(i) \
175 + IF_M(i) { \
176 + acc##i = __riscv_vfmacc_vf_f32m2(acc##i, *(a##i + k + 2), b, vl); \
177 + }
178 + ROWS_APPLY(STEP_K2)
179 + #undef STEP_K2
180 + }
181 + // k + 3
182 + {
183 + fixed_vfloat32m2_t b =
184 + load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 3) * ldb);
185 + #define STEP_K3(i) \
186 + IF_M(i) { \
187 + acc##i = __riscv_vfmacc_vf_f32m2(acc##i, *(a##i + k + 3), b, vl); \
188 + }
189 + ROWS_APPLY(STEP_K3)
190 + #undef STEP_K3
191 + }
192 + }
193 +
194 + // K tail
195 + for (; k < K; ++k) {
196 + fixed_vfloat32m2_t b = load_row8_B_as_f32<kv_cache_t>(B + (int64_t)k * ldb);
197 + #define TAIL_ROW(i) \
198 + IF_M(i) { acc##i = __riscv_vfmacc_vf_f32m2(acc##i, *(a##i + k), b, vl); }
199 + ROWS_APPLY(TAIL_ROW)
200 + #undef TAIL_ROW
201 + }
202 +
203 + // store accumulators to C
204 + #define STORE_ROW(i) \
205 + IF_M(i) { __riscv_vse32_v_f32m2(C + (i) * ldc, acc##i, vl); }
206 + ROWS_APPLY(STORE_ROW)
207 + #undef STORE_ROW
208 +
209 + #undef ROWS_APPLY
210 + #undef IF_M
211 +}
212 +
213 +// ============================================================================
214 +// Macro kernel: dispatch M tiles of {8,4,2,1}, step N by 8
215 +// ============================================================================
216 +
217 +template <int32_t N, typename kv_cache_t>
218 +FORCE_INLINE void gemm_macro_rvv_fma_Mx8_Ku4(const float* __restrict A,
219 + const kv_cache_t* __restrict B,
220 + float* __restrict C, int32_t M,
221 + int32_t K, int64_t lda,
222 + int64_t ldb, int64_t ldc,
223 + bool accumulate) {
224 + static_assert(N % 8 == 0, "N must be a multiple of 8");
225 + for (int32_t m = 0; m < M;) {
226 + int32_t mb = (M - m >= 8) ? 8 : (M - m >= 4) ? 4 : (M - m >= 2) ? 2 : 1;
227 + const float* Ab = A + m * lda;
228 + float* Cb = C + m * ldc;
229 +
230 + for (int32_t n = 0; n < N; n += 8) {
231 + const kv_cache_t* Bn = B + n;
232 + float* Cn = Cb + n;
233 + switch (mb) {
234 + case 8:
235 + gemm_micro_rvv_fma_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
236 + K, accumulate);
237 + break;
238 + case 4:
239 + gemm_micro_rvv_fma_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
240 + K, accumulate);
241 + break;
242 + case 2:
243 + gemm_micro_rvv_fma_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
244 + K, accumulate);
245 + break;
246 + default:
247 + gemm_micro_rvv_fma_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
248 + K, accumulate);
249 + break;
250 + }
251 + }
252 + m += mb;
253 + }
254 +}
255 +
256 +// ============================================================================
257 +// TileGemm wrapper — plugs into AttentionMainLoop
258 +// ============================================================================
259 +
260 +template <typename kv_cache_t>
261 +class TileGemmRVV {
262 + public:
263 + template <AttentionGemmPhase phase, int32_t k_size>
264 + FORCE_INLINE static void gemm(const int32_t m_size,
265 + float* __restrict__ a_tile,
266 + kv_cache_t* __restrict__ b_tile,
267 + float* __restrict__ c_tile, const int64_t lda,
268 + const int64_t ldb, const int64_t ldc,
269 + const int32_t block_size,
270 + const int32_t dynamic_k_size,
271 + const bool accum_c) {
272 + if constexpr (phase == AttentionGemmPhase::QK) {
273 + gemm_macro_rvv_fma_Mx8_Ku4<BLOCK_SIZE_ALIGNMENT, kv_cache_t>(
274 + a_tile, b_tile, c_tile, m_size, k_size, lda, ldb, ldc, accum_c);
275 + } else {
276 + gemm_macro_rvv_fma_Mx8_Ku4<HEAD_SIZE_ALIGNMENT, kv_cache_t>(
277 + a_tile, b_tile, c_tile, m_size, dynamic_k_size, lda, ldb, ldc,
278 + accum_c);
279 + }
280 + }
281 +};
282 +
283 +} // namespace
284 +
285 +// ============================================================================
286 +// AttentionImpl<ISA::RVV> — mirrors ISA::NEON specialization
287 +// ============================================================================
288 +
289 +template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
290 +class AttentionImpl<ISA::RVV, scalar_t, head_dim, kv_cache_scalar_t> {
291 + public:
292 + using query_t = scalar_t;
293 + using q_buffer_t = float;
294 + using kv_cache_t = scalar_t;
295 + using logits_buffer_t = float;
296 + using partial_output_buffer_t = float;
297 + using prob_buffer_t = float;
298 +
299 + constexpr static int64_t BlockSizeAlignment = BLOCK_SIZE_ALIGNMENT;
300 + constexpr static int64_t HeadDimAlignment = HEAD_SIZE_ALIGNMENT;
301 + constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER;
302 + constexpr static int64_t HeadDim = head_dim;
303 + constexpr static ISA ISAType = ISA::RVV;
304 + constexpr static bool scale_on_logits = false;
305 +
306 + static_assert(HeadDim % HeadDimAlignment == 0);
307 + static_assert(HeadDimAlignment % 8 == 0);
308 + static_assert(BlockSizeAlignment % 8 == 0);
309 +
310 + public:
311 + template <template <typename tile_gemm_t> typename attention>
312 + FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
313 + attention<TileGemmRVV<kv_cache_t>> attention_iteration;
314 + attention_iteration(CPU_ATTENTION_PARAMS);
315 + }
316 +
317 + constexpr static int64_t k_cache_token_group_stride(
318 + const int32_t block_size) {
319 + return BlockSizeAlignment;
320 + }
321 +
322 + constexpr static int64_t v_cache_token_group_stride(
323 + const int32_t block_size) {
324 + return head_dim * BlockSizeAlignment;
325 + }
326 +
327 + constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
328 + return HeadDimAlignment;
329 + }
330 +
331 + static void copy_q_heads_tile(scalar_t* __restrict__ src,
332 + float* __restrict__ q_buffer,
333 + const int32_t q_num,
334 + const int32_t q_heads_per_kv,
335 + const int64_t q_num_stride,
336 + const int64_t q_head_stride, float scale) {
337 + static_assert(head_dim % 16 == 0);
338 + constexpr int32_t unroll_size = head_dim / 16;
339 + using load_vec_t = typename VecTypeTrait<scalar_t>::vec_t;
340 +
341 + vec_op::FP32Vec16 scale_vec(scale);
342 + for (int32_t q_num_idx = 0; q_num_idx < q_num; ++q_num_idx) {
343 + for (int32_t q_head_idx = 0; q_head_idx < q_heads_per_kv; ++q_head_idx) {
344 + scalar_t* __restrict__ curr_q =
345 + src + q_num_idx * q_num_stride + q_head_idx * q_head_stride;
346 + float* __restrict__ curr_q_buffer =
347 + q_buffer + q_num_idx * q_heads_per_kv * head_dim +
348 + q_head_idx * head_dim;
349 +
350 + vec_op::unroll_loop<int32_t, unroll_size>([&](int32_t i) {
351 + load_vec_t vec(curr_q);
352 + vec_op::FP32Vec16 fp32_vec(vec);
353 + fp32_vec = fp32_vec * scale_vec;
354 + fp32_vec.save(curr_q_buffer);
355 +
356 + curr_q += 16;
357 + curr_q_buffer += 16;
358 + });
359 + }
360 + }
361 + }
362 +
363 + static void reshape_and_cache(
364 + const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
365 + scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
366 + const int64_t* __restrict__ slot_mapping, const int64_t token_num,
367 + const int64_t key_token_num_stride, const int64_t value_token_num_stride,
368 + const int64_t head_num, const int64_t key_head_num_stride,
369 + const int64_t value_head_num_stride, const int64_t num_blocks,
370 + const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
371 + const int64_t block_size, const int64_t block_size_stride,
372 + const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
373 + #pragma omp parallel for collapse(2)
374 + for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
375 + for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
376 + const int64_t pos = slot_mapping[token_idx];
377 + if (pos < 0) {
378 + continue;
379 + }
380 +
381 + const int64_t block_idx = pos / block_size;
382 + const int64_t block_offset = pos % block_size;
383 + {
384 + // Write Key (transpose to column-major: [head_dim, block_size])
385 + const scalar_t* key_start_ptr = key +
386 + token_idx * key_token_num_stride +
387 + head_idx * key_head_num_stride;
388 + scalar_t* key_cache_start_ptr =
389 + key_cache + block_idx * num_blocks_stride +
390 + head_idx * cache_head_num_stride + block_offset;
391 +
392 + // Strided vector store for efficient transpose.
393 + // Load contiguous key elements, store with stride = block_size.
394 + {
395 + const ptrdiff_t byte_stride = block_size * sizeof(scalar_t);
396 + int64_t i = 0;
397 + for (; i < head_dim;) {
398 + size_t vl;
399 + if constexpr (std::is_same_v<scalar_t, float>) {
400 + vl = __riscv_vsetvl_e32m2(head_dim - i);
401 + vfloat32m2_t v = __riscv_vle32_v_f32m2(
402 + reinterpret_cast<const float*>(key_start_ptr + i), vl);
403 + __riscv_vsse32_v_f32m2(
404 + reinterpret_cast<float*>(key_cache_start_ptr +
405 + i * block_size),
406 + byte_stride, v, vl);
407 + } else {
408 + // Half and BFloat16 are both 16-bit types
409 + vl = __riscv_vsetvl_e16m1(head_dim - i);
410 + vuint16m1_t v = __riscv_vle16_v_u16m1(
411 + reinterpret_cast<const uint16_t*>(key_start_ptr + i), vl);
412 + __riscv_vsse16_v_u16m1(
413 + reinterpret_cast<uint16_t*>(key_cache_start_ptr +
414 + i * block_size),
415 + byte_stride, v, vl);
416 + }
417 + i += vl;
418 + }
419 + }
420 + }
421 + {
422 + // Write Value (row-major: [block_size, head_dim])
423 + const scalar_t* value_start_ptr = value +
424 + token_idx * value_token_num_stride +
425 + head_idx * value_head_num_stride;
426 + scalar_t* value_cache_start_ptr =
427 + value_cache + block_idx * num_blocks_stride +
428 + head_idx * cache_head_num_stride + block_offset * head_dim;
429 + std::memcpy(value_cache_start_ptr, value_start_ptr,
430 + sizeof(scalar_t) * head_dim);
431 + }
432 + }
433 + }
434 + }
435 +};
436 +
437 +} // namespace cpu_attention
438 +
439 + #undef BLOCK_SIZE_ALIGNMENT
440 + #undef HEAD_SIZE_ALIGNMENT
441 + #undef MAX_Q_HEAD_NUM_PER_ITER
442 +
443 +#endif // __riscv_v_min_vlen == 128
444 +
445 +#endif // CPU_ATTN_RVV_HPP
csrc/cpu/cpu_types_riscv_defs.hpp
+1 -1
@@ -72,7 +72,7 @@ typedef RVVTYPE(vuint16, LMUL_512, _t) fixed_u16x32_t
72 __attribute__((riscv_rvv_vector_bits(512)));
73
74 // bfloat16
75 -#ifdef RISCV_BF16_SUPPORT
75 +#ifdef __riscv_zvfbfmin
76 typedef RVVTYPE(vbfloat16, LMUL_128, _t) fixed_bf16x8_t
77 __attribute__((riscv_rvv_vector_bits(128)));
78 typedef RVVTYPE(vbfloat16, LMUL_256, _t) fixed_bf16x16_t
csrc/cpu/cpu_types_riscv_impl.hpp
+29 -6
@@ -15,8 +15,17 @@
15 #include <torch/all.h>
16 namespace vec_op {
17
18 -// BFloat16 is always supported on RISC-V: natively when RISCV_BF16_SUPPORT
19 -// is defined, otherwise via the FP32-simulation fallback path.
18 +// FP8 KV cache is not supported on RISC-V. These tag types and the
19 +// corresponding BF16Vec32 stub constructors below exist solely so that
20 +// templates referencing vec_op::fp8_*_tag in their bodies (e.g. in
21 +// cpu_attn_vec.hpp) compile under GCC's -Wtemplate-body lookup. The
22 +// stubs are never instantiated by CPU_ATTN_DISPATCH on __riscv.
23 +struct fp8_e4m3_tag {};
24 +struct fp8_e5m2_tag {};
25 +
26 +// BFloat16 is always supported on RISC-V: natively when __riscv_zvfbfmin
27 +// is defined (compiler-provided when -march includes zvfbfmin), otherwise
28 +// via the FP32-simulation fallback path.
29 #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
30 AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
31 AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
@@ -106,7 +115,7 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
115 // BF16 Implementation
116 // ============================================================================
117
109 -#ifdef RISCV_BF16_SUPPORT
118 +#ifdef __riscv_zvfbfmin
119
120 FORCE_INLINE fixed_u16x8_t bf16_to_u16(fixed_bf16x8_t v) {
121 return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_128, _u16, LMUL_128)(v);
@@ -185,6 +194,13 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
194
195 explicit BF16Vec32(fixed_bf16x32_t data) : reg(data) {};
196
197 + // FP8 KV cache stubs: never instantiated on RISC-V (CPU_ATTN_DISPATCH
198 + // omits FP8 cases on __riscv); exist only so name lookup succeeds.
199 + explicit BF16Vec32(const uint8_t* ptr, fp8_e4m3_tag)
200 + : BF16Vec32(static_cast<const void*>(ptr)) {}
201 + explicit BF16Vec32(const uint8_t* ptr, fp8_e5m2_tag)
202 + : BF16Vec32(static_cast<const void*>(ptr)) {}
203 +
204 explicit BF16Vec32(const BF16Vec8& v) {
205 fixed_u16x8_t u16_val = bf16_to_u16(v.reg);
206 fixed_u16x32_t u16_combined =
@@ -323,6 +339,13 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
339 reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp, 32);
340 }
341
342 + // FP8 KV cache stubs: never instantiated on RISC-V (CPU_ATTN_DISPATCH
343 + // omits FP8 cases on __riscv); exist only so name lookup succeeds.
344 + explicit BF16Vec32(const uint8_t* ptr, fp8_e4m3_tag)
345 + : BF16Vec32(static_cast<const void*>(ptr)) {}
346 + explicit BF16Vec32(const uint8_t* ptr, fp8_e5m2_tag)
347 + : BF16Vec32(static_cast<const void*>(ptr)) {}
348 +
349 explicit BF16Vec32(const BF16Vec8& v) {
350 float tmp_small[8];
351 RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp_small, v.reg_fp32, 8);
@@ -410,7 +433,7 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
433 explicit FP32Vec8(fixed_fp16x8_t v)
434 : reg(RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(v, VEC_ELEM_NUM)) {};
435
413 -#ifdef RISCV_BF16_SUPPORT
436 +#ifdef __riscv_zvfbfmin
437 explicit FP32Vec8(fixed_bf16x8_t v)
438 : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(v, VEC_ELEM_NUM)) {};
439 explicit FP32Vec8(const BF16Vec8& v)
@@ -608,7 +631,7 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
631 explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
632 explicit FP32Vec16(const FP16Vec16& v);
633
611 -#ifdef RISCV_BF16_SUPPORT
634 +#ifdef __riscv_zvfbfmin
635 explicit FP32Vec16(fixed_bf16x16_t v)
636 : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_512)(v, VEC_ELEM_NUM)) {};
637 explicit FP32Vec16(const BF16Vec16& v)
@@ -868,7 +891,7 @@ inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
891 acc = acc.fma(a, b);
892 }
893
871 -#ifdef RISCV_BF16_SUPPORT
894 +#ifdef __riscv_zvfbfmin
895 template <>
896 inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
897 *ptr = static_cast<__bf16>(v);
csrc/cpu/generate_cpu_attn_dispatch.py
+31 -4
@@ -20,7 +20,8 @@ ISA_TYPES = {
20 "VEC16": 2,
21 "NEON": 3,
22 "VXE": 4,
23 - "VSX": 5,
23 + "RVV": 5,
24 + "VSX": 6,
25 }
26
27 # KV cache index: 0 = auto (same as scalar_t), 1 = fp8_e4m3, 2 = fp8_e5m2
@@ -38,7 +39,7 @@ KV_CACHE_CPP_TYPES = {
39 }
40
41 # ISAs supported for head_dims divisible by 32
41 -ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE", "VSX"]
42 +ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE", "RVV", "VSX"]
43
44 # ISAs supported for head_dims divisible by 16 only
45 ISA_FOR_16 = ["VEC16"]
@@ -149,6 +150,15 @@ def generate_header_file() -> str:
150 #include "cpu_attn_vxe.hpp"
151 #endif
152
153 +// cpu_attn_rvv.hpp is hardcoded to VLEN==128 (m1/m2 intrinsics, vl=8) and
154 +// itself includes <riscv_vector.h>, which is unavailable on scalar
155 +// (-march=rv64gc) builds. Gate the include the same way as the dispatch
156 +// macro below, so non-128 / scalar RISC-V builds skip it entirely.
157 +#if defined(__riscv) && defined(__riscv_v_min_vlen) && \
158 + __riscv_v_min_vlen == 128
159 + #include "cpu_attn_rvv.hpp"
160 +#endif
161 +
162 #ifdef __powerpc__
163 #include "cpu_attn_vsx.hpp"
164 #endif
@@ -212,6 +222,23 @@ def generate_header_file() -> str:
222 ["VXE", "VEC", "VEC16"],
223 fp8=False,
224 )
225 + # RISC-V with RVV. cpu_attn_rvv.hpp is hardcoded to VLEN==128
226 + # (riscv_rvv_vector_bits(128) typedefs + vl=8 m1/m2 intrinsics), so
227 + # we split the dispatch into two top-level branches: VLEN==128 builds
228 + # get the full RVV+VEC+VEC16 case set, other VLEN builds get a
229 + # VEC/VEC16-only fallback. Preprocessor directives cannot appear
230 + # inside a #define body, so this duplication is necessary.
231 + header += _macro_block(
232 + "#elif defined(__riscv) && defined(__riscv_v_min_vlen) "
233 + "&& __riscv_v_min_vlen == 128",
234 + ["RVV", "VEC", "VEC16"],
235 + fp8=False,
236 + )
237 + header += _macro_block(
238 + "#elif defined(__riscv)",
239 + ["VEC", "VEC16"],
240 + fp8=False,
241 + )
242 header += _macro_block(
243 "#elif defined(__powerpc__)",
244 ["VSX", "VEC", "VEC16"],
@@ -233,8 +260,8 @@ def generate_header_file() -> str:
260 fp8=False,
261 )
262 header += (
236 - "#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / "
237 - "__s390x__ / __powerpc__ */\n\n"
263 + "#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ /"
264 + " __riscv / __powerpc__ */\n\n"
265 "#endif // CPU_ATTN_DISPATCH_GENERATED_H\n"
266 )
267
tests/kernels/attention/test_cpu_attn.py
+54 -9
@@ -49,15 +49,14 @@ def get_attn_isa(
49 block_size: int | None = None,
50 dtype: torch.dtype | None = None,
51 ):
52 - if block_size and dtype:
53 - return _get_attn_isa(dtype, block_size)
54 - else:
55 - if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
56 - return "neon"
57 - elif torch.cpu._is_amx_tile_supported():
58 - return "amx"
59 - else:
60 - return "vec"
52 + # Delegate to _get_attn_isa so the fallback path applies the same arch
53 + # gating (e.g. RISC-V RVV is only chosen when the build's hardcoded
54 + # VLEN=128 kernel is actually present; on VLEN=256 / scalar hosts it
55 + # correctly falls through to vec/vec16).
56 + return _get_attn_isa(
57 + dtype if dtype is not None else torch.bfloat16,
58 + block_size if block_size else 32,
59 + )
60
61
62 # rand number generation takes too much time, cache rand tensors
@@ -583,6 +582,52 @@ def test_varlen_with_paged_kv_normal_neon(
582 )
583
584
585 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3"])
586 +@pytest.mark.parametrize("seq_lens", SEQ_LENS)
587 +@pytest.mark.parametrize("num_heads", NUM_HEADS)
588 +@pytest.mark.parametrize("head_size", HEAD_SIZES)
589 +@pytest.mark.parametrize("block_size", [96, 128])
590 +@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
591 +@pytest.mark.parametrize("dtype", QTYPES)
592 +@pytest.mark.parametrize("soft_cap", [None])
593 +@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
594 +@pytest.mark.parametrize("use_alibi", [False])
595 +@pytest.mark.parametrize("use_sink", [False])
596 +@pytest.mark.parametrize("isa", ["rvv"])
597 +@pytest.mark.skipif(
598 + current_platform.get_cpu_architecture() != CpuArchEnum.RISCV,
599 + reason="Not a RISC-V CPU.",
600 +)
601 +def test_varlen_with_paged_kv_normal_rvv(
602 + seq_lens: list[tuple[int, int]],
603 + num_heads: tuple[int, int],
604 + head_size: int,
605 + sliding_window: int | None,
606 + dtype: torch.dtype,
607 + block_size: int,
608 + soft_cap: float | None,
609 + num_blocks: int,
610 + use_alibi: bool,
611 + use_sink: bool,
612 + isa: str,
613 + kv_cache_dtype: str,
614 +) -> None:
615 + varlen_with_paged_kv(
616 + seq_lens=seq_lens,
617 + num_heads=num_heads,
618 + head_size=head_size,
619 + sliding_window=sliding_window,
620 + dtype=dtype,
621 + block_size=block_size,
622 + soft_cap=soft_cap,
623 + num_blocks=num_blocks,
624 + use_alibi=use_alibi,
625 + use_sink=use_sink,
626 + isa=isa,
627 + kv_cache_dtype=kv_cache_dtype,
628 + )
629 +
630 +
631 @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3"])
632 @pytest.mark.parametrize("seq_lens", SEQ_LENS)
633 @pytest.mark.parametrize("num_heads", NUM_HEADS)
vllm/v1/attention/backends/cpu_attn.py
+28
@@ -1,5 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import functools
4 from dataclasses import dataclass
5 from typing import TYPE_CHECKING, ClassVar
6
@@ -34,6 +35,7 @@ _CPU_ARCH_PREFER_MIXED_BATCH = (
35 CpuArchEnum.X86,
36 CpuArchEnum.ARM,
37 CpuArchEnum.S390X,
38 + CpuArchEnum.RISCV,
39 CpuArchEnum.POWERPC,
40 )
41
@@ -506,6 +508,29 @@ def _make_sliding_window_bias(
508 return attn_biases
509
510
511 +@functools.lru_cache(maxsize=1)
512 +def _riscv_supports_rvv_vlen128() -> bool:
513 + """Whether the C++ RVV attention path (hardcoded to VLEN==128) is usable.
514 +
515 + The kernel in csrc/cpu/cpu_attn_rvv.hpp uses riscv_rvv_vector_bits(128)
516 + typedefs and m1/m2 intrinsics with vl=8; CMake's auto-detect picks the
517 + largest zvl<N>b advertised by /proc/cpuinfo, so the binary contains the
518 + RVV path only when the build host advertised exactly zvl128b. Mirror
519 + that here so the Python dispatch doesn't request ISA::RVV on builds
520 + where it wasn't compiled in (would TORCH_CHECK at first attention call).
521 + """
522 + try:
523 + with open("/proc/cpuinfo") as f:
524 + cpuinfo = f.read()
525 + except OSError:
526 + return False
527 + if "zvl128b" not in cpuinfo:
528 + return False
529 + # CMake auto-detect picks the largest advertised VLEN; if the host
530 + # advertises zvl256b or higher, the build skipped the RVV-128 path.
531 + return all(f"zvl{n}b" not in cpuinfo for n in (256, 512, 1024))
532 +
533 +
534 def _get_attn_isa(
535 dtype: torch.dtype,
536 block_size: int,
@@ -523,6 +548,7 @@ def _get_attn_isa(
548 arch = current_platform.get_cpu_architecture()
549 supports_arm = arch == CpuArchEnum.ARM
550 supports_vxe = arch == CpuArchEnum.S390X
551 + supports_riscv = arch == CpuArchEnum.RISCV
552 supports_vsx = arch == CpuArchEnum.POWERPC
553 supports_avx512 = torch.cpu._is_avx512_supported()
554 if fp8_kv and not supports_amx and not supports_avx512:
@@ -535,6 +561,8 @@ def _get_attn_isa(
561 if supports_arm:
562 # support ARM NEON FMLA and BFMMLA (bf16) for block size 32
563 return "neon"
564 + elif supports_riscv and _riscv_supports_rvv_vlen128():
565 + return "rvv"
566 elif supports_vxe:
567 return "vxe"
568 elif supports_vsx: