Replay compact fixed state for hybrid decoding

[Tianlei Wu (tianleiwu)]

Summary

Replace dense recurrent checkpoint tensors with compact state-update replay for dynamically batched hybrid models. The engine now owns fixed convolution and GatedDeltaNet state beside paged KV cache, captures packed per-token updates, and replays only an accepted prefix into inactive state banks before atomically publishing both cache types.

Stack

• Stacked on Export compact GatedDeltaNet state updates for Qwen3.5/3.8 #2453, which exports compact Qwen3.5/3.8 value and capsule outputs.
• Depends on [CUDA] Add compact GatedDeltaNet operator onnxruntime#32282 and [CUDA] Add compact VarlenCausalConv state updates onnxruntime#32290.
• MTP scheduling and drafter orchestration remain outside this PR.

Key Changes

• Parse explicit paged_kv and fixed decoder state groups, including compact convolution value and GDN capsule update bindings.
• Add a fixed-state pool with per-request ownership, double-buffered publication, compact prefix replay, and CPU/CUDA provider implementations.
• Add composite paged/fixed reservations with validation and preparation before the allocation-free publication boundary.
• Bind packed variable-length inputs and fixed state together through HybridDecoderIO, including optional packed position_ids.
• Resolve paged cache tensor names from explicit manifest bindings while retaining legacy name-template synthesis for models without state_groups.
• Add synthetic hybrid/composite fixtures and focused manifest, reservation, replay, transaction, and end-to-end Engine coverage.

Memory Impact

The runtime no longer requires dense per-token recurrent checkpoint outputs. Speculative rollback data is bounded by state_update_capacity and represented as compact convolution values or packed GDN transition capsules.

Testing

• test/python/builder/test_decoder_state_groups.py: 24 passed
• CPU Debug engine_unit_tests: 256 passed
• CUDA Debug engine_unit_tests on H200: 261 passed
• Paged and hybrid Engine Python tests across CPU/CUDA: 32 passed
• CUDA compact convolution/GDN partial-prefix replay: passed

Checklist

• Tests added and updated
• CPU and CUDA replay paths validated
• No MTP/drafter scheduling included
• Existing paged-only Engine behavior preserved

[Tianlei Wu (tianleiwu)]

Stack: 1/6 → #2455 → #2456 → #2457 → #2458 → #2459

[Copilot]

Pull request overview

This PR updates ONNX Runtime GenAI’s dynamic-batching Engine to support hybrid decoding with both paged KV cache and fixed recurrent state, replacing dense fixed-state checkpoints with compact per-token state-update capture + partial-prefix replay. It extends the decoder-state manifest schema (explicit bindings + fixed state groups + optional state_update metadata) and wires those resources through cache reservation, decoder IO binding, and device-specific replay (CPU/CUDA), with new synthetic fixtures and tests.

Changes:

• Extend decoder state_groups to include explicit bindings for paged KV and fixed state, plus optional compact state_update capture/replay contract.
• Add fixed-state pool support for dynamic batching, including double-buffered publication and compact prefix replay via DeviceInterface::ReplayStateUpdates (CPU + CUDA implementations).
• Add hybrid decoder IO path (HybridDecoderIO) to bind fixed state (+ optional packed position_ids) alongside packed paged inputs, plus new synthetic composite fixtures and tests.

Reviewed changes

Copilot reviewed 36 out of 38 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	New Python integration test for deterministic packed paged+fixed Engine behavior.
test/python/create/create_synthetic_paged_model.py	Updates synthetic paged config to use legacy templates + explicit paged_kv bindings.
test/python/create/create_synthetic_hybrid_model.py	Updates synthetic hybrid fixture to use unified fixed state groups + compact state_update outputs.
test/python/create/create_synthetic_composite_model.py	New synthetic composite (paged+fixed) model + config used by integration tests.
test/python/builder/test_decoder_state_groups.py	Updates/extends builder tests for bindings and compact state_update metadata.
test/models/engine/synthetic-paged/genai_config.json	Updates checked-in synthetic paged config to include explicit paged_kv bindings.
test/models/engine/synthetic-hybrid/genai_config.json	Updates checked-in synthetic hybrid config for unified fixed groups + state_update metadata.
test/models/engine/synthetic-composite/genai_config.json	New checked-in composite config (paged+fixed + state_update).
test/cpp/engine/paged_key_value_cache_tests.cpp	Adjusts tests to reflect fixed groups now being parsed but rejected by paged-only manager/session.
test/cpp/engine/model_state_manifest_tests.cpp	Updates manifest tests for new bindings schema and dynamic-engine fixed support.
test/cpp/engine/fixed_state_pool_tests.cpp	Extends fixed-state pool tests for compact capture gating and partial-prefix replay (CPU/CUDA).
test/cpp/engine/dynamic_batching_config_tests.cpp	Updates config parsing tests for unified fixed kind + explicit bindings.
src/smartptrs.h	Adds ReplayStateUpdates ABI (version bump) and replay descriptor types.
src/python/py/models/builders/qwen.py	Adds Qwen-specific state_groups emission including compact state_update metadata.
src/python/py/models/builders/base.py	Updates base builder state_groups emission to use explicit bindings and paged_kv helper.
src/models/model_state_manifest.h	Adds helpers for querying manifest state-group kinds (fixed presence).
src/models/model_state_manifest.cpp	Reworks validation to be binding-driven; validates state_update session contract and shapes.
src/ep/cuda/model_kernels.cu	Adds CUDA kernel to replay compact state updates on device.
src/ep/cuda/kernels.h	Exposes LaunchReplayStateUpdates entrypoint.
src/ep/cuda/interface.cpp	Implements DeviceInterface::ReplayStateUpdates using H2D descriptor copy + kernel launch.
src/ep/cpu/interface.cpp	Implements DeviceInterface::ReplayStateUpdates for causal-conv and GDN replay on CPU.
src/engine/paged_key_value_cache.cpp	Resolves paged KV names from manifest bindings (with legacy fallback).
src/engine/fixed_state_pool.h	Extends pool API/bindings for state_update capture metadata and prefix-commit support.
src/engine/fixed_state_pool.cpp	Implements compact capture tensors, prefix commit bookkeeping, and replay via DeviceInterface.
src/engine/execution_context.h	Threads fixed-state reservation/bindings/staging bytes through execution context.
src/engine/engine.cpp	Ensures cache reservation PrepareCommit runs before crossing the transaction boundary.
src/engine/decoders/varlen_decoder_io.h	Adds packed position_ids preparation hook for varlen execution.
src/engine/decoders/varlen_decoder_io.cpp	Implements packed (rank-1) position_ids input construction for varlen batches.
src/engine/decoders/simple_decoder.h	Tracks whether the model has fixed state groups to select IO path.
src/engine/decoders/simple_decoder.cpp	Routes dynamic batching to HybridDecoderIO when fixed groups exist; adjusts CUDA-graph eligibility.
src/engine/decoders/hybrid_decoder_io.h	New DecoderIO wrapper combining varlen IO + fixed-state bindings.
src/engine/decoders/hybrid_decoder_io.cpp	Binds fixed state tensors (and optional state_update tensors) into the ORT run.
src/engine/cache_manager.h	Extends step reservation interface to carry fixed-state handles/bindings + prepare/validate stages.
src/engine/cache_manager.cpp	Adds composite reservation (paged + fixed), fixed-state pool ownership, and joint commit/release flow.
src/config.h	Extends decoder state-group schema: unified fixed kind, explicit bindings, and state_update metadata.
src/config.cpp	Parses new state_groups schema including bindings and state_update blocks (with validation).

---

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