refactor(engine): unify native + ORT session-state policy behind SessionStore

[justinchuby]

What & why

The engine kept two parallel session implementations selected by
if self.decode_backend == EngineDecodeBackend::Native, one over
self.native_sessions and one over self.sessions. The policy around a
session — the "session {id} not found" lookup, the rewind bound-checks and their
exact error strings, the checkpoint arithmetic, and truncate/clear — was
backend-independent but duplicated verbatim across the two arms (the
checked_sub bound check and its message in rewind_session_by, the
position > current check in rewind_session_to, etc.). This is the same class
of drift risk that #1170 closed for the KV length layer with KvPrefixStore.

This PR extracts that policy into one place and routes both backends through it,
following the KvPrefixStore precedent rather than inventing a new pattern.

The seam

New module crates/onnx-genai-engine/src/engine/session_state.rs:

• SessionStore trait — mirrors KvPrefixStore. Backends implement
  logical_len (via SessionLen), validate_rewind, rewind, reset,
  close, with identical method names/signatures on both sides, exactly as
  the NativeKvStore/OrtKvStore adapters already do.
• Free policy functions require_session, checkpoint, token_count,
  rewind_by, rewind_to, reset, close own everything that is
  backend-independent: the lookup + not-found text, both rewind bound-check
  strings, SessionPosition::new(len) checkpoint arithmetic, and the dispatch
  into the single backend truncation method.
• Adapters in runtime.rs: OrtSessions/OrtSessionsRef (always compiled)
  and NativeSessions/NativeSessionsRef (#[cfg(feature = "native-backend")])
  implement SessionStore. The six public methods (checkpoint_session,
  rewind_session_by, rewind_session_to, reset_session, close_session,
  session_token_count) now select an adapter and call session_state::*. The
  require_ort_backend(...) guard stays in the ORT arm.

Error text is now authored once. Rendered strings are byte-identical to the
originals (only internal variable names differ, e.g. {id} vs {session_id}),
so behaviour is unchanged. No test asserted these strings before; the ones that
assert backend rewind errors (failed_rewind_of_*) are untouched.

Asymmetries I deliberately kept (and why)

Forcing these into false symmetry would be worse than the small duplication that
remains, so — mirroring how KvPrefixStore's reset is legitimately different
per backend — they stay backend-specific behind the trait:

• validate_rewind: ORT validates the target and draft/speculative state
  and paged-KV page ownership (validate_target_state_rewind_to_len) so a
  rejected rewind leaves the session untouched; the three failed_rewind_of_*
  tests pin this. Native's validate_rewind is a deliberate Ok(()) — it always
  admits and clamps against its in-process decoder length. These are genuinely
  different objects: ORT owns paged KV + a draft runner; native owns one
  persistent in-process decoder.
• rewind/reset/close mechanism: ORT does
  validate → scheduler.complete → remove → rewind_target_state_to_len →
  reinsert; native truncates its token vec, rewinds/resets the single persistent
  decoder (gated on native_active_session) and touches LRU last_access.
• create_session is left backend-specific and is not routed through the
  seam: it is object construction, not policy, and the two session value types
  (EngineSession vs NativeSessionState) are genuinely different. Unifying it
  would be a forced abstraction with no shared policy to host.
• Token truncate for ORT is intentionally not lifted into the shared
  policy. It lives inside rewind_target_state_to_len (kv_bridge.rs), which is
  shared with the speculative-decode hot path (speculative/mod.rs). Lifting
  it would fracture that path for no gain. The shared policy therefore owns the
  logical-length bound check (truly common) and delegates the truncation.

DRY guard (mandatory, per #1170)

Two-part, mirroring ReusedPrefix + KV_REWIND_CALLERS:

1. Compile-time CheckedPosition newtype with a private field. It is only
   ever constructed by session_state::rewind_to after the bound check, and it
   is the required argument to SessionStore::rewind/validate_rewind. A future
   backend therefore cannot rewind KV without going through the one bound check
   — there is no other way to obtain the token.
2. Test-time tripwire the_rewind_bound_check_lives_only_in_the_shared_policy
   scans the crate sources for the "cannot rewind session" fragment and fails
   if it appears anywhere except engine/session_state.rs. A third backend that
   open-codes a copy turns the suite red. (fork_session uses the distinct
   "cannot fork session", so there is no collision.)

Performance

These are per-turn / per-request session ops, not the per-token decode hot
path (generate_incremental_with_callback). The trait dispatch happens once per
rewind/checkpoint/reset, off the hot path, so the dynamic call cannot show up in
decode throughput. No hot-path code changed; no per-token dispatch was
introduced. (Restructured, not measured — because there is nothing on the hot
path to measure.)

Falsification (the important part)

I broke the shared policy deliberately — inverted the one bound check in
session_state::rewind_to (requested > current → requested < current) — and
confirmed both backends go red from that single line:

• ORT (default backend): failed_rewind_of_windowed_evicted_position…,
  failed_rewind_of_ort_owned_kv…, failed_rewind_of_runner_backed_state… — all
  3 FAILED (they now hit the bound check before their backend validation).
• Native (--features native-backend):
  native_session_rewind_by_truncates_logical_length FAILED with
  cannot rewind session 1 to token 3; current length is 5.

Reverting the one line restored both to green. A break in the shared path turns
both backends red, which is the evidence that the extraction actually unified
them rather than leaving two look-alikes. I added
native_session_rewind_by_truncates_logical_length because there was previously
no integration test exercising a valid explicit rewind on the native
backend (the existing native "rewind" test drives prefix-reuse, a different
path) — this closes that gap and gives the native half of the falsification.

Verification

• cargo test -p onnx-genai-engine --lib → 414 passed, 0 failed, 1 ignored.
• cargo test -p onnx-genai-engine --features native-backend → all targets pass
  except the pre-existing pipeline_load_rejects_kv_page_pool_before_fixed_state_when_host_budget_cannot_fit_either
  (a host-RAM KV-budget test — confirmed failing identically on a stashed
  baseline; it does not touch session code).
• cargo clippy -p onnx-genai-engine --all-targets -- -D warnings and the same
  with --features native-backend → clean.
• CUDA (RTX 4060, driver 591.55): cargo test -p onnx-genai-engine --features native-backend,cuda --test native_engine -- --test-threads=1 ran on
  the GPU (not ignored): 16 passed, 1 failed, 0 ignored. Passing includes the
  new native_session_rewind_by_truncates_logical_length,
  native_sub4_cuda_fallback_generates_coherent_decode, and
  qwen15b_native_decode_locks_accurate_near_tie_token (real CUDA decode). The 1
  failure, engine_native_cuda_matches_cpu_tokens ("full CUDA mask reservation
  size overflow" during native decoder load), is a pre-existing CUDA
  mask-workspace sizing issue on this 8 GB card — confirmed failing identically
  on a stashed baseline, unrelated to session state.

Behaviour note

close_session/reset_session now run the existence check (via
require_session) before scheduler.complete. Previously ORT close called
scheduler.complete before its remove-based existence check. This is not
observable through any error (scheduler.complete on a missing id is benign) and
reset already checked existence first.

Do not merge — please review.

[justinchuby]

Reviewed and merged. This closes the last structural piece of the owner's standing directive that native and ORT decode must share one set of state-management logic rather than being treated separately.

Why this was worth doing. Before this, runtime.rs carried two parallel session implementations selected by ~15 if decode_backend == Native branches, operating over two different maps (native_sessions vs sessions). The duplication was not merely stylistic: in rewind_session_by the checked_sub bound check and its error string were reproduced verbatim in both arms. Two copies of a policy is two places for it to drift, and a third backend would have made a third.

The asymmetries kept are the right ones. I specifically did not want a forced symmetry here. ORT must validate and rewind draft/speculative state and paged-KV page ownership; native rewinds a persistent in-process decoder. Keeping validate_rewind as a real implementation on one side and a documented Ok(()) no-op on the other is honest, and keeping ORT's token truncation inside rewind_target_state_to_len (shared with the speculative-decode path) rather than dragging it into the shared policy is the correct seam: the shared code owns the bound check, which is genuinely common, not the truncation, which is not. This mirrors the KvPrefixStore precedent where reset legitimately differs per backend.

Falsification — the evidence I actually cared about. Inverting the single shared bound check turns both backends red (3 ORT failed_rewind_of_* tests plus the new native native_session_rewind_by_truncates_logical_length). If breaking the shared path had reddened only one backend, the extraction would have unified nothing and I would have rejected it.

I also falsified the DRY guard itself, because a guard that cannot fail is worse than no guard — we have shipped one of those before. Appending a function to runtime.rs that open-codes the policy string makes the_rewind_bound_check_lives_only_in_the_shared_policy fail at session_state.rs:394 with the offender listed; removing it returns 9 passed / 0 failed. Note the guard also self-checks that its source scan found more than 10 files, so it cannot pass by silently scanning nothing — that is the exact failure mode we hit previously.

The CheckedPosition newtype is the better half of the guard: with a private field constructible only inside session_state, a backend cannot reach rewind without having passed the one bound check. That is a compile-time property rather than a convention.

Verification: cargo test --lib 414 passed / 0 failed; --features native-backend green except a pre-existing host-RAM KV-budget failure (confirmed identical on a stashed baseline); clippy -D warnings clean on both feature sets; CUDA on the RTX 4060 16 passed / 1 failed, that failure being the pre-existing mask reservation size overflow in native_decode/cuda.rs — independently reproduced by a second agent today and now under separate investigation. Per-turn operations only, so no per-token dispatch was introduced on the decode hot path.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 80.52%. Comparing base (dcf9c0f) to head (cb837b6).
⚠️ Report is 15 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1255      +/-   ##
==========================================
+ Coverage   79.85%   80.52%   +0.67%     
==========================================
  Files         358      360       +2     
  Lines      155466   158790    +3324     
  Branches   155466   158790    +3324     
==========================================
+ Hits       124144   127872    +3728     
+ Misses      26720    26291     -429     
- Partials     4602     4627      +25     

Flag	Coverage Δ
mlas	85.22% <ø> (?)
offline	80.43% <ø> (+0.58%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.
see 9 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

🔴 Benchmark Regression Detected

Comparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).

ℹ️ Absolute times are informational only — they vary with runner load. The % change column is the reliable signal because both sides ran under identical conditions.

Status	Scenario	Base	PR	Change
🔴	gather/medium_f32_threads=1-internal/32768	4.00 µs	7.71 µs	+92.5%
🔴	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	409.27 µs	596.70 µs	+45.8%
🔴	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	70.55 µs	102.19 µs	+44.8%
🔴	matmul/medium_generic_bf16_threads=8/32x512x512	424.63 µs	565.78 µs	+33.2%
🔴	matmul/small_generic_f16_threads=8/1x256x256	35.40 µs	46.34 µs	+30.9%
⚠️	matmul/small_generic_f32_threads=8/1x256x256	55.03 µs	64.53 µs	+17.3%
⚠️	matmul/small_generic_f16_threads=1/1x256x256	38.29 µs	44.33 µs	+15.8%
⚠️	tokenization/encode_tokens_per_second	348.63 µs	401.15 µs	+15.1%
✅	sampling_latency/greedy_per_token	2.98 µs	3.40 µs	+14.2%
✅	tokenization/decode_tokens_per_second	5.88 ms	6.63 ms	+12.7%
✅	sampling_latency/top_p_per_token	361.36 µs	402.52 µs	+11.4%
✅	gather/medium_f16_threads=1-internal/32768	2.45 µs	2.70 µs	+10.0%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	2.00 ms	2.19 ms	+9.6%
✅	sampling_latency/min_p_per_token	195.72 µs	214.45 µs	+9.6%
✅	gather/small_bf16_threads=1-internal/4096	534.4 ns	585.4 ns	+9.5%
✅	matmul/large_generic_f32_threads=8/32x1024x1024	4.82 ms	5.22 ms	+8.4%
✅	sampling_latency/top_k_per_token	52.42 µs	56.48 µs	+7.7%
✅	add/small_bf16_threads=1-internal/1024	481.0 ns	512.8 ns	+6.6%
✅	grammar_masking/llguidance_compute_mask/32	78.43 µs	82.40 µs	+5.1%
✅	kv_cache/alloc_dealloc_pages	39.66 µs	41.06 µs	+3.5%
✅	matmul/small_generic_bf16_threads=1/1x256x256	33.29 µs	33.82 µs	+1.6%
✅	matmul/large_generic_f16_threads=1/32x1024x1024	74.72 µs	75.64 µs	+1.2%
✅	gather/medium_bf16_threads=1-internal/32768	2.44 µs	2.47 µs	+1.2%
✅	logit_processing/seven_processor_chain_per_step	334.07 µs	334.46 µs	+0.1%
✅	matmul/large_generic_f32_threads=1/32x1024x1024	9.72 ms	9.67 ms	-0.5%
✅	gather/small_f32_threads=1-internal/4096	747.9 ns	726.0 ns	-2.9%
✅	add/medium_f16_threads=1-internal/262144	116.13 µs	112.34 µs	-3.3%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	586.65 µs	563.90 µs	-3.9%
✅	qwen3_sampling_processors/top_k_partial_selection	148.73 µs	142.82 µs	-4.0%
✅	qwen3_sampling_processors/top_k_top_p_fast	717.93 µs	689.20 µs	-4.0%
✅	add/medium_bf16_threads=1-internal/262144	116.75 µs	111.95 µs	-4.1%
✅	add/medium_f32_threads=1-internal/262144	27.27 µs	26.09 µs	-4.3%
✅	gather/large_bf16_threads=1-internal/131072	17.53 µs	16.73 µs	-4.6%
✅	gather/small_f16_threads=1-internal/4096	548.8 ns	518.1 ns	-5.6%
✅	gather/large_f16_threads=1-internal/131072	16.93 µs	15.91 µs	-6.1%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	564.39 µs	528.12 µs	-6.4%
✅	matmul/small_generic_bf16_threads=8/1x256x256	37.69 µs	35.16 µs	-6.7%
✅	add/small_f16_threads=1-internal/1024	511.2 ns	475.5 ns	-7.0%
✅	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	6.57 ms	6.10 ms	-7.3%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	2.09 ms	1.94 ms	-7.3%
✅	matmul/small_generic_f32_threads=1/1x256x256	44.18 µs	40.21 µs	-9.0%
✅	add/large_f32_threads=1-internal/4194304	873.61 µs	793.70 µs	-9.1%
✅	matmul/medium_generic_f32_threads=1/32x512x512	2.56 ms	2.27 ms	-11.3%
✅	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	77.21 µs	66.86 µs	-13.4%
✅	reduce_mean/large_f32_threads=1-internal/262144	1.22 ms	1.06 ms	-13.8%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	4.13 ms	3.56 ms	-13.8%
🟢	reduce_mean/medium_f32_threads=1-internal/65536	326.23 µs	269.29 µs	-17.5%
🟢	add/small_f32_threads=1-internal/1024	268.5 ns	220.8 ns	-17.7%
🟢	matmul/large_generic_f16_threads=8/32x1024x1024	99.52 µs	81.24 µs	-18.4%
🟢	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	698.87 µs	566.53 µs	-18.9%
🟢	reduce_mean/small_f32_threads=1-internal/4096	20.41 µs	16.28 µs	-20.2%
🟢	matmul/medium_generic_f16_threads=8/32x512x512	40.27 µs	31.51 µs	-21.8%
🟢	add/large_f16_threads=1-internal/4194304	2.26 ms	1.72 ms	-23.8%
🟢	add/large_bf16_threads=1-internal/4194304	2.20 ms	1.64 ms	-25.4%
🟢	matmul/medium_generic_f32_threads=8/32x512x512	1.80 ms	1.26 ms	-29.8%
🟢	matmul/medium_generic_f16_threads=1/32x512x512	48.70 µs	32.96 µs	-32.3%
🟢	matmul/large_generic_bf16_threads=8/32x1024x1024	2.19 ms	1.34 ms	-38.6%
🟢	gather/large_f32_threads=1-internal/131072	58.69 µs	35.98 µs	-38.7%
🟢	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	82.10 µs	48.85 µs	-40.5%

Visual flags: ⚠️ ≥ 15% slower, 🔴 ≥ 30% slower — calibrated against measured runner noise (~27% worst-case on multi-threaded matmul)

Host info

CPU: Apple M1 (Virtual)
Cores: 3
OS: Darwin 25.5.0 arm64
Rust: rustc 1.97.1 (8bab26f4f 2026-07-14)
Load avg: { 3.71 3.58 5.94 }

What this cannot catch

• Regressions in code paths not covered by these benchmarks (e.g., end-to-end decode with a real model)
• Sub-threshold regressions that compound over multiple PRs
• Performance changes that only manifest under GPU execution
• Latency changes in the ORT integration path (these benchmarks exercise the native Rust kernels)