[CUDA] Gracefully disable MatMulNBits fpA_intB fast path when it cannot serve prefill (narrow-N int4)

[justinchuby]

Description

The MatMulNBits CUDA fpA_intB fast path (onnxruntime_USE_FPA_INTB_GEMM=ON, enabled at runtime via ORT_FPA_INTB_GEMM) serves prefill shapes (m >= 16) exclusively through the CUTLASS weight-only GEMM, because the fused CUDA GEMV kernel only supports small m and throws unsupported m (contrib_ops/cuda/llm/fpA_intB_gemv/dispatcher.h) otherwise.

For narrow-N int4 models, the CUTLASS heuristic can fail to find a valid GEMM tile for one or more m buckets (the profiler logs Have not found any valid GEMM config for shape ...). At prefill this leaves the fast path with no usable tactic, so the node either:

• getBestConfig() returns nullopt → fails with No valid fpA_intB MatMulNBits tactic for M=..., or
• only the small-m CUDA GEMV kernel "survived" → the GEMV dispatcher throws unsupported m.

Because PrePack drops the original quantized weights, there is no runtime fallback once the fast path is committed — the model cannot run at all.

Reported in #29691.

Fix

At construction, after RunGemmProfile, verify that every profiled m >= 16 bucket has a usable CUTLASS GEMM tactic (getBestConfig returns a value that is not the small-m-only CUDA GEMV kernel). If any bucket cannot be served, disable the fast path — but only when the model did not ship pre-prepacked weights (those have their own contract enforced just below). Disabling has_fpA_intB_gemm_ before PrePack skips prepacking, so the original quantized weights are retained and the default MatMulNBits kernel runs correctly instead of failing at prefill.

Tactics are profiled per rounded power-of-two m bucket, so probing powers of two covers the full range. For well-shaped (large-N) models every bucket has a valid GEMM tactic, so the fast path stays enabled and behavior is unchanged.

Single-file change (contrib_ops/cuda/quantization/matmul_nbits.h). The runtime nullopt guard already present in matmul_nbits.cc is left as-is.

Validation

Built ONNX Runtime from source with --use_cuda --cmake_extra_defines onnxruntime_USE_FPA_INTB_GEMM=ON, CMAKE_CUDA_ARCHITECTURES=90a (Hopper/H200), CUDA 12.9. Repro model: Qwen2.5-0.5B-Instruct int4, block_size=128, symmetric (3-input), narrow-N attention projections.

	ORT_FPA_INTB_GEMM=1 (before)	ORT_FPA_INTB_GEMM=1 (after)	ORT_FPA_INTB_GEMM=0 (default)
prefill	❌ crash unsupported m	✅ runs	✅ runs
decoded text	—	identical to default	baseline

After the fix, ORT_FPA_INTB_GEMM=1 produces bit-identical output to the default kernel, confirming a clean fallback.

Note: local validation was performed against a from-source build of the equivalent code on SM90a with a single narrow-N int4 model. I was not able to run the full CUDA CI matrix locally; a maintainer review of the gating heuristic and a large-N regression check would be appreciated.

[Copilot]

Pull request overview

This PR hardens the CUDA MatMulNBits fpA_intB fast path selection so that if CUTLASS can’t provide a valid weight-only GEMM tactic for any prefill-relevant m >= 16 bucket (common in narrow‑N int4 models), the fpA_intB path is disabled early (before PrePack) to preserve original quantized weights and allow the default MatMulNBits kernel to run instead of failing at runtime.

Changes:

• After RunGemmProfile, probes profiled power-of-two m buckets (m >= 16) and disables the fpA_intB fast path if any bucket has no usable tactic.
• Applies this disablement only when weights were not shipped as pre-prepacked (weight_prepacked == 0), preserving the offline-prepacked contract.

[tianleiwu]

FpA_IntB is an opt-in feature, and we shall not silently fallback to non FPA_INTB when user explicitly enable FPA_INTB. Recommend solution is like:

If a CUTLASS tactic is missing for any $M \ge 16$ bucket, raise an error such as:

fpA_intB MatMulNBits is unsupported for this shape: no CUTLASS prefill tactic for N=224, K=896, int4, block_size=128. Disable ORT_FPA_INTB_GEMM to use the standard MatMulNBits kernel.