Current released-runtime behavior
ORT GenAI 0.15.2 already discovers more heterogeneous decoder state than this issue previously credited:
DefaultKeyValueCache discovers sparse/nonconsecutive KV layer indices from graph inputs.
RecurrentState discovers past_key_values.<layer>.conv_state inputs and derives a companion past_key_values.<layer>.recurrent_state name.
The blocker is therefore not a blanket inability to describe mixed KV and recurrent slots.
Precise remaining gaps
- Recurrent-state naming contract: Mobius hybrid exports currently use
past_key_values.<layer>.ssm_state, while ORT GenAI 0.15.2 derives and binds past_key_values.<layer>.recurrent_state from each discovered conv_state input. Nemotron-H/Jamba packages cannot bind that state without renaming one side or making the runtime suffix configurable.
- Beam reorder:
DecoderOnly_State::UpdateInputsOutputs forwards beam indices to the KV cache update, but calls RecurrentState::Update() without beam indices. RecurrentState::Update() only swaps input/output buffers, so recurrent state is not reordered when beams are reordered.
- Nonzero rewind:
RecurrentState::RewindTo(index) rejects every index != 0. Stateful rollback to a retained nonzero sequence position is therefore unsupported.
- Auxiliary graph topology remains separate work: MTP/NextN graphs still need explicit ownership, conditioning, and shared-weight relationships rather than being overloaded into KV metadata.
Tagged source evidence:
Acceptance criteria
- A mixed attention + Mamba/Mamba2 package binds sparse KV, convolution, and recurrent inputs/outputs with exact names, shapes, dtypes, and global layer indices.
- The recurrent suffix is compatible with Mobius
ssm_state exports, either by a documented Mobius rename or a configurable/runtime-supported alias.
- Recurrent state applies the same beam-index reorder semantics as KV state, with a test whose beam permutation changes the expected state rows.
- Recurrent state supports nonzero rewind/rollback, with a test that rewinds to a retained index and then reproduces deterministic decode.
- Prefill, cached decode, replay, reorder, and rollback pass for a tiny mixed-state package without treating recurrent state as KV cache.
- A package with two MTP/NextN heads can represent graph ownership, conditioning outputs, and shared-weight relationships, or fails before durable package output with an actionable error.
- Existing homogeneous decoder packages remain backward compatible.
Mobius continues to fail closed for released-runtime claims until these exact behaviors are proven.
Current released-runtime behavior
ORT GenAI 0.15.2 already discovers more heterogeneous decoder state than this issue previously credited:
DefaultKeyValueCachediscovers sparse/nonconsecutive KV layer indices from graph inputs.RecurrentStatediscoverspast_key_values.<layer>.conv_stateinputs and derives a companionpast_key_values.<layer>.recurrent_statename.The blocker is therefore not a blanket inability to describe mixed KV and recurrent slots.
Precise remaining gaps
past_key_values.<layer>.ssm_state, while ORT GenAI 0.15.2 derives and bindspast_key_values.<layer>.recurrent_statefrom each discoveredconv_stateinput. Nemotron-H/Jamba packages cannot bind that state without renaming one side or making the runtime suffix configurable.DecoderOnly_State::UpdateInputsOutputsforwards beam indices to the KV cache update, but callsRecurrentState::Update()without beam indices.RecurrentState::Update()only swaps input/output buffers, so recurrent state is not reordered when beams are reordered.RecurrentState::RewindTo(index)rejects everyindex != 0. Stateful rollback to a retained nonzero sequence position is therefore unsupported.Tagged source evidence:
recurrent_state.cppat v0.15.2recurrent_state.hat v0.15.2decoder_only.cppat v0.15.2kv_cache.cppat v0.15.2Acceptance criteria
ssm_stateexports, either by a documented Mobius rename or a configurable/runtime-supported alias.Mobius continues to fail closed for released-runtime claims until these exact behaviors are proven.