onnxruntime
e0ad0718 - [CUDA] Enable native SM90, block_size=32, and fused bias for fpA_intB MatMulNBits (#29585)

Commit
16 days ago
[CUDA] Enable native SM90, block_size=32, and fused bias for fpA_intB MatMulNBits (#29585) ### Description Extends the CUDA fpA_intB weight-only `MatMulNBits` path (FP16/BF16 int4/int8) in three ways. All are gated behind the existing `ORT_FPA_INTB_GEMM` opt-in; default behavior is unchanged. **1. Native SM90 (Hopper) mixed-GEMM kernel** - `cutlass_extensions` `GemmFpAIntB::operator()`: reuse the SM80 (Ampere) mixed-GEMM kernel for `__CUDA_ARCH__ >= 900` (Hopper/Blackwell) instead of stubbing it out, so the SM80-compat dispatch produces a real kernel on Hopper (mirrors `MoeFCGemm::operator()`). - `fpA_intB_launcher_sm90.inl`: gate the host-callable launcher on `COMPILE_HOPPER_TMA_GEMMS` only (not `__CUDA_ARCH__` / `__NV_SASS_VERSION__`, which are unset in the host pass), so the native SM90 TMA/WGMMA launcher symbol is emitted instead of the "recompile with 90a" stub. - `CutlassFpAIntBGemmRunnerInterface`: add `setArch(int)` and `setUseSm90Native(bool)`; the runner routes SM90 to `sm90_dispatch` when native, else the SM80 compat path. - `matmul_nbits`: `FpAIntBPackingSmForKernel()` returns 90 for `weight_prepacked=2` on an SM90 device (else 80); `InitGemmProfiler` opts into the native kernel (sm==90) or forces the runner to SM80 (compat on Hopper) so tactic enumeration and workspace sizing match the dispatched kernel; `GemmIdCore` gains an `sm` field so SM80-compat and SM90-native tactics for the same shape are not confused; the GEMV is launched with the packing arch (not the raw device SM) to match the packed interleave layout. - `weight_prepacked=2` (SM90 layout) is now accepted (was reserved/rejected): requires an SM90 device and `block_size ∈ {64, 128}`. Contrib-op docs updated accordingly. **2. `block_size=32`** - Relax the group-size gates in the GEMV dispatcher, the CUTLASS fine-grained `can_implement` / kernel launcher, and the `MatMulNBits` eligibility check. The SM80 fine-grained kernel supports group size 32 natively. **3. Fused bias** - Drop the `!has_bias_` gate: a fused bias (input 5) is already supported by the GEMV, the SM80/SM90 CUTLASS epilogue, and the tactic profiler, so bias-bearing nodes (e.g. gpt-oss `qkv_proj`/`o_proj`) become eligible. ### Motivation and Context Enables the fpA_intB weight-only path for more shapes and for Hopper-native execution, unblocking int4/int8 models with fused bias and `block_size=32` on H200. Builds on #29499 (prepacked fpA_intB weights). ### Testing - Built with `USE_FPA_INTB_GEMM=ON` (arch 80;90) on H200. - C++ `onnxruntime_provider_test --gtest_filter='*MatMulNBits*:*FpAIntB*'`: 65 passed, 1 skipped. - `test_op_matmulnbits_prepacked_cuda.py`: 7 passed (int4/int8, SM80/SM90, bias parity vs runtime prepack).
Author
Parents
Loading