microsoft / microsoft/onnxruntime-genai

DirectML Memory Accumulation with Repeated OgaGenerator Creation

Open
#1,620 0 comments 0 reactions 0 assignees View on GitHub
ep:DML
Dominant language
C++
Stars
1.1k
Forks
354
Avg merge
2d 16h
Merged PRs (30d)
85

Description

# DirectML Memory Accumulation with Repeated OgaGenerator Creation

## Issue
I have a program using ort-genai in which GPU memory apparently accumulates ~1GB per inference iteration when using DirectML execution provider, despite apparently proper `OgaGenerator` cleanup. Memory grows from 6GB to 14GB over 8 iterations, suggesting KV cache memory is not being released.

I’m looking to understand KV Cache management expectations with the 0.8.3 release of the library. I’m curious if there is any documentation on the topic that I am missing? Perhaps, current known limitations that I am missing? Or, more memory control on the future roadmap? I understand in dealing with pre-1.0 releases many things may change.

## Environment
- **OS**: Windows 11
- **Hardware**: Intel Core Ultra 9 185H, NVIDIA RTX 4060 (8GB VRAM), 32GB RAM
- **ONNX Runtime GenAI**: v0.8.3
- **Execution Provider**: DirectML
- **Model**: Phi-3.5-mini-instruct-onnx (int4-awq) from HuggingFace (https://huggingface.co/microsoft/Phi-3.5-mini-instruct-onnx)
- **Use Case**: MLPerf client integration with repeated benchmark iterations

## Detail:
I’m in progress on a fork of mlperf-client project (v0.6) (https://github.com/mlcommons/mlperf_client). I'm extending it to work on more arbitrary models. To me, this most naturally means using onnxruntime-genai->onnxruntime->dml to get more ubiquitous targeting. I see ort-genai staged but unused within that project likely for future release. So I’m using your stack.

On my fork, I have Phi-3.5-mini-instruct-onnx (int4-awq) running through the onnxruntime-genai stack on CPU+IntelGPU+NVGPU… However, it may not complete due to increases in memory through the frame work.

Through my understanding and LLM based analysis of the onnxruntime-genai repo, I believe KV Cache sticks around for multiple inferences within the life of an OgaModel. This creates unused burden in my usecase. I believe this flow to be typical enough. Here is some extracted/simplified code flow from my fork.:
```cpp
// Setup phase (once per test session)
OgaModel* model = OgaModel::Create(model_path);

// Init phase - params created once and reused
OgaGeneratorParams* params = OgaGeneratorParams::Create(*model);
params->SetSearchOption("max_length", max_length);
params->SetSearchOption("top_k", top_k);
params->SetSearchOption("temperature", temperature);

for (int iteration = 0; iteration < 8; iteration++) {
// Prepare phase - create fresh generator using existing params
OgaGenerator* generator = OgaGenerator::Create(*model, *params);

// Run phase - actual inference
OgaSequences* sequences = OgaSequences::Create();
sequences->Append(input_tokens, token_count);
generator->AppendTokenSequences(*sequences);

// Generate tokens until done
while (!generator->IsDone()) {
generator->GenerateNextToken();
}

// Reset phase - cleanup generator only
delete generator; // Generator destroyed
delete sequences; // Sequences destroyed
// params remain alive for next iteration

// Memory issue: ~1GB increase per iteration
}

// Deinit phase - final cleanup
delete params; // Only deleted at very end
```

Per Issue #538 (https://github.com/microsoft/onnxruntime-genai/issues/538), I understand the creation of the generator object to correlate with the KVCache creation…

How might I free the cache?

1. **Generator Lifecycle**: Does `OgaGenerator` destruction guarantee DirectML GPU KV cache memory release when `OgaGeneratorParams` remains alive across iterations?

2. **KV Cache Scope**: Is KV cache memory tied to the Generator, the Params, or the underlying Model/Session?

### Speculative Solutions

So far in my crawling the ort-genai repo the best speculation I see is would be:

1.
Rewinding. Does rewinding to 0 trigger cleanup? So far as I can tell only onnxruntime-genai internal datastructures are wiped… but I haven’t observed a clear indication of a KV Cache reset/freed down the stack as a consequence. Ex:

```
void Phi35Inference::Reset() {
BaseInference::Reset();

if (oga_bench_generator_) {
try {
oga_bench_generator_->RewindToLength(0); // Clear KV data structures
} catch (const std::exception& e) {
logger_(LogLevel::kWarning, "RewindToLength failed: " + std::string(e.what()));
}
}

oga_bench_generator_.reset();
oga_benchmark_sequences_.reset();

}
```

2.
session.use_env_allocators as a suggestion toward an internal possibility for onnxruntime-genai only... Is this arena based allocation sharing available as a resolution through some Oga* API call?:

```
Memory consumption can be reduced between multiple sessions by configuring shared arena-based
allocation, and you can share allocators between sessions to reduce memory usage.
Implement shared allocator usage:
cpp// Create shared allocator once
OrtEnv* env;
OrtCreateEnv(ORT_LOGGING_LEVEL_WARNING, "genai", &env);
// Register shared allocator with env
OrtCreateAndRegisterAllocator(env, allocator_info, arena_cfg);

// For each model creation
OrtSessionOptions* session_options;
OrtCreateSessionOptions(&session_options);
OrtAddSessionConfigEntry(session_options, "session.use_env_allocators", "1");
```

3.
Could using these session options may help? Are additional DirectML-specific session options needed for proper memory management? Could you all explain why?:
```
"session_options": {
"enable_cpu_mem_arena": false,
"enable_mem_pattern": false,
"execution_mode": "sequential",
```

## Request

Is this expected behavior, or should `OgaGenerator` destruction release DirectML GPU memory? If this is a known limitation, guidance on workarounds would be appreciated.

I’m hoping a full repro isn’t warranted, and I'm just missing something (either doc or expectation/understanding)...

Thanks for your effort on the project. Great work so far! Any guidance is appreciated!

Thanks,
-MichaelC

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by reproducing the repeated OgaGenerator::Create/delete loop with DirectML on ONNX Runtime GenAI v0.8.3, keeping OgaGeneratorParams alive as shown. Read issue #538 and trace the OgaGenerator, OgaGeneratorParams, and model/session lifecycles. Done means documenting whether generator destruction releases the DirectML KV cache and identifying any supported workaround or limitation.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
backend, machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.