microsoft / microsoft/onnxruntime-genai
Eliminate per-step allocations from Engine-hosted MTP chaining
- 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