Repository navigation
perf(executor): eliminate per-layer GQA shape-input D2H in eager decode (+35% eager tok/s) - #1427
Merged
Merged
Conversation
… present-KV in eager decode Native eager decode was pinned at a launch-bound floor (~300 tok/s on clean dense 1.5b GQA exports vs ORT eager 475) by blocking scalar D2H readbacks in the executor's dynamic-shape resolver. `resolve_node_outputs` / `dynamic_output_shapes` read GroupQueryAttention inputs 5 (`seqlens_k`) and 6 (`total_sequence_length`) back to host per layer (~56 D2H/token) to size the present-KV outputs, each forcing a host wait on the prior kernel. For the native decode path, present-KV is bound aliased in-place to past-KV at fixed physical capacity (see `kernel_input_uses_physical_capacity`), so the present sequence extent equals `past_key[2]` — already host-known from the input shape — making the readback redundant. `seqlens_k` is never used in GQA shape inference at all. This skips materializing GQA integer shape-inputs, but ONLY when the present-KV output is externally (persistently) capacity-bound; a growing `past ⧺ current` GQA cache keeps the readback (fail-closed, correctness + generality preserved). `dynamic_output_shapes` sizes present from `past_key[2]` when the total value is absent. Keys only on structural shape conditions (GQA op + capacity-bound present) — no model-family, head-size, or head-count branch — so it lifts qwen/llama/glm/phi/deepseek eager decode uniformly. Gated by `ONNX_GENAI_GQA_SHAPE_ONDEVICE` (default-on; disable via 0/false/off/no), mirroring #1383's `ONNX_GENAI_DEFER_EAGER_SYNC` escape-hatch pattern. Measured (deepseek-r1-distill-qwen-1.5b-int4-cuda, GPU H200, eager, greedy, medians-of-5): short-ctx 298.12 -> 402.85 tok/s (+35.1%), deep ~2600-ctx 255.83 -> 367.31 tok/s (+43.6%); shape-input D2H/step 1392 -> 104. Byte-identical: DeepSeek-V2-Lite, qwen3-0.6b, and GLM-4-9B golden decode locks all pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
pushed a commit
that referenced
this pull request
Aug 19, 2026
Three bypassed merges left `main` failing both required checks: * #1427 introduced `map_or(true, ..)` in `executor/dispatch.rs`, which clippy rejects under `-D warnings` (`this map_or can be simplified`). That fails the cross-compile gate, so the aarch64 pass never even runs. Replaced with the equivalent `is_none_or`. * #1427 and #1418 also left `dispatch.rs`, `gather_block_quantized.rs` and `gpt_oss_20b_decode_lock.rs` unformatted. `cargo fmt --all -- --check` runs in *both* required jobs while the cross-compile gate runs only in `Rust quality`, so while `main` is broken this way no single-fix PR can go green: a fmt-only PR still fails the cross-compile step, and this PR's own aarch64 fix still fails fmt. Only a branch carrying both can pass, which is why the repair is folded in here. It drops out on rebase once fixed upstream. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This was referenced Aug 19, 2026
justinchuby
pushed a commit
that referenced
this pull request
Aug 19, 2026
Three bypassed merges left `main` failing both required checks: * #1427 introduced `map_or(true, ..)` in `executor/dispatch.rs`, which clippy rejects under `-D warnings` (`this map_or can be simplified`). That fails the cross-compile gate, so the aarch64 pass never even runs. Replaced with the equivalent `is_none_or`. * #1427 and #1418 also left `dispatch.rs`, `gather_block_quantized.rs` and `gpt_oss_20b_decode_lock.rs` unformatted. `cargo fmt --all -- --check` runs in *both* required jobs while the cross-compile gate runs only in `Rust quality`, so while `main` is broken this way no single-fix PR can go green: a fmt-only PR still fails the cross-compile step, and this PR's own aarch64 fix still fails fmt. Only a branch carrying both can pass, which is why the repair is folded in here. It drops out on rebase once fixed upstream. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
pushed a commit
that referenced
this pull request
Aug 19, 2026
Bypassed merges left `main` failing both required checks: * #1427 introduced `map_or(true, ..)` in `executor/dispatch.rs`, which clippy rejects under `-D warnings` (`this map_or can be simplified`). That aborts the cross-compile gate before its aarch64 pass runs, which is why the gate could not report the aarch64 defect this PR fixes. Replaced with the equivalent `is_none_or`. * #1427 and #1418 also left `dispatch.rs`, `gather_block_quantized.rs` and `gpt_oss_20b_decode_lock.rs` unformatted. `cargo fmt --all -- --check` runs in *both* required jobs while the cross-compile gate runs only in `Rust quality`, so while `main` is broken this way no single-fix PR can go green: a fmt-only PR still fails the cross-compile step, and this PR's own fix still fails fmt. Only a branch carrying both can pass. Drops out on rebase once fixed upstream. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
pushed a commit
that referenced
this pull request
Aug 19, 2026
Bypassed merges left `main` failing both required checks: * #1427 introduced `map_or(true, ..)` in `executor/dispatch.rs`, which clippy rejects under `-D warnings` (`this map_or can be simplified`). That aborts the cross-compile gate before its aarch64 pass runs, which is why the gate could not report the aarch64 defect this PR fixes. Replaced with the equivalent `is_none_or`. * #1427 and #1418 also left `dispatch.rs`, `gather_block_quantized.rs` and `gpt_oss_20b_decode_lock.rs` unformatted. `cargo fmt --all -- --check` runs in *both* required jobs while the cross-compile gate runs only in `Rust quality`, so while `main` is broken this way no single-fix PR can go green: a fmt-only PR still fails the cross-compile step, and this PR's own fix still fails fmt. Only a branch carrying both can pass. Drops out on rebase once fixed upstream. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
🔴 Benchmark Regression DetectedComparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).
Visual flags: Host infoWhat this cannot catch
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1427 +/- ##
===========================================
+ Coverage 80.26% 82.93% +2.67%
===========================================
Files 362 12 -350
Lines 157159 5584 -151575
Branches 157159 5584 -151575
===========================================
- Hits 126137 4631 -121506
+ Misses 26406 760 -25646
+ Partials 4616 193 -4423
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 19, 2026
…1382) ## What this fixes Validating the merged scheduler/fmt wave, I found `#1363` had dropped the `allow(dead_code)` guards on three prefill fan-out symbols, breaking the aarch64 cross-compile gate. **#1443 has since fixed that** — but by `#[cfg]`-ing the three items to `target_arch = "x86_64"`, which removes them outright. That trades one break for two others. ### 1. 🔴 `aarch64 + feature = "mlas"` no longer compiles The three symbols have two callers, gated on **different** things: | call site | enclosing fn | its gate | | --- | --- | --- | | `matmul_nbits.rs:2457` | `run_mlas_shards` | `#[cfg(feature = "mlas")]` | | `matmul_nbits.rs:6914-6915` | `borrowed_affine_int4_matmul_prefill` | `#[cfg(target_arch = "x86_64")]` | Off x86 the second caller disappears, but the **first does not** — it is arch-independent. So on `aarch64 + mlas`, which is Apple Silicon, the caller is compiled and its callee is not: ``` error[E0425]: cannot find function `prefill_fan_out` in this scope --> crates/onnx-runtime-ep-cpu/src/kernels/matmul_nbits.rs:2457:27 error[E0425]: cannot find function `prefill_fan_out` in this scope --> crates/onnx-runtime-ep-cpu/src/kernels/matmul_nbits.rs:6914:8 error[E0425]: cannot find function `prefill_column_grain` in this scope --> crates/onnx-runtime-ep-cpu/src/kernels/matmul_nbits.rs:6915:48 ``` **How that was produced.** MLAS's vendored sources do not cross-build to aarch64 in this container (`arm_neon.h: inlining failed in call to always_inline vaddq_f16 — target specific option mismatch`, an mlas-sys/toolchain issue unrelated to this PR), so instead I reproduced the *exact cfg resolution* on the host: on `main`, retarget the three item gates from `x86_64` to a third arch so they are absent, leave every caller alone, and build the lib with MLAS on — ``` sed -i '395s/x86_64/s390x/; 403s/x86_64/s390x/; 543s/x86_64/s390x/' \ crates/onnx-runtime-ep-cpu/src/kernels/matmul_nbits.rs cargo check --locked -p onnx-runtime-ep-cpu --features mlas --lib # exit 101 ``` That is precisely the configuration `aarch64 + mlas` produces. The `:2457` error is the load-bearing one — that call site is `feature`-gated only, so it is present on aarch64 for real. **Why no gate caught it.** `check_cross_compile.sh`'s aarch64 pass builds **default features**, and `mlas` is off by default — the same blind spot #1443 was written under. The configuration *is* built elsewhere: the `rust-coverage` job's macOS-arm64 leg runs with `RUSTFLAGS: -D warnings` (`ci.yml` L434, L439) and builds `cargo build -p onnx-runtime-ep-cpu-plugin --features mlas` (L501-503), whose `mlas` feature forwards to `onnx-runtime-ep-cpu/mlas`. That is the shipped-wheel build path, so `main` as it stands also breaks the macOS-arm64 MLAS wheel at release time. ### 2. 🔴 The #1363 fan-out policy stopped being tested on aarch64 Gating the items forced gating their tests, so #1443 also put `#[cfg(target_arch = "x86_64")]` on six unit tests. All six are pure functions of explicit literal arguments — `prefill_fan_out(WIDE_PREFILL_MACS - 1, 16, 32)`, `prefill_column_grain(8, 1024, 3072)` — with nothing architecture-specific in them. They are now simply not compiled off x86, so the policy that #1363 rewrote has no aarch64 coverage. ## The fix `cfg_attr(.., allow(dead_code))` instead of `cfg`. The item always exists, so whichever caller survives can reach it; the lint is silenced only where **no** caller exists. The six tests are ungated and run everywhere again. `prefill_tile_grain` in this same file already uses this idiom (`not(feature = "mlas")`) — that is the shape #1363 deleted. The predicate is **per-symbol**, because the caller sets differ: | symbol | callers | predicate | | --- | --- | --- | | `WIDE_PREFILL_MACS`, `prefill_fan_out` | `run_mlas_shards` **and** `borrowed_affine_int4_matmul_prefill` | `not(any(feature = "mlas", target_arch = "x86_64"))` | | `prefill_column_grain` | `borrowed_affine_int4_matmul_prefill` **only** — `run_mlas_shards` takes `prefill_tile_grain` instead | `not(target_arch = "x86_64")` | Giving `prefill_column_grain` the union predicate would leave the lint live on `aarch64 + mlas`, where it has no caller — converting #1443's `E0425` into a `never used` error in the same configuration. Review caught exactly that in the first draft of this branch; the four-way probe below is the regression check for it. ### Four-config probe of the predicates `.validation-worktrees/cfgprobe/probe.rs` reproduces the two items, the two callers and their gates with `mlas`/`x86` standing in for the real cfgs, compiled under `-D warnings`: ``` === union predicate on prefill_column_grain (wrong) === PASS [aarch64 default] FAIL [--cfg mlas] <- error: function `prefill_column_grain` is never used PASS [--cfg x86] PASS [--cfg mlas --cfg x86] === per-symbol predicates (this PR) === PASS [aarch64 default] PASS [--cfg mlas] PASS [--cfg x86] PASS [--cfg mlas --cfg x86] ``` The probe is sharp, not vacuous: it fails on exactly the configuration that is wrong, and only that one. ## Second commit: unbreaking `main`'s required lane `main` currently fails **both** required checks, from merges landed past queued checks: | defect | source | breaks | | --- | --- | --- | | `map_or(true, ..)` in `executor/dispatch.rs` — clippy `this map_or can be simplified` under `-D warnings` | #1427 | `Rust quality` — and it aborts the cross-compile gate *before* its aarch64 pass, which is why the gate never reported defect 1 | | `dispatch.rs`, `gather_block_quantized.rs`, `gpt_oss_20b_decode_lock.rs` unformatted | #1427, #1418 | `Fast (Linux x86_64)` **and** `Rust quality` | `cargo fmt --all -- --check` runs in *both* required jobs (`ci.yml` L162, L276) while `check_cross_compile.sh` runs only in `Rust quality` (L401), and PR checks run against `merge(base, head)`. So while `main` is broken this way a fmt-only PR still fails the cross-compile step and this PR alone still fails fmt — only a branch carrying both can go green. It is mechanical (`cargo fmt --all`, plus `map_or(true, f)` → `is_none_or(f)`, identical on `Option`) and `git rebase` drops it once fixed upstream. ## Verification at `9fb04f5b5` (base `main` `81f99ff42`) | check | step | `main` | this branch | | --- | --- | --- | --- | | `Fast` + `Rust quality` | `cargo fmt --all -- --check` | **FAIL** (4 diffs / 3 files) | **pass** | | `Rust quality` | `bash scripts/check_cross_compile.sh` | **FAIL** exit 1 | **pass** exit 0, `scope: full offline set (aarch64 cross toolchain present)` | | `Rust quality` | 30-crate `cargo clippy --locked --all-targets … -- -D warnings` | **FAIL** exit 1 | **pass** exit 0 | | `Rust quality` | 9 guard scripts | pass | **9/9 pass** | | aarch64 | `cargo clippy --target aarch64-unknown-linux-gnu --all-targets -p onnx-runtime-ep-cpu -- -D warnings` | pass | **pass** (now *with* the 6 tests compiled) | | `aarch64 + mlas` cfg resolution | `cargo check -p onnx-runtime-ep-cpu --features mlas --lib`, items absent | **FAIL** exit 101, 3 × E0425 | **pass** — `cfg_attr` never removes the item, so E0425 cannot occur | | all 4 `(mlas on/off) x (x86 / non-x86)` | `rustc -D warnings` cfg probe | — | **4/4 pass** | | tests | `cargo test -p onnx-runtime-ep-cpu --lib` | — | **1447 passed / 0 failed**; the 8 prefill policy tests pass | ## A note on the gate that found this `scripts/check_cross_compile.sh` **false-passes locally** without an aarch64 cross toolchain: at L191-194 it silently swaps `CRATES_FULL` → `CRATES_NO_FFI`, dropping `onnx-runtime-ep-cpu` — the crate the gate exists for — and still exits 0 with a ✓. The "REDUCED SCOPE" note prints *below* the checkmark. **Read the scope note, not the exit code**; only `scope: full offline set (aarch64 cross toolchain present)` means anything. On Actions it `exit 2`s instead (L178-190), and `ci.yml` L396-399 installs `gcc-aarch64-linux-gnu` + `libc6-dev-arm64-cross` before invoking it, so the fail-loud coverage is intact — this is a local-only trap. All results above were produced with the toolchain installed, at full scope. Two of the three defects in this PR would have been caught by the required checks had they been allowed to run. ## Process No admin bypass, no ruleset bypass, no merge with checks queued or failing. Auto-merge has been armed since 2026-08-19T04:55:37Z and merges only once `Fast (Linux x86_64)` and `Rust quality` are green. Every CI run in this repo is currently `queued` with zero in progress, so the required contexts have not been created yet. Waiting. --------- Co-authored-by: Pris <pris@squad.local> 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
Raises the native eager decode "launch-bound floor" by removing the dominant per-token blocking scalar D2H. On clean dense GQA exports native eager sat at ~300 tok/s vs ORT eager 475; this lever lifts it to ~403 tok/s short-ctx (+35.1%) / ~367 deep-ctx (+43.6%), byte-identical.
Root cause
The floor is not the GQA kernel readback (that is already skipped in steady-state eager by the kernel's
warmed_signatureself-warming). The real cost is the executor's dynamic-shape resolver:resolve_node_outputs→dynamic_output_shapesreads GroupQueryAttention inputs 5 (seqlens_k) and 6 (total_sequence_length) back to host per layer (~56 blocking 4-byte D2H/token) to size the present-KV outputs. Each forces a host stall on the prior kernel.Fix (executor-only, 2 files)
dynamic_shapes.rs(GQA branch):present_sequencefalls back topast_key[2]whentotal_sequence_lengthis absent. Native binds present-KV aliased in-place to past-KV at fixed physical capacity (kernel_input_uses_physical_capacity), sopresent == max(past_key[2], total) == past_key[2]always — host-known from the input shape.dispatch.rs(resolve_node_outputs): skip materializing GQA integer shape-inputs only when the present-KV output (outputs[1]) is externally capacity-bound (external.outputs). A growingpast ⧺ currentGQA keeps the readback → fail-closed, correctness + generality preserved.seqlens_kis never used in GQA shape inference (pure waste).Generality
Keys ONLY on structural shape conditions (GQA op + capacity-bound present output). No model-family / head-size / head-count branch. Lifts qwen / llama / glm / phi / deepseek eager decode uniformly (RULES.md §2).
Escape hatch
ONNX_GENAI_GQA_SHAPE_ONDEVICE(default-on; disable via0/false/off/no), mirroring #1383'sONNX_GENAI_DEFER_EAGER_SYNC.Measurements
Export
deepseek-r1-distill-qwen-1.5b-int4-cuda(1536h/28L GQA), H200 pinned idle,ONNX_GENAI_CUDA_GRAPH=0eager, greedy, medians-of-5,profile_native --steady:=0, old path)vs ORT eager 475 (matched io-binding onnxruntime-genai harness, same export): native 0.85× short — did NOT cross 475 (honest gap −15% short / −23% deep). Residual is per-launch host dispatch:
cuLaunchKernelunchanged at ~219/token — closing it needs kernel fusion / launch-batching, a separate larger lever.Correctness — golden locks (byte-identical to pre-change, default-ON,
--test-threads=1)ONNX_GENAI_GQA_SHAPE_ONDEVICE=0— qwen3 lock still PASS (old readback path restored)Notes
Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com