onnxruntime

d2d1e8fe - [CUDA] Add tinygemm2 as a GEMM auto-tune candidate for small-M MatMul on SM 9.0+ (#32901)

[CUDA] Add tinygemm2 as a GEMM auto-tune candidate for small-M MatMul on SM 9.0+ (#32901) ## Description Adds TensorRT-LLM's `tinygemm2` kernel as a third candidate for the CUDA fp16/bf16 `MatMul` auto-tuner introduced in #32876 (`ep.cuda.enable_gemm_auto_tune=1`). On H200 it makes decode-sized projections up to 2.2x faster than cuBLAS under CUDA graph replay, and the tuner keeps cuBLAS wherever tinygemm2 is slower. Without the session option, behavior is unchanged. **Stacked on #32876** (base branch `tlwu/cuda-small-m-decode-gemv`); only the last commit belongs to this PR. ## Summary of Changes | File | Change | |------|--------| | `onnxruntime/core/providers/cuda/math/tinygemm2_kernel.cuh` | Kernel adapted from TensorRT-LLM `cpp/tensorrt_llm/kernels/tinygemm2` (Apache 2.0, header retained). A CTA computes 16 output columns x 8 rows; 4 compute warps run `mma.sync.m16n8k16` with fp32 accumulation, 4 warps TMA-load B and 4 load A through a 16-stage mbarrier pipeline. | | `onnxruntime/core/providers/cuda/math/tinygemm2.{h,cu}` | Eligibility, per-device capability probe, tensor-map encoding and launch with programmatic dependent launch. | | `onnxruntime/core/providers/cuda/math/matmul.{h,cc}` | Eligible candidates (small-N GEMV, tinygemm2) are collected into a mask; the tuner times cuBLAS against every eligible candidate. `ORT_ENABLE_SMALL_N_GEMV` keeps forcing only the small-N GEMV. | | `onnxruntime/core/providers/cuda/math/gemm_auto_tuner.{h,cc}` | `GemmKernel::kTinyGemm2`; the tuning key includes the eligible-candidate mask, since operand alignment can change it. | | `include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h` | Documents the new candidate. | Changes from the TensorRT-LLM kernel: - **B is read in the ONNX `MatMul` `[K, N]` layout in place.** TensorRT-LLM expects `[N, K]` weights. The B tile is loaded as 64 K-rows x 16 columns with a 32B TMA swizzle and read with `ldmatrix.x4.trans`, so no transposed weight copy (and no extra memory) is needed. - fp16 in addition to bf16 (only the MMA instruction and output conversion differ). - No bias; raw `mbarrier` PTX instead of `cuda::barrier`; the compute-warp reduction uses a named barrier for its 128 threads instead of a `__syncthreads()` that the loader warps never reach. Eligibility: `1 <= M <= 64`, `N * K <= 32M` elements, `N % 8 == 0` and `K % 8 == 0` (TMA row pitch), 16-byte aligned A and B, no transposes, `alpha == 1`. Device support is probed once per device: the kernel's PTX version must be >= 9.0 (a binary without SM 9.0+ code, or JIT from older PTX, has an empty kernel body) and its ~195 KiB of shared memory must fit `sharedMemPerBlockOptin` (so e.g. SM 12.0 consumer Blackwell is excluded). The kernel is launched with programmatic dependent launch (CUDA >= 12.3): the weight loads overlap the preceding kernel, and only the activation loads wait for it via `cudaGridDependencySynchronize()`, which also orders the output stores. ## Testing - Standalone check of the ported kernel against a float64 reference for fp16 and bf16, M in {1,3,8,9,17,64} x N in {8,16,24,48,136,1000} x K in {8,64,136,720,1032,5120}, plus 4x2880x2880 and 64x4096x5120: all within output rounding. `compute-sanitizer` memcheck, synccheck and racecheck report 0 errors. - CUDA internal tests (`onnxruntime_providers_cuda_ut`, H200): new `TinyGemm2Test.Eligibility` and `TinyGemm2Test.MatchesReference` (partial row/column tiles, K below one tile and not a multiple of the 1024-element K loop, NaN-poisoned output), plus the existing `GemmAutoTunerTest.*` / `MatMulSmallNGemv*` suites with tinygemm2-eligible operator shapes (N = 2880 and 1032) added: 17 passed. - `onnxruntime_provider_test --gtest_filter='*MatMul*:*Gemm*'` (1003 tests) passes with and without `ORT_CUDA_GEMM_AUTO_TUNE=1`; Python CUDA graph tests pass. - CUDA plugin EP build: auto-tune selects tinygemm2 and CUDA graph capture/replay (with the programmatic launch edges) produces correct results. - Compiled for sm_86 and sm_90a in both the in-tree and plugin builds. H200, CUDA graph replay of enough MatMul copies that weights stream from DRAM, K = 5120, per-MatMul us, best of two sessions per mode: | dtype | M | N | cuBLAS | auto-tune | choice | speedup | |---|---:|---:|---:|---:|---|---:| | bf16 | 1 | 48 | 6.24 | 3.35 | tinygemm2 | 1.87x | | bf16 | 8 | 48 | 6.99 | 3.38 | tinygemm2 | 2.07x | | bf16 | 64 | 48 | 7.05 | 3.41 | tinygemm2 | 2.07x | | bf16 | 8 | 1024 | 8.06 | 4.73 | tinygemm2 | 1.70x | | bf16 | 16 | 1024 | 8.16 | 8.17 | cuBLAS | 1.00x | | bf16 | 8 | 2048 | 10.99 | 8.09 | tinygemm2 | 1.36x | | bf16 | 8 | 5120 | 19.23 | 19.30 | cuBLAS | 1.00x | | fp16 | 16 | 48 | 7.29 | 3.35 | tinygemm2 | 2.18x | | fp16 | 16 | 1024 | 8.10 | 6.94 | tinygemm2 | 1.17x | | fp16 | 8 | 2048 | 10.75 | 7.82 | tinygemm2 | 1.37x | | fp16 | 64 | 1024 | 8.51 | 8.51 | cuBLAS | 1.00x | Auto-tune was never slower than cuBLAS across M in {1,4,8,16,64} x N in {48,1024,2048,5120} for both dtypes. The isolated tuner timing does not capture the programmatic-launch overlap, so near ties (e.g. bf16 16x1024) conservatively stay on cuBLAS. ## Motivation and Context - Stage 2 of the GEMM auto-tuner in #32876, which was designed as a candidate table so kernels like this can be added without changing the dispatch policy. - tinygemm2 targets exactly the latency-bound decode/speculative-verify shapes (a handful of rows against a small-to-medium weight) where cuBLAS spends most of its time on launch/ramp and tile under-utilization. - Follow-ups: FlashInfer TGV GEMM for SM 10.x, cuBLASLt heuristic algorithms as candidates, and a persistent tuning cache (#29588). ## Checklist - [x] Tests added/updated - [x] No breaking changes (only active with `ep.cuda.enable_gemm_auto_tune=1` on SM 9.0+) - [x] Documentation updated (session option key comment) --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>