Configure MTP heads for paged Engine decoding - #2458
Closed
Tianlei Wu (tianleiwu) wants to merge 3 commits into
Closed
Configure MTP heads for paged Engine decoding#2458Tianlei Wu (tianleiwu) wants to merge 3 commits into
Tianlei Wu (tianleiwu) wants to merge 3 commits into
Conversation
The head inherits every build option from the main model, so exporting a paged model with enable_mtp already built a paged head -- except that its hidden-state values were still declared as dense [batch, sequence, hidden]. That shape contradicts the packed [num_tokens, hidden] rows the rest of a paged graph carries, so the head could not be fed from a paged main model. Declare them with hidden_state_shape(), which is the helper every other hidden state in the builder already uses. A dense export is unchanged (the helper returns the same three dimensions); a paged export now produces an mtp.onnx whose contract matches text.onnx exactly, plus hidden_states in and hidden_states_out out. That lets the paged Engine drive the head with the same varlen decoder and cache machinery as the main model, instead of needing a separate dense KV cache.
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.
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.
This was referenced Aug 23, 2026
Contributor
Author
Contributor
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
model.mtpinto an internal paged decoder configurationStack
PR 5/6. Depends on
tlwu/20260823/gdn-paged-spec-api.Validation
engine_unit_tests322/322 passed