Skip to content

Close remaining heterogeneous recurrent-state runtime gaps #605

Description

@justinchuby

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

  1. 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.
  2. 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.
  3. Nonzero rewind: RecurrentState::RewindTo(index) rejects every index != 0. Stateful rollback to a retained nonzero sequence position is therefore unsupported.
  4. 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

  1. A mixed attention + Mamba/Mamba2 package binds sparse KV, convolution, and recurrent inputs/outputs with exact names, shapes, dtypes, and global layer indices.
  2. The recurrent suffix is compatible with Mobius ssm_state exports, either by a documented Mobius rename or a configurable/runtime-supported alias.
  3. Recurrent state applies the same beam-index reorder semantics as KV state, with a test whose beam permutation changes the expected state rows.
  4. Recurrent state supports nonzero rewind/rollback, with a test that rewinds to a retained index and then reproduces deterministic decode.
  5. Prefill, cached decode, replay, reorder, and rollback pass for a tiny mixed-state package without treating recurrent state as KV cache.
  6. 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.
  7. Existing homogeneous decoder packages remain backward compatible.

Mobius continues to fail closed for released-runtime claims until these exact behaviors are proven.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions