microsoft / microsoft/onnxruntime-genai
Generator.rewind_to(0) fails for multimodal models (gemma4, phi4mm)
- Dominant language
- C++
- Stars
- 1.1k
- Forks
- 354
- Avg merge
- 2d 16h
- Merged PRs (30d)
- 85
Description
## Description
`Generator.rewind_to(0)` fails with an attention mask / KV cache inconsistency error for multimodal model types (`gemma4`, `phi4mm`). This prevents Generator reuse for batch loglikelihood evaluation, forcing a new Generator per sample.
## Reproduction
```python
import onnxruntime_genai as og
model = og.Model(og.Config("path/to/gemma4-e2b-it-onnx"))
params = og.GeneratorParams(model)
params.set_search_options(max_length=4096, past_present_share_buffer=False)
gen = og.Generator(model, params)
gen.append_tokens([[2, 100, 200, 300, 400]])
logits = gen.get_output("logits") # Works
gen.rewind_to(0) # Appears to succeed
gen.append_tokens([[2, 101, 201]]) # CRASHES
```
**Error:**
```
RuntimeError: inconsistent total_sequence_length (between attn_mask and past_key and past_value)
```
## Root Causes
Three issues:
### 1. `MultiModalPipelineState::RewindTo` is a no-op (inherits base `State::RewindTo` which does nothing)
`MultiModalPipelineState` does not override `RewindTo`, so the decoder's KV cache and attention mask are never reset. Fix: delegate to `decoder_state_->RewindTo(index)`.
### 2. `DecoderState::RewindTo` is also a no-op (same reason)
`DecoderState` (used inside multimodal pipeline) does not override `RewindTo`. Fix: implement it to delegate to `position_inputs_`, `kv_cache_`, and `recurrent_state_`.
### 3. `DefaultKeyValueCache::RewindTo(0)` fails bounds check
Line 411: `shape_[2] <= static_cast(index)` evaluates to `0 <= 0` = true, throwing an error even for `index=0`. This is because `shape_[2]` starts at 0 before any KV cache updates. Fix: skip bounds check when `index == 0`.
## Performance Motivation
Generator reuse via `rewind_to(0)` would eliminate the need to create a new `og.Generator` per sample in lm-eval benchmarks. Current benchmark throughput:
- **With new Generator per call (ortgenai)**: 9.5 req/s
- **Direct ORT (no Generator overhead)**: 142 req/s
- **PyTorch HF**: 185 req/s
## Fix
PR forthcoming with the three fixes above. Tested on Gemma4 E2B-IT with CUDA EP — 5 consecutive rewind_to(0) + append_tokens cycles succeed.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by tracing MultiModalPipelineState::RewindTo and DecoderState::RewindTo, then inspect DefaultKeyValueCache::RewindTo around line 411. Reproduce the Python example with Gemma4 or Phi4MM, and verify five rewind_to(0) plus append_tokens cycles complete without the attention-mask or KV-cache inconsistency error.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- ai, backend
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 58/100