feat(spec): frontier-anchored multi-layered MTP - #1080
Merged
Conversation
Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
Add the third decode-window repair design to the multi-depth MTP drafter: every depth re-runs the SAME k rows per request, ending at the committed frontier (valid_cache_lengths + accept, computed once before depth 0). Every row then holds a committed token or a fresh draft, so rejected-tail garbage rows are structurally impossible and stale per-depth KV/sconv entries are rewritten in place with exact values. Depth 0 re-runs purely to regenerate chain hiddens, splicing the drafter stash (last k-1 committed tokens + the target hiddens one position behind them) with the round's verify hiddens at the accept boundary; the sample gather becomes static (last row) and the row count stays bs*k, the standard decode shape. Backend side, the window keeps the plain k-row catch-up shape and only its anchor moves: enter_draft_frontier_window rebuilds the conv metadata at the frontier and enter_draft_frontier re-anchors the inner backend's seq_lens and grouped write locations. The frontier is accept-dependent and accept is only known inside the captured graph, so both are plain tensor ops recorded at capture and recomputed per replay (the lookback variant's host-side fill cannot serve this). Arming and per-round hooks raise on refusal instead of falling back. INKLING_MTP_DECODE_LOOKBACK becomes a three-way knob parsed in one place (0 = unrepaired window, 1 = lookback rows, 2 = frontier window, the new default); the conv ring carries the lookback depth only in mode 1. Validated on GB200: positional-oracle unit tests for the id/hidden composition, including equality with the lookback window on all shared positions; conv-metadata and grouped-loc hook tests; cache plan/setup suites; and real-weight engine A/B runs (mode 1 vs 2, CUDA graphs on) with identical outputs and verify-round counts, plus full-model accept-length verification. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
Drop the two superseded decode-window designs: the unrepaired verify-anchored window and the lookback variant (D repair rows, widened loc-stack sizing, per-round fallback chain). The frontier-anchored window is now the only decode mechanism and runs uniformly at every spec_num_steps; EXTEND rounds always prime the drafter stash. The INKLING_MTP_DECODE_LOOKBACK knob, the model-side mode attribute, the conv ring's lookback depth, the mixin's second loc stack and the configure/arm handshake all go away — the drafter calls the re-anchor hooks directly, and mis-wiring fails on the attribute access itself. Also sweeps the residue of the removed designs: the wrapper's orphaned advance_draft_forward_metadata (classic per-step-loop hook; only eagle-family backends implement it now), three unread draft_* model class flags, _decode_shifted_ids' design-1 default window, stale kernel-doc wording, and the lookback-named stash buffers and helpers (renamed to stash naming). The drafter-owned draft_seq_lens_buf stays: the CUDA-graph wrapper aliases it into every draft metadata init. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
Cleanup pass over the frontier-only Inkling MTP drafter and shared drafter helpers: - Merge _run_first_step/_run_extend_depth_catchup into a single _run_extend_depths loop: depth 0 is not special, every depth runs the same ragged rows with shifted inputs and chained hiddens. Drop the MIXED-batch skeleton; MIXED is rejected at the Inkling backend's metadata init. - Remove never-binding guards: speculative_num_draft_tokens getattr/max fallbacks, the base-class override_num_extends(0) no-op, the next_tokens broadcast initialization, and clamps whose bounds have no producer (req_pool_indices clamp_min(0) — CUDA-graph padding uses the reserved slot 0, never PAD_SLOT_ID). - Revert PR lightseekorg#574's accept_lengths clamp and helper: its producer (_cap_accept_to_context_len, DFLASH/DSPARK-gated) was removed in lightseekorg#974 in favor of the spec_context_pad physical extent, so verify's [1, spec_num_tokens] contract holds on every path. Inline the gather arithmetic back at the call sites (pre-lightseekorg#574 form) in mtp/eagle/dflash. - Unify token-value dtype to int32 end to end (stash tokens included); int64 remains only in the index domain. - Align nvtx markers with method names. Verified: test_drafter_accept_indexing, test_inkling_mtp_conv_state, test_cache_memory_plan, test_cache_setup — 121 passed + 36 subtests. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
…ution to fill Decode rounds (Inkling MTP): - The depth-0 frontier window IS the last k committed (token, hidden) pairs, so the cross-round stash is its tail [:, 1:] — replace the separate accept-indexed re-gathers (_update_stash_decode / _committed_tail_update) with two slice copies; randomized equivalence verified against the old primitive before deletion. - Depths >= 1 are the depth-0 window rolled one left with the previous step's draft appended: one cat per depth replaces the per-depth three-source splice (_decode_shifted_ids deleted, _frontier_shifted_ids is depth-0 only). - _run_frontier_window renamed _run_decode_depths, mirroring _run_extend_depths; it derives slot/v/accept itself (_decode_slices inlined; the accept clamp and int64 cast were never-binding — verify emits [1, k] and both consumers promote against int64 arange). MM pad substitution moves from the drafters to the fill: - fill_input_buffers rewrites the drafter-only shift-1 buffer's media hash ids in place (substitute_mm_pad_, once per chunk); every id the drafters consume is sampled or substituted in-vocab by construction, so the per-depth substitute+clamp (_prepare_draft_input_ids) is deleted from both mtp and eagle. The -1 final-chunk placeholder is preserved (the shift-1 buffer stays unclamped). - wire_drafter targets InputBuffers and fails loudly at startup for a multimodal-active target with no resolvable substitutes, instead of clamping media ids to garbage per token. - maybe_substitute_mm_pad (out-of-place + legacy scalar branch) deleted; substitute_mm_pad_ is the sole substitution utility. Verified: test_drafter_accept_indexing, test_inkling_mtp_conv_state, test_cache_memory_plan, test_cache_setup, test_multimodal_pad_ids — 123 passed + 36 subtests. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
…nd decode window - Delete InklingConvMetadata.is_decode: the sconv kernels are unified behind one entry point selected by num_extends, so the field had writers everywhere and zero readers. - Frontier hooks renamed enter_draft_frontier_window/enter_draft_frontier -> update_draft_forward_metadata (both levels; the per-round absolute counterpart of advance_draft_forward_metadata). Strip never-binding [:bs] slices and a double int32 cast, drop the then-unused bs param, and rebuild the conv metadata with dataclasses.replace (only seq_lens changes). Docstrings cut to the contract. - Decode window: per-depth ids are now a rolling window (cat of the previous window's tail and the new draft) instead of strided slices of a scratch buffer — depth 0 becomes a free view; the last-row sample gather is precomputed at init (padded_gather_ids_buf); the frontier/step_positions math drops redundant casts via promotion; single-use locals inlined; dp_size demoted to an __init__ local (no external readers). - Trim mtp.py docstrings/comments to constraints the code cannot state; fix the stale base_model_output shape comment. Verified: test_drafter_accept_indexing, test_inkling_mtp_conv_state, test_cache_memory_plan, test_cache_setup, test_multimodal_pad_ids — 123 passed + 36 subtests. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
Ruff F821: the else branch read logits_output as a variable leaked from the previous loop iteration. Seed prev_hidden with the depth-0 splice before the loop and publish each depth's hiddens right after its forward — same dataflow, no cross-iteration leak. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
Member
Author
|
@lightseek-bot Looks like
Locally hit the same error, downgrading to Failures are unrelated to this PR. Merging. |
Contributor
|
Hi @syuoni Let's remove Claude as co-author next time when merging. |
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
Makes the frontier-anchored decode window the only multi-depth MTP drafting
mechanism and cleans the path end to end.
Core change
Every draft depth
0..steps-1runs the same k rows per request, ending atthe committed frontier (
frontier = vc + accept). Rows re-run committedpositions, so KV/sconv slots written last round from since-rejected drafts are
rewritten in place with exact values — no repair passes, no per-depth metadata
advance. All round geometry (positions, cache write locations, seq_lens,
sample gather) is
frontier ± const, recomputed per CUDA-graph replay fromacceptalone with pure tensor math; no host syncs.Supporting pieces:
keyed by req pool slot) reconstructs the window's pre-verify history; it is
the tail
[:, 1:]of the depth-0 window, so the roll is two slice copies.with its draft appended (one
cat); only depth 0 is accept-dependent.update_draft_forward_metadata(frontier)re-anchors convmetadata and grouped write locs (the per-round absolute counterpart of
advance_draft_forward_metadata).configure/arm handshakes are removed; EXTEND rounds always prime the stash.
Cleanups riding along
fill_input_buffers(once per chunk on the drafter-only shift-1 buffer): drafters only ever see
in-vocab ids; per-depth substitute+clamp deleted from MTP and EAGLE;
unconfigured multimodal targets now fail loudly at startup.
producer died in feat(cache): unify draft KV into the target's cache pool — one plan, … #974 — redundant casts/slices, dead
is_decodeconvmetadata field, dead
bsparams, single-use locals).constraints the code can't state.