[MLAS] Add RISC-V RVV backend for MLAS QNBitGemm (MatMulNBits)

[velonica0]

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.

[velonica0]

Hi, @hariharans29
Could you please take a look at this PR? Thank you for your help.

[Copilot]

Pull request overview

Adds a RISC-V Vector (RVV) backend to MLAS’s QNBitGemm (MatMulNBits) so n-bit quantized GEMM can execute on rv64gcv (and optionally Zvfh for fp16), enabling vectorized kernels and associated packing/workspace plumbing.

Changes:

• Introduces RVV implementations for QNBitGemm packing, per-GEMM workspace sizing, and SQ4/SQ8 compute kernels, plus optional Zvfh fp16 dequant + GEMM.
• Wires the new RVV QNBitGemm dispatch into MLAS platform initialization and exports the dispatch symbol.
• Extends MLAS unit tests to validate RISC-V SQ8 prepack layout expectations and adds an RVV fp16 end-to-end test.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
cmake/onnxruntime_mlas.cmake	Adds the new RVV and Zvfh source files to the MLAS build for RISC-V.
onnxruntime/core/mlas/lib/mlasi.h	Declares the RVV QNBitGemm dispatch for platform wiring.
onnxruntime/core/mlas/lib/platform.cpp	Assigns the RVV QNBitGemm dispatch when MLAS_USE_RVV is enabled on RISC-V.
onnxruntime/core/mlas/lib/riscv64/qnbitgemm_kernel_rvv.cpp	Implements RVV QNBitGemm packing/workspace helpers and SQ4/SQ8 kernels; defines the dispatch object.
onnxruntime/core/mlas/lib/riscv64/hqnbitgemm_kernel_rvv.cpp	Implements Zvfh fp16 dequant + fp16 GEMM kernel for HQNBIT_CompFp16 paths.
onnxruntime/test/mlas/unittest/test_sq8bitgemm.cpp	Adds MLAS_TARGET_RISCV64-specific reference packing/layout checks for SQ8 prepack tests.
onnxruntime/test/mlas/unittest/test_qnbitgemm_rvv_fp16.cpp	Adds RISC-V Zvfh fp16 end-to-end tests for HQ4/HQ8 paths.

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

[hariharans29]

Review: PR #29537 — [MLAS] Add RISC-V RVV backend for MLAS QNBitGemm (MatMulNBits)

Approve pending CI green on commit 4408d47. Well-scoped PR that turns MlasIsQNBitGemmAvailable from false to true on RISC-V for the compute types the MatMulNBits operator actually selects among — SQNBIT_CompFp32 (4-bit weights), SQNBIT_CompInt8 (4-bit and 8-bit weights), and, gated on Zvfh, HQNBIT_CompFp16 (4-bit and 8-bit weights). Same shape as the sibling LASX dispatch: private packed-B/block-sum/dequant-buffer layouts consumed only by these kernels, portable pack/workspace helpers shared with the other backends, and clean feature-macro gating in the dispatch object so a build without MLAS_USE_RVV (or without MLAS_USE_RVV_ZVFH) leaves those compute pointers null and falls back to the generic path. 7 files, +1,641/−1.

What looks correct

1. Dispatch object is coherent. qnbitgemm_kernel_rvv.cpp constructs MlasSQNBitGemmDispatchRvv in three well-scoped tiers: portable pack/workspace helpers always, MLAS_USE_RVV-gated compute for SQ4 CompFp32 / SQ4 CompInt8 / SQ8 CompInt8, and MLAS_USE_RVV_ZVFH-gated pack + compute for HQ4/HQ8 CompFp16. platform.cpp assigns this->QNBitGemmDispatch = &MlasSQNBitGemmDispatchRvv; unconditionally in the RVV branch, which is safe because the dispatch itself self-gates the per-ISA compute pointers. mlasi.h adds the extern const declaration next to the other backends' declarations. Nothing surprising.

2. SubBlkLen=16 packed-B layout is consistent across all Q4 consumers. RvvSQ4BitGemmPackQuantBData transforms each 16-nibble sub-block | v0 v1 | v2 v3 | ... | vE vF | (8 bytes) into | v0 v8 | v1 v9 | ... | v7 vF | — element b in the low nibble of byte b, element b+8 in the high nibble of the same byte. All four Q4 consumers — ComputeColumnDot_CompFp32 (M=1 GEMV), Q4BitBlkDequantBForSgemm_CompFp32_Impl (dequant-for-SGEMM), SQ4BitGemmKernel_CompInt8_Impl via UnpackQbCentered, and hqnbitgemm_kernel_rvv.cpp's HQ4BitBlkDequantBForHgemm_CompFp16_Impl via DequantSubblockToFp16 — read the same layout using the same low/high nibble split (b & 0x0F for the first 8 elements, b >> 4 for the next 8). Layout invariant is enforced by the fact that the packer and consumers all live in this dispatch's private cone.

3. Q4BitBlkDequantBForSgemm_CompFp32_Impl correctly targets the MlasSgemmCopyPackB output layout. PackWidth = 16 column panels, K rows per panel, element(k, n_local) at FpData[panel * PackWidth * CountK + k * PackWidth + n_local]. Uses __riscv_vsse32_v_f32m1(dst + off * PackWidth, DstColStride, v, vl) with DstColStride = PackWidth * sizeof(float) to scatter one dequant sub-block down column nl of the current panel — that's the right strided-store idiom for a scalable-vector arch. Partial panels are pre-zeroed via ZeroFloats(panel_base, PackWidth * CountK) before the column-fill, so any n_local >= n_cols slot reads 0.0f as SGEMM expects. Correct.

4. SQ4BitGemmKernel_CompInt8_Impl design is right. The comment on the tile choice — "unpack B once per block into scratch, reuse across MTILE=8 rows, K-reduction at LMUL=4" — matches the code exactly. For each n, for each row-tile [m0, m0+MTILE), for each K-block, for each UNPACK_CHUNK=128 chunk within the block:

   1. Call UnpackQbCentered a few times to fill qbc[UNPACK_CHUNK] with centered int8 (nibble - offset).
   2. Inner loop over the 8 tile rows: vwmul.vv i8 -> i16 at e8m4 (128 int8 lanes at VLEN=256), vwredsum.vs i16 -> i32 accumulates into isum[mi].
   3. After all blocks, acc[mi] += a_scale * b_scale * float(isum[mi]), correctly factoring the two scale multiplies outside the integer dot.

   UNPACK_CHUNK = 128 matching e8m4 VLMAX at VLEN≥256 is the tightest possible fit for a stack int8_t qbc[128] scratch that stays register-file-resident during the tile loop. Nice.

5. SQ8BitGemmKernel_BlkSum_CompInt8_Impl implements the BlkSum identity correctly.

   C += aScale * bScale * sum_i(qa_i * qbRaw_i) - ABlockSum * (bScale * bZeroPoint)
      = sum_i (aScale * qa_i) * (bScale * (qbRaw_i - bZeroPoint))
   

   ABlockSum = aScale * sum_i(qa_i) is computed in RvvQuantizeARowComputeBlkSum_CompInt8_Impl (correct: vsext.vf2 i8→i16 then vwredsum i16→i32, times the block scale). QuantBBlkSum = bScale * bZeroPoint is written by RvvSQ8BitGemmPackQuantBDataAndBlkSum in two places — at scale-packing time with the default zp=128 when zero points are absent, and at zp-packing time with the actual zp when they arrive. Both cover the MatMulNBits 3-call PrePack protocol.

   The kernel uses vwmulsu.vv i8×u8 → i16 at e8m4 — that's the right widening-multiply for signed A × unsigned B raw bytes without pre-subtracting the zero point on B side, since the zp contribution is folded into QuantBBlkSum and subtracted at block granularity via asum_row[b] * bblksum_col[b].

6. PrePack 3-call protocol handled correctly in RvvSQ8BitGemmPackQuantBDataAndBlkSum. The MlasTrySimpleParallel body:

   • Weight-data call (QuantBDataBegin != nullptr): memcpy the raw 8-bit bytes into PackedData.
   • Scale call (QuantBScaleBegin != nullptr): copy scale, and if !HasZeroPoint also write BlkSum[i] = scale * 128.0f as the default.
   • Zero-point call (QuantBZPBegin != nullptr): rewrite BlkSum[i] = PackedScale[i] * float(zp).

   This is the same 3-call handshake the NEON and LASX dispatches use, and it correctly handles all call orderings the operator may deliver.

7. RvvSQ4BitGemmPackQuantBData is thread-parallel over blocks. MlasTrySimpleParallel(ThreadPool, Iterations = N * BlockCountK, ...) parallelizes across (n, k_blk) pairs, with each iteration writing an independent 8-byte-per-SubBlk region — no false-sharing. Correct.

8. HQ4/HQ8 CompFp16 (Zvfh) path is well-decomposed. hqnbitgemm_kernel_rvv.cpp:

   • DequantSubblockToFp16 widens u8 → u16 → u32 → f32, subtracts offset, multiplies by scale, vfncvt.f.f.w f32 → f16, stores. Consistent with the fp32 sibling.
   • HQ4BitBlkDequantBForHgemm_CompFp16_Impl writes column-major B[n*ldb + k] and zero-pads [CountK, ldb), matching what the fp16 GEMM expects.
   • HQ8BitBlkDequantBForHgemm_CompFp16_Impl reads plain byte-per-weight (no nibble split), correct for 8-bit.
   • RvvHQ4BitGemmKernel_CompFp16 uses a 4-row main tile with named vfloat32m2_t acc0..acc3 because RVV vector types are sizeless (cannot be array elements) — documented in a comment. B is widened f16 → f32 once and reused across 4 rows via vfmacc.vv at LMUL=2. R1 tail for odd trailing rows. Reduction via vfredusum back to f32m1, cast to fp16 for the C store. Accumulation is done in fp32 for accuracy — the design note is honored.

9. RvvQuantizeARow_CompInt8_Impl symmetric-int8 quantization is textbook: vfabs.v → vfredmax.vs for per-block amax, scale = amax/127.0f, inv_scale = (scale != 0.0f) ? 1/scale : 0.0f, then vfmul.vf(inv_scale) → vfcvt.x.f → vmax.vx(-127) → vmin.vx(127) → vncvt.x.x.w twice (i32→i16→i8) to store i8. Trailing tail zero-pad. Good — the != 0.0f guard on inv_scale is the standard bit-exact protection against the all-zero A-block case.

10. The RISC-V branch in test_sq8bitgemm.cpp reflects the RVV writer's layout. The x86 branch expects packed-B in a tile-swizzled layout with idx = (((n) / 16) * BlkCount + k) * 16 + (n) % 16 and the ARM64 branch expects its own layout. The new RISC-V branch expects plain [N][ldb] for B, plain [N][BlockCountK] for scale, and plain [N][BlockCountK] for block-sums (idx = n * BlockCountK + k). Those match what RvvSQ8BitGemmPackQuantBDataAndBlkSum actually writes. The prepack self-consistency check is now valid on RISC-V.

11. test_qnbitgemm_rvv_fp16.cpp end-to-end tests cover the Zvfh path. 16 direct tests total: 8 HQ4 + 8 HQ8, sweeping {M, N, K, BlkLen, HasZp, HasBias} including corner cases the design specifically has to handle (K=88, BlkLen=32 — K not a multiple of BlkLen; K=4096, BlkLen=128 — large K; N=17, BlkLen=64 and N=33, BlkLen=64 — N tail). The reference is a straightforward fp32-accumulate GEMM over the packed-and-then-dispatched-dequant B, which is the correct oracle because both the reference and the kernel see the same fp16 dequant. Tolerance maxrel < 6e-3 is loose enough for fp16 downcast noise on the reduction and tight enough to catch dequant errors. Good.

Non-blocking observations

1. RvvQ8BitGemmPackQuantBDataSize slight over-allocation for BlkSum. The sizing MlasDivRoundup(N, 16) * BlockCountK * 16 * sizeof(float) matches the NEON/LASX/AVX-512 convention that reserves a [N/16 × BlockCountK × 16] padded layout, but the RVV writer stores plain [N][BlockCountK]. This over-allocates by (round_up(N, 16) - N) * BlockCountK * sizeof(float) for N not a multiple of 16. Harmless (the writer touches less than the reserved region and the PackedQuantBDataStruct cursor arithmetic still lands on the right offset), and staying with the standard formula keeps the sizing consistent with the sibling backends. Could be tightened to N * BlockCountK * sizeof(float) in a follow-up if desired, but it's not worth churning for.

2. HQ8 test at test_qnbitgemm_rvv_fp16.cpp line 142 calls dispatch->HQ4BitGemmKernel_CompFp16 even though the widths are 8-bit. This is deliberate — after the HQ4/HQ8 dequant, the downstream fp16 GEMM consumes an identical [N][ldb] fp16 buffer, and the AVX-512/NEON siblings share the same fp16 kernel between HQ4 and HQ8 for the same reason (the dispatch struct doesn't even have a separate HQ8BitGemmKernel_CompFp16 entry). A one-line comment at that call site — // shared HQ4/HQ8 fp16 GEMM after dequant, dispatch has one entry — would remove a potential double-take.

3. SQ4BitGemmKernel_CompInt8 and SQ8BitGemmKernel_BlkSum_CompInt8 have no dedicated RVV unit tests. Coverage rests on the existing generic MLAS suite (SQNBitGemm*, SQ8Bit*, BlockQ4*, BlockQ8*), which the PR body reports at 13,526/13,526 passing on the K3 board. Those tests do exercise the RVV path once the RISC-V platform dispatch is selected, so this is effectively fine, but a small dedicated file mirroring test_qnbitgemm_rvv_fp16.cpp for the CompInt8 paths (SQ4 and SQ8, with/without zero points, one shape per BlkLen) would tighten local pre-submit turnaround for future RVV kernel changes. Non-blocking.

4. No CI leg on rv64gcv. As with the LASX and CANN dispatches, kernel correctness rests on the author's local hardware validation (K3 board, rv64gcv_zvfh, VLEN=256). That's the right posture for a hand-optimized ISA-specific backend that adds no shared surface — the packed-B / block-sum / dequant-buffer layouts are private to MlasSQNBitGemmDispatchRvv, so a bug here cannot silently corrupt any other arch's dispatch. Just worth noting explicitly that the merge signal comes from the author's local results plus code review, not from CI.

5. ZeroFloats(panel_base, PackWidth * CountK) in Q4BitBlkDequantBForSgemm_CompFp32_Impl zero-fills the entire partial panel before filling n_cols < PackWidth columns. Correct but slightly wasteful — most of the ZeroFloats bytes will be immediately overwritten. Could be tightened to zero only the unused-column strides (for nl in [n_cols, PackWidth): ZeroFloats(panel_base + nl, CountK) with stride PackWidth), but a partial panel occurs at most once per GEMM, so this is imperceptible. Style nit.

6. assert(BlkLen >= 16 && BlkLen % 16 == 0) in RvvSQ4BitGemmPackQuantBData. Debug-only. Existing MLAS kernels are inconsistent about whether kernel-internal preconditions use assert or MLAS_ASSERT, so this doesn't stand out as drift; just carrying forward the same nit I raised on the sibling AVX2 W2 PR.

7. CompFp32 M=1 GEMV kernel is n-major (for n) with a nested for k_blk. For typical LLM shapes (N large, M=1) this is fine — the outer loop is the natural parallelization axis if a future threadpool split were added. No action needed.

8. UNPACK_CHUNK = 128 bound in SQ4BitGemmKernel_CompInt8_Impl is documented as "≤ e8m4 VLMAX at VLEN≥256." For VLEN=128 hardware (rv64gcv minimum), e8m4 VLMAX = 64, so the __riscv_vsetvl_e8m4(clen - off) inside the row loop will return 64 and iterate twice per 128-element chunk. Correct — vsetvl handles the VLEN-clamp — but on VLEN=128 the double iteration slightly reduces the amortization benefit of the MTILE=8 tile. Not actionable; the author benchmarked on VLEN=256 hardware and that's what the PR is optimized for.

9. RvvHQ4BitGemmKernel_CompFp16 reinterprets fp16 loads via u16m1. __riscv_vreinterpret_v_u16m1_f16m1(__riscv_vle16_v_u16m1(...)) is the standard RVV idiom because there's no direct fp16 vle intrinsic in the rv64gcv_zvfh spec I'm aware of. Correct, and worth keeping as-is.

10. SQ4BitGemmKernel_CompInt8_Impl isum promotion path. is = __riscv_vmv_s_x_i32m1(0, 1) seed, vwredsum_vs_i16m8_i32m1 accumulates the widened products. __riscv_vmv_x_s_i32m1_i32(is) extracts the scalar. Standard RVV reduce-into-lane-0 idiom. Correct.

CI

Latest commit 4408d47 shows 1/1 checks OK, and the preceding 17e4b47 was also 1/1 — only a subset of pipelines ran on these last two commits, which matches the microsoft-side CI infrastructure issues that have been in evidence on other in-flight PRs this week. The initial functional commit 1226f35 reached 68/85. Nothing in the two doc-only follow-up commits should regress those pipelines. Recommend confirming the full x86 / Linux / macOS legs (which the RVV code doesn't touch at all — the RVV files are only compiled under MLAS_TARGET_RISCV64 and MLAS_USE_RVV, and the test file is gated on MLAS_TARGET_RISCV64 && MLAS_USE_RVV_ZVFH) are still green on 4408d47 once the pipelines are back up.

Summary

Clean, tightly scoped MLAS backend PR that closes the RISC-V gap for MatMulNBits. Design mirrors the LoongArch LASX sibling (private packed layouts, shared portable helpers, feature-macro-gated compute). RVV kernel implementations look correct at every level I checked — nibble extraction and layout invariance across the four Q4 consumers, tail handling in the sub-block dequant and in the K-quantize row, MTILE-based B-unpack reuse for CompInt8, BlkSum identity for SQ8, R4 fp16 tile with named accumulators (the correct workaround for RVV's sizeless vector types), LMUL=4 K-reduction for the CompInt8 kernels tied explicitly to the widening-multiply throughput lever the author identified in profiling. Test coverage is thorough for the Zvfh path (16 end-to-end tests, corner cases included) and rests on the generic MLAS suite for CompInt8 — same posture as the other arch-specific backends without CI legs. Two very small comment-fix commits since the initial submission address the two Copilot AI passes and the lintrunner clang-format finding.

[hariharans29]

If you have resolved the Copilot comments in the code - can you please manually resolve the comments visible on the PR ?

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[velonica0]

If you have resolved the Copilot comments in the code - can you please manually resolve the comments visible on the PR ?

Okay, I've resolved. Thank you.