Fix compact state configuration and staging

[Tianlei Wu (tianleiwu)]

Description

This follow-up addresses the two deferred review comments from #2454. It reapplies the decoder-I/O-owned state binding schema from #2463 and preallocates compact state-update staging so speculative fixed-state replay does not allocate device buffers after scheduling.

Summary of Changes

State Manifest Configuration

• Move reusable state binding templates to decoder inputs and outputs, while state groups retain layer membership and per-group update geometry.
• Update fixed convolution, fixed recurrent, and paged KV consumers to use the revised manifest schema.
• Align model-builder output, synthetic model configurations, and manifest/configuration tests with the schema from Move state bindings to decoder IO config #2463.

Fixed-State Staging

• Allocate capacity-sized compact value, capsule, capture-count, and active buffers when FixedStatePool is constructed.
• Create exact-batch tensor views over the preallocated buffers in Reserve() and retain their backing storage for the reservation lifetime.
• Include compact-update buffers in persistent memory accounting and verify address reuse across different reservation batch sizes.

Testing

• lintrunner -a
• cmake --build build/cu130/Release --target engine_unit_tests -j $(nproc)
• build/cu130/Release/engine_unit_tests --gtest_color=no --gtest_filter='FixedStatePoolTest.*:CudaFixedStatePoolTest.*:DecoderStateGroupsConfigTest.*:ModelStateManifestTest.*:PagedKeyValueCacheTest.*' (65 passed)
• python3 -m pytest test/python/builder/test_decoder_state_groups.py -q --tb=short (25 passed)
• git diff --check origin/main...HEAD

Checklist

• Tests added/updated
• No breaking changes
• Documentation not required; this aligns the implementation with the already-reviewed Move state bindings to decoder IO config #2463 schema

[Copilot]

Pull request overview

This PR aligns the engine/model-builder “decoder state manifest” schema with the decoder-IO-owned binding-template design (state group entries are now primarily kind + layer membership + per-group state_update geometry), and updates fixed-state staging so compact state replay/capture can reuse preallocated buffers rather than allocating after scheduling.

Changes:

• Updates the state group schema: moves binding templates to model.decoder.inputs / model.decoder.outputs, splits fixed state kinds into fixed_conv and fixed_recurrent, and simplifies state_update fields to capacity/key-head geometry.
• Updates Python builders, synthetic model generators/configs, and C++/Python tests to match the revised schema.
• Preallocates compact state-update staging in FixedStatePool and uses tensor views in Reserve(), updating persistent memory accounting and address-reuse tests.

Reviewed changes

Copilot reviewed 20 out of 20 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
test/python/test_onnxruntime_genai_hybrid_engine.py	Updates expected fixed state group kinds in the hybrid fixture config assertions.
test/python/create/create_synthetic_paged_model.py	Updates synthetic paged config to use decoder IO templates and removes per-group KV bindings.
test/python/create/create_synthetic_hybrid_model.py	Updates synthetic hybrid config to use decoder IO templates; renames fixed group kinds; adjusts state_update schema usage.
test/python/create/create_synthetic_composite_model.py	Updates composite synthetic config and decoder IO templates; removes per-group bindings; updates state_update control placement.
test/python/builder/test_decoder_state_groups.py	Updates builder expectations for state group kinds and removes assertions for deprecated per-group bindings.
test/models/engine/synthetic-paged/genai_config.json	Updates stored synthetic paged model config to decoder-IO template schema.
test/models/engine/synthetic-hybrid/genai_config.json	Updates stored synthetic hybrid model config to decoder-IO template schema and fixed kind split.
test/models/engine/synthetic-composite/genai_config.json	Updates stored synthetic composite config to decoder-IO template schema and fixed kind split.
test/cpp/engine/paged_key_value_cache_tests.cpp	Updates manifest tests to use FixedConv kind and removes state binding construction that no longer exists.
test/cpp/engine/model_state_manifest_tests.cpp	Updates manifest construction/validation tests to use decoder IO templates and fixed kind split.
test/cpp/engine/fixed_state_pool_tests.cpp	Extends tests to cover reuse of preallocated compact state-update staging across reservation batch sizes and updates byte accounting expectations.
test/cpp/engine/dynamic_batching_config_tests.cpp	Updates JSON parsing/overlay tests to the new schema and fixed kind split.
src/python/py/models/builders/qwen.py	Updates Qwen builder emitted state group kinds and state_update payload shape to the new schema.
src/python/py/models/builders/base.py	Updates common builder to emit state groups without per-group bindings and adjusts paged-kv helper signature.
src/models/model_state_manifest.h	Stores decoder inputs/outputs in the manifest for binding resolution at validation time.
src/models/model_state_manifest.cpp	Resolves bindings from decoder IO templates, splits fixed kinds, derives state_update kind from group kind, and updates config/session validation accordingly.
src/engine/paged_key_value_cache.cpp	Resolves paged KV tensor names from decoder IO templates instead of per-group bindings.
src/engine/fixed_state_pool.cpp	Uses decoder IO templates for fixed state bindings; preallocates state_update staging and creates per-reservation tensor views over it.
src/config.h	Updates config schema types: removes per-group StateBinding, replaces Fixed with FixedConv/FixedRecurrent, simplifies StateUpdate payload.
src/config.cpp	Updates JSON parsing for new kinds/state_update fields and removes parsing for deprecated binding/state_update fields.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Baiju Meswani (baijumeswani)]

Looks good.