Repository navigation
perf(mtp): per-slot executor host capture state (foundation for two-slot verify replay) - #1654
Merged
Merged
Conversation
Make the executor's host-side captured-graph bookkeeping per-slot so the M=1
Primary decode graph and the M=k+1 Verify speculative-verify graph can be
captured and replayed independently on ONE executor, instead of sharing a single
signature/schedule that the M=1 decode clobbered every verify step (the
replays=0 MTP blocker).
onnx-runtime-ep-api: DeviceGraphSlot::{COUNT, index()} (Primary=0, Verify=1).
onnx-runtime-session:
- New SlotCaptureState struct holding the 8 host capture fields
(device_graph_signature, capture_schedule, capture_segmentation,
capture_cf_shapes, capture_warm_signature, capture_warm_shapes,
capture_warm_seeded, capture_quarantine_ops). Executor now holds one per slot
in `slot_capture: [SlotCaptureState; COUNT]`, indexed by graph_slot via
cap()/cap_mut(). buffer_shapes/if_last_predicate stay per-run scratch.
- set_graph_slot is now a pure retarget (NO reset of the other slot) — switching
slots keeps each slot's installed graph, the precondition for coexistence.
- All capture/replay/seed sites in bindings/run/capture/mod/tests route through
cap()/cap_mut(). Primary maps to index 0, so greedy (Primary-only) bookkeeping
is byte-identical to the historical single-field layout.
onnx-genai-engine:
- Drop the enable_decode_inline() gate on configure_verify_capture: verify
capture now arms on sibling-less GDN models too (both graphs coexist safely).
- configure_verify_capture no longer permanently retargets the main exec to
Verify; it only pins the StepScoped workspace. run_verify_captured flips the
main exec to Verify around the verify forward and back to Primary after, so
M=1 decode captures into Primary and the verify into Verify.
- Remove the dead set_retain_decode_graph_across_spec/retain_decode_graph_across_spec
accessors (superseded by verify-capture). Refresh stale comments.
Compiles (session + engine native-backend). Tests/GPU validation follow.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Proves Primary (greedy's only slot) host capture state survives a switch to Verify and back — the invariant that keeps greedy byte-identical while MTP flips the executor to Verify around each verify forward. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…by shared interior buffer arena Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #1654 +/- ##
==========================================
- Coverage 81.41% 81.02% -0.39%
==========================================
Files 384 384
Lines 180337 180358 +21
Branches 180337 180358 +21
==========================================
- Hits 146813 146141 -672
- Misses 28582 29272 +690
- Partials 4942 4945 +3
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 21, 2026
The rustfmt breakage is main's, not this branch's: a clean `origin/main` worktree at 9fd0526 fails `cargo fmt --all -- --check` with exit 1 at the same four sites (capture.rs:585, run.rs:670, tests.rs:4979 and :4990), introduced by #1652/#1654. Carried here only because this branch cannot be green while its base is red. The three reformatted files are byte-identical to #1655, which fixes the same four sites upstream, so that merge will not conflict. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: c80f8522-983c-47f7-8241-2155a823aabe
justinchuby
added a commit
that referenced
this pull request
Aug 21, 2026
…d again) (#1655) Sixth formatting break on `main` today, from #1654 — a few hours after #1653 fixed the fifth. Four sites in `onnx-runtime-session/src/executor/{capture.rs,run.rs,tests.rs}`, all the same shape: a `self.cap_mut().field.insert(..)` chain rustfmt wants split across lines. Straight `cargo fmt --all`, no hand edits. **Verified inert** per file: each before/after pair is identical after stripping whitespace and trailing commas. Reproduced on a pristine `origin/main` worktree first, so it is main's breakage, not my branch's. Because `Rust quality` runs sequentially, this red `Check formatting` step is again **skipping every other lint in that job** — clippy, publish order, dispatch-manifest lints, feature-gate coverage. Those have effectively not run on main since #1654 landed. This is the sixth today (#1637/#1641, #1644, #1648, #1647/#1648, #1652, #1654) and my fifth repair PR. The structural cause and the two candidate fixes — a merge queue, or running `Rust quality` on `main` post-merge — are written up in #1653. Nothing about the individual changes is careless; required checks simply run on each PR's merge ref and never on the actual post-merge `main`, so sequentially-green PRs can still leave the tip red. cc @justinchuby Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 21, 2026
…imary + Verify) (#1658) ## Summary Gives the fixed **M=k+1 speculative-verify** forward its OWN executor sibling with a **private interior device-buffer arena**, unblocking verify CUDA-graph capture+replay, then closes **Blocker B** so the M=1 Primary decode graph also replays. **Both `cuda_graph` slots now replay on one executor** — the campaign's core proof — GPU-validated on the real qwen38-27b-int4 hybrid (GDN+GQA) MTP artifact. Branches off `39885ffc5` (per-slot executor host capture state, #1654). ## What changed - **Verify-dedicated sibling** (`build_verify_sibling`, `Session::verify_exec` + verify-sibling run/capture/replay/reset methods): a structural clone of the main graph on the `Verify` slot, sharing only immutable weights/EP + the persistent external KV/recurrent-state bindings. Its interior `[1,M]` scratch is independent of the M=1 decode's `[1,1]` scratch, so the interleaved base decode can no longer resize/clobber the M=2 verify graph (the old `Slice [1,2] vs [1,1]` decline → eager fallback). - **Race fix (CUDA-700):** the sibling PINS its StepScoped workspace (it only ever runs the fixed M=K shape), so its captured graph's baked scratch pointer is never freed back to the shared EP arena between replays. Reset the verify sibling's graph only at the generation boundary. - **Blocker B:** re-advance accepted tokens one at a time (M=1, state-equivalent, byte-identical) instead of a batched M=num_accepted forward; and enable `retain_decode_graph_across_spec` when verify-capture arms so the contents-only KV roll-back no longer invalidates the Primary graph. No hardcoded dims (M derived from `num_speculative_tokens+1`). Greedy path inert (no `verify_width`, never hits a non-zero rewind → byte-identical). ## Validation (H200, k=1, 83.3% accept, median of 5, ORT 1.28 cuda13, int4 block-32) **Both slots replay, fallbacks=0, invalidations→~0 (every one of 5 runs):** ``` cuda_graph: captures=4 replays=504 fallbacks=0 invalidations=3 cuda_graph_verify: captures=4 replays=184 fallbacks=0 invalidations=0 ``` - Token-identical across all 5 runs (deterministic), no NaN, coherent text. - MTP throughput: **median 20.63 tok/s** (up from 15.6 pre-fix). - engine lib **579/0** (`native-backend`), ep-cuda `graph::tests` **8/8**, recurrent-commit + `native_verify_logits_require_restored_recurrent_state` byte-identity oracles pass. ## Honest perf finding (not a speedup) Graphing the verify was **necessary but not sufficient**. Approach-A (snapshot→restore→**re-advance**) redundantly re-runs the accepted tokens on commit, and MTP runs verify(M=2) + per-token re-advance forwards per step for ~2.67 tokens/step — strictly more device work per accepted token than greedy's single graphed M=1 forward on this launch-bound decode. So MTP (20.63 tok/s) still trails the campaign-established greedy baseline (~56 tok/s; corroborated here by Primary `replays=504` exactly matching the greedy run). The greedy baseline artifact is out-of-sandbox this session. **Next lever:** move to approach-B — thread `num_accepted` through the verify to SELECT the committed recurrent state from per-position post-states instead of re-scanning, removing the redundant re-advance forwards. The two-slot graphed executor here is the prerequisite. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- 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
1 participant
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.
Per-slot executor host capture state (foundation for two-slot MTP verify replay)
Makes the
Executor's host-side CUDA-graph capture state per graph slot so the M=1Primarydecode graph and the M=k+1
Verifyspeculative-verify graph can coexist on one executor.Changes (coordinator's 4-item scope + test)
DeviceGraphSlot::{COUNT=2, index()}(Primary=0, Verify=1) in ep-api.SlotCaptureState;Executorholdsslot_capture: [SlotCaptureState; 2]withcap()/cap_mut()indexed bygraph_slot.buffer_shapes/if_last_predicate/control_flow_output_valuesstay shared per-run scratch.set_graph_slotis now a non-resetting pure retarget.enable_decode_inline()gate onconfigure_verify_capture; main exec stays onPrimary for M=1 decode,
run_verify_capturedflips to Verify around the verify forward.set_retain_decode_graph_across_spec/retain_decode_graph_across_specaccessors.set_graph_slot_is_non_resetting_and_per_slot_isolated(greedy uses onlyPrimary=index0; Primary marker survives a Verify round-trip, no cross-slot bleed).
Byte-identity / gates
Primary(index 0) → bookkeeping byte-identical to the historicalsingle-field layout (proven by the new test).
--features native-backend: 579 passed / 0 failed (greedy inert, oraclenative_verify_logits_require_restored_recurrent_stategreen).graph::tests--features cuda,cuda-13000: 8/8.83.3% acceptance, 14.67 tok/s; verify now ARMS+DECLINES gracefully (was inert in MTP verify-capture infra (option-c): capture-safe, gated on decode-inline sibling #1652).
Remaining gap (GPU-proven, honest — no speedup yet)
Per-slot host state is necessary but not sufficient for verify replays>0. The Verify capture
declines because the interior device
buffersarena is SHARED: the M=1 decode JIT-sizes interiorscratch to
[1,1], so the M=2 verify capture errorsSlice output shape [1,2], expected [1,1]→graceful eager fallback (
cuda_graph_verify: captures=0 replays=0 fallbacks=1). Next lever = averify-dedicated device buffer arena / verify executor so interior scratch is independent of the
interleaved M=1 decode. Blocker B (separate/pre-existing): Primary M=1 decode is invalidated by
the variable-M
commit_recurrent_state_to_acceptedre-advance sharing Primary bindings.Validated by unit tests + GPU no-regression; end-to-end replays>0 pending the verify-dedicated arena.
Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com