[CUDA] MatMul: add an opt-in split-K GEMV for small-N fp16 shapes

[Tianlei Wu (tianleiwu)]

Description

Decode-time projections in hybrid / MoE LLMs are GEMVs with an N that is too small to fill a cuBLAS tile kernel. On the Qwen3.6-35B-A3B NVFP4 decode loop, router, linear-attention, and shared-expert gate projections repeatedly launch with M <= 8, N <= 1024, and K >= 128.

This PR adds an experimental split-K GEMV kernel for that corner of the shape space. The path remains off by default because real-model measurements are currently slower than cuBLAS.

Summary of Changes

Kernel and dispatch

• Adds a row-major FP16 split-K GEMV for M <= 8, N <= 1024, and K >= 128.
• Gates dispatch with ORT_ENABLE_SMALL_N_GEMV=1; all ineligible layouts, shapes, transpose modes, and alpha values fall through to cuBLAS.
• Reads the opt-in setting once when each MatMul kernel instance is created, keeping environment parsing out of Compute().
• Compares leading dimensions in int64_t so large shape values are not narrowed before eligibility checks.

Cross-block reduction

• Uses a grid of (ceil(N / 32), split_k) so K slices can run across multiple SMs.
• Publishes FP32 partials through volatile global workspace accesses before __threadfence() and the completion atomic, following CUDA's last-block reduction visibility pattern.
• Reduces partials in fixed slice order for deterministic output.
• Uses per-invocation scratch for both workspace and completion counters. The caller clears counters before every launch; the kernel does not reset them.

Tests

• Direct-kernel tests cover every M specialization, N = 1, N = 32, N = 1024, uneven K splits, and repeated launches with explicit counter initialization.
• Inputs vary along K and results are checked against an FP32 CPU reference.
• Operator-level MatMul tests enable the feature and cover eligible M = 8, N = 1 and M = 8, N = 1024 dispatches plus an ineligible K = 127 cuBLAS fallback.

Current Status

The path is default-off because it is currently slower than cuBLAS in the real model. On H200 (SM90) for the target decode shapes:

	us / call	us / decode step
cuBLAS	5.0	671
this kernel	10.6	1491

A standalone microbenchmark had suggested approximately 2.1 us/call after subtracting the back-to-back launch floor. That result kept the small weight matrix resident in H200's L2, while the model reads cold weights. Keeping the implementation behind an opt-in flag preserves it for follow-up work that addresses that cold-read cost.

Validation

• Compiled matmul.cc, matmul_small_n_gemv.cu, and both small-N GEMV test translation units with the CUDA 13.0 Release build.
• Ran scoped lintrunner checks for all changed source and test files.
• Full onnxruntime_test_all linking is currently blocked by an unrelated -Werror=unused-variable in contrib_ops/cuda/bert/paged_attention.cc on the rebased main branch.

Checklist

• No behavior change by default
• Cache-safe cross-block workspace publication
• Deterministic fixed-order reduction
• Operator-level enabled dispatch and fallback coverage
• Enabled by default; blocked on closing the real-model performance gap against cuBLAS

[Copilot]

Pull request overview

Adds an experimental, opt-in CUDA fast path for decode-time MatMul shapes that behave like small‑N GEMVs (fp16, row-major, no transpose), using a split‑K kernel to increase SM utilization. The path is gated behind ORT_ENABLE_SMALL_N_GEMV and falls back to the existing cuBLAS implementation by default.

Changes:

• Introduces SmallNGemvSplitKKernel and host-side helpers (eligibility, split‑K heuristic, workspace/counter sizing, launch).
• Adds a guarded fp16 fast path in MatMul<T>::ComputeDefault before the cuBLAS GEMM call.
• Adds a CUDA unit test that directly launches the kernel and checks counter reset behavior across repeated launches.

Reviewed changes

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

File	Description
onnxruntime/core/providers/cuda/math/matmul.cc	Adds env-var gate and routes eligible fp16 MatMul into the new small‑N GEMV split‑K kernel.
onnxruntime/core/providers/cuda/math/matmul_small_n_gemv.h	Declares eligibility + sizing helpers and the launch entry point for the split‑K GEMV kernel.
onnxruntime/core/providers/cuda/math/matmul_small_n_gemv.cu	Implements the split‑K kernel and dispatch for M=1..8, including deterministic slice-ordered reduction.
onnxruntime/test/providers/cuda/test_cases/matmul_small_n_gemv_test.cc	Adds direct-kernel unit tests covering M variants, tile boundaries, and counter reuse/reset semantics.

Suppressed comments (1)

onnxruntime/core/providers/cuda/math/matmul.cc:348

• The kernel resets the arrival counter internally (matmul_small_n_gemv.cu sets counter[blockIdx.x] = 0), and the unit test relies on that by launching twice without re-memsetting. Here the MatMul fast-path still does a cudaMemsetAsync on every launch, which adds overhead on the microsecond-scale GEMV path and defeats the intended "clear once, then reuse" design described in the PR.

Consider caching a counter buffer across launches (or otherwise ensuring one-time initialization) so the per-call memset can be removed, or alternatively drop the in-kernel reset if you intend the counter to be per-launch scratch only.

        const size_t counter_elements = SmallNGemvCounterElements(n);
        auto counter = GetScratchBuffer<unsigned int>(counter_elements, this->GetComputeStream(ctx));
        CUDA_RETURN_IF_ERROR(cudaMemsetAsync(counter.get(), 0, counter_elements * sizeof(unsigned int), Stream(ctx)));
        auto workspace = GetScratchBuffer<float>(SmallNGemvWorkspaceElements(m, n, k), this->GetComputeStream(ctx));

[Ti-Tai Wang (titaiwangms)]

Thanks for putting this behind an opt-in flag and documenting the real-model benchmark result. The split-K decomposition, bounds, fixed-order reduction, build integration, scratch allocation, and stream ordering all look reasonable. However, I think this needs changes before merge.

Blocking issue: cross-block workspace visibility

In matmul_small_n_gemv.cu, blocks publish ordinary stores to ws, execute __threadfence(), increment the completion counter, and then the last block reads all workspace partials.

CUDA's documented last-block reduction pattern states that a memory fence orders operations but does not, by itself, guarantee visibility to other blocks. The communicated buffer must use volatile/cache-bypassing accesses so the final block cannot consume stale L1 data. Here ws is a plain float* __restrict__.

This is particularly relevant for small shapes such as M=1, N=1, where all split-K partials share a cache line. The result could be an architecture- or scheduling-dependent silent wrong answer.

Please make workspace publication and consumption explicitly cache-safe, for example using the documented volatile pattern, cache-global load/store intrinsics, or a separate reduction kernel. An A100 stress run passing does not resolve the specification issue because its cache behavior may not expose the hazard.

Test coverage

The added tests call LaunchSmallNGemv() directly, so they do not cover the externally relevant path in MatMul<T>::ComputeDefault:

• ORT_ENABLE_SMALL_N_GEMV enabled with an eligible shape;
• ineligible shapes falling back to cuBLAS;
• the layout/transpose/alpha guards;
• scratch-buffer sizing and stream integration.

Please add at least one operator-level MatMul test that exercises the enabled dispatch and one fallback case. The current test inputs are also constant along K, which makes some K-indexing and split-boundary errors invisible. A random FP16 case checked against a higher-precision CPU reference would provide stronger coverage, especially for M=8, N=1 and M=8, N=1024.

Counter contract and PR description

The production path clears the counter before every launch, while the kernel also resets it and the test relies on that reset for a second launch. The header nevertheless says callers must provide a zeroed counter. Please choose and document one ownership/reset contract; if kernel-side reset is retained, use an explicitly atomic reset.

The PR description also appears stale: it mentions a cached counter and mutex in matmul.h, but the current diff allocates a scratch counter and clears it for each invocation.

Verdict: request changes. The overall integration is clean and the feature is safely default-off, but the cross-block publication protocol should be made CUDA-memory-model safe and the actual MatMul dispatch should be covered before merge.

[Tianlei Wu (tianleiwu)]

Addressed the blocking feedback in e268362:

• workspace publication/consumption now uses volatile global accesses around the fence + completion atomic protocol;
• counter ownership is caller-only: scratch counters are zeroed before each launch and the kernel no longer resets them;
• direct-kernel inputs vary along K and are checked against an FP32 CPU reference;
• operator-level tests cover enabled eligible dispatch (M=8, N=1/1024) and an ineligible K=127 cuBLAS fallback;
• the PR description now reflects the scratch counter design and current tests.

The four touched production/test translation units compile with the CUDA 13.0 Release build, and scoped lintrunner checks pass. A full onnxruntime_test_all link is currently blocked by an unrelated -Werror=unused-variable in contrib_ops/cuda/bert/paged_attention.cc on the rebased main branch.