[Chi Lo (chilo-ms)]

Description

Adds the K1 KV-cache declaration contract and session integration for CPU/GPU memory planning, following the merged S0 protocol (#33841).

• Strict versioned JSON parsing for cache tensor pairs and a bounded workload/resident-instance envelope. Version 1 remains dense-only; version 2 also supports explicit paging.
• Graph resolution for dense FP32/FP16/BF16 BNSH com.microsoft::GroupQueryAttention pairs and unquantized FP16/BF16 SEPARATE PagedAttention pools with shared int32 block tables. Checked backing bounds include whole pools, unique tables, alias policy, lifetime, and resident multiplicity without speculative prefix-sharing credit.
• Dedicated public APIs: GetDeclaredStateSupport positively negotiates supported declaration schema versions; SessionOptionsSetDeclaredStateContract parses/copies JSON before session creation. Corresponding C++ wrappers are included. APIs are appended without changing shipped slots or bumping ORT_API_VERSION.
• Declaration transport is not subject to the 8192-byte session-config value limit. Invalid replacements preserve the previous declaration; clones and sessions retain immutable snapshots after caller buffers/options are released. Embedded model runtime options do not discard the caller's declaration.
• Session integration: resolve before AOT capability discovery, refresh after pre-partition transforms, and refresh after partitioning/runtime optimization. SessionState retains resolved byte/alias/geometry/lifetime facts and checked backing bounds with session-scoped graph identity and advancing revisions. Failed refreshes invalidate current metadata and fail initialization; unsupported subgraphs or transformed-away cache wiring do not produce stale/partial facts.
• The S0 alias/projection helpers preserve existing pins/policies and reject conflicting or stale graph facts. Projection and session resolution do not rewrite model shapes or placement.
• Nine focused API/session regression cases cover support negotiation, invalid parameters, large buffers, clone/replacement atomicity, discovery ordering, post-optimization refresh, legacy behavior, embedded model options, and unsupported subgraphs.
• Review follow-up rejects quantized dense GQA attributes, one-sided dense cache pairs, mismatched past_present K/V capacity shapes, and shared tables whose common block-ID range cannot cover all resident pools. Six regressions also preserve stateless GQA, disjoint-capacity bounds, declaration-order independence, and valid mixed-lifetime pools.
• Fixes the new setter's non-MSVC build failure by consistently using the existing _In_reads_ annotation for its char buffer in public/internal declarations and implementation; the API signature and byte-length semantics are unchanged.

See the contract/workflow document.

Scope boundary: support negotiation advertises declaration transport, resolution, and retention only, not budget admission or bind-time enforcement. This PR does not adapt C1/C2 estimators, charge declared backing to the existing accountant, pass S0 queries into EP callbacks, select/commit a CPU/GPU cut, allocate/move KV, enforce resident-instance/Run-time envelopes, or wire GenAI/model-package producers. session.declared_state is not registered or consumed; use the dedicated API. follow_consumer remains the only location policy. Paged declarations do not add CPU kernels or prove provider/kernel eligibility. Sessions without a dedicated declaration retain legacy behavior.

Motivation and Context

Related to the K-track in #33372. The planner needs bounded external-state requirements before owner allocation, independent of whether GenAI, an external harness, or ORT owns KV.

This extension makes K1 an explicit session feature rather than a helper-only contract or an arbitrarily accepted configuration key. It supplies stable resolved facts for the remaining C1/C2/M2/R3/I1/I2 integration. GenAI's separate dense per-cache colocation work is microsoft/onnxruntime-genai#2959; placement discovery/colocation alone is not admission.

Complete future handshake: owner declares before session creation -> ORT resolves and commits a budget-feasible placement -> owner queries each cache's required location -> allocates/binds colocated backing -> ORT enforces the binding contract before execution. Only transport, resolution/retention, bounds, and projection are implemented here.

KV-colocation experiment: performance and transfer savings

The reported GPT-OSS-20B v1 mixed CPU/GPU experiment kept GQA layers 12-15 on CPU, with MoE and projections on GPU. For the same mixed-placement cut:

KV capacity (tokens)	GPU-resident KV (ms/token)	Colocated KV (ms/token)
256	14.04	13.20
2,048	20.02	12.42
8,192	32.92	11.75
32,768	84.91	12.49

At 32,768 capacity, colocation gave a 6.8x speedup over the same CPU cut with GPU-resident KV, eliminating 512 MiB of KV transfers per token and moving 256 MiB of KV backing off GPU. GPU-only execution remained faster at 10.39 ms/token: colocation reduces offload's transfer cost; it does not establish a performance advantage over GPU-only execution.

Scope of the result: synchronized capacity measurements with a 99-token prompt and 64 decode steps, not 32K active-context inference. The two KV layouts produced bit-identical logits and unforced generation for the same mixed-placement cut. CPU-versus-GPU compute differed in one top-1 prediction out of 65, so cross-device compute equivalence is not certified. This is separate prototype evidence, not K1/GenAI integration validation or an admission guarantee.

Validation

Native build/runtime validation is delegated to CI, as requested. The previous helper-only head completed its checks; those results are not claimed for this session-integration extension. No local native build/test pass, CUDA execution, placement/admission execution, or real-model K1 validation is claimed.

CI on the initial session-integration head found the undefined _In_reads_bytes_ annotation on non-MSVC compilers. This follow-up replaces it with the portable existing annotation; native results for the new head remain pending CI.

Local checks completed in an isolated environment before committing:

$files = @(git diff --name-only -- '*.cc' '*.h' | ForEach-Object { $_ -replace '/', '\' })
lintrunner -v @files
git --no-pager diff --check

Results: lintrunner passed for all five C++ files changed by this review follow-up; diff whitespace checks passed. The initial session-integration changes separately passed lint for 13 C++ files.

Focused native filter for onnxruntime_test_all:

DeclaredStateParserTest.*:DeclaredStateSessionApiTest.*:DeclaredStateTest.*:CapabilityPlanningTest.*

This selects 40 declared-state/API/session cases plus 10 S0 protocol cases. New-head CI results are pending; a push starts fresh validation without re-running old-head workflows.

[Copilot]

🟡 Changes recommended

Unresolved validation gaps accept unsupported dense GQA cache configurations and incompatible shared block-pool capacities, risking runtime failures or aliasing of independent state.

4 open issues

• Reject quantized attributes for dense GQA caches · New
• Check shared-table coverage across the common block-ID range · New
• Require both dense K/V cache pairs when either is declared · New
• Require matching capacities for dense K/V cache pairs · New

What changed in this PR

Adds the internal K1 KV-cache declaration contract on top of S0 capability planning, preparing ONNX Runtime for bounded CPU/GPU memory planning without enabling session integration or allocation.

Changes:

• Adds versioned parsing, graph resolution, and checked backing bounds for dense and paged caches.
• Extends planning queries with cache capacity shapes, workload bounds, and input/output alias facts.
• Adds validation tests and documents the contract, proposed owner workflow, and integration boundaries.

File	Description
onnxruntime/​test/​framework/​declared_state_test.cc	Tests parsing, resolution, backing bounds, and atomic projection.
onnxruntime/​test/​framework/​capability_planning_test.cc	Tests alias validation and query snapshot binding.
onnxruntime/​core/​framework/​declared_state.h	Defines declaration contracts and helper interfaces.
onnxruntime/​core/​framework/​declared_state.cc	Implements parsing, resolution, backing calculations, and projection.
onnxruntime/​core/​framework/​capability_planning.h	Adds per-port input/output alias constraints.
onnxruntime/​core/​framework/​capability_planning.cc	Validates alias values and duplicate pairs.
docs/​annotated_partitioning/​declared_state_contract.md	Documents K1 schemas, accounting, workflow, and scope.
docs/​annotated_partitioning/​capability_planning_interfaces.md	Documents alias semantics and K1 projection into S0.

_{Review effort: Balanced}

[Chi Lo (chilo-ms)]

Pushed 7849a38 to address the CI failures:

• Corrected the ScopedEnvironmentVariables map initializer that selected the iterator-range constructor and broke compilation across GCC, MSVC, Clang, Android, and WebAssembly.
• Excluded unused K1 declaration/session ownership state from minimal builds, where the APIs remain unsupported. This targets the Android minimal binary's 672-byte budget overrun without increasing the size threshold. Full-build behavior and public API stubs are unchanged.

Native compilation and the resulting binary size still need confirmation from the new-head CI; no local native build was run, as requested.

[Chi Lo (chilo-ms)]

Follow-up in 41afd00: fixed-message argument validation now calls noexcept status creation directly, avoiding temporary strings in minimal-build K1 stubs. Exception translation remains around supported contract parsing/ownership. API errors, public slots, and the Android size threshold are unchanged; CI must confirm the new size.

The Windows release failure at 7849a38 was test_arm64_spectre: its first CMake configure probe timed out after 30 seconds; the other 27 cases passed. No unrelated test change or old-head rerun was made—the new push starts a fresh attempt.

[Chi Lo (chilo-ms)]

Pushed 8847cef for the last measured minimal-size blocker. Removing unsupported ownership state and temporary-string allocation saved 576 section bytes; the required public API slots and nonthrowing NOT_IMPLEMENTED stubs remain. CI now measures 1,611,898 bytes, so the Android threshold is rebaselined from 1,611,802 by exactly 96 bytes, without extra headroom. The latest passing main baseline is 1,611,570 bytes.

The other failures on 41afd00 were infrastructure/setup failures: Windows NuGet installation hit TLS EPROTO, and Detox failed terminating the not-yet-installed iOS simulator app before tests. No unrelated source changes or old-head reruns were made. Fresh-head CI will validate this adjustment.