[Tianlei Wu (tianleiwu)]

Description

Validate GEMM auto-tuning candidates using L2-conditioned CUDA graph replay before caching the selected kernel in graph-enabled sessions. This addresses cases where stream-based timing selects a candidate that is slower during subsequent graph replay, while retaining the existing default-cuBLAS speedup margin and a legacy-policy opt-out.

Stacked on #32901, with base branch tlwu/cuda-gemm-tinygemm2. Only commit 18e327f8a7 belongs to this PR; it is a clean cherry-pick of c3914243a960865ff2d7bdc70ac2661561db5a2d. Retarget to main after #32901 merges.

Summary of Changes

• Capture interleaved candidate measurements, including L2 eviction and activation warming outside the event-timed regions. Warm and replay the graph, synchronize completion, and apply the existing 5% margin before inserting the winner into the cache.
• Use external event-record nodes to measure replayed work; scope capture, graph, and executable cleanup to tuning. Reject timing inside an existing capture without ending the caller's capture.
• Separate stream-timed and replay-timed cache keys. Cache-hit execution remains unchanged.
• Propagate the plugin's enable_cuda_graph setting to MatMul kernel configuration; use the existing provider setting for traditional CUDA EP routing.
• Preserve stream timing for non-graph sessions and expose ORT_CUDA_GEMM_GRAPH_REPLAY_TUNING=0 to restore legacy timing in graph-enabled sessions. Unset or 1 enables replay validation when both CUDA graphs and GEMM auto-tuning are enabled.
• Add regression tests for replay selection, a stream-winner/replay-loser case, cache separation, capture-failure cleanup, existing-capture ownership, and L2/hot-buffer replay.

Existing auto-tuning defaults and forced-kernel precedence are unchanged. Verbose logs identify timing=stream or timing=cuda_graph_replay.

Testing

The implementation patch was built and tested locally in the isolated current-main + #32901 integration worktree on H200. The submitted cherry-pick has the same stable patch ID; no new build was performed on the publication branch.

• CUDA plugin build and focused GemmAutoTunerTest.*: 12/12 passed, with no skipped or zero-match tests. Tests were rerun after formatting and through the archived build harness.
• 21 real-weight trials: 96 unchanged decoder FP16 MatMuls, K=5120/N=48, M=1/8/16 with graph replay, plus graph-OFF controls; three fresh processes per setting. Candidate/core mappings, finite outputs, and bitwise replay stability verified. TinyGemm2 remained selected with unchanged output hashes across policies.
• Six restricted M=8 cuBLAS/small-N trials, excluding TinyGemm2 without changing production eligibility: replay policy selected cuBLAS in 3/3 trials; legacy policy selected small-N in one of three. Subsequent hot replay measured approximately 6.212 us cuBLAS versus 6.640 us small-N.
• GenAI smoke: 512-context, batch one, seven drafts, 512 generated tokens; active speculation and zero failures, disables, or fallback. Loaded-library hashes and unchanged model configuration verified.
• Clean cherry-pick, exact seven-file scope, whitespace checks, and lintrunner -a passed on this branch.

Cost and Limitations

Replay validation added approximately 10 ms to first-run startup in the fixed-input diagnostic; no steady-state speedup is claimed where the selected candidate is unchanged. L2 conditioning approximates decoder cache state, but does not model every surrounding graph or cache history.

Local numerical measurements used real FP16 weights and synthetic activations. Traditional CUDA EP build validation, BF16 operator coverage, and broader model accuracy/acceptance validation remain to be completed before promotion. The installed runtime and serving models were not replaced.

Checklist

• Tests added/updated
• No public API changes
• Tuning contract and opt-out documented
• Local CUDA plugin validation completed
• Traditional CUDA EP and BF16 validation completed

[David Fan (jiafatom)]

Please resolve conflicts.

[Copilot]

Copilot review overview

🔵 Needs a closer look

CUDA graph timing affects kernel selection, while traditional CUDA EP and BF16 validation remain incomplete.

Review effort: Balanced
Findings: None

What changed in this PR

Adds CUDA graph replay-based validation for GEMM autotuning in graph-enabled sessions.

Changes:

• Measures and caches kernels using L2-conditioned CUDA graph replay.
• Separates stream and graph-replay cache entries and adds an opt-out.
• Propagates graph configuration to plugin kernels and adds regression tests.

File	Description
onnxruntime/​test/​providers/​cuda/​test_cases/​gemm_auto_tuner_test.cc	Tests replay selection, cleanup, capture ownership, and cache separation.
onnxruntime/​core/​providers/​cuda/​plugin/​cuda_kernel_adapter.h	Exposes graph enablement to plugin kernels.
onnxruntime/​core/​providers/​cuda/​plugin/​cuda_ep.cc	Propagates the plugin graph setting.
onnxruntime/​core/​providers/​cuda/​math/​matmul.h	Stores replay-tuning policy in MatMul kernels.
onnxruntime/​core/​providers/​cuda/​math/​matmul.cc	Resolves replay policy and selects its cache key.
onnxruntime/​core/​providers/​cuda/​math/​gemm_auto_tuner.h	Extends tuning keys and timing interfaces.
onnxruntime/​core/​providers/​cuda/​math/​gemm_auto_tuner.cc	Implements graph capture, replay timing, and cleanup.

---

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