Skip to content

Configure MTP models for paged runtime decoding - #2494

Open
Tianlei Wu (tianleiwu) wants to merge 11 commits into
tlwu/20260828/runtime-spec-apifrom
tlwu/20260828/runtime-mtp-config
Open

Configure MTP models for paged runtime decoding#2494
Tianlei Wu (tianleiwu) wants to merge 11 commits into
tlwu/20260828/runtime-spec-apifrom
tlwu/20260828/runtime-mtp-config

Conversation

@tianleiwu

@tianleiwu Tianlei Wu (tianleiwu) commented Aug 28, 2026

Copy link
Copy Markdown
Contributor

Summary

Project exported model.mtp metadata into the paged decoder configuration used by the dynamic batch engine. Also bind the optional packed hidden_states decoder output that chained MTP drafting consumes.

This is a runtime-only change and does not depend on an ONNX Runtime package bump or include model-builder changes.

Stack

Depends on #2493.

Changes

  • Parse and validate the MTP model configuration for paged decoding.
  • Derive the internal MTP decoder configuration from the target decoder settings.
  • Bind the optional packed hidden_states output in VarlenDecoderIO.
  • Add focused configuration coverage for valid and invalid MTP projections.

Testing

  • cmake --build build/Linux/Debug
  • MTP configuration and speculative unit tests: 52 passed.
  • Packed decoder and speculative engine tests: 17 passed.
  • Full engine_unit_tests: 286 passed.

This was referenced Aug 28, 2026
Closed
Open
@tianleiwu
Tianlei Wu (tianleiwu) force-pushed the tlwu/20260828/runtime-mtp-config branch from 89e252b to 7bf9f3b Compare August 29, 2026 00:02
@tianleiwu
Tianlei Wu (tianleiwu) force-pushed the tlwu/20260828/runtime-mtp-config branch from 7bf9f3b to 401cce5 Compare August 29, 2026 08:51
@tianleiwu
Tianlei Wu (tianleiwu) force-pushed the tlwu/20260828/runtime-mtp-config branch from 401cce5 to 81ef6a1 Compare August 29, 2026 17:03
@tianleiwu
Tianlei Wu (tianleiwu) marked this pull request as ready for review August 29, 2026 17:04
Copilot AI lite review requested due to automatic review settings August 29, 2026 17:04
@tianleiwu
Tianlei Wu (tianleiwu) requested a review from a team as a code owner August 29, 2026 17:04
Copilot AI reviewed Aug 29, 2026
View reviewed changes

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR wires Multi-Token Prediction (MTP) model metadata into the paged/runtime decoding path so the dynamic-batch engine can run an internal “MTP head” session, and it optionally binds a packed hidden_states decoder output that the MTP drafting loop consumes.

Changes:

  • Add CreateMtpDecoderConfig to project model.mtp metadata into an internal decoder config suitable for paged decoding.
  • Extend VarlenDecoderIO / CUDA-graph buffer holders to optionally bind an exported hidden_states output ([packed tokens, hidden_size]).
  • Add a targeted unit test validating that the projection preserves the paged-attention contract while stripping main-model fixed/sliding state.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.

Show a summary per file
File Description
src/config.h Declares CreateMtpDecoderConfig for projecting model.mtp into a decoder-config view.
src/config.cpp Implements CreateMtpDecoderConfig and resets main-model-only decoder state for the MTP projection.
src/engine/decoders/varlen_decoder_io.h Documents and exposes optional packed hidden-states output access on VarlenDecoderIO.
src/engine/decoders/varlen_decoder_io.cpp Allocates/binds optional hidden_states output tensors (including CUDA-graph persistent buffers).
test/cpp/mtp_config_test.cpp Adds a projection-focused unit test for paged MTP decoder configuration.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment thread src/config.cpp Outdated
Comment thread src/engine/decoders/varlen_decoder_io.h
Comment thread test/cpp/mtp_config_test.cpp Outdated
@tianleiwu
Tianlei Wu (tianleiwu) force-pushed the tlwu/20260828/runtime-mtp-config branch from 81ef6a1 to fc08ffe Compare August 29, 2026 22:02
@tianleiwu
Tianlei Wu (tianleiwu) force-pushed the tlwu/20260828/runtime-mtp-config branch from fc08ffe to 13d5696 Compare August 30, 2026 06:41
@tianleiwu
Plumb speculative checkpoint capture from the step plan to the model IO
960e3d3 
RequestStepPlan carries a draft count, the composite planner turns any nonzero
count into a checkpoint-capturing fixed reservation (and sizes its staging for
it), and the hybrid decoder binds the checkpoints output only on such a step, so
the operators skip writing it entirely otherwise. Nothing sets a draft count yet,
so every step still plans and reserves exactly what it did before.

(cherry picked from commit d4abc79)
@tianleiwu
Verify speculative drafts inside a dynamic engine step
ef5bd51 
(cherry picked from commit 0ff9d6d)
@tianleiwu
fix(engine): address speculative verification review feedback
2034abd
@tianleiwu
fix(engine): preserve budgets after row reorder
d6a6985
@tianleiwu
Expose speculative draft proposals through the C, C++ and Python APIs
3849604
@tianleiwu
Fix draft block-table planning and the sampled-token boundary on a ve…
bcd3041 
…rify step
@tianleiwu
Verify drafts with a batched on-device argmax instead of a host copy …
f8cd282 
…per row
@tianleiwu
fix(engine): harden speculative request bindings
7d69e0c
@tianleiwu
Bind the paged decoder's hidden-states output when the model exposes one
d018552 
A model exported with include_hidden_states emits a [num_tokens, hidden_size]
row per packed token alongside its logits. The Engine never bound it, so the
paged path had no way to feed an MTP draft head, which consumes exactly that
row for the last accepted token.

Bind it whenever the config names it and the session actually has it, so every
model exported without it runs a byte-identical step. The rows are sized by the
packed token count rather than by the logits row count, because a pruned LM head
emits only one logits row per request while the hidden states stay per token;
PrepareLogits and PrepareHiddenStates now share a TokenCount helper so the two
cannot drift apart. VarlenGraphBuffers gains a matching static buffer so a
captured step keeps a stable output address.

315 engine unit tests pass, and the tiny paged fixture (now exported with
include_hidden_states) still reports exact ragged-batch parity.
@tianleiwu
Project model.mtp into an internal paged decoder config
87ef4d2 
The paged Engine should host mtp.onnx directly rather than requiring callers to
rewrite and load a second genai_config.json. CreateMtpDecoderConfig copies the
main config so device, batching and packed-attention input names remain aligned,
then replaces the filename/session options, dimensions, I/O names and shared
initializers from model.mtp.

The head is one full-attention layer, so the projection explicitly drops the
main hybrid model's fixed state groups, convolution metadata and sliding-window
state. A focused test verifies both the inherited paged contract and the removed
main-model state; all six MtpConfigTest cases pass.
@tianleiwu
fix(engine): validate MTP decoder projection
413ad82
@tianleiwu
Tianlei Wu (tianleiwu) force-pushed the tlwu/20260828/runtime-mtp-config branch from 13d5696 to 413ad82 Compare August 30, 2026 20:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

Copilot code review Copilot
Copilot review effort
Applies to this pull request for everyone.Learn more about Copilot code review.
Copilot left review comments

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@tianleiwu