[CPU][RISC-V] Add RVV path for W4A8 INT4 GEMM (#45269)
Signed-off-by: wcy <233313160abc@gmail.com> Co-authored-by: lyd1992 <liuyudong@iscas.ac.cn> Co-authored-by: OpenAI Codex <codex@openai.com>
wcy committed
Jun 25, 2026 at 16:18 UTC
638b1a99ccd7b47f45f9a773456f77aeb66d4a1b
5 files changed
+157
-7
cmake/cpu_extension.cmake
+4
-3
@@ -166,12 +166,13 @@ elseif (S390_FOUND)
166
"-mtune=native")
167
elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
168
message(STATUS "RISC-V detected")
169
- if(DEFINED VLLM_RVV_VLEN AND NOT VLLM_RVV_VLEN GREATER 0)
169
+ if(DEFINED VLLM_RVV_VLEN AND VLLM_RVV_VLEN LESS 0)
170
message(FATAL_ERROR
171
- "VLLM_RVV_VLEN must be a positive integer; got '${VLLM_RVV_VLEN}'")
171
+ "VLLM_RVV_VLEN must be zero or a positive integer; got '${VLLM_RVV_VLEN}'")
172
endif()
173
# VLLM_RVV_VLEN selects the target VLEN. Auto-detected from /proc/cpuinfo
174
- # by default; override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256.
174
+ # by default; set -DVLLM_RVV_VLEN=0 to force scalar RISC-V build.
175
+ # Override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256 for RVV.
176
if(NOT DEFINED VLLM_RVV_VLEN)
177
# Auto-detect: find the largest zvl<N>b in /proc/cpuinfo isa line.
178
if(EXISTS /proc/cpuinfo)
csrc/cpu/cpu_types_riscv_defs.hpp
+18
-3
@@ -3,13 +3,17 @@
3
4
// VLEN-to-LMUL mapping for RISC-V Vector extension.
5
//
6
-// LMUL_<N> expands to the LMUL suffix giving N total bits of vector data:
7
-// VLEN=128: LMUL_128=m1, LMUL_256=m2, LMUL_512=m4, LMUL_1024=m8
8
-// VLEN=256: LMUL_128=mf2, LMUL_256=m1, LMUL_512=m2, LMUL_1024=m4
6
+// LMUL_<N> expands to the LMUL suffix giving N total bits of vector data.
7
+// LMUL_64 is used by 8-lane int8/uint8 vectors.
8
+// VLEN=128:
9
+// LMUL_64=mf2, LMUL_128=m1, LMUL_256=m2, LMUL_512=m4, LMUL_1024=m8
10
+// VLEN=256:
11
+// LMUL_64=mf4, LMUL_128=mf2, LMUL_256=m1, LMUL_512=m2, LMUL_1024=m4
12
13
#include <riscv_vector.h>
14
15
#if __riscv_v_min_vlen == 128
16
+ #define LMUL_64 mf2
17
#define LMUL_128 m1
18
#define LMUL_256 m2
19
#define LMUL_512 m4
@@ -17,6 +21,7 @@
21
#define BOOL_256 b16
22
#define BOOL_512 b8
23
#elif __riscv_v_min_vlen == 256
24
+ #define LMUL_64 mf4
25
#define LMUL_128 mf2
26
#define LMUL_256 m1
27
#define LMUL_512 m2
@@ -41,6 +46,16 @@
46
47
// ---- Semantic fixed-vector typedefs (named by element count) ----
48
49
+// uint8 / int8
50
+typedef RVVTYPE(vuint8, LMUL_64, _t) fixed_u8x8_t
51
+ __attribute__((riscv_rvv_vector_bits(64)));
52
+typedef RVVTYPE(vint8, LMUL_64, _t) fixed_i8x8_t
53
+ __attribute__((riscv_rvv_vector_bits(64)));
54
+
55
+// int16
56
+typedef RVVTYPE(vint16, LMUL_128, _t) fixed_i16x8_t
57
+ __attribute__((riscv_rvv_vector_bits(128)));
58
+
59
// float16
60
typedef RVVTYPE(vfloat16, LMUL_128, _t) fixed_fp16x8_t
61
__attribute__((riscv_rvv_vector_bits(128)));
csrc/cpu/sgl-kernels/gemm_int4.cpp
+124
@@ -285,6 +285,125 @@ inline int32_t load_uint4_vnni(const uint8_t* __restrict__ B, int64_t k, int64_t
285
return (n_group % 2 == 0) ? (packed & 0x0f) : ((packed >> 4) & 0x0f);
286
}
287
288
+#if defined(CPU_CAPABILITY_RVV)
289
+template <int64_t N, int64_t ldb, int group>
290
+inline fixed_i8x8_t load_uint4_as_int8_rvv(const uint8_t* __restrict__ B, int64_t k) {
291
+ constexpr int64_t n_group_size = 8;
292
+ constexpr int64_t vnni_size = 4;
293
+ static_assert(N == 32);
294
+ static_assert(ldb == N / 2);
295
+ static_assert(group >= 0 && group < N / n_group_size);
296
+
297
+ // Unpack: gather 8 packed int4 values from the VNNI4 layout.
298
+ const int64_t ki = k % vnni_size;
299
+ const int64_t k_base = k - ki;
300
+ constexpr int64_t packed_group = group / 2;
301
+ const uint8_t* packed_ptr = B + k_base * ldb + packed_group * n_group_size * vnni_size + ki;
302
+
303
+ fixed_u8x8_t packed = RVVI(__riscv_vlse8_v_u8, LMUL_64)(packed_ptr, vnni_size, n_group_size);
304
+ if constexpr (group % 2 == 1) {
305
+ packed = RVVI(__riscv_vsrl_vx_u8, LMUL_64)(packed, 4, n_group_size);
306
+ }
307
+ fixed_u8x8_t nibbles = RVVI(__riscv_vand_vx_u8, LMUL_64)(packed, 0x0f, n_group_size);
308
+ return RVVI4(__riscv_vreinterpret_v_u8, LMUL_64, _i8, LMUL_64)(nibbles);
309
+}
310
+
311
+inline fixed_i32x8_t gemm_accum_uint8_int8_rvv(fixed_i32x8_t acc, uint8_t a, fixed_i8x8_t b) {
312
+ constexpr int64_t vl = 8;
313
+ fixed_i16x8_t b_i16 = RVVI(__riscv_vsext_vf2_i16, LMUL_128)(b, vl);
314
+ return RVVI(__riscv_vwmacc_vx_i32, LMUL_256)(acc, static_cast<int16_t>(a), b_i16, vl);
315
+}
316
+
317
+template <int64_t N, int64_t ldb, int group>
318
+inline fixed_i32x8_t gemm_accum_uint4_rvv(
319
+ fixed_i32x8_t acc,
320
+ const uint8_t* __restrict__ B,
321
+ const int8_t* __restrict__ qzeros_b,
322
+ uint8_t a,
323
+ int64_t k) {
324
+ constexpr int64_t n_group_size = 8;
325
+ fixed_i8x8_t b = load_uint4_as_int8_rvv<N, ldb, group>(B, k);
326
+ fixed_i8x8_t qzeros =
327
+ RVVI(__riscv_vle8_v_i8, LMUL_64)(qzeros_b + group * n_group_size, n_group_size);
328
+ b = RVVI(__riscv_vsub_vv_i8, LMUL_64)(b, qzeros, n_group_size);
329
+ return gemm_accum_uint8_int8_rvv(acc, a, b);
330
+}
331
+
332
+template <int group>
333
+inline void _dequant_and_store_rvv(
334
+ float* __restrict__ C,
335
+ fixed_i32x8_t acc,
336
+ const float* __restrict__ scales_a,
337
+ const int32_t* __restrict__ qzeros_a,
338
+ const float* __restrict__ scales_b,
339
+ const int32_t* __restrict__ compensation,
340
+ int64_t m,
341
+ int64_t ldc) {
342
+ constexpr int64_t n_group_size = 8;
343
+ constexpr int64_t n = group * n_group_size;
344
+ constexpr int64_t vl = n_group_size;
345
+
346
+ // Dequant compensation: remove activation zero-point contribution.
347
+ fixed_i32x8_t comp = RVVI(__riscv_vle32_v_i32, LMUL_256)(compensation + n, vl);
348
+ fixed_i32x8_t zp_comp = RVVI(__riscv_vmul_vx_i32, LMUL_256)(comp, qzeros_a[m], vl);
349
+ acc = RVVI(__riscv_vsub_vv_i32, LMUL_256)(acc, zp_comp, vl);
350
+
351
+ // Scale: convert int32 accumulators to fp32 and apply activation/weight scales.
352
+ fixed_fp32x8_t acc_f = RVVI(__riscv_vfcvt_f_x_v_f32, LMUL_256)(acc, vl);
353
+ acc_f = RVVI(__riscv_vfmul_vf_f32, LMUL_256)(acc_f, scales_a[m], vl);
354
+ fixed_fp32x8_t scale_b = RVVI(__riscv_vle32_v_f32, LMUL_256)(scales_b + n, vl);
355
+ acc_f = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(acc_f, scale_b, vl);
356
+
357
+ // Store: accumulate into the float scratch buffer that already holds bias/zero.
358
+ float* c_ptr = C + m * ldc + n;
359
+ fixed_fp32x8_t c_old = RVVI(__riscv_vle32_v_f32, LMUL_256)(c_ptr, vl);
360
+ fixed_fp32x8_t c_new = RVVI(__riscv_vfadd_vv_f32, LMUL_256)(c_old, acc_f, vl);
361
+ RVVI(__riscv_vse32_v_f32, LMUL_256)(c_ptr, c_new, vl);
362
+}
363
+
364
+template <int64_t N, int64_t ldb>
365
+void _dequant_gemm_accum_rvv(
366
+ float* __restrict__ C,
367
+ const uint8_t* __restrict__ A,
368
+ const float* __restrict__ scales_a,
369
+ const int32_t* __restrict__ qzeros_a,
370
+ const uint8_t* __restrict__ B,
371
+ const float* __restrict__ scales_b,
372
+ const int8_t* __restrict__ qzeros_b,
373
+ const int32_t* __restrict__ compensation,
374
+ int64_t M,
375
+ int64_t K,
376
+ int64_t lda,
377
+ int64_t ldc) {
378
+ static_assert(N == 32);
379
+ static_assert(ldb == N / 2);
380
+ constexpr int64_t vl = 8;
381
+
382
+ // Accumulate one C row over the 32-column block.
383
+ for (int64_t m = 0; m < M; ++m) {
384
+ fixed_i32x8_t acc0 = RVVI(__riscv_vmv_v_x_i32, LMUL_256)(0, vl);
385
+ fixed_i32x8_t acc1 = RVVI(__riscv_vmv_v_x_i32, LMUL_256)(0, vl);
386
+ fixed_i32x8_t acc2 = RVVI(__riscv_vmv_v_x_i32, LMUL_256)(0, vl);
387
+ fixed_i32x8_t acc3 = RVVI(__riscv_vmv_v_x_i32, LMUL_256)(0, vl);
388
+ // A[m][k] @ B[k][0:32] -> acc[m][0:32]
389
+ for (int64_t k = 0; k < K; ++k) {
390
+ // GEMM K step: one scalar activation updates four 8-column RVV tiles.
391
+ const uint8_t a = A[m * lda + k];
392
+ acc0 = gemm_accum_uint4_rvv<N, ldb, 0>(acc0, B, qzeros_b, a, k);
393
+ acc1 = gemm_accum_uint4_rvv<N, ldb, 1>(acc1, B, qzeros_b, a, k);
394
+ acc2 = gemm_accum_uint4_rvv<N, ldb, 2>(acc2, B, qzeros_b, a, k);
395
+ acc3 = gemm_accum_uint4_rvv<N, ldb, 3>(acc3, B, qzeros_b, a, k);
396
+ }
397
+
398
+ // Dequant/scale/store each 8-column group back into C.
399
+ _dequant_and_store_rvv<0>(C, acc0, scales_a, qzeros_a, scales_b, compensation, m, ldc);
400
+ _dequant_and_store_rvv<1>(C, acc1, scales_a, qzeros_a, scales_b, compensation, m, ldc);
401
+ _dequant_and_store_rvv<2>(C, acc2, scales_a, qzeros_a, scales_b, compensation, m, ldc);
402
+ _dequant_and_store_rvv<3>(C, acc3, scales_a, qzeros_a, scales_b, compensation, m, ldc);
403
+ }
404
+}
405
+#endif
406
+
407
template <int64_t N, int64_t ldb, bool sym_quant_act>
408
void _dequant_gemm_accum(
409
float* C,
@@ -336,6 +455,11 @@ void _dequant_gemm_accum(
455
_dequant_and_store<true, N, sym_quant_act>(
456
C, C_i32, scales_a, qzeros_a, scales_b, compensation, M, N /*ldi*/, ldc, 1 /*ldsa*/);
457
} else
458
+#elif defined(CPU_CAPABILITY_RVV)
459
+ if constexpr (!sym_quant_act && N == BLOCK_N && ldb == BLOCK_N / 2) {
460
+ _dequant_gemm_accum_rvv<N, ldb>(C, A, scales_a, qzeros_a, B, scales_b, qzeros_b, compensation, M, K, lda, ldc);
461
+ return;
462
+ } else
463
#endif
464
{
465
for (int64_t m = 0; m < M; ++m) {
csrc/cpu/sgl-kernels/vec.h
+8
@@ -9,11 +9,19 @@
9
#define CPU_CAPABILITY_AVX512
10
#endif
11
12
+#if defined(__riscv_v_min_vlen) && (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256)
13
+#define CPU_CAPABILITY_RVV
14
+#endif
15
+
16
#include <ATen/cpu/vec/functional.h>
17
#include <ATen/cpu/vec/vec.h>
18
#if defined(CPU_CAPABILITY_AVX512)
19
#include <immintrin.h>
20
#endif
21
+
22
+#if defined(CPU_CAPABILITY_RVV)
23
+#include "../cpu_types_riscv_defs.hpp"
24
+#endif
25
namespace {
26
27
using namespace at::vec;
vllm/model_executor/kernels/linear/mixed_precision/cpu.py
+3
-1
@@ -168,11 +168,13 @@ class CPUWNA16LinearKernel(MPLinearKernel):
168
if zp.output_dim == 0:
169
zp.data = zp.t().contiguous()
170
171
+ supports_amx = torch.cpu._is_amx_tile_supported()
172
+ supports_riscv = current_platform.get_cpu_architecture() == CpuArchEnum.RISCV
173
layer.use_w4a8 = (
174
envs.VLLM_CPU_INT4_W4A8
175
and not self.config.has_g_idx
176
and self.config.act_type == torch.bfloat16
175
- and torch.cpu._is_amx_tile_supported()
177
+ and (supports_amx or supports_riscv)
178
)
179
# layer.use_w4a8 = False
180
# AWQ format will be converted to GPTQ format in `AutoAWQMarlinLinearMethod`