onnxruntime
d3fab37b - [MLAS] Add RISC-V RVV backend for MLAS QNBitGemm (MatMulNBits) (#29537)

Commit
40 days ago
[MLAS] Add RISC-V RVV backend for MLAS QNBitGemm (MatMulNBits) (#29537) ## Motivation MLAS has hand-optimized QNBitGemm kernels for x86 (AVX2/AVX512), ARM (NEON), and LoongArch, but none for RISC-V — on RISC-V the dispatch was null, so MlasIsQNBitGemmAvailable returned false and MatMulNBits fell back to the generic scalar path, leaving the RISC-V Vector unit idle for the hottest kernel in quantized-LLM inference. This PR adds a native RISC-V Vector (RVV) implementation of the QNBitGemm dispatch so weight-quantized matmul runs on the vector unit. ## What's implemented All compute modes the operator selects among, for rv64gcv: - 4-bit weights × fp32 activations (SQNBIT_CompFp32) — M=1 GEMV + dequant-to-SGEMM - 4-bit weights × int8 activations (SQNBIT_CompInt8) — quantize-A + int8×int4 kernel - 8-bit weights × int8 activations (SQNBIT_CompInt8, BlkSum path) - 4-/8-bit weights × fp16 activations (HQNBIT_CompFp16) — requires the Zvfh extension (rv64gcv_zvfh, MLAS_USE_RVV_ZVFH) plus the packing / quantization / workspace-sizing plumbing each mode needs (including the 3-call PrePack protocol). ## Design notes - The packed-B / block-sum / dequant-buffer layouts are private to this dispatch (produced and consumed only by these kernels), so plain layouts are used — natural for RVV's scalable vectors. - The SQ8 prepack unit test asserts a per-arch B layout (#ifdef ARM64 / #else x86); added a MLAS_TARGET_RISCV64 branch describing the RVV layout, following the existing per-arch pattern. - fp16 kernels accumulate in fp32 (fp16 → vfwcvt → fp32 FMA → fp16) for accuracy. ## Files New: - onnxruntime/core/mlas/lib/riscv64/qnbitgemm_kernel_rvv.cpp — SQ4/SQ8 kernels, packing, dispatch object - onnxruntime/core/mlas/lib/riscv64/hqnbitgemm_kernel_rvv.cpp — fp16 (Zvfh) dequant + GEMM - onnxruntime/test/mlas/unittest/test_qnbitgemm_rvv_fp16.cpp — RISC-V fp16 e2e test (HQ4 + HQ8) Modified: - cmake/onnxruntime_mlas.cmake — add the two sources (RVV / Zvfh source lists + - onnxruntime/core/mlas/lib/mlasi.h — extern decl of the dispatch - onnxruntime/core/mlas/lib/platform.cpp — assign QNBitGemmDispatch on RISC-V - onnxruntime/test/mlas/unittest/test_sq8bitgemm.cpp — RISC-V layout branch in the SQ8 prepack test ## Performance (RVV, VLEN=256, N=256, K=2048) The K-reduction runs at LMUL=4; profiling showed the kernels are widening-multip), so cutting per-chunk instruction overhead is the lever. - SQ4 CompInt8 (4-bit): ~3.4× vs the naive per-sub-block version (2.2 → 7.6 GOP/s) - SQ8 CompInt8 (8-bit): ~1.2× - fp16 GEMM tiled to reuse the B vfwcvt across rows ## Testing Built and run on real RISC-V hardware (K3 board, rv64gcv_zvfh, VLEN=256) via onnxruntime_mlas_test: - SQNBitGemm* : SQ8Bit* : BlockQ4* : BlockQ8* : RvvFp16* → 13,526 / 13,526 pass - Broader suite run passed with 0 failures as far as it ran (Activation → all QGemm); the run was only bounded by the very slow threaded large-GEMM stress tests, unrelated to this change. ## Build Enable fp16 with -Donnxruntime_USE_RVV_ZVFH=ON. The 4-/8-bit int8/fp32 paths build with the existing onnxruntime_USE_RVV.
Author
Parents
Loading