[Tianlei Wu (tianleiwu)]

Description

The fpA_intB MatMulNBits tactic profiler ran 5 warmup and 10 timed launches for every candidate tactic, no matter how long a launch took. For large GEMMs this makes profiling very expensive. On H200, profiling Qwen3.8-27B's lm_head (N=248320, K=5120, INT4) at M=4096 covers 60 valid tactics with a median launch time of 50 ms. That is about 48 s for one GEMM shape.

The cost matters because unprofiled M buckets are profiled lazily inside Compute(). The default initial sweep stops at M=2048, so the first run with M=4096 (for example, a GenAI Engine with max_scheduled_tokens=4096 doing batched prefill) stalls. In an 8-request, 512-token-prompt batch the first step took 71–98 s. The same step takes 1–2 s once M=4096 is profiled. That looked like a hang in end-to-end serving tests.

This change makes tactic timing adaptive:

• Run one untimed launch, then time one warm launch as a probe.
• If the probe takes 2 ms or less, average 10 launches, the same timed count as before. Small GEMMs are not affected.
• For costlier tactics:
  • If the probe is already more than 1.5x the best average so far, skip the tactic, since it cannot win.
  • Otherwise, average as many launches as fit in a 20 ms budget, with at least one.

Summary of Changes

File	Change
onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h	Add GetProfileTimedRuns() (the timing policy). profileTacticForProblem() now probes one warm launch, prunes or budgets the timed runs, and uses RAII for its CUDA events.
onnxruntime/test/providers/cuda/test_cases/matmul_nbits_workspace_test.cc	Unit test for the policy: cheap tactics keep the full average, expensive tactics that can still win are averaged within budget, and expensive tactics that are clearly slower are skipped.

Testing

• MatMulNBitsWorkspace.* CUDA EP internal tests: 26/26 pass, including the new TacticProfilerTimedRunsPruneOnlyExpensiveSlowTactics.
• onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py: 4 passed, 9 skipped because the weight packer is not available in the local build.
• H200 A/B, same build with and without this change.

Standalone MatMulNBits node, lm_head shape (N=248320, K=5120, block 32, INT4):

	Before	After
First run at M=4096 (lazy profile)	48.7–50.8 s	12.1–15.3 s
Selected tactic at M=4096	91	91
Steady-state GEMM at M=4096	29.0–29.3 ms	29.0–29.1 ms

Tactic selection across the 8 Qwen3.8-27B GEMM shapes and 13 M buckets (1…4096):

• 102 of 104 (M, N, K) problems picked the same tactic.
• The two differences were noise between separate runs:
  • One reproduced identically when rerun.
  • The other was 0.1% apart.
• Using the per-tactic times recorded during profiling, the modeled total profiling GPU time for these shapes drops from 135 s to 47 s.

End to end: Qwen3.8-27B INT4 + DFlash2, GenAI Engine, batch 8 × 512-token prompts, H200:

Scenario	Before	After
First step with max_scheduled_tokens=4096 (lazy M=4096 profile)	98.1 s	20.8 s
Process wall time, max_scheduled_tokens=512 (load + run)	104.6 s	56.2 s
Prefill step, max_scheduled_tokens=512	1.295 s	1.298 s
128-token speculative decode	2.322 s	2.332 s

• The two runs of the 128-token speculative decode accepted the same number of drafts (843/843).
• The 4-way concurrent runs above were on separate GPUs, so the absolute wall times include some host contention.

Motivation and Context

Reported from a GenAI Engine probe that appeared to stall indefinitely with max_scheduled_tokens=4096. Root cause:

• The first 4096-row step triggered lazy fpA_intB profiling for every distinct GEMM shape, and most of that time went to the vocabulary projection.
• Session-construction profiling of the default buckets (up to M=2048) had the same per-tactic cost. That is why model load time also drops.

Deployments that know their largest M can still set ep.cuda.fpa_intb_profile_m so all profiling happens at session creation. This change makes both paths considerably cheaper.

[Copilot]

Copilot review overview

🟢 Approval recommended

The focused implementation is consistent with its stated policy and has appropriate boundary coverage.

Review effort: Balanced
Findings: None

What changed in this PR

Bounds CUDA fpA_intB GEMM tactic-profiling latency while preserving full sampling for inexpensive tactics.

Changes:

• Adds probe-based pruning and a 20 ms timing budget.
• Uses RAII for CUDA events and adds policy unit tests.

File	Description
onnxruntime/​contrib_ops/​cuda/​llm/​gemm_profiler.h	Implements adaptive tactic timing.
onnxruntime/​test/​providers/​cuda/​test_cases/​matmul_nbits_workspace_test.cc	Tests timing, budgeting, and pruning thresholds.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.