microsoft / microsoft/onnxruntime-genai

Generator.rewind_to(0) fails for multimodal models (gemma4, phi4mm)

Open
#2,140 0 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.