Repository navigation
fix(cuda): replace QMoE's incidental trailing sync with explicit scratch/teardown drains - #1788
Merged
Merged
Conversation
…tch/teardown drains QMoEKernel::execute()'s trailing self.runtime.synchronize() call was originally suspected (per #1777) of causing observed host_us convergence with device execution time for grouped M>1 shapes. Investigation this cycle found that suspicion was already moot: CudaRuntime::synchronize() has been a deferred no-op by default since #1383 (predating #1777), gated by defer_eager_sync, so the call was already inert in production. An A/B test confirms removing/keeping it changes nothing measurable. The real invariants that call accidentally protected are now handled explicitly, using the runtime's existing unconditional-barrier primitive (drain_for_unmap, the same one elementwise.rs's BroadcastMetadataCache already uses) rather than a second synchronization authority: - ScratchPool::ensure(): drains before freeing a growing slot's old pointer (drain now precedes the replacement alloc, so a drain failure cannot leak the fresh allocation). - Drop for QMoEKernel: best-effort drain before freeing scratch, since execute() no longer implicitly guarantees prior in-flight kernels have retired by the time a later call (or teardown) runs. self.runtime.synchronize() itself is kept (not deleted) at the end of execute(), unconditionally on the non-capturing path: in the default configuration it remains a no-op, but keeping it preserves the ONNX_GENAI_DEFER_EAGER_SYNC=0 debug escape hatch for QMoE, consistent with every other kernel in this EP (e.g. MatMulNBitsKernel::run). A second, related bug was found and fixed in the benchmark harness (qmoe_gpu.rs, from #1777): median_us()'s and setup_gemv_bench()'s own inter-rep/post-warmup runtime.synchronize() calls were also no-ops, letting consecutive reps' async-enqueued kernels pile up in the CUDA launch queue until cuLaunchKernel itself blocked on a full queue -- a normal consequence of a genuinely memory-bandwidth-bound workload (3.5-25% of peak HBM bandwidth) pushed through a bounded async pipeline, not a QMoE synchronization defect. This is the actual mechanism behind #1777's host_us=median_us observation for grouped shapes. Both call sites now use drain_for_unmap (a real barrier) between reps/after warm-up; host_us for grouped M>1 drops from ~672-7150us to a clean ~60-67us, while median_us (GPU event time) is unchanged, confirming this is a measurement-methodology fix, not a compute-time change. Stale doc comments describing the original (falsified) theory are corrected. New regression tests (qmoe_gpu.rs): - qmoe_scratch_pool_regrows_and_shrinks_across_calls_matches_cpu: one kernel instance cycled through growing/shrinking row counts with consecutive calls enqueued back-to-back (no intervening drain), so a missing drain_for_unmap in ScratchPool::ensure's growth path has a genuine chance to corrupt data; checked against the CPU oracle. - qmoe_execute_does_not_block_host_for_device_completion: asserts host-side execute() return time stays a small fraction of measured GPU completion time for a genuinely slow grouped case; fails if an unconditional blocking sync is reintroduced. - qmoe_drop_after_inflight_launch_leaves_runtime_usable: drops a kernel immediately after an async call with no readback, then verifies the runtime/stream remain healthy for independent subsequent work. Verified on an idle A100 (device 1): 37/37 targeted CUDA tests pass (34 existing + 3 new), 3 independent runs of the merged #1777 bandwidth probe show consistent before/after numbers, cargo clippy -D warnings and cargo fmt are clean, and an independent review pass found and confirmed resolution of three issues (two regression tests that could not have caught their target bug due to forced drains, a leak on an error path, and a doc-comment/debug-escape-hatch inconsistency) before this commit. No paging/residency, per-expert identity, or exporter changes are included in this PR. Refs: #1777, #82 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1788 +/- ##
==========================================
+ Coverage 80.62% 80.79% +0.17%
==========================================
Files 411 411
Lines 199558 199979 +421
Branches 199558 199979 +421
==========================================
+ Hits 160884 161578 +694
+ Misses 33214 32942 -272
+ Partials 5460 5459 -1
Flags with carried forward coverage won't be shown. Click here to find out more. 🚀 New features to boost your workflow:
|
justinchuby
added a commit
that referenced
this pull request
Aug 23, 2026
…afe, close coverage gaps (#82) (#1800) ## Summary Working as Sebastian (Performance Engineer). Follow-up to #1777/#1788 for issue #82. This cycle's task was to determine whether grouped QMoE (M>=2) CUDA-graph capture eligibility could be made safe, **if and only if a property-level invariant can be proven** — no M/model allowlists, no copying the unrelated MatMulNBits fix. **Finding: it already is.** Grouped-path (M>=2) QMoE CUDA-graph capture is already implemented and already safe in current main. No functional change to `qmoe.rs` was needed. This PR is purely additive test/benchmark coverage that proves and documents the invariant, closes specific test-coverage gaps, and extends the existing bandwidth benchmark with capture+replay timing. ## The capability invariant 1. The session executor's per-shape `KernelKey` kernel-instance cache ("Chew's guarantee") never reuses a kernel instance for a shape it wasn't compiled/warmed for, so `capture_support()`'s row-count-agnostic `warmed` boolean is honest in practice for every M. 2. Every grouped-path launch grid (`launch_grouping`/`launch_gather`/`launch_grouped_linear`/`launch_linear`) is sized from host-static values, never a device-computed per-expert count — no launch-shape instability, no host dependency on device state. 3. `ScratchPool::ensure(.., capturing)` already rejects (loud `Err`) any capacity growth attempted while capturing — a violated per-shape-kernel invariant cannot silently reallocate/corrupt a captured graph. 4. The only host readback in `execute()` and its trailing `synchronize()` are both already gated `if !capturing`. No new global synchronization/allocation authority is introduced; all of the above are existing #1777/#1788-era mechanisms. ## What this PR adds - `qmoe_64experts_top6_m2/m4_capture_replay_reresolves_changed_router_probs`: explicit M=2/M=4 grouped capture+replay correctness (M=8 already had coverage), changing router logits between capture and replay. - `qmoe_grouped_capture_support_denied_before_warmup_with_reason`: missing-runtime-hook / warm-up-precondition taxonomy. - `qmoe_grouped_capture_rejects_shape_growth_without_corrupting_kernel`: dynamic-allocation/growth taxonomy, direct `ScratchPool::ensure` exercise. - `qmoe_grouped_capture_never_syncs_even_with_deferral_disabled`: hidden-sync/drain regression test — forces a REAL blocking `synchronize()` so a capture-time sync would provably corrupt/fail the recording. - `qmoe_grouped_capture_repeated_replays_and_consecutive_capture_cycles_stay_correct`: consecutive graphs, repeated replay, basic device-memory accounting stability. - `qmoe_grouped_eager_execute_after_capture_replay_still_correct` / `qmoe_grouped_drop_kernel_after_capture_replay_leaves_runtime_usable`: lifecycle/teardown safety. - `qmoe_grouped_capture_replay_bandwidth_probe`: extends the #1777/#1788 bandwidth harness with capture+replay timing across grouped M={2,4,8} for both DeepSeek-V2-Lite and GLM-5.2 shapes. ## Measurements (idle A100, `CUDA_VISIBLE_DEVICES=1`, 3 independent process runs, reps=25, batch=16) | shape | M | capture_us (one-time) | replay_us (GPU) | eager_us (GPU) | host launch-gap cut | GB/s (dedup) | |---|---|---|---|---|---|---| | deepseek-v2-lite | 2 | ~70-72 | 656.6 | 674.6-674.8 | 96.1-96.2% | 118.6 | | deepseek-v2-lite | 4 | ~72-74 | 1226.7-1227.0 | 1244.3-1245.6 | 96.1-96.2% | 116.3-116.4 | | deepseek-v2-lite | 8 | ~71-80 | 3289.2-3297.2 | 3307.4-3316.4 | 95.9-96.1% | 72.8-73.0 | | glm-5.2 | 2 | ~71-81 | 4058.4-4058.6 | 4076.3-4077.8 | 95.3-96.0% | 111.6 | | glm-5.2 | 4 | ~71-84 | 5106.0-5106.0 | 5124.9-5125.7 | 95.8-96.0% | 177.4 | | glm-5.2 | 8 | ~72-83 | 7135.0-7135.7 | 7153.3-7157.6 | 95.3-95.8% | 253.9-254.0 | Replay GPU-event time is within noise of eager time at every configuration (this workload is memory-bandwidth-bound, not launch-bound, at these M values) — capture's measured benefit here is **host dispatch time**, cut ~95-96% (from ~60-67us to ~2-3us per launch), not GPU throughput. Variance across 3 independent process runs is <0.3% on `replay_median_us`. Correctness gated by CPU oracle + byte-identical-vs-eager check on a reduced-expert-count proxy case before any timing is trusted, at every (shape, M). Full 46-test suite (`cargo test -p onnx-runtime-ep-cuda --release --features gpu-tests --test qmoe_gpu`) passes, 0 failures, ~455-465s wall on this A100. `cargo fmt`/`cargo clippy` clean on the changed file (scoped to `--test qmoe_gpu`; full-workspace clippy hits a pre-existing, unrelated compile break in `weight_paging.rs`, not touched by this PR). ## Review An independent code-review agent found: 1. Four tests whose capture-window assertions could panic between `begin_graph_capture`/`end_graph_capture` without aborting the capture or freeing buffers on that path. 2. One test that leaked its `GemvBenchSetup` buffers via a bare `drop`. Both fixed in this PR: added a `CaptureGuard` RAII helper (aborts capture if dropped while still armed) and a `GemvBenchSetup::free_without_readback` method (mirrors the existing `PendingExecution::free_without_readback`). Full suite re-verified passing after the fix. ## Scope Test/benchmark-only — no production code in `qmoe.rs`/`runtime.rs`/`capture.rs` changed. No paging/residency, exporter, or BlockQuantizedMoE changes (per instruction). ## Next gate for #82 Per the ordering established in the prior cycle's follow-up: **BlockQuantizedMoE paging+prefetch wiring** is next — the sound benchmark (#1777/#1788) and the capture-safety proof (this PR) are both now in place as prerequisites/instrumentation for that work. This PR does not close #82; it removes one more item from that issue's dependency chain. Per user directive: no CI wait — this PR is validated by local A100 tests + independent review; CI is asynchronous supplementary evidence. Merging on local gates. 🤖 Generated by Sebastian (Performance Engineer), Squad autonomous cycle. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 24, 2026
…ndary window (#1810 Slice 7A) (#1922) ## Slice 7A — inert, default-OFF, producer-only expert-route telemetry (QMoE/BQMoE) Closes #1810 (Slice 7A). **Draft. Do not merge.** Crate-internal / test-only / default-OFF. > **Cycle-22 revision by Sebastian (Performance Engineer), revision owner.** > The original author is under reviewer-protocol lockout and did not advise, pair, or contribute to this revision. Reworked independently from the PR diff, the approved Slice-6 design doc, and Roy's Cycle-22 review. ### What changed vs the rejected HEAD (`8ec9d8230`) Roy's Cycle-22 NO-GO was correct and is **addressed at the design level**: the rejected design launched a **separate reset/epoch kernel on every execute/replay**, which contradicts the approved Slice-6 **coarse-boundary window** contract (`docs/memory/EXPERT_ROUTE_TELEMETRY_SLICE6_DESIGN.md` §2.3/§3 — the epoch is bumped and the record consumed at a *coarse safe boundary*, **not per replay**). That per-call reset is why the fused path measured **2.526 us / 2.18 % -> NO-GO**. This revision makes the fused route kernels **accumulate-only** and moves reset/epoch-advance to an explicit, host-ordered **coarse boundary**. The earlier NO-GO was scoped to the wrong (per-call reset) design point; the boundary-window producer now measures **GO** (see below), independently reproduced. ### Invariant / API (producer-only boundary-window contract) - **`arm(request, device, experts)`** - validates identity, allocates the stable-VA record via the existing runtime allocator, opens **window 1** (stamps identity, `epoch = 1`, zeroes bitmap/counters). **No kernel is compiled or launched.** - **execute / replay** - the fused `route_telemetry_mark_row` inside `qmoe_route` / `bqmoe_route` **only accumulates** into the stable record: `atomicOr` the routed bit into the bitmap, `atomicAdd` the in-range count, `atomicOr` the sticky **poison** bit on an out-of-range id, `atomicOr` the sticky **overflow** bit on count saturation. The **epoch is fixed for the whole window**; every eager call and captured replay in the window accumulates the routed-expert **union + count**. **No reset/epoch kernel in the captured graph; no host sync/alloc/drain/VMM on this path.** - **`reset_route_telemetry_boundary()`** - the **only** place the window advances. **Rejected while the EP stream is capturing/replaying** (fail closed, checked *before* any drain); otherwise drains prior stream work via the existing `drain_for_unmap` authority, bumps the **host-side** epoch, and re-stamps identity + zeroes the record so the next window starts empty with **no stale carryover**. Allocates nothing, moves no pointer, touches no PMM/VMM/cache/global coordinator. - **Fail-closed** - overflow **saturates** into the sticky overflow bit (never wraps into a smaller "success"); out-of-range routes **poison** without touching the bitmap; snapshot/validate run on the host at a boundary against an already-copied record. - **Isolation / stability** - multi-request / multi-device / multiple kernel instances use distinct records; the bitmap **VA is stable** across windows and captures. Epoch is a host counter, so the record footprint **drops the former 4-byte device epoch buffer** (`footprint = 4*ceil(experts/32) + 24`). - **Reuses only existing authorities** - `alloc_raw`/`free_raw`, `htod`/`dtoh`, `is_capturing`, `drain_for_unmap`, `ordinal`. No new coordinator/allocator/cache/PMM/VMM and **no host sync in steady-state replay**. Scope unchanged: **default-OFF, crate-internal, test-only**. Ordinary inference is **byte-identical** and has **zero overhead when disarmed**. No lifecycle/policy/consume wiring in this slice. ### Tests (rewritten to the new semantics, same commit) QMoE (`tests/qmoe_gpu.rs`) and BQMoE (`tests/block_quantized_moe_gpu.rs`) `mod route_telemetry`: - off/on output **byte-identity**; - eager calls **accumulate the CPU-oracle union + count** within a window at a **fixed epoch**, then a boundary reset opens an empty **epoch-2** window; - **>=3 graph replays** with different routes accumulate the union at a **fixed epoch** and a **stable VA** (QMoE - proves no per-replay reset); - boundary reset **increments the epoch and starts an empty window**; - boundary reset **rejected during an active capture** (epoch unchanged); - capacity / device mismatch stays **inert** and never fails inference; - multi-instance **request/device isolation**, footprint, teardown/accounting; - **QMoE and BQMoE routes for M in {1, 2, 4, 8}** match the oracle. Host unit tests cover **poison/overflow fail-closed** on the validator; the #1884 probe harness continues to cover device-level poison/overflow. **Results** (idle A100, GPU 5, `--test-threads=1`): QMoE `route_telemetry` 10 passed; BQMoE 8 passed; telemetry host unit tests 7 passed; #1884 probe functional regressions 6 passed; **full QMoE suite 56 passed**, **full BQMoE suite 13 passed** (no regressions to #1788/#1800/#1884/#1854). ### Measurement - G1 gate (`microbench_fused_route_telemetry_g1_gate`) Idle **pinned A100**, serialized off/on back-to-back, whole QMoE layer **CUDA-graph captured** and measured via `replay_graph` (host-enqueue gaps excluded), **BATCH = 512, RUNS = 5**, ~8 s clock ramp + first-shape recheck, GPU-event and host-enqueue reported separately. Gate denominator is a **realistic** DeepSeek-V2-Lite decode layer (64 experts, top_k 6, hidden 2048); tiny synthetic is informational only. **n = 3 independent runs:** | Run | route+layer overhead (GPU) | % of layer | host-enqueue delta | epoch | count | output | |-----|---------------------------|-----------|----------------|-------|-------|--------| | 1 | 0.976 us | 0.84 % | -0.016 us | fixed = 1 | 30726 | identical | | 2 | 0.948 us | 0.82 % | -0.006 us | fixed = 1 | 30726 | identical | | 3 | 1.054 us | 0.91 % | -0.008 us | fixed = 1 | 30726 | identical | **Range: 0.948-1.054 us, 0.82-0.91 %. Clock drift < 0.5 %.** The epoch stays **fixed at 1** across the whole BATCH x RUNS replay window (count accumulates to 30726) - a moving epoch would have proven a forbidden per-replay reset survived. **G1 GATE (<= 2 us AND <= 2 %): GO.** Independently reproduced vs Roy's boundary-only experiment (0.884 us / 0.76 %); same order, same verdict - not a reused number. Telemetry remains **default-OFF** regardless of the GO. ### Review Independent review (excluding the locked-out author and the revision owner): **APPROVE** - no blocking or substantive findings; confirmed accumulate-only execute path, reject-before-drain boundary reset, fail-closed overflow/poison, stable VA, isolation, footprint, and that the tests assert the *new* semantics (they would fail under the rejected per-replay-reset design). One MINOR note: the device-side poison branch of the fused mark is defensive and unreachable-by-construction from the production route kernel (Roy's prior approved framing) - covered by the host validator unit tests and the #1884 probe harness. **Roy's explicit final re-review is requested** before this leaves draft. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 24, 2026
… coarse residency (#1810 Slice 7B) (#1971) ## #1810 Slice 7B — boundary-time route-telemetry **consumer** **Status: DRAFT — awaiting independent review. Do not merge.** Closes the loop opened by the merged Slice-6/7A **producer** (expert-route telemetry, PR #1922 `e1ec495ee`) and the merged Slice-4/5 **coarse-boundary residency lifecycle** (PR #1854). This is the smallest production seam the Slice-6 design §8 specifies: > expose a boundary-time consumer that produces a per-expert desired-set, and feed that set to the **existing** Slice 4/5 coarse-boundary plan application (`coarse_residency.rs`) as its policy input — with **no** new allocator, **no** id→slot rewrite, and every mapping change still owned by PMM/VMM. ### What it does New module `crates/onnx-runtime-ep-cuda/src/route_residency.rs`: `consume_route_window_at_boundary(...)`: 1. **Gate** — default-off via the existing `COARSE_RESIDENCY_ENABLE_ENV` (`coarse_residency_profile_enabled()`). When off (shipped default) it returns `Disabled` *before* reading the snapshot or touching any allocator. 2. **Safe boundary** — re-reads the existing `CudaWeightResidency::resize_safe_point` and fails closed with `RejectedNotSafeBoundary { reason }` if a graph is capturing/replaying, an admission is in flight, a deferred release has not settled, execution is multi-device, or a routed-residency guard is live. Coarse-boundary only — never a per-token remap. 3. **Validate** — runs the producer's own `consume_and_validate` on the already-completed window snapshot; fail-closed (→ `WholeBank { reason }`) on poison / overflow / stale epoch / foreign request / foreign device, or when no in-range experts were recorded. 4. **Plan** — turns the routed-expert union into a desired **hot set** and asks the already-validated `StaticProfileResidencyPolicy` to shape a `ResidencyPlan` (the design's "record → desired set" `RouteObserverPolicy` role — **reused, not duplicated**, so exactly one validated policy emits `PerExpertCandidate`). 5. **Apply** — hands the plan to the existing `CudaWeightResidency::apply_coarse_residency_plan`, which remains the **sole** authority that maps / unmaps / accounts / quarantines / rolls back through PMM/VMM. `RouteWindowConsumeOutcome` carries the exact reason on every non-applied path — **no silent fallback**. ### Invariants held by construction - **Allocates nothing**, opens no stream, owns no VA, copies no device bytes — pure host glue between two existing authorities. - **No remap during capture/replay**; coarse-boundary only, never per-token. - **No new host sync** in steady state (only the producer's already-taken snapshot and the existing transition primitive's drain). - **Default-off & byte-identical** when disarmed (two independent default-off switches: telemetry disarmed *and* this gate off). - **Same-device fail-closed** and **PMM/VMM remains the sole mapping/accounting/quarantine/rollback authority**. ### Tests — `tests/route_residency_consume_gpu.rs` (6 GPU tests, all passing on an idle A100) - `disabled_gate_is_structural_no_op` — off path is a structural no-op (byte-identical). - `route_window_hot_set_transitions_cold_experts` — routed hot-set stays resident, cold set tiers to host, bytes identical. - `expert_group_transitions_atomically_from_window` — atomic expert-group transition driven from a window. - `active_capture_and_multi_device_reject_consume` — active-capture and multi-device boundaries reject the consume (fail-closed). - `foreign_identity_and_defective_windows_fail_closed` — foreign request/device and poison/overflow/stale windows fail closed to whole-bank. - `injected_fault_rolls_back_consumer_transition` — injected driver fault (`fail_nth(Unmap, 3)` over an isolated + merged cold range) rolls back range-precisely and quarantines; `rollback_count == 1`, `values_touched == 0`, `committed_values` empty, content bit-identical (mirrors the proven `coarse_residency_plan_gpu.rs` rollback fixture). Run: ``` env -u ONNX_GENAI_WEIGHT_OFFLOAD_COARSE_RESIDENCY_ENABLE CUDA_VISIBLE_DEVICES=<idle> ONNX_GENAI_CUDA_DEVICE=0 \ cargo test -p onnx-runtime-ep-cuda --features cuda,cuda-13000,gpu-tests --release \ --test route_residency_consume_gpu -- --ignored --test-threads=1 ``` ### Honest scope Like `coarse_residency::apply_residency_plan_at_boundary` when it shipped (Slice 5), this consumer has **no live decode-loop call site yet** — wiring it into a running session's request boundary is the next slice. It ships here as the production seam: reachable, default-off, and proven by the GPU tests. `cargo fmt` applied; `cargo clippy` on the crate lib is clean for the new file (pre-existing unrelated warnings/errors in `optimizer.rs` / `standard_attention.rs` are out of scope). ### Constraints honored Did not touch PagedAttention or IQ1 fusion files. Branch `squad/1810-slice7b-telemetry-residency-consume` based on latest `main` (`011fbb284`, includes #1922 `e1ec495ee`). #1788/#1800/#1884/#1854 lifecycle regressions preserved (reused, not forked). Refs #1810. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Working as Sebastian (Performance Engineer). Follow-up to #1777 for issue #82: investigates and resolves
QMoEKernel::execute()'s unconditional eager-pathruntime.synchronize()call.Headline correction to the record:
CudaRuntime::synchronize()has been a deferred no-op by default since #1383 (merged before #1777 was authored), gated bydefer_eager_sync. This means the trailing sync inexecute()was already inert in production before this PR. An A/B test (remove the call vs. keep it) shows zero measurable performance difference — confirming this was never a live blocking point in the default configuration. #1777's causal claim that this call causedhost_us≈median_usconvergence for grouped M>1 shapes is falsified by this A/B test.The real invariants that call accidentally protected (only in configurations where the debug flag forces synchronous behavior, or in any future default change) are now handled explicitly with the runtime's existing unconditional-barrier primitive (
drain_for_unmap, already used byelementwise.rs'sBroadcastMetadataCache) rather than a second synchronization authority:ScratchPool::ensure(): drains before freeing a growing slot's old pointer (order fixed so a drain failure can't leak the fresh allocation — mirrors theelementwise.rsprecedent).Drop for QMoEKernel: best-effort drain before freeing scratch at teardown.self.runtime.synchronize()itself is kept, not deleted, at the end ofexecute(): it's a no-op by default, but keeping it preserves theONNX_GENAI_DEFER_EAGER_SYNC=0debug escape hatch for QMoE, consistent with every other kernel in this EP (e.g.MatMulNBitsKernel::run).A second, related bug was found and fixed in the #1777 benchmark harness itself (
qmoe_gpu.rs): its own inter-rep/post-warmupruntime.synchronize()calls were also no-ops, letting consecutive reps' async-enqueued kernels pile up in the CUDA launch queue untilcuLaunchKernelitself blocked on a full queue. This queue-saturation effect — not a QMoE sync defect — is the actual mechanism behind #1777'shost_us≈median_usobservation for grouped shapes. Fixed by usingdrain_for_unmap(a real barrier) between reps/after warm-up. Stale doc comments describing the original (falsified) theory are corrected.Measurements (idle A100,
CUDA_VISIBLE_DEVICES=1, merged #1777 probe, 25 reps)3 independent "after" runs are consistent within noise;
median_us(true GPU event time) is unchanged across all runs — confirming this is a pure measurement-methodology fix to the harness, not a compute-time change. No tok/s claims are made.Tests
New regression tests in
qmoe_gpu.rs(in the same commit):qmoe_scratch_pool_regrows_and_shrinks_across_calls_matches_cpu: one long-lived kernel instance cycled through growing/shrinking row counts ([1,8,8,2,8,4,16,1,16,2]), with consecutive calls enqueued back-to-back (no intervening drain) so a missingdrain_for_unmapin the growth path has a genuine chance to corrupt data; every step checked against the CPU oracle.qmoe_execute_does_not_block_host_for_device_completion: asserts host-sideexecute()return time stays a small fraction of measured GPU completion time; fails if an unconditional blocking sync is reintroduced.qmoe_drop_after_inflight_launch_leaves_runtime_usable: drops a kernel immediately after an async call with no readback, then verifies the runtime/stream remain healthy for independent subsequent work (both a fresh QMoE kernel and a plain htod/dtoh round trip).Verified on idle A100 (device 1):
cargo test -p onnx-runtime-ep-cuda --release --features gpu-tests --test qmoe_gpu→ 37/37 passed (34 existing + 3 new).cargo clippy -p onnx-runtime-ep-cuda --release --features gpu-tests --test qmoe_gpu -- -D warningsandcargo fmt --checkare clean.Review
An independent review pass (code-review agent) found 3 substantive issues in an earlier revision of this diff, all fixed and confirmed resolved by a follow-up review pass before this commit:
dtoh) immediately after every call, making them structurally unable to expose the bugs they claimed to guard against — fixed by splitting "enqueue" from "read back" (PendingExecution/enqueue_with_existing_kernel) so consecutive calls genuinely overlap.drain_for_unmap()error path inScratchPool::ensure()(alloc happened before the fallible drain) — fixed by reordering to drain-before-alloc, mirroringelementwise.rs's existing precedent.MatMulNBitsKernelwhile actually diverging from it (having fully removed the debug-escape-hatch-respecting sync call) — fixed by reinstatingself.runtime.synchronize()and correcting the comment.Scope
No paging/residency, per-expert-identity, or exporter changes are included in this PR (explicitly out of scope per issue #82's follow-up plan).
Next gate for #82
With this merged, the remaining #82 dependency chain is: (1) this PR — QMoE
execute()async-safety + benchmark methodology fix [this PR], (2) the batch-M≥2 capture-eligibility fix (still open — QMoE grouped paths remain capture-ineligible; the "repeat-launch + CUDA-event" pattern used by the #1777 probe is the correct proxy until captured replay is supported for M>1), then (3) wiringBlockQuantizedMoEpaging+prefetch (explicitly deferred, not touched here). This PR does not change capture eligibility; it only fixes the async-safety and measurement-methodology issues blocking trustworthy baselines for that next step.Per explicit user directive, this PR proceeds on local A100 validation (34+3 targeted tests, 3 independent benchmark runs, clippy/fmt, independent review) without waiting on remote CI, which remains supplementary asynchronous evidence.
Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com