@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c68c55d43e
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, 03:00 UTC

[CPU][RISC-V] Add VLEN=256 support to RVV attention kernels (#42943)

Signed-off-by: velonica0 <like@mail.nankai.edu.cn> Signed-off-by: velonica0 <47554626+velonica0@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>

velonica0 committed May 21, 2026 at 19:50 UTC c68c55d43e504745dbfc2d46b552e80acb74d4b9
4 files changed +91 -128
csrc/cpu/cpu_attn_rvv.hpp
+68 -101
@@ -4,17 +4,18 @@
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
7 +// RVV attention kernel using VLEN-agnostic RVVI() macros from
8 +// cpu_types_riscv_defs.hpp. The Mx8 tile GEMM uses 8 FP32 elements
9 +// per vector (LMUL_256 bits of FP32 data), which maps to:
10 +// VLEN=128: m2 (256 bits = 8 x FP32)
11 +// VLEN=256: m1 (256 bits = 8 x FP32)
12 +// Only VLEN=128 and VLEN=256 are supported; other VLENs (512, 1024)
13 +// and scalar RISC-V builds fall back to VEC/VEC16.
14 +#if defined(__riscv_v_min_vlen) && \
15 + (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256)
16
17 #include "cpu_attn_impl.hpp"
18 + #include "cpu_types_riscv_defs.hpp"
19 #include <riscv_vector.h>
20 #include <type_traits>
21
@@ -22,73 +23,50 @@ namespace cpu_attention {
23
24 namespace {
25
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 -
26 #define BLOCK_SIZE_ALIGNMENT 32
27 #define HEAD_SIZE_ALIGNMENT 32
28 #define MAX_Q_HEAD_NUM_PER_ITER 16
29
30 // ============================================================================
47 -// B-matrix row loading: load 8 elements as FP32 (using m2 LMUL at VLEN=128)
31 +// B-matrix row loading: load 8 elements as FP32
32 // ============================================================================
33
34 template <typename kv_cache_t>
51 -FORCE_INLINE fixed_vfloat32m2_t load_row8_B_as_f32(const kv_cache_t* p);
35 +FORCE_INLINE fixed_fp32x8_t load_row8_B_as_f32(const kv_cache_t* p);
36
37 template <>
54 -FORCE_INLINE fixed_vfloat32m2_t load_row8_B_as_f32<float>(const float* p) {
55 - return __riscv_vle32_v_f32m2(p, 8);
38 +FORCE_INLINE fixed_fp32x8_t load_row8_B_as_f32<float>(const float* p) {
39 + return RVVI(__riscv_vle32_v_f32, LMUL_256)(p, 8);
40 }
41
42 template <>
59 -FORCE_INLINE fixed_vfloat32m2_t
60 -load_row8_B_as_f32<c10::Half>(const c10::Half* p) {
43 +FORCE_INLINE fixed_fp32x8_t load_row8_B_as_f32<c10::Half>(const c10::Half* p) {
44 #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);
45 + fixed_fp16x8_t h = RVVI(__riscv_vle16_v_f16, LMUL_128)(
46 + reinterpret_cast<const _Float16*>(p), 8);
47 + return RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(h, 8);
48 #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.
49 alignas(16) float tmp[8];
50 for (int i = 0; i < 8; ++i) {
51 tmp[i] = static_cast<float>(p[i]);
52 }
74 - return __riscv_vle32_v_f32m2(tmp, 8);
53 + return RVVI(__riscv_vle32_v_f32, LMUL_256)(tmp, 8);
54 #endif
55 }
56
57 template <>
79 -FORCE_INLINE fixed_vfloat32m2_t
58 +FORCE_INLINE fixed_fp32x8_t
59 load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p) {
60 #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);
61 + fixed_bf16x8_t bf = RVVI(__riscv_vle16_v_bf16, LMUL_128)(
62 + reinterpret_cast<const __bf16*>(p), 8);
63 + return RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(bf, 8);
64 #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);
65 + fixed_u16x8_t raw = RVVI(__riscv_vle16_v_u16, LMUL_128)(
66 + reinterpret_cast<const uint16_t*>(p), 8);
67 + fixed_u32x8_t wide = RVVI(__riscv_vzext_vf2_u32, LMUL_256)(raw, 8);
68 + fixed_u32x8_t shifted = RVVI(__riscv_vsll_vx_u32, LMUL_256)(wide, 16, 8);
69 + return RVVI4(__riscv_vreinterpret_v_u32, LMUL_256, _f32, LMUL_256)(shifted);
70 #endif
71 }
72
@@ -96,14 +74,12 @@ load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p) {
74 // Micro kernel: Mx8 tile, K unrolled by 4, RVV scalar-broadcast FMA
75 // ============================================================================
76 //
99 -// NEON uses vfmaq_laneq_f32 (lane-indexed FMA from a preloaded A vector).
77 // 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.
78 +// use vfmacc_vf (scalar * vector + accumulator).
79 //
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.
80 +// The 8-column tile uses LMUL_256 bits of FP32 data:
81 +// VLEN=128: m2 (2 regs per accumulator), M=8 => 18 of 32 regs
82 +// VLEN=256: m1 (1 reg per accumulator), M=8 => 9 of 32 regs
83
84 template <int32_t M, typename kv_cache_t>
85 FORCE_INLINE void gemm_micro_rvv_fma_Mx8_Ku4(
@@ -115,94 +91,90 @@ FORCE_INLINE void gemm_micro_rvv_fma_Mx8_Ku4(
91
92 constexpr size_t vl = 8;
93
118 - // helpers for per-M codegen
94 #define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
95 #define IF_M(i) if constexpr (M > (i))
96
122 - // A row base pointers
97 #define DECL_A(i) const float* a##i = A + (i) * lda;
98 ROWS_APPLY(DECL_A)
99 #undef DECL_A
100
127 - // declare one m2 accumulator per row
128 - #define DECL_ACC(i) fixed_vfloat32m2_t acc##i;
101 + #define DECL_ACC(i) fixed_fp32x8_t acc##i;
102 ROWS_APPLY(DECL_ACC)
103 #undef DECL_ACC
104
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 - } \
105 + #define INIT_ACC(i) \
106 + IF_M(i) { \
107 + if (accumulate) { \
108 + acc##i = RVVI(__riscv_vle32_v_f32, LMUL_256)(C + (i) * ldc, vl); \
109 + } else { \
110 + acc##i = RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.f, vl); \
111 + } \
112 }
113 ROWS_APPLY(INIT_ACC)
114 #undef INIT_ACC
115
116 int32_t k = 0;
117
146 - // K unrolled by 4
118 for (; k + 3 < K; k += 4) {
148 - // k + 0
119 {
150 - fixed_vfloat32m2_t b =
120 + fixed_fp32x8_t b =
121 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); \
122 + #define STEP_K0(i) \
123 + IF_M(i) { \
124 + acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)(acc##i, *(a##i + k + 0), \
125 + b, vl); \
126 }
127 ROWS_APPLY(STEP_K0)
128 #undef STEP_K0
129 }
159 - // k + 1
130 {
161 - fixed_vfloat32m2_t b =
131 + fixed_fp32x8_t b =
132 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); \
133 + #define STEP_K1(i) \
134 + IF_M(i) { \
135 + acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)(acc##i, *(a##i + k + 1), \
136 + b, vl); \
137 }
138 ROWS_APPLY(STEP_K1)
139 #undef STEP_K1
140 }
170 - // k + 2
141 {
172 - fixed_vfloat32m2_t b =
142 + fixed_fp32x8_t b =
143 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); \
144 + #define STEP_K2(i) \
145 + IF_M(i) { \
146 + acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)(acc##i, *(a##i + k + 2), \
147 + b, vl); \
148 }
149 ROWS_APPLY(STEP_K2)
150 #undef STEP_K2
151 }
181 - // k + 3
152 {
183 - fixed_vfloat32m2_t b =
153 + fixed_fp32x8_t b =
154 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); \
155 + #define STEP_K3(i) \
156 + IF_M(i) { \
157 + acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)(acc##i, *(a##i + k + 3), \
158 + b, vl); \
159 }
160 ROWS_APPLY(STEP_K3)
161 #undef STEP_K3
162 }
163 }
164
194 - // K tail
165 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); }
166 + fixed_fp32x8_t b = load_row8_B_as_f32<kv_cache_t>(B + (int64_t)k * ldb);
167 + #define TAIL_ROW(i) \
168 + IF_M(i) { \
169 + acc##i = \
170 + RVVI(__riscv_vfmacc_vf_f32, LMUL_256)(acc##i, *(a##i + k), b, vl); \
171 + }
172 ROWS_APPLY(TAIL_ROW)
173 #undef TAIL_ROW
174 }
175
203 - // store accumulators to C
176 #define STORE_ROW(i) \
205 - IF_M(i) { __riscv_vse32_v_f32m2(C + (i) * ldc, acc##i, vl); }
177 + IF_M(i) { RVVI(__riscv_vse32_v_f32, LMUL_256)(C + (i) * ldc, acc##i, vl); }
178 ROWS_APPLY(STORE_ROW)
179 #undef STORE_ROW
180
@@ -381,7 +353,6 @@ class AttentionImpl<ISA::RVV, scalar_t, head_dim, kv_cache_scalar_t> {
353 const int64_t block_idx = pos / block_size;
354 const int64_t block_offset = pos % block_size;
355 {
384 - // Write Key (transpose to column-major: [head_dim, block_size])
356 const scalar_t* key_start_ptr = key +
357 token_idx * key_token_num_stride +
358 head_idx * key_head_num_stride;
@@ -389,8 +360,6 @@ class AttentionImpl<ISA::RVV, scalar_t, head_dim, kv_cache_scalar_t> {
360 key_cache + block_idx * num_blocks_stride +
361 head_idx * cache_head_num_stride + block_offset;
362
392 - // Strided vector store for efficient transpose.
393 - // Load contiguous key elements, store with stride = block_size.
363 {
364 const ptrdiff_t byte_stride = block_size * sizeof(scalar_t);
365 int64_t i = 0;
@@ -405,7 +374,6 @@ class AttentionImpl<ISA::RVV, scalar_t, head_dim, kv_cache_scalar_t> {
374 i * block_size),
375 byte_stride, v, vl);
376 } else {
408 - // Half and BFloat16 are both 16-bit types
377 vl = __riscv_vsetvl_e16m1(head_dim - i);
378 vuint16m1_t v = __riscv_vle16_v_u16m1(
379 reinterpret_cast<const uint16_t*>(key_start_ptr + i), vl);
@@ -419,7 +387,6 @@ class AttentionImpl<ISA::RVV, scalar_t, head_dim, kv_cache_scalar_t> {
387 }
388 }
389 {
422 - // Write Value (row-major: [block_size, head_dim])
390 const scalar_t* value_start_ptr = value +
391 token_idx * value_token_num_stride +
392 head_idx * value_head_num_stride;
@@ -440,6 +407,6 @@ class AttentionImpl<ISA::RVV, scalar_t, head_dim, kv_cache_scalar_t> {
407 #undef HEAD_SIZE_ALIGNMENT
408 #undef MAX_Q_HEAD_NUM_PER_ITER
409
443 -#endif // __riscv_v_min_vlen == 128
410 +#endif // __riscv_v_min_vlen == 128 || 256
411
412 #endif // CPU_ATTN_RVV_HPP
csrc/cpu/cpu_types_riscv_defs.hpp
+4
@@ -71,6 +71,10 @@ typedef RVVTYPE(vuint16, LMUL_256, _t) fixed_u16x16_t
71 typedef RVVTYPE(vuint16, LMUL_512, _t) fixed_u16x32_t
72 __attribute__((riscv_rvv_vector_bits(512)));
73
74 +// uint32
75 +typedef RVVTYPE(vuint32, LMUL_256, _t) fixed_u32x8_t
76 + __attribute__((riscv_rvv_vector_bits(256)));
77 +
78 // bfloat16
79 #ifdef __riscv_zvfbfmin
80 typedef RVVTYPE(vbfloat16, LMUL_128, _t) fixed_bf16x8_t
csrc/cpu/generate_cpu_attn_dispatch.py
+7 -12
@@ -150,12 +150,10 @@ 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.
153 +// cpu_attn_rvv.hpp supports VLEN=128 and VLEN=256 via RVVI() macros.
154 +// Other VLENs and scalar RISC-V builds skip it entirely.
155 #if defined(__riscv) && defined(__riscv_v_min_vlen) && \
158 - __riscv_v_min_vlen == 128
156 + (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256)
157 #include "cpu_attn_rvv.hpp"
158 #endif
159
@@ -222,15 +220,12 @@ def generate_header_file() -> str:
220 ["VXE", "VEC", "VEC16"],
221 fp8=False,
222 )
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.
223 + # RISC-V with RVV. cpu_attn_rvv.hpp supports VLEN=128 and VLEN=256
224 + # via RVVI() macros. Builds with a supported VLEN get
225 + # RVV+VEC+VEC16; other RISC-V builds fall back to VEC/VEC16 only.
226 header += _macro_block(
227 "#elif defined(__riscv) && defined(__riscv_v_min_vlen) "
233 - "&& __riscv_v_min_vlen == 128",
228 + "&& (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256)",
229 ["RVV", "VEC", "VEC16"],
230 fp8=False,
231 )
vllm/v1/attention/backends/cpu_attn.py
+12 -15
@@ -514,26 +514,23 @@ def _make_sliding_window_bias(
514
515
516 @functools.lru_cache(maxsize=1)
517 -def _riscv_supports_rvv_vlen128() -> bool:
518 - """Whether the C++ RVV attention path (hardcoded to VLEN==128) is usable.
519 -
520 - The kernel in csrc/cpu/cpu_attn_rvv.hpp uses riscv_rvv_vector_bits(128)
521 - typedefs and m1/m2 intrinsics with vl=8; CMake's auto-detect picks the
522 - largest zvl<N>b advertised by /proc/cpuinfo, so the binary contains the
523 - RVV path only when the build host advertised exactly zvl128b. Mirror
524 - that here so the Python dispatch doesn't request ISA::RVV on builds
525 - where it wasn't compiled in (would TORCH_CHECK at first attention call).
517 +def _riscv_supports_rvv() -> bool:
518 + """Whether the C++ RVV attention path is usable.
519 +
520 + The kernel in csrc/cpu/cpu_attn_rvv.hpp uses VLEN-agnostic RVVI()
521 + macros and supports VLEN=128 and VLEN=256. CMake auto-detects the
522 + largest zvl<N>b from /proc/cpuinfo and passes it via -mrvv-vector-bits.
523 + The RVV path is compiled whenever __riscv_v_min_vlen is defined, so
524 + we check that at least one supported zvl<N>b is advertised.
525 """
526 try:
527 with open("/proc/cpuinfo") as f:
528 cpuinfo = f.read()
529 except OSError:
530 return False
532 - if "zvl128b" not in cpuinfo:
533 - return False
534 - # CMake auto-detect picks the largest advertised VLEN; if the host
535 - # advertises zvl256b or higher, the build skipped the RVV-128 path.
536 - return all(f"zvl{n}b" not in cpuinfo for n in (256, 512, 1024))
531 + return any(f"zvl{n}b" in cpuinfo for n in (128, 256)) and all(
532 + f"zvl{n}b" not in cpuinfo for n in (512, 1024)
533 + )
534
535
536 def _get_attn_isa(
@@ -566,7 +563,7 @@ def _get_attn_isa(
563 if supports_arm:
564 # support ARM NEON FMLA and BFMMLA (bf16) for block size 32
565 return "neon"
569 - elif supports_riscv and _riscv_supports_rvv_vlen128():
566 + elif supports_riscv and _riscv_supports_rvv():
567 return "rvv"
568 elif supports_vxe:
569 return "vxe"