Repository navigation
perf(mtp): approach-B (skip re-advance on full accept) + device-to-device recurrent snapshot — 20.68→34.34 tok/s, token-identical - #1663
Merged
Conversation
…tate is committed) Per spec step the driver did THREE lots of target work: (a) M=1 base decode, (b) M=K verify, (c) a snapshot->restore->re-advance that re-ran EVERY accepted token through the full model just to rebuild the destructive GDN/conv recurrent state. Lot (c) is pure redundancy on a full accept: the eager verify forward already advanced both the KV (base+K) and the recurrent/conv state by exactly those K==accepted tokens, and the committed length equals the current length. When `accepted == draft.len()`, skip the KV rewind, the snapshot restore, AND the accepted-token re-advance entirely — the verify's own post-state IS the committed state. Only a partial accept still needs the snapshot->restore-> re-advance rebuild to the shorter prefix. The snapshot is still taken before the verify (acceptance is only known after it); it is simply unused on the fast path. Per-step target-forward budget (128-tok run, 48 verify steps, 83.3% accept, 32/48 full-accept): before: 48 base(M=1) + 80 re-advance(M=1) + 48 verify(M=2) [Primary replays=504] after: 48 base(M=1) + 16 re-advance(M=1) + 48 verify(M=2) [Primary replays=248] => 64 redundant re-advance forwards/generation eliminated. GPU-validated (H200, qwen38-27b-int4-mtp, median-of-5, ORT 1.28 cuda13, int4 block-32, --release bench-native,native-cuda,cuda-13000): throughput 20.68 -> 26.73 tok/s (+29%), TOKEN-IDENTICAL (md5 be7ed565 across all 5 runs == pre-change baseline), no NaN. Both slots still replay: cuda_graph: captures=4 replays=248 fallbacks=0 invalidations=3 cuda_graph_verify: captures=4 replays=184 fallbacks=0 invalidations=0 Engine lib suite 579/0. Token-identity empirically confirms the M=K verify post-state equals K sequential M=1 advances for the accepted prefix. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
InferenceSession gained a verify_exec field in #1658 but the two executor/tests.rs InferenceSession initializers were not updated, leaving `cargo test -p onnx-runtime-session --features cuda` red on main. Add verify_exec: None to both. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…rip) The speculative recurrent-state snapshot copied all GDN/conv fixed-state bindings device->HOST every spec step (snapshot_fixed_states via read_bytes_range), then wrote them back on partial-accept. That PCIe round-trip (~10ms/step) was the dominant remaining per-step overhead and is unused on the ~67% full-accept steps. Replace it with a stream-ordered device->device snapshot into per-binding device scratch: - Add ExecutionProvider::copy_device_to_device (default errors; CUDA overrides with cuMemcpyDtoDAsync on the EP stream, so the snapshot is ordered ahead of the verify forward that overwrites the state with no host sync). - DeviceIoBinding::snapshot_device_into / restore_device_from + allocator() accessor; re-export DeviceBuffer from onnx-runtime-session. - DecodeCudaState grows a lazily-allocated fixed_state_snapshot_scratch arena (sized from fixed_state_binding_range — no hardcoded dims) and snapshot_fixed_states_device / restore_fixed_states_device replace the host-bytes variants. RecurrentStateSnapshot carries a device_scratch flag. Byte-identical (D2D is an exact copy). Inert for greedy (only runs under spec-decode with recurrent state). GPU: MTP token-identical to baseline (md5 be7ed565, median-of-5), both slots replay (Primary=248, Verify=184, fallbacks=0), no NaN. Throughput 26.73 -> 34.34 tok/s median-of-5. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…(20.68→34.34 tok/s) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
# Conflicts: # crates/onnx-genai-engine/src/native_decode/mod.rs
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1663 +/- ##
===========================================
- Coverage 82.11% 81.40% -0.71%
===========================================
Files 12 384 +372
Lines 5473 180716 +175243
Branches 5473 180716 +175243
===========================================
+ Hits 4494 147112 +142618
- Misses 780 28662 +27882
- Partials 199 4942 +4743
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
…rify int4-GEMM kernel (#1673) ## Summary Investigated the coordinator's proposed **EAGLE-style base-decode fusion** (fuse out the redundant separate M=1 base decode) as the final structural lever to push MTP self-spec past the ~56 tok/s greedy baseline. **Measurement refuted the hypothesis: the fusion is throughput-negative and was NOT implemented** (per the standing "STOP and report rather than ship a regression" directive). This PR lands the measured analysis + an env-gated diagnostic, not a perf change. ## What's in the diff - **Env-gated per-phase profiler** in `native_speculative.rs` (`ONNX_GENAI_PROFILE_SPEC_PHASES=1`): times base decode / propose / verify / commit, prints cumulative + per-call ms at generation end. Inert unless the env var is set; lives only in the spec-decode path → **greedy byte-identical**. - **Decision note** `.squad/decisions/inbox/gaff-mtp-base-decode-fusion-regresses-verify-kernel-is-the-lever.md`. ## Measured per-step decomposition Steady, 48 outer steps, 128-tok run, k=1, 83.3% accept, **2.67 tok/step**, H200, int4 block-32, ORT 1.28 cuda13, both graph slots replaying (fallbacks=0): | phase | ms/step | share | |---|--:|--:| | base decode M=1 | 18.3 | 23% | | propose (MTP) | 5.9 | 7% | | verify M=2 | 49.2 | 62% | | commit | 6.2 | 8% | | **total** | **79.6** | 100% | → 29.8 ms/tok → **34.2 tok/s** (matches origin/main #1663). ## Why the fusion regresses The base decode is only 23% of step time but contributes ~0.87 of the 2.67 tokens/step (it folds the prior bonus and produces `h_bonus`, which the MTP head requires). Fusing it out keeps the M=2 verify (~49ms) but drops tokens/step to ~1.83 → **~30 tok/s**. The premise "base decode is redundant, tokens/step stays ~1.8" is empirically wrong (measured 2.67). ## The real lever Both phases are CUDA-graph **replayed** (fallbacks=0), so the 18.3 vs 49.2 gap is pure kernel compute. The replayed **M=2 verify (49.2ms) costs 2.7× the M=1 base and MORE than 2 sequential M=1 forwards (36.6ms)** — inverting speculation's core premise. Root cause: the int4 `MatMulNBits` GEMV(M=1)→GEMM(M≥2) cliff. Making the M=2..K verify a **batched-GEMV int4 kernel** (~M=1 cost) → per-step 48.4/2.67 = 18.1 ms/tok ≈ **55 tok/s**, crossing greedy with zero token loss. This is a CUDA/EP kernel change, not a loop-structure change. ## Validation - Token-identity md5 `be7ed565` == greedy (48/48). - Engine lib **582 passed** (`--features native-backend`), 0 warnings. - Both graph slots replay (Primary 248, Verify 184), fallbacks=0. **Caveat:** GEMV→GEMM attribution is inferred (impossible negative-overhead linear model + M=2 > 2×M=1 inversion); nsys/ncu are sandbox-blocked. Follow-up: `ONNX_GENAI_PROFILE_OPS=1` to confirm `MatMulNBits` dominates the verify, then attack that kernel. 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.
Summary
Two GPU-validated, token-identical speedups on the real Qwen3.8-27B int4 hybrid (GDN+GQA) MTP self-speculative decode (k=1, 83.3% accept, 2.67 tok/step). Branch off origin/main `66e19c2b5`. Median-of-5 on an idle H200: 20.68 → 34.34 tok/s (+66%).
1. approach-B — eliminate the redundant re-advance
On full accept (~67% of steps), the eager M=K verify forward already advanced KV + GDN/conv recurrent state to exactly the committed length, so the snapshot→restore→re-advance commit is pure redundancy and is skipped. Partial accept still rebuilds by re-running the accepted prefix (GDN exposes no intermediate per-position states). 20.68 → 26.73 tok/s (+29%); Primary graph replays 504 → 248.
2. device-to-device recurrent snapshot
The per-step recurrent snapshot copied all fixed-state (GDN SSM + conv1d) bindings device→HOST every spec step (~10ms PCIe round-trip), unused on full-accept steps. Replaced with a stream-ordered device→device snapshot into a per-binding device scratch arena:
26.73 → 34.34 tok/s.
Also fixes a pre-existing red on main: `InferenceSession` gained `verify_exec` in #1658 but two `executor/tests.rs` initializers were never updated, breaking `cargo test -p onnx-runtime-session --features cuda`.
Per-step target-forward budget (48 verify steps)
D2D snapshot does not change the forward count; it removes the per-step PCIe D2H state stage.
Validation (all green)
Honest remaining gap
Still below the ~56 tok/s greedy baseline. The structural lever left is fusing out the separate M=1 base decode (per-step cost → a single M=2 verify forward). Blocked by a circular dependency (MTP head needs the hidden after the last committed token, only produced by the base decode); solvable via EAGLE-style frontier-hidden reseeding while staying token-identical, but the speedup hinges on preserving acceptance rate through the changed seeding — deferred as the next, riskier turn. approach-B + D2D are landed and durable regardless. Validated by unit tests + GPU token-identity; end-to-end campaign number reported honestly.
Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com