[Model] Add Inkling multi-depth MTP support [5/N] - #48768
Conversation
mergify
Bot
added
ci/build
rust
new-model
WoosukKwon
changed the title
| self.prefill_cudagraph_manager: SpeculatorCudaGraphManager | None = None | ||
| self.decode_cudagraph_manager: SpeculatorCudaGraphManager | None = None | ||
|
|
||
| self.num_spec_prefill_steps = ( |
There was a problem hiding this comment.
I really don't like the idea of folding multi-module MTP into the autoregressive speculator implementation.
There was a problem hiding this comment.
@benchislett This is a great input. cc @TheEpicDolphin
The MTP integration might be a bit messy now. We will split this PR into smaller ones, and merge them one by one; MTP will be the last one since it involves the deepest changes and nontrivial design questions. Please share any alternative ideas that you might have.
There was a problem hiding this comment.
I'm considering splitting it into it's own separate class that inherits from DraftModelSpeculator. The multi-module MTP implementation has become complex enough that i think it warrants separation
There was a problem hiding this comment.
Here's the draft PR: #48892. Will try to have it ready for review later today/tomorrow.
|
Hey — we're doing a 0-day build of this PR on H200s (SM90) and hit a crash in Problem: The The same mismatch exists in the execution path ( Additionally, Workaround we're carrying downstream:
Wanted to flag in case you want to fix it upstream in the tml-fa4 interface or gate the Inkling model to SM100+ only. Happy to share the patch if useful. |
SM90 (H200) compatibility reportWe got Inkling-NVFP4 serving on 8x H200 (SM90) using this PR merged onto our v0.25.0 base. Two issues came up that needed patching — sharing in case it's useful for the PR or for anyone else running on Hopper. 1. tml_fa4
|
|
@dougbtv Thanks for reporting the issue. Yeah we made some last minutes change and they were not rigorously tested on Hopper GPUs. I think we should fall back to normal FA4 score_mod (instead of tml-fa4) for Hopper. |
|
This pull request has merge conflicts that must be resolved before it can be |
WoosukKwon
changed the title
WoosukKwon
changed the title
WoosukKwon
changed the title
Extend the merged Inkling MTP=1 implementation to multiple checkpoint depths with multi-step speculative decoding and KV-cache plumbing. Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <mozf@inferact.ai> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
WoosukKwon
changed the title
|
This pull request has merge conflicts that must be resolved before it can be |
|
This pull request has merge conflicts that must be resolved before it can be |
Ports the upstream mlx-vlm inkling_mtp drafter semantics (one dense block per draft depth, vllm-project/vllm#48768 equivalent) onto the omlx chain cycle: the history fold runs block 0 on the persistent per-block cache list, chain step j runs block min(j, depth-1) on the head-clone copy, and the fold consumes pre-norm trunk hidden via a new per-model marker. Verify rollback reslices the short-conv states from padded-input stashes and trims the KV subs, replacing upstream's replay forward. oQ preserve_mtp maps model.mtp.layers.* through a sanitize hook (head floors at 4-bit), the bpw estimator and config normalization understand the nested mtp_config block, and prompt-priming capture is wired through the inner model. Measured on the oQ4e-mtp output the adaptive controller parks at depth 0 under prefix-cached serving (0.98-1.00x, outputs valid); a warmed head accepts 1.32 tokens/cycle at depth 4 on repetitive text with agreement decaying past depth 2, so the toggle stays opt-in.
The first cut copied mlx-vlm's inkling_mtp drafter, which only folds block 0 and leaves depth blocks 1..7 with empty caches every round; measured per-depth agreement collapsed to 0.68/0.45/0.19/0.00 and the controller always parked. The reference semantics (vllm-project/vllm#48768) keep every block stateful over the accepted sequence: a cycle is k chained passes over one uniform token window (pass j = block j, hidden chained, token stream left-shifted), each fold trims provisional draft rows back to w0 = min_j valid_j and refolds the gap from a small ring of recent pairs, and the embedding side is double-normed through the backbone embed_norm. Rewritten that way (plus a begin-cycle hook so the pass counter lives on the cache list, no more per-cycle head clone, rolling conv-input stashes for cross-cycle rewinds, sliding-window prompt capture folded at activation, and a capture gate that skips return_hidden forwards - the v1 primed=0 bug), per-depth agreement is 0.87/0.76/0.67/0.61 with the same head weights and the server sustains 2.1-2.6 accepted tokens per cycle at 63-75% draft accept. Net decode is still 0.95-0.98x on this hardware because a fine-grained MoE verify pass costs 2.1x a plain step (per-token expert routing) plus 0.9x for the 8-pass head, so the adaptive controller correctly parks; the toggle stays opt-in. Unit oracles pin the trim/refold cycle against one-shot refolds and the runtime against the standalone reference engine.
* fix: mixed CacheList store/restore integrity across cache subsystems Composite CacheList layers mixing a sliceable KVCache with a non-sliceable ArraysCache were never exercised before. The prefix cache stored such layers as cumulative boundary snapshots but restored the KVCache sub by concatenating per-block data, duplicating the sequence. Restore now mirrors the store-side layer-level decision. Also: SSD compat signatures stamp CacheList sub composition (plus a per-block signature gate for hot-cache loads), TurboQuant excludes composite layers with KVCache members until a TQ sub-state serialization exists, scheduler emits sub_class_names, prefill boundary snapshots are extracted eagerly instead of aliasing live caches, and ArraysCache.extract gets the None-slot guard its siblings already had. * feat: Inkling Small support (vendored model, serving, sliding-window slice) Vendors the inkling package from mlx-vlm PR #1756 head plus the switch_layers/mlp/activations shims missing at the pin, with marked local changes: batched right-padded prefill correctness (conv_mask wiring, lengths-aware conv state writes, per-sequence key masking and log-tau positions, BatchKVCache-safe buffer-coordinate offsets) and an opt-out sliding-window K/V slice that keeps the 512-token layers O(window) instead of O(S). Adds a torch-free port of the transformers 5.14 Inkling processors, NVFP4 hf_quant_config translation, discovery registration, a channel output parser (content_thinking to think tags, invoke_tool_json to tool_calls, end_sampling stop), and prefill-guard pricing for the materialized banded attention bias. * feat: oQ/oQe wiring for Inkling (streaming sanitize, sconv skip, estimator) Name-driven rules mostly carry over; the additions are an explicit skip for the depthwise sconv conv weights (nn.Conv1d has no to_quantized), an inkling branch in the calibration layer walk, embed-norm-aware embedding for use_embed_norm checkpoints, class-method resolution on the VLM sanitize proxy plus a guard for packages without a VisionModel export, and a bpw/size estimator that prices post-sanitize names through the discovered plan (raw source names like experts.w13_weight priced 97% of the checkpoint as fp16 passthrough). * feat: Inkling Lightning MTP runtime (single-checkpoint multi-depth head) Ports the upstream mlx-vlm inkling_mtp drafter semantics (one dense block per draft depth, vllm-project/vllm#48768 equivalent) onto the omlx chain cycle: the history fold runs block 0 on the persistent per-block cache list, chain step j runs block min(j, depth-1) on the head-clone copy, and the fold consumes pre-norm trunk hidden via a new per-model marker. Verify rollback reslices the short-conv states from padded-input stashes and trims the KV subs, replacing upstream's replay forward. oQ preserve_mtp maps model.mtp.layers.* through a sanitize hook (head floors at 4-bit), the bpw estimator and config normalization understand the nested mtp_config block, and prompt-priming capture is wired through the inner model. Measured on the oQ4e-mtp output the adaptive controller parks at depth 0 under prefix-cached serving (0.98-1.00x, outputs valid); a warmed head accepts 1.32 tokens/cycle at depth 4 on repetitive text with agreement decaying past depth 2, so the toggle stays opt-in. * fix: rewrite Inkling MTP cycle to vLLM per-depth-stateful semantics The first cut copied mlx-vlm's inkling_mtp drafter, which only folds block 0 and leaves depth blocks 1..7 with empty caches every round; measured per-depth agreement collapsed to 0.68/0.45/0.19/0.00 and the controller always parked. The reference semantics (vllm-project/vllm#48768) keep every block stateful over the accepted sequence: a cycle is k chained passes over one uniform token window (pass j = block j, hidden chained, token stream left-shifted), each fold trims provisional draft rows back to w0 = min_j valid_j and refolds the gap from a small ring of recent pairs, and the embedding side is double-normed through the backbone embed_norm. Rewritten that way (plus a begin-cycle hook so the pass counter lives on the cache list, no more per-cycle head clone, rolling conv-input stashes for cross-cycle rewinds, sliding-window prompt capture folded at activation, and a capture gate that skips return_hidden forwards - the v1 primed=0 bug), per-depth agreement is 0.87/0.76/0.67/0.61 with the same head weights and the server sustains 2.1-2.6 accepted tokens per cycle at 63-75% draft accept. Net decode is still 0.95-0.98x on this hardware because a fine-grained MoE verify pass costs 2.1x a plain step (per-token expert routing) plus 0.9x for the 8-pass head, so the adaptive controller correctly parks; the toggle stays opt-in. Unit oracles pin the trim/refold cycle against one-shot refolds and the runtime against the standalone reference engine. * perf: fix Inkling MTP cycle cost (lazy-graph accumulation, mask kernel recompiles) Two cycle-cost defects were hiding the acceptance win. The head's bookkeeping arrays (rolling conv stashes, refold ring, per-block KV) were never on any eval path, so their lazy graphs deepened across cycles; a sync-free async_eval at fold time cuts a depth-4 head cycle from 59 ms to 8.2 ms. And the banded rel-position mask kernel templated S/LQ/Q_OFF into the Metal source, recompiling a pipeline every time the cache grew (every decode step on the seven global layers, every MTP verify and head pass); those are runtime params now, which also benefits the plain decode path. Matrix after the fix: code_4k 1.34x, code_16k 1.25x, story_4k 1.26x, story_16k 1.17x decode with the controller never parking (512-token generations sustain 1.9-2.3 accepted tokens per cycle at 58-70% draft accept).
|
Closing this as #48892 implemented this last remaining feature |
4 participants
Summary
This draft now contains the multi-depth MTP work that remains after the following model slices landed:
Additional independent changes were extracted into draft PRs:
The standalone QKVR preparation benchmark was removed.
The remaining model changes:
The core model, LoRA support, MTP=1 foundation, breakable CUDA graphs, and Hopper attention support are no longer introduced by this diff.
Duplicate-work check
This does not duplicate another open PR. It is rebased directly onto current
mainand builds on the merged slices #48799, #48822, #48858, and #48869. The changes shown here are specifically the delta from one draft depth to multi-depth MTP. LoRA is isolated in #48884, and the unrelated infrastructure changes are isolated in #48880 and #48881.Validation
.venv/bin/python -m pytest tests/config/test_speculative_draft_hf_overrides.py tests/models/inkling/test_mtp_input_fusion.py -v: 38 passedgit diff --name-only -z origin/main...HEAD | xargs -0 .venv/bin/pre-commit run --files: passed, including Python 3.10 mypy and Rust formattinggit diff --check origin/main...HEAD: passedModel evaluations
The source PR reports the following evaluations for the complete Inkling implementation. These were not rerun while extracting the multi-depth MTP delta.
AI assistance
AI assistance (OpenAI Codex) was used to split and rebase the original branch, preserve contributor attribution, isolate LoRA from multi-depth MTP, run validation, and update this draft PR. The human submitter must review every changed line, understand the change end-to-end, and confirm the model evaluations before marking this ready for review.