microsoft / microsoft/onnxruntime-genai

Eliminate per-step allocations from Engine-hosted MTP chaining

Open
#2,517 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

## Context

Engine-hosted MTP currently creates packed hidden-state storage, per-feedback-stage device tensors, and stage-local request objects in the drafted-token loop. This can reduce speculative-decoding gains, especially at batch size 1.

Follow-up to microsoft/onnxruntime-genai#2495 review comment https://github.com/microsoft/onnxruntime-genai/pull/2495#discussion_r3910718953 and the DFlash2 integration in microsoft/onnxruntime-genai#2496.

## Scope

- Reuse scratch buffers sized for the largest supported batch and draft count for packed hidden states and feedback tensors.
- Remove or pool avoidable stage-local `ScheduledRequests` and related host allocations.
- Preserve transaction rollback, dynamic batching, and request-lifecycle behavior.
- Exercise the CUDA MTP chaining path with correctness tests.

## Performance qualification

Benchmark the same MTP head before and after the change for:

- Draft counts: K=1, 2, 4, 7
- Batch sizes: B=1, 4, 8

Report end-to-end decode latency/throughput and per-step allocation counts. Include ordinary decoding and the current Engine-hosted MTP implementation as baselines. The optimized path should perform no avoidable dynamic allocations in the steady-state token loop and should not regress any tested B/K point.

Contributor guide

No contributing guide indexed for this repository

Research direction

Trace the Engine-hosted MTP drafted-token loop and its packed hidden-state, feedback-tensor, and stage-local request allocations, using the linked MTP integration and review discussion for context. Add reusable scratch storage and pooling while preserving rollback, dynamic batching, and request lifecycles. Validate with CUDA MTP correctness tests and benchmark K=1, 2, 4, 7 across B=1, 4, 8 against ordinary decoding and the current implementation.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
ai, performance
Issue type
Refactor
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.