Expose speculative draft proposals through Engine APIs - #2457
Closed
Tianlei Wu (tianleiwu) wants to merge 35 commits into
Closed
Expose speculative draft proposals through Engine APIs#2457Tianlei Wu (tianleiwu) wants to merge 35 commits into
Tianlei Wu (tianleiwu) wants to merge 35 commits into
Conversation
ORT captures a CUDA graph by re-running the model inside a single user-visible Run() until the EP reports capture complete (InferenceSession::RunImpl recursion, bounded by kMaxGraphCaptureRunAttempts and driven by the CUDA EP's min_num_runs_before_cuda_graph_capture_ = 2, so three executions). Those internal runs re-feed identical inputs, so idempotent in-place writes such as the KV-cache append are unaffected, but the recurrent state is an accumulator and was advanced three times on the first Run of each captured shape. That is the real mechanism behind the bias this code previously attributed to a per-step effect on replay; the comment is corrected accordingly. Undo it instead of avoiding it: let the capture happen on a throwaway Run, restore the state and replay once. The backup lives in the CPU mirror of the wrapped live buffers, so it costs no device memory, and it runs once per captured graph id. Verified bit-exact against the double-buffered path on Qwen3.8-27B MTP: identical generated_token_sha256 across runs, and 24/24 byte-identical completions on a paired GPQA sample. Decode is 3-5% faster in 5/5 repeats because sharing collapses the two captured graph variants into one and drops the per-step buffer swap. Peak device memory is unchanged -- the avoided allocation becomes arena headroom rather than a lower high-water mark. Off by default; enable with ORTGENAI_SHARE_RECURRENT_STATE_UNDER_GRAPH_CAPTURE=1.
The cherry-picked commits carried `self._recurrent_state_window`, which this branch renamed to `self._state_window` (extra option `recurrent_state_window` -> `state_window`). The patches applied cleanly because the surrounding context was unchanged, so nothing flagged it; the two remaining references would have raised AttributeError for every Qwen3.5/3.8 build.
The dense Qwen3.8-27B-NVFP4 checkpoint ships in the compressed-tensors format,
which the ModelOpt loader rejected outright. That format names the NVFP4 payload
`weight_packed`, stores the reciprocal of the global scale as
`weight_global_scale`, and quantizes FP8 weights per channel instead of per
tensor. Separately, `Qwen3_5ForConditionalGeneration` routed to
`Qwen35TextModel`, which carries no MTP wiring, so `enable_mtp` was unavailable
for dense models.
Loader (quantized_model.py, builders/base.py):
* dispatch `compressed-tensors` to `ModeloptModel`, and let `make_matmul_op`
take the native-quant path for it so pre-quantized weights are preserved
rather than dequantized
* fall back to `weight_packed`, and derive `weight_scale_2` from
`reciprocal(weight_global_scale)`
* accept scalar or per-channel FP8 weight scales via
`_validate_positive_weight_scale`; the emitter already handled [N, 1]
* build a dense MLP when a layer has no `mlp.gate.weight`
Builder (builders/qwen.py, builders/qwen_mtp.py, builder.py, __init__.py):
* extract the MTP wiring out of `Qwen35MoeTextModel` into `_init_mtp`,
`_make_mtp_head`, `_save_mtp_head` and a `_mtp_head_class()` hook, so it can
be reused without inheriting the MoE layer builder
* add `Qwen35DenseTextModel` and `Qwen35DenseMtpHead`, and route
`Qwen3_5ForConditionalGeneration` to the former
* add `_make_mtp_decoder_layer()` and `_mtp_mlp_modules()` hooks so the dense
head builds `Qwen3_5DecoderLayer` with a dense MLP
The old `Qwen35NativeQuantTextModel` is deliberately not reintroduced: native
quantization is now dispatched generically in `make_matmul_op` from the tensors
`ModeloptModel` already carries, so a dedicated checkpoint-reading class is no
longer needed.
Also fixes MTP weight loading against current safetensors, whose `safe_open` is
no longer iterable -- this broke the existing Qwen3.6 MoE head too -- and lets
the example drive both sessions from one packaged directory through
`Config.overlay`.
Verified on H200 against Qwen3.8-27B-NVFP4. The 64-layer export reproduces the
reference graph exactly (233 FP8 + 168 FP4 + 96 FP16 MatMul, 16 GQA, 3,815 MB of
shared initializers), with 48 GatedDeltaNet in place of 48 LinearAttention. MTP
self-speculative decoding is token-identical to greedy over 3 prompts x 80
tokens and runs 1.69x faster (42.3 vs 25.0 tok/s at a ~70% accept rate).
Building a main model with kv_cache_quant_type=int8_per_channel plus enable_mtp failed outright: the head inherits kv_cache_scale_file and then rejects it, because calibrated scales are per main-model layer (16 full-attention layers here) while the head is a single layer: ValueError: kv_cache_scale_file must provide 1 (per layer) or 1 (per KV layer) scales, got k=16 v=16 Unless the file carries an explicit `mtp` section calibrated for the head, drop the KV quantization options from the head's inherited extra_options so it keeps an unquantized KV cache. That reproduces the pre-refactor behaviour, where the combination was expressed as mtp_kv_cache_quant_type=none. Routing around this with mtp_quant_config was not equivalent: it also disables the native-ModelOpt MTP path, and the requantization path then tripped over compressed-tensors FP8 lm_head weights, which carry per-channel rather than per-tensor scales. Accept both there, matching the main loader's _validate_positive_weight_scale. Together these keep the head's lm_head as native FP8 (0.85 GiB of MTP external data instead of 3.22 GiB) while the main model uses INT8 KV.
The greedy N>1 MTP step pays an extra main forward on every partially accepted round, plus a SnapshotState() (2*num_layers device copies) on every step. The sampling path in the same file already avoids both: it skips the snapshot when the model has a croppable recurrent-state window, and its partial-accept branch commits straight out of the wide verify with no replay. The greedy replay exists to make the bonus token decode-consistent, but the branch it guards already documents that it is not lossless in practice: the cropped state is itself derived from the wide batched verify, so greedy near-ties flip regardless. The result is that the common case pays for a guarantee it does not provide. Restore the direct arena commit behind ORTGENAI_MTP_FAST_COMMIT (default off, so current behaviour is unchanged). Measured on the Qwen3.8-27B INT8-KV MTP model, 2048-token prompt, 512 generated tokens, max_draft_tokens=3, median of 3: fast commit off 45.5 tok/s 419 target forwards 1.22 tokens/forward fast commit on 60.8 tok/s 222 target forwards 2.31 tokens/forward Tokens per target forward returns to the 2.415 recorded for this model before the MTP runtime was upstreamed, and the forward count halves. Generated tokens differ from the replay path on near-ties, which is why this is opt-in pending a quality gate. max_draft_tokens=1 is unaffected.
The direct arena commit added in the previous change was gated behind ORTGENAI_MTP_FAST_COMMIT and defaulted off. Remove the environment variable and make the path unconditional for any main model that exposes a croppable recurrent-state window, which is the same condition GenerateStepMultiSample already uses. That also removes the greedy partial-accept crop-and-replay branch, which is now unreachable: it required CanCropRecurrentState() and only ever existed to make the bonus token decode-consistent. Its own note recorded that it does not achieve that, because the state it crops to is itself derived from the wide batched verify, so greedy near-ties flip either way. The remaining snapshot and replay fallback is unchanged and still covers models without a state window. Qwen3.8-27B INT8-KV MTP, 2048-token prompt, 2048 generated tokens, max_draft_tokens=3: before 62.70 tok/s 1608 target forwards 1.27 tokens/forward after 94.56 tok/s 857 target forwards 2.39 tokens/forward max_draft_tokens=1 is unaffected. 254 unit tests pass.
Avoid pinned-input destructor fences and collapse verify-row argmax readback to one synchronization so chained MTP work can remain queued on the shared CUDA stream.
Allow one MTP generator to serve repeated requests while retaining allocated state and captured graphs. Reset also restores the configured sampling seed so benchmark replays remain deterministic.
Use the graph-capture save/restore fixup instead of double buffering by default, reducing recurrent-state memory and graph variants. Keep the environment override as a diagnostic fallback.
(cherry picked from commit 0061d8e)
(cherry picked from commit 7824ad9)
(cherry picked from commit 33ac793)
(cherry picked from commit ee26531)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a8d618bb-0feb-4d05-a918-93b6accde1b1 (cherry picked from commit 618980b)
(cherry picked from commit 2ea2993)
(cherry picked from commit 9f1495a)
(cherry picked from commit 59c0843)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a8d618bb-0feb-4d05-a918-93b6accde1b1 (cherry picked from commit 9b73842)
(cherry picked from commit 9071d5e)
(cherry picked from commit 53dcd9e)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a8d618bb-0feb-4d05-a918-93b6accde1b1 (cherry picked from commit 9b3d165)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: a8d618bb-0feb-4d05-a918-93b6accde1b1 (cherry picked from commit beb17d8)
Cherry-picking PR #2442 onto the RC2 GatedDeltaNet branch left three defects that only the combined tree can hit: - Both branches define a `make_gated_delta_net` with a different contract. The dense one keeps the operator's fused qwen/sigmoid gates, checkpoint window and optional `cu_seqlens`; the packed one requires `cu_seqlens` and pre-casts the gates. Keep both and name the packed one `make_varlen_gated_delta_net`, matching the existing `make_varlen_*` spelling. - `_make_linear_attention` returned through the dense gated_delta_net path before reaching the packed branch, so a paged build never emitted the packed graph, and it referenced a variable the packed rename had removed. - `_setup_hybrid_cache_io` carried both branches' state-dtype rules. V-major float32 state is what GatedDeltaNet requires, which is every paged build and the dense gated_delta_net export, so the two conditions collapse into one. The builder state-group tests predate main's windowed paged ring (#2358), whose `use_windowed_paged_kv_cache` the stub must now carry.
These two emitters hardcoded the dense [batch, sequence, features] annotation, so a paged build gave every native-quantized MatMul a rank-3 value shape against a packed rank-2 activation. With prune_lm_head the LM head then declared logits as [batch_size, batch_size, vocab] instead of [batch_size, vocab]. Every other emitter already routes through hidden_state_shape.
The packed varlen operators already produce the per-token state series a speculative decoder needs to roll back a partially accepted draft (VarlenCausalConvWithState.prefix_states, GatedDeltaNet.checkpoints), but the builder hardcoded both windows to zero and rejected state_window outright on the paged path. Emit them as a separate third output per operator instead of widening the committed state the way the dense export does, so the Engine's fixed-state bank contract is unchanged and a model exported with a window still runs unmodified when the outputs are not fetched. genai_config carries the window as checkpoint_count plus a checkpoint_alignment, because the conv writes slot j after local token j while GatedDeltaNet writes slot W-1 after the last token.
A speculative step runs more tokens than it ends up keeping. The operators already publish the state after each of those tokens, so rolling back a rejected draft only needs the pool to stage a different source row: FixedStateReservation now binds the checkpoint series when the step asks for it, and CommitPrefix selects the slot for the accepted prefix instead of the step's final state. The slot depends on the group's alignment, which is why the manifest carries it. Capture is opt-in per reservation, so a non-speculative step allocates exactly the staging it did before and every existing caller is unchanged.
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.
This was referenced Aug 23, 2026
Contributor
Author
This was referenced Aug 23, 2026
Tianlei Wu (tianleiwu)
force-pushed
the
tlwu/20260823/gdn-paged-spec-verify
branch
from
August 28, 2026 09:57
847509f to
0711e91
Compare
Contributor
Author
|
Superseded by #2493, which ports this runtime API and device-argmax work onto the compact replay stack without an ONNX Runtime version bump. |
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
2 participants
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
Stack
PR 4/6. Depends on
tlwu/20260823/gdn-paged-spec-verify.Validation
engine_unit_tests322/322 passed