Enable static-cache + Flash-attention path (runtime-gated, ready for ORT #28958) - #364
Merged
Conversation
The build_from_module opset lowering downgraded the default-domain opset from 24 to 23 for every non-default EP when MOBIUS_ORT_LOWER_OPSET_FOR_EP is set. nonpad_kv_seqlen (Attention input #6) and TensorScatter are opset-24-only, so declaring opset 23 on a static-cache graph is invalid and silently strips the static-cache Flash-attention path. Add _graph_requires_opset24() and skip lowering for any sub-model that contains a TensorScatter node or an Attention consuming a non-empty input #6. Standard-op-only sub-models still lower as before, so there is no memcpy-storm regression for the common path. Defensive fix: the flag defaults OFF. Add CPU-authorable unit tests in src/mobius/_builder_test.py. Agent-signed-off: Developer (71529949) [claude-opus-4.8 via copilot] Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The comment above the opset 24->23 lowering claimed the flag 'defaults to True', but ort_lower_opset_for_ep defaults to False (_flags.py:112-113, _env_bool MOBIUS_ORT_LOWER_OPSET_FOR_EP default False). Correct the comment so the documented behavior matches reality. Agent-signed-off: Developer (71529949) [claude-opus-4.8 via copilot] Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Re-adds the end-to-end static-cache (TensorScatter + ONNX-domain Attention) CUDA coverage dropped with the aborted PR #340, adapted to the current maskless graph on main (is_causal=1 + nonpad_kv_seqlen, no attn_mask/past_key). tests/static_cache_flash_e2e_test.py builds a tiny fp16 static-cache model, runs prefill (S_q>1) and decode (S_q=1) on the CUDA EP asserting finite outputs and in-place cache advance, and asserts the Attention runs on the Flash kernel (not MEA/unfused). Flash dispatch is asserted via the kernel's own VERBOSE log ("ONNX Attention: using Flash Attention" from cuda/llm/attention.cc), captured by redirecting the native stderr fd and setting set_default_logger_severity(0) -- this build exposes neither CUPTI Kernel-cat profiling events nor the per-op AttentionKernelDebugInfo, so the VERBOSE log is the authoritative channel. tests/_static_cache_support.py is a shared helper: a functional capability probe (static_cache_cuda_supported) that skips when ORT lacks microsoft/onnxruntime#28958, a quick_build_enabled gate that skips the Flash assertion under onnxruntime_QUICK_BUILD (Flash compiled for head_dim 128 only there), and the dispatch-log capture/parse helpers. Verified end-to-end on an A100 (ORT 1.28.0 contains #28958): both tests pass, the Flash assertion skips under onnxruntime_QUICK_BUILD, and the file is collectable without GPU. Agent-signed-off: Developer (ed516084) [claude-opus-4.8 via copilot] Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Add tests/static_cache_parity_test.py guarding the maskless static-cache attention combination (is_causal=1 + nonpad_kv_seqlen + TensorScatter) that mobius emits and that onnx#8068 / onnxruntime#28958 enable on CUDA Flash. Gate: reuses the shared functional probe static_cache_cuda_supported() from tests/_static_cache_support.py (also used by the e2e Flash-dispatch test) so both static-cache CUDA tests share one capability check. Pre-#28958 ORT raises NOT_IMPLEMENTED -> skip; post-#28958 ORT runs. Also skips when no CUDA EP is registered (TensorScatter + external-cache Attention are CUDA-only). Tests (all @pytest.mark.integration): - test_static_vs_dynamic_vs_hf_decode: builds Qwen2.5-0.5B twice from the same weights (static + dynamic cache); asserts token-id and last-token logit parity across static / dynamic / HuggingFace over N greedy decode steps on CUDA f32. - test_chunked_prefill_structurally_empty_rows_are_zero: q_seq>1 chunked prefill with nonpad_kv_seqlen < q_seq; asserts structurally-empty bottom-right rows are EXACTLY 0 (LaunchZeroFullyMaskedRows guard), all finite, valid rows non-zero. Exact valid-row values are NOT asserted because upstream's bottom-right numpy parity reference is deferred (ort#28958 2E). - test_prefill_chunk_causal_triangle_matches_reference: q_seq>1 prefill with nonpad==q_seq verifies the kernel reproduces the causal triangle against a NumPy SDPA reference (the well-defined bottom-right math). Validated: 3 passed on A100 / ORT 1.28.0; skips cleanly on CPU-only. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang <titaiwang@microsoft.com>
…sting.ort_capabilities The functional 'does this ORT run the maskless is_causal=1 + nonpad_kv_seqlen + TensorScatter static-cache graph on CUDA (needs onnxruntime#28958)' probe was hand-rolled in two test modules. Move the single canonical implementation to src/mobius/_testing/ort_capabilities.py as supports_static_cache_flash() — a fail-closed (any exception -> False), lru_cached probe built on a minimal self-contained TensorScatter + opset-24 Attention graph with no tests/ deps. - New: src/mobius/_testing/ort_capabilities.py (CUDA_AVAILABLE, supports_static_cache_flash()). - tests/_static_cache_support.py: delete the duplicate probe + its prefill feeds; keep only the e2e-specific model builder + Flash-dispatch capture helpers; re-export the canonical probe under the old name as a temporary migration bridge. - tests/static_cache_flash_e2e_test.py: import the canonical probe. Verified on A100 (ORT 1.28.0 source build w/ #28958): probe returns True and all 5 static-cache tests pass (2 e2e Flash-dispatch + 3 parity). Agent-signed-off: Developer (ed516084) [claude-opus-4.8 via copilot] Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…_capabilities Swap the capability-probe import from the transitional tests/_static_cache_support bridge to the canonical mobius._testing.ort_capabilities.supports_static_cache_flash. Import-only change: no test logic altered. Docstring repointed to the canonical module. All 3 tests still pass on A100/ORT 1.28.0. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang <titaiwang@microsoft.com>
Convenience wrapper over supports_static_cache_flash() that returns None when the static-cache Flash path runs, else a skip message distinguishing the two causes (no CUDA EP vs. an ORT build predating microsoft/onnxruntime#28958). Lets tests skip via pytest.skip(reason) inside a fixture/helper instead of CUDA_AVAILABLE-branching at the call site. Verified on A100: returns None (path supported); isolated ruff check + format clean. Agent-signed-off: Developer (ed516084) [claude-opus-4.8 via copilot] Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Adopt mobius._testing.ort_capabilities.static_cache_flash_skip_reason() in _require_static_cache_attention(), replacing the local two-branch CUDA-EP / onnxruntime#28958 skip text with the single canonical source. Drops the now-unused CUDA_AVAILABLE import. Skip behavior is identical; all 3 tests pass on A100/ORT 1.28.0. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang <titaiwang@microsoft.com>
…o probe bridge Per the tightened consolidation DoD, the e2e Flash-dispatch test no longer depends on tests/_static_cache_support.py: its e2e-only helpers (tiny-model builder, cache feeds, and the VERBOSE Flash-dispatch capture) are inlined, and it imports the canonical probe from mobius._testing.ort_capabilities (CUDA_AVAILABLE, supports_static_cache_flash). tests/_static_cache_support.py is reduced to a minimal back-compat bridge that only re-exports the canonical probe under its historical name, so the not-yet-migrated parity test keeps collecting. The bridge is deleted in a follow-up once the parity test imports the canonical module directly. Verified on A100 (ORT 1.28.0 w/ #28958): e2e 2/2 and parity 3/3 pass; isolated ruff check + format clean. Agent-signed-off: Developer (ed516084) [claude-opus-4.8 via copilot] Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Delete tests/_static_cache_support.py now that the capability probe lives solely in mobius._testing.ort_capabilities and both the parity and e2e tests import it directly (parity via static_cache_flash_skip_reason; e2e via supports_static_cache_flash with its fixture helpers inlined). No file imports the shim anymore. Parity 3/3 pass and e2e collects clean on A100/ORT 1.28.0. Single canonical import path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang <titaiwang@microsoft.com>
Strengthen the nonpad<q_seq chunked-prefill test: replace reliance on the near-vacuous non-zero check with an authoritative convex-combination invariant. Softmax attention output is a convex combination of the attended V rows, so each surviving output element must lie within the [min, max] hull of V over the valid key window [0, nonpad), per head and per channel. This catches a wrong attention-frontier bug (attending padded keys / mean-of-all-V) without depending on the upstream-deferred exact valid-row reference (ort#28958 §2E). The exact-math causal-triangle test (nonpad==q_seq) is unchanged. Also move the mobius._testing.ort_capabilities import into the first-party isort block. All 3 parity tests pass on A100/ORT 1.28.0. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang <titaiwang@microsoft.com>
Two robustness improvements to the conditional opset 24->23 lowering (decision semantics unchanged): 1. _graph_requires_opset24 now scans with ir.traversal.RecursiveGraphIterator instead of only top-level nodes, so a TensorScatter or Attention input #6 (nonpad_kv_seqlen) nested inside a future If/Loop/Scan subgraph is still detected. Docstring updated to drop the shallow-scan assumption. 2. The lowering loop is extracted into _apply_opset_lowering(pkg, ep), which owns the flag/EP/opset call-site gates and is now driven directly by the tests instead of a re-implemented decision. Added a nested-subgraph test and a mixed-package test (one static-cache sub-model that must keep opset 24 + one standard sub-model that must lower to 23), plus default-EP and flag-disabled gate tests. Agent-signed-off: Developer (71529949) [claude-opus-4.8 via copilot] Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
M1 (ort_capabilities.py): make supports_static_cache_flash() fail-closed but
not fail-silent. Introduce a cached _probe_static_cache_flash() returning a
_ProbeOutcome enum (SUPPORTED/NO_CUDA/NEEDS_FIX/PROBE_ERROR). Expected ORT
kernel rejects (NotImplemented/Fail from session.run, pre-#28958) are logged at
debug -> NEEDS_FIX; unexpected build/serialize/session errors and a vanished
CUDA EP are logged at warning with exc_info -> PROBE_ERROR. static_cache_flash_
skip_reason() now reports the true cause and no longer misattributes a broken
probe to needs-#28958.
M2 (static_cache_flash_e2e_test.py): add _flash_capable_gpu() (torch compute
capability >= (8,0)) and skipif-guard the Flash-dispatch test so pre-SM80 CUDA
(T4/Volta) skips cleanly instead of hard-failing the =={'flash'} assertion;
update the docstring accordingly.
nits: hoist prefill_len=4 to module constant _PREFILL_LEN; annotate config:
object params as ArchitectureConfig.
Validated on A100 (SM 8.0): 5/5 static-cache tests pass (43s), ruff clean.
Agent-signed-off: Developer (ed516084) [claude-opus-4.8 via copilot]
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…e check Part (b) of the expanded scope: the structured-outcome enum (part a) cannot catch a SILENT CPU fallback because no exception is raised. ORT always appends the CPU EP as implicit fallback, and the CPU opset-24 Attention kernel runs the maskless is_causal=1 + nonpad + TensorScatter graph WITHOUT erroring (with the historical wrong top-left-causal values on pre-fix builds). So if a CUDA build declines the node at GetCapability, session.run() succeeds on CPU and the probe would fail-OPEN -> report SUPPORTED with zero CUDA coverage. Close it with a deterministic known-answer (value-based, not log parsing): identical keys/query -> uniform softmax over the two valid KV positions; distinct per-position value tags (1.0, 3.0) -> the correct bottom-right output is their mean (2.0). A wrong top-left kernel lets the single decode query attend only KV slot 0 -> 1.0, which _probe_output_is_correct() rejects, classifying the run as NEEDS_FIX rather than SUPPORTED. Probe geometry switched to the decode regime (S_q=1, two valid KV slots) so bottom-right vs top-left genuinely differ. Validated on A100 (SM 8.0): correct ref (2.0) accepted, top-left (1.0) rejected, A100 still SUPPORTED; 5/5 static-cache tests pass (46s), ruff clean. Agent-signed-off: Developer (ed516084) [claude-opus-4.8 via copilot] Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…guard Confirm the core contract for the new S_q=1 probe geometry: a pre-#28958 ORT still returns False. Source-grounded in CUDA llm/attention.cc Attention<T>:: ComputeInternal, guard `causal_cross_no_past = is_causal && (q_seq != total_seq) && (past == 0)` then `if (causal_cross_no_past && nonpad_kv_seqlen != nullptr) return NOT_IMPLEMENTED`. For S_q=1 over total_kv>=2, q_seq != total_seq holds and there is no S_q==1 decode fast-path that bypasses it (S_q=1 is the case the reject message explicitly names), so a pre-fix build RAISES at run from ComputeInternal -> expected reject -> NEEDS_FIX -> False. The known-answer value check remains the second defense for non-raising wrong-value fallbacks. Comment-only; no behavior change. 5/5 static-cache tests pass on A100 (45s), ruff clean. Agent-signed-off: Developer (ed516084) [claude-opus-4.8 via copilot] Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Comment/docstring only, zero logic change: 1. ort_capabilities.py _probe_static_cache_flash docstring: record the one residual assumption — the probe validates a CUDA EP + output correctness but does NOT explicitly assert the Attention node was placed on CUDA; safe today because the pre-#28958 CPU fallback yields the wrong top-left value the value-check rejects, would only mis-pass if a future CPU kernel became bottom-right-correct AND CUDA declined the node. 2. e2e test module docstring: note the new SM>=8.0 (_flash_capable_gpu) skip on the Flash-dispatch test alongside the QUICK_BUILD skip, so the doc matches behavior. 3. ort_capabilities.py: note _PROBE_QUERY_FILL is intentionally equal to _PROBE_KEY_FILL (identical K/Q -> uniform softmax the known-answer relies on). 5/5 static-cache tests pass on A100 (44s), ruff clean. Agent-signed-off: Developer (ed516084) [claude-opus-4.8 via copilot] Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Performance Comparison
|
|
The author of this PR, titaiwangms, is not an activated member of this organization on Codecov. |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds ORT capability gating plus CUDA integration tests to validate mobius’s static-cache (TensorScatter KV) attention path, and updates opset-lowering logic so the static-cache Flash-attention form remains valid when lowering opset declarations for older EPs.
Changes:
- Add CUDA integration coverage for static-cache correctness (parity vs dynamic-cache ONNX + HF) and for Flash-kernel dispatch.
- Introduce a cached, functional ORT capability probe (
supports_static_cache_flash/static_cache_flash_skip_reason) to skip tests until ORT includes microsoft/onnxruntime#28958 behavior. - Refactor opset 24→23 lowering so sub-models that require opset-24-only semantics (TensorScatter or Attention input #6
nonpad_kv_seqlen) are not lowered; add unit tests for this scan/lowering behavior.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/static_cache_parity_test.py | Integration parity + edge-case tests for static-cache Attention behavior (incl. bottom-right causal corner cases). |
| tests/static_cache_flash_e2e_test.py | CUDA e2e regression + Flash-dispatch verification via ORT VERBOSE logs. |
| src/mobius/_testing/ort_capabilities.py | Adds functional CUDA capability probe to gate static-cache Flash tests until ORT supports the kernel combo. |
| src/mobius/_builder.py | Moves opset-lowering into helpers and prevents lowering for graphs requiring opset-24-only semantics. |
| src/mobius/_builder_test.py | Unit tests for _graph_requires_opset24 and per-submodel opset-lowering behavior. |
Contributor
Author
Review summaryVerdict: sound, mergeable. No Critical issues across a five-reviewer pass (readability / correctness / adversarial / spec-deep / cross-module). The math and spec are independently verified: opset-24 🟠 MajorM1 — The probe's
M2 — Duplicated opset-lowering logic has drifted. 🟡 Minor
🔵 Nits
❓ Open questions (grounded)
✅ PraiseThe known-answer probe design (identical Q/K → uniform softmax → mean-of-value-tags, which cleanly separates top-left Synthesized from a five-agent review (readability, correctness, adversarial, spec-deep, cross-module). GPU-gated tests were not executed here; M1/M2 and the open questions are the load-bearing items. |
titaiwangms
added a commit
that referenced
this pull request
Jun 19, 2026
…er -> after-scatter) (#365) ## What `examples/static_cache_generation.py` fed `nonpad_kv_seqlen` as the **before-scatter** KV count. A separate `nonpad_kv_seqlen` accumulator stayed always-equal to `write_indices` (both initialized to 0, both incremented only *after* `session.run`), so the example passed `nonpad == write_indices`. The spec-correct value is `write_indices + cur_seq_len` — the valid KV count **after** the current chunk is scattered into the cache, per the bottom-right `is_causal` contract (onnx/onnx#8068). ## Impact Wrong in **every** phase: - **Prefill**: passed `0` instead of `prompt_len`, so the kernel saw ~0 valid keys and produced wrong logits. - **Decode**: off-by-one on every step. On the CPU EP the maskless graph **runs-but-wrong silently** (no error raised). This is **example-only** — it is the canonical usage reference and does **not** affect library/graph code or PR #364. ## Fix - Feed `nonpad_kv_seqlen = write_indices + cur_seq_len` at the run inputs. - Remove the now-dead redundant `nonpad_kv_seqlen` accumulator (it merely duplicated `write_indices`). - Add an explanatory inline comment + docstring note so the canonical example teaches the correct **after-scatter** pattern. ## Verified - Matches the authoritative reference in `tests/static_cache_parity_test.py` (prefill nonpad = prompt_len = write_indices(0) + query_len, lines 322-327; decode nonpad = valid_len + 1 = write_indices(valid_len) + 1, lines 361-362) and the consumer `src/mobius/components/_attention.py:169-183` (the scatter precedes the maskless `is_causal` Attention's `nonpad` read). - Ran on CPU with `Qwen/Qwen2.5-0.5B` -> coherent output: `The capital of France is` -> ` Paris. It is the largest city in`. Pre-fix the all-zero prefill nonpad silently produced wrong logits. - Triple-reviewed (code / critical / readability all ship). ## Reference - onnx/onnx#8068 (causal bottom-right errata). - Relates to the static-cache + Flash enablement in PR #364 — this is a separate, example-only follow-up. --------- Signed-off-by: titaiwang <titaiwang@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Three grounded review fixes on the static-cache probe + opset lowering:
(2) ort_capabilities.py: the inner reject handler matched (NotImplemented,
Fail) by type alone, so a genuine ORT Fail (CUDA OOM, kernel bug, install
drift) was silently misclassified as 'needs onnxruntime#28958' -> whole
suite skips, masking a regression. Add _is_expected_pre28958_reject():
NotImplemented is always the confirmed pre-#28958 path; a Fail (or the
defensive RuntimeError fallback) is accepted only when its message carries
the reject signature (nonpad_kv_seqlen / tensorscatter), else it returns
PROBE_ERROR and logs at WARNING so real failures stay loud. Value-based
known-answer check preserved. Adds CPU-only unit tests for the classifier
and the output check.
(1) _builder.py _apply_opset_lowering: the gate guarded only EP != 'default',
so it also fired for EP == 'cpu', diverging from ort_inference. Also skip
'cpu' (CPU EP already has opset-24 kernels). Adds a cpu-skip unit test.
(M2) ort_inference.py _should_lower_opset: only scanned top-level nodes and
lacked the Attention input-#6 (nonpad_kv_seqlen) check. Delegate to
_builder._graph_requires_opset24 (recursive + input-#6) so both
opset-24-detection paths agree.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: titaiwang <titaiwang@microsoft.com>
…ew fixes Extend coverage for each behavior added in d0e96f5, per review: (2) ort_capabilities.py: extract _classify_run_error(exc) -> _ProbeOutcome (the caught session.run error -> NEEDS_FIX vs PROBE_ERROR mapping, with the debug/warning logging colocated) so the outcome is unit-testable without CUDA. New tests assert NotImplemented -> NEEDS_FIX, a signature-matching Fail -> NEEDS_FIX, a genuine Fail (CUDA OOM, no signature) -> PROBE_ERROR (not a silent 'needs #28958' skip), and the RuntimeError fallback -> PROBE_ERROR. (M2) add ort_inference_test.py: _should_lower_opset with a nested If-subgraph Attention(nonpad_kv_seqlen input #6) -> False, proving the recursion + input-#6 alignment now fires where the old top-level-only scan missed it; plus top-level TensorScatter, standard-graph-lowers, cpu-skip, flag-disabled, and opset<=max guards. (1) cpu-skip already covered by _builder_test.test_apply_opset_lowering_skipped_for_cpu_ep. Behavior unchanged; 28 unit tests pass, static-cache GPU tests remain fail-closed-skip on pre-#28958 ORT. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang <titaiwang@microsoft.com>
…t.py Every sibling file under src/mobius/_testing/ carries the standard 2-line MIT license header; this newly-added test file was missing it. Header-only change, no logic touched. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang <titaiwang@microsoft.com>
CI lint (ruff 0.15.14 with --config=pyproject.toml) flagged three errors in the static-cache files that the isolated local ruff run did not surface: - LOG014 (ort_capabilities.py): `exc_info=True` sits outside an except handler after _classify_run_error was extracted from the except arm. Pass the explicitly-captured exception (`exc_info=exc`) instead — semantically correct regardless of handler context and preserves the traceback in the warning log. - SIM112 (static_cache_flash_e2e_test.py): capitalize the CI-exported env var to ONNXRUNTIME_QUICK_BUILD. The lowercase onnxruntime_QUICK_BUILD references that name the actual ORT cmake build flag are left unchanged. - SIM115 (static_cache_flash_e2e_test.py): open the stderr-capture tempfile.TemporaryFile via a `with` context manager. Behavior-preserving: 28 unit tests pass, static-cache GPU trio still fail-closed-skips on pre-#28958 ORT. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang <titaiwang@microsoft.com>
The ② narrowing (a signature-less ORT Fail must surface as PROBE_ERROR, not be misclassified as the pre-#28958 NEEDS_FIX reject) was only tested at the pure _classify_run_error classifier level. Add a small end-to-end test that mocks ort.InferenceSession so session.run raises a signature-less Fail (CUDA OOM) and asserts _probe_static_cache_flash returns PROBE_ERROR AND logs a WARNING - i.e. a genuine failure stays loud instead of silently skipping the whole suite as "needs #28958". A contrasting case asserts the real NotImplemented reject still maps to NEEDS_FIX end-to-end. Quality improvement only: this is NOT the codecov/patch fix. codecov/patch is red due to a Codecov org account-activation gap (the PR author is not an activated Codecov member), which only a maintainer can resolve at app.codecov.io/members/gh/onnxruntime - adding tests cannot turn it green. Production code untouched; CPU-only (no CUDA needed). 14 ort_capabilities unit tests pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang <titaiwang@microsoft.com>
titaiwangms
force-pushed
the
static-cache-flash-enablement
branch
from
June 19, 2026 23:52
c39a5b8 to
6e19723
Compare
justinchuby
reviewed
Jun 22, 2026
justinchuby
approved these changes
Jun 22, 2026
justinchuby
reviewed
Jun 22, 2026
justinchuby
left a comment
justinchuby
left a comment
Member
There was a problem hiding this comment.
Review verdict: 🟡 minor-issues (non-blocking). Thorough, correctness-first PR — the opset-24 retention logic, the fail-closed/fail-open probe design, and the parity tests are all rigorous and well-reasoned. I traced the probe + tests against the production components/_attention.py static branch (Attention input #6 = nonpad_kv_seqlen, is_causal=1, no attn_mask/past_key) and they faithfully replicate it. A few non-blocking observations:
1. Description ↔ code mismatch (no fix needed in code). Change-summary item (e) says "deleted tests/_static_cache_support.py", but the diff contains no deletion and that file exists on neither main nor the head ref (gh api .../contents/tests/_static_cache_support.py → 404 on both). The end state is correct (no shim), but the "deleted" claim is misleading — worth correcting the PR body so reviewers aren't hunting for a phantom removal.
2. Nice latent-bug fix. Routing the build-side _apply_opset_lowering through the recursive _graph_requires_opset24 (and adding the per-submodel retention) closes a real hole: the previous build_from_module lowering had no TensorScatter/nonpad guard at all, so with MOBIUS_ORT_LOWER_OPSET_FOR_EP=1 on a non-default EP it would have emitted an invalid opset-23 graph carrying TensorScatter. Latent only (flag default is False), but a genuine correctness improvement, and the build/inference gates now agree.
3. Probe value-mismatch classification (inline) — a wrong-but-not-top-left output is silently NEEDS_FIX; see inline note.
4. Logger-severity restore nit (inline).
Nothing here blocks merge; the gating story (probe-gated skips, self-enabling once ORT #28958 is pinned) is sound and CI-safe. 👍
justinchuby
reviewed
Jun 22, 2026
justinchuby
reviewed
Jun 22, 2026
Three non-blocking review nits from justinchuby: 1. Rename `_apply_opset_lowering` -> `_maybe_apply_opset_lowering` (it gates on the flag/EP and may not lower), updating the build_from_module call site and the _builder_test import, calls, and test names. 2. ort_capabilities: in the silent-wrong-output branch, distinguish the *expected* pre-#28958 top-left fallback (mean ~= first value tag, 1.0) from an *unexpected* wrong value. The known fallback stays at debug; any other wrong value is now logged at WARNING so a genuine post-#28958 regression (or ORT drift) is not silently bucketed as the known fallback. Adds the `_PROBE_TOPLEFT_OUTPUT` constant naming that signature. 3. captured_attention_dispatch: restore the ORT default logger severity to the value that was set before the context (captured via a module-level shadow, since ORT exposes a setter but no getter) instead of a hardcoded 2. Lint clean (real pyproject config); _builder_test + ort_capabilities_test pass (24 passed, 2 CUDA-only e2e skipped on CPU). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang <titaiwang@microsoft.com>
Address Copilot review comments on PR #364: - tests/static_cache_flash_e2e_test.py: the module-level `skipif` evaluated `supports_static_cache_flash()` eagerly at import/collection time, which runs a CUDA build + serialize + session.run. Replace it with an autouse fixture that calls `static_cache_flash_skip_reason()` at test setup, so importing the module and `pytest --collect-only` are side-effect-free. The probe is lru-cached (runs at most once) and the helper reports the true skip cause (no CUDA EP, missing microsoft/onnxruntime#28958, or an unexpected probe failure). The remaining Copilot comments are already satisfied at HEAD: the opset-lowering gate already excludes the CPU and default EPs (_builder.py), and the probe only treats an ORT `Fail` as the expected pre-#28958 reject when its message carries the reject signature, routing generic failures (OOM, drift) to PROBE_ERROR (ort_capabilities.py). The _builder.py:216 comment was a PR-description note, not a code defect. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: titaiwang <titaiwang@microsoft.com>
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
titaiwangms
commented
Jun 22, 2026
titaiwangms
left a comment
titaiwangms
left a comment
Contributor
Author
There was a problem hiding this comment.
Review summary — static-cache Flash enablement / probe / CI
Reviewed at d95460c by a multi-model review team plus independent verification. Clean, well-gated infrastructure; no Critical/Major findings specific to this PR. Mergeable.
Verified correct ✅
- Recursive opset-24 guard (
_graph_requires_opset24+_maybe_apply_opset_lowering): correctly detectsTensorScatterandAttentionwith a non-empty input #6 (nonpad_kv_seqlen) insideIf/Loop/Scanbodies, decides per sub-model, and errs conservatively — an empty-name optional placeholder only keeps opset 24, never wrongly lowers to 23. The default-off flag means no shipped export changes. - Capability probe (
supports_static_cache_flash): genuinely fail-closed. The known-answer value check (2.0) closes the latent CPU-fallback fail-open where a CUDA build that declines the node would silently run on CPU with wrong (top-left) values; the_ProbeOutcomeenum cleanly separates the expected pre-#28958 reject from unexpected probe errors (logged withexc_info). Strong unit-test coverage of the classifier. - New GPU tests skip automatically unless the installed ORT can actually run the path, so CI stays green today and flips green automatically once an ORT release with microsoft/onnxruntime#28958 is installed — no code change needed.
Notes (no action required for this PR)
- This PR enables the maskless
is_causal=1path; the bias-aware sibling (#367) and the genai consumer (microsoft/onnxruntime-genai#2235) edit parallel, non-conflicting structures — confirmed compatible across branches. - The recursive guard is future-proof but currently only matches
TensorScatter/Attention-input#6; if other opset-24-only default-domain ops are introduced into these graphs later, the predicate set will need extending (a lowered-to-23 graph would otherwise be silently invalid).
Praise: the probe design (functional runtime check + known-answer value gate, not a version-string check) is exactly the right pattern for a capability that depends on an unreleased ORT kernel, and the fail-closed-but-loud behavior is well thought through.
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.
What
mobius
mainalready emits the correct masklessis_causal=1+nonpad_kv_seqlen+TensorScatterstatic-cache decoder graph. This PR adds the runtime enablement, verification, and CI for that path on the CUDA Flash-attention kernel — it is not graph surgery. Nothing in the emitted graph is reverted or rewritten.The path becomes runnable once an ONNX Runtime build containing microsoft/onnxruntime#28958 is installed (Flash eligibility widening for the bottom-right-causal errata onnx/onnx#8068).
Issues
tests/static_cache_parity_test.py: static vs dynamic vs HuggingFace, chunked-prefill zero-guard, V convex-hull valid-row invariant, exact-trianglenonpad == q_seq).is_causal=1+nonpad_kv_seqlenstatic-cache graph once upstream lands"). The ORT pin bump + node-count rebaseline + ONNX 1.22 pin + skill docs remain deferred until an official PyPI ORT release with Fix Attention is_causal bottom-right alignment for external KV cache (onnx#8068, #28904) microsoft/onnxruntime#28958, so Emit maskless is_causal=1 + nonpad_kv_seqlen static-cache graph once upstream lands #345 stays open (intentionallyPart of, notCloses).Changes
_builder.pyvia_graph_requires_opset24(a recursive subgraph scan) +_apply_opset_lowering, so graphs carryingTensorScatteror theAttentionnonpad_kv_seqleninput correctly stay at opset 24. Flag-gated, default off. (_builder_test.pyexercises the real branch.)src/mobius/_testing/ort_capabilities.py—supports_static_cache_flash()is a functional, fail-closed-but-loud runtime probe (not a version-string check). It builds a minimalTensorScatter+ masklessAttentiongraph and runs it on CUDA. A known-answer value check closes a latent CPU-fallback fail-open: because ORT implicitly appends the CPU EP, a CUDA build that declines the node would silently run on CPU with wrong (top-left) values; the probe's deterministic reference (2.0) rejects that →NEEDS_FIX→False. A structured_ProbeOutcomeenum distinguishes the expected pre-#28958 reject from unexpected probe errors (logged withexc_info).tests/static_cache_parity_test.py— static vs dynamic vs HuggingFace, chunked-prefill zero-guard, V convex-hull valid-row invariant, exact-trianglenonpad == q_seq.tests/static_cache_flash_e2e_test.py— asserts the ONNX-domainAttentionactually routes to Flash (via VERBOSE dispatch capture), gated on SM ≥ 8.0 (_flash_capable_gpu) andonnxruntime_QUICK_BUILD.tests/_static_cache_support.py; one canonical probe module, no shim.Gating
All new GPU tests skip automatically unless the installed ORT can actually run the path (probe-gated). CI stays green today and flips green automatically once an official ORT release containing microsoft/onnxruntime#28958 is installed — zero code change needed to enable.
Verified
causal_cross_no_past = is_causal && (q_seq != total_seq) && (past == 0)inAttention<T>::ComputeInternal(llm/attention.cc) raisesNOT_IMPLEMENTEDfor theS_q=1decode shape (no fast-path bypass); (2) empirically — a realonnxruntime-gpu==1.27.0isolated venv raisesNotImplemented→supports_static_cache_flash() == False.Explicitly out of scope (intentionally held, separate follow-ups)
examples/static_cache_generation.pynonpad-before-scatter check (verify-example-nonpad).RUF067ruff version-skew inpyproject.toml.References
Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com