onnxruntime
a1fc71e7 - [CUDA] Fix QMoE profiler cross-stream race and CUDA-graph-capture safety (#29584)

Commit
6 days ago
[CUDA] Fix QMoE profiler cross-stream race and CUDA-graph-capture safety (#29584) ### Description Three related fixes in the QMoE (`contrib_ops/cuda/moe`) weight-only GEMM profiling path that can cause a sticky `CUDA 700` (illegal memory access) surfacing at a later MoE kernel launch: - **`moe_kernels.cu`** — `GemmProfilerBackend::runProfiler` now keys its workspace layout on `mSM >= 90`, the same key used by `getWorkspaceSize` / `prepareRouting` / `prepareQuantParams` / `prepareTmaWsInputs`. Previously it keyed on the per-tactic `is_tma_warp_specialized`, which diverges on SM90 when a non-TMA (Ampere-fallback) tactic is profiled (e.g. INT4 mixed-input GEMM tiles). That divergence drops the `tma_ws_input` region and shifts every subsequent sub-buffer offset, so the profiler reads `expert_first_token_offset` and the GEMM inputs from the wrong (random-filled) locations → OOB read in the grouped GEMM. - **`moe_gemm_profiler`** — `profileTactics`/`runProfiling` accept an optional `timing_stream`; the profiler runs on the caller-supplied compute stream when provided, so its kernels are strictly ordered with surrounding compute-stream work and share the temp allocator's stream context. Profiling on a private side stream races with the stream-aware temp arena, which can hand the same scratch block to a later compute-stream allocation (e.g. the real MoE workspace) while the profiler's grouped-GEMM kernels are still in flight. - **`moe_quantization.cc`** — Skip profiling while the compute stream is being captured into a CUDA graph (profiling launches kernels, records/synchronizes events, and allocates/frees scratch — all illegal during capture) and fall back to a config cached from an earlier non-capturing run, or the runner's default tactic. ### Motivation and Context These are latent correctness/stability bugs in the QMoE tactic profiler that manifest on SM90 (H200) and under CUDA graph capture. They are independent of the fpA_intB MatMulNBits work and are split out as a standalone fix. ### Testing - Built with `USE_FPA_INTB_GEMM=ON` (arch 80;90) on H200. - `test_qmoe_cuda.py` passes (SwiGLU int4/int8, FP16/BF16, multiple batch/seq shapes).
Author
Parents
Loading