[BBaditi]

Description

Each QMoE/MoE kernel instance owns a MoeGemmProfiler with a private (GemmId, M bucket) cache. A model whose MoE layers all have the same shape therefore profiles the identical problem once per layer for every new M bucket. Qwen3-30B-A3B has 48 such QMoE nodes.

This change shares successful profiling results process-wide. The key holds everything a profiling run reads:

• the runner type;
• runner state that changes the candidate tactics (SM80 FP4 family) or the profiled GEMM2 epilogue (deterministic Hopper reduce);
• the device;
• every GemmProfilerBackend::init parameter;
• the GemmId and the M bucket.

This is the in-process part of the QMoE cross-node dedup planned in docs/contrib_ops/cuda/gemm_profiler_cache.md (section 8, phase 2). It follows the process-wide profile map that MatMulNBits shares through GemmPluginProfilerManager (matmul_nbits.cc), with two differences:

• Coordination is per key. A short map lock hands out the key's entry, and only that entry's mutex is held while it is profiled. Concurrent first uses of one key profile it once, while other problems, M buckets and devices profile in parallel.
• Only successful results are shared, as the design doc requires for the MatMulNBits cache. A run that returns no config (for example after a transient allocation failure) is kept by that profiler instance as before, and the next profiler retries instead of reusing the failure.

The cache is intentionally leaked, like the MatMulNBits tactic cache registry, because a CUDA EP owned by another static can be destroyed after this library's statics. ORT_DISABLE_MOE_SHARED_TACTIC_CACHE=1 restores per-instance profiling and is documented in gemm_profiler_cache.md. MoeGemmProfiler::ProfilingRunCountForTesting() is a read-only counter that only the tests use.

Tests (moe_gemm_profiler_shared_cache_test.cc):

• IdenticalProblemsAreProfiledOncePerProcess and RunnerStateThatChangesProfilingIsPartOfTheKey: sharing works, and runner state that changes tactics keeps results apart.
• FailedProfilingIsNotShared: a failed run is retried by the next profiler.
• ConcurrentFirstUsesOfOneKeyProfileOnce: four threads that reach one key profile it once.
• OtherKeysProfileWhileOneIsInFlight: a second problem profiles to completion while the first is held inside its workspace allocation.

FailedProfilingIsNotShared and OtherKeysProfileWhileOneIsInFlight fail against an earlier, unpublished draft of this change that held one global lock while profiling and cached failed results. With sharing forced off, the four sharing tests fail.

Measured on an A100 80GB, main (6596e0c) vs this change:

• Qwen3-30B-A3B-Instruct-2507 through ORT GenAI, 32,640-token prompt, builds run in A-B-B-A order:
  • Cold time to first token is 82,465 and 83,704 ms on main vs 5,524 and 5,630 ms with this change.
  • Warm time to first token changes by -0.21% (95% CI -0.26% to -0.17%, 10 requests).
  • Generated tokens are identical, and peak memory is the same.
  • Decode measured -0.65% (95% CI -1.12% to -0.18%). Decode at 8 expanded rows runs the GEMV path, which uses no profiled tactic, so I read this as run-to-run drift rather than an effect of the change.
• Multi-turn chat in one session, with prompts of 2,048, 8,192, 16,384 and 32,640 tokens. Time to first token for turns 1 and 2 is 798 and 2,347 ms instead of 22,922 and 79,956 ms. Turn 4 completes; on main it runs out of memory, because each node's profiling request leaves a 1 GiB arena region (fixed separately in [CUDA] Reserve MoE GEMM profiler scratch outside the arena #34167).
• 48 same-shape QMoE nodes built from layer 0 of that model, first run with profiling at 2,048 rows: 22,916 ms to 811 ms (10 interleaved pairs), with identical outputs.

Motivation and Context

On Qwen3-30B-A3B-Instruct-2507 at a 32,640-token prompt on an A100, profiling dominated the first request: 82.5 s cold time to first token versus about 3.4 s warm. Most of that was the same two GEMMs being profiled 48 times.

Limits:

• Results are shared only inside one process, with no disk persistence. That is still future work in the design doc.
• A key that never yields a config is profiled again by every node, as each node did before.
• Tested on SM80 Linux. CI covers SM86, Windows and the plugin EP build.

[Copilot]

🟡 Changes recommended

The tests reuse process-wide cache entries across repeated runs, causing assertion failures and a 60-second stall, and lack exception-safe worker cleanup, allowing an allocation failure to abort the test process.

2 open issues

• Use fresh cache keys for each repeated test invocation · New
• Join the worker before rethrowing profiling exceptions · New

What changed in this PR

This PR shares successful CUDA MoE/QMoE GEMM profiling results across identical problems within a process, reducing repeated profiling across model layers. It implements in-process deduplication without adding disk persistence.

Changes:

• Adds a shared tactic cache keyed by runner type and state, device, problem parameters, and M bucket, with per-key locking and retries after failed profiling.
• Adds tests for reuse, runner-state separation, failure retries, and concurrent profiling.
• Documents the environment variable that restores per-instance profiling.

File	Description
onnxruntime/​test/​contrib_ops/​cuda_kernels/​moe_gemm_profiler_shared_cache_test.cc	Tests cache sharing, failure retries, and concurrency.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_kernels.h	Exposes runner state needed to distinguish cache keys.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_gemm_profiler.h	Declares the profiling-run counter used by tests.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_gemm_profiler.cc	Implements shared caching, per-key coordination, and profiling counts.
docs/​contrib_ops/​cuda/​gemm_profiler_cache.md	Documents in-process sharing and its opt-out setting.

_{Review effort: Balanced}

[Tianlei Wu (tianleiwu)]

Nice cold-start improvement (~82.5s → ~5.5s at 32K tokens). The two-level locking design (SharedTacticCache::mutex_ for map lookup plus per-entry SharedTacticEntry::mutex while profiling), fast-path bypass via the per-instance config_cache_ on warm runs, and inclusion of tactic/epilogue-relevant runner state (useSm80Fp4(), use_deterministic_hopper_reduce_) are all well reasoned.

+1 to the two existing inline comments on moe_gemm_profiler_shared_cache_test.cc regarding --gtest_repeat cache isolation and exception-safe held.join() in OtherKeysProfileWhileOneIsInFlight. Left three additional suggestions/nitpicks inline on SharedTacticKey representation, built-in CUDAAllocator OOM exception handling, and the env-var reference table in gemm_profiler_cache.md.

[Copilot]

🔵 Needs a closer look

Process-wide CUDA tactic reuse changes cross-thread synchronization, but runtime validation was unavailable and GEMM1-only tests leave SM90 GEMM2 epilogues and multi-device cache isolation awaiting GPU validation.

No open issues 🎉

2 resolved since last review

• Join the worker before rethrowing profiling exceptions
• Use fresh cache keys for each repeated test invocation

_{Review effort: Balanced}

[Tianlei Wu (tianleiwu)]

Reviewed a5b018a, including the profiler/backend inputs, runner state, MoE/QMoE callers, per-key synchronization, failure retries, and the updated tests. The earlier review feedback is addressed, and I found no new production correctness issue. One test-registration gap is noted inline.

Validation: source review and current-head CI inspection. Linux CUDA CI passed, and all five new tests passed twice in Windows CUDA CI (the dedicated internal-test step and the subsequent provider-test run). I did not run the tests locally because the available build artifacts predate this change.