Conversation
Execute the checkpoint's shared MTP head through the existing engine and scheduler. Verify up to four candidates, commit matching Conv/GDN states without target replay, and reclaim request resources on stop or failure. Reuse column-parallel vocabulary projection, Marlin FP8 packing and PagedCompiler draft capture. Bound recurrent rows independently from KV pages and make exact-prompt snapshots opt-in. Wire MTP through the CLI, benchmark and service entrypoints; retain three core MTP test modules. Validated NVIDIA A6000 TP1/TP2, graph recapture, forced acceptance lengths, real 27B FP8 batching/cancellation, and ordinary Qwen2 pre-transpose. Runtime support: InfiniTensor/InfiniCore#1565; graphs also require #1560.
big-hip
marked this pull request as draft
September 19, 2026 13:22
Keep ordinary short-prefill dispatch, preserve live KV during graph recapture, and defer cleanup until a timed-out worker exits. Retain opt-in MTP acceleration features and cover ordinary projection, recapture, shutdown and remote-release contracts.
big-hip
marked this pull request as ready for review
September 19, 2026 15:57
Reuse Decode Attention with per-request causal lengths for packed short verification. Align NVIDIA batched gate projections with checkpointed recurrence to avoid shape-dependent BF16 state drift after cancellation and admission. Add request-isolation and causal-boundary coverage to the existing MTP execution tests. Validate the real TP2 RTX 5090 cancel/re-admit reproducer for K=1/2/4, exact Conv/GDN state comparisons, and ordinary eager/graph regressions.
39 tasks
Keep the MTP PR focused on greedy speculation, accepted-state checkpoints and scheduler integration. Remove the separate TP1 snapshot cache, its public budget option and invalidation hooks, and its dedicated tests. MTP requires prefix caching to be disabled. Reuse the existing full-prompt model input directly; retain core acceptance, lifecycle, packed execution and ordinary-graph regressions.
Remove the unused worker and engine state-export methods with their Python binding after prompt snapshots were deferred. Restore the upstream KV-cache binding and drop capture tracing logs. Retain checkpoint execution, request ownership and ordinary graph capture. Validate the native rebuild, CPU contracts and all four CLI entrypoints.
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
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
Add opt-in text-only greedy decoding through Qwen's built-in MTP head. Generate 1–4 candidates, verify with the target model, and select the accepted Conv/GDN checkpoint without target replay after rejection.
--disable-prefix-caching; no separate prompt snapshot cache or MTP graph capture is included.csrc/models/qwen3_5/, shared recurrence inqwen3_next, and orchestration inllm/model_runner/mtp_runner.py. No weights, research scripts or profiler artifacts enter this PR.Motivation
The existing model path discards MTP weights. Hybrid speculation must commit both Conv and GDN state at the accepted token boundary; attention KV rollback alone is insufficient.
Ordinary Qwen inference also uses vocabulary-parallel projection, corrected norm/tied-weight loading and the configured recurrent-state precision. FP32 doubles the recurrent-state component versus BF16, not total model memory; its capacity is independently configurable. Ordinary short Prefill retains chunked GDN dispatch. NVIDIA batched Decode uses the same per-token gate projection shape as verification to avoid shape-dependent BF16 state drift, so some previous batched outputs change. Shared graph recapture and shutdown fixes preserve active state and in-flight ownership. Performance is not certified unchanged for every model/backend.
Type of Change
feat— opt-in model-head speculation and service integration.fix— required weight-layout, state-ownership and lifecycle correctness.Test Results of Involved Models on Supported Platforms (Please attach screenshots)
Current paired heads: InfiniLM
31ef4290+ InfiniCorea3ac4df4, verified on this server's two RTX A6000 48 GiB cards (2026-09-23). The small FP8-block/Marlin hybrid fixture passes TP1: 66 tests, one TP2-only skip; TP2: 67 tests, no skips, including K1/K2/K4 checkpoints, cancellation/re-admission, request isolation and ordinary graph recapture. These totals include CPU tests. Sampled test-process peaks: 370 MiB for TP1; 468 MiB/device for TP2. Core graph/cast checks (13) and paged Prefill cases (88) also pass. Logs, fixture conditions, binary hashes and commands.The initial attempt loaded a September 19 Core library and aborted in graph allocation. Rebuilding/installing the current #1566 dependency resolved that reproducer with unchanged LM code and test assertions. The four CLI entrypoint checks and local native build also pass. No additional implementation change was needed.
Limits: these are small-model correctness checks under shared GPU occupancy. Full 27B current-head regression and uncontended throughput were not rerun because remaining memory is insufficient. The real-model results below are historical, not current-head passes. Joint integration with #573/#575 also remains unverified.
Historical model/entrypoint evidence:
78d19f74: ordinary graph vs K2, 3 prompts × 3 repeats × 64 outputs, exact token equality78d19f74: original token-21 cancel/re-admit mismatch resolved for K1/K2/K4; 384/384 Conv/GDN tensors and both hidden vectors match in the same-state probeSource/build conditions and A6000 checks. 5090 reproducer, resolved numerical issue, exact-state comparison, benchmark and memory evidence. The tiny boundary test also passes the old binary; the real 27B reproducer establishes the original failure.
Benchmark / Performance Impact
Revision
78d19f74: 2×RTX 5090 (32607 MiB/device), SM120, driver 610.43.02, CUDA 13.3, PyTorch2.13.0a0+8145d630e8.nv26.06; SYS topology without CUDA P2P. Qwen3.8-27B-FP8, Marlin W8A16 (FP8 storage/BF16 compute), FP32 recurrent state, text greedy, TP2/PP1, B1, 80×64-token KV pages, prefix reuse off. Prefill eager; 64 output tokens, EOS ignored for fixed-work timing, three repeats after warmup. Median rate(N-1)/(wall-TTFT)includes CPU/scheduler, excludes loading/Prefill.All outputs match exactly. These measurements predate the two scope cleanups; neither removed feature was enabled in this comparison. They are not fresh measurements or gains caused by deleting code. No native W8A8 integration or validated speed advantage over vLLM is claimed.
100 ms sampled device peaks (GPU0/GPU1): ordinary graph B1 23768/23782 MiB, K2 B1 24384/24392 MiB. K4 B2 with additional checkpoint rows peaks at 25586/25586 MiB. Brief peaks may be missed. A separate same-machine MTP-off B2 gate-fix regression (63/1023-token prompts, 32 outputs/request, 9 state rows, three repeats) measured ordinary eager 93.50→100.04 tok/s, ordinary graph 92.66→99.34 tok/s; this does not establish performance for every batch/model.
Notes for Reviewers
InfiniLM-v0.2.9c. Runtime prerequisite: issue/1565 fix(runtime): support Qwen MTP and recurrent graph replay InfiniCore#1566. No dependency on Mamba scan #1562; generic linear/capture fixes overlap feat(mamba2): integrate model loading and recurrent inference #575 and require merge-order resolution.README.md, recurrent input forwarding). Once an earlier PR lands, resolve against that release head while retaining both features; the independent-PR checks are not a joint-integration validation.CI / ChatOps
Current
31ef4290: fork format and Ruff passed. These workflows do not run accelerator tests. Upstream PR checks require maintainer approval. Current small-model GPU checks pass as reported above; full 27B current-head regression and joint PR integration remain pending.Checklist
Title, Branch, and Commits
feat(nvidia): …,fix(cuda/gemm): …).<type>/xxx-yyyy-zzzzwhere<type>matches the PR title's Conventional Commits type and words are joined with hyphens (seeCONTRIBUTING.md§Branches).CONTRIBUTING.md§Pull Requests).InfiniLM-v0.2.9c(270feb3e).fixup!/squash!/wipcommits remain.Scope and Design
CONTRIBUTING.md§Code/General).printf/std::cout/print(...)left behind, orTODOwithout an owner and issue link.General Code Hygiene (applies to all languages)
CONTRIBUTING.md§Code/General).CONTRIBUTING.md§Code/General).the `seqlens_k` tensor) (CONTRIBUTING.md§Code/General).CONTRIBUTING.md§Code/General).CONTRIBUTING.md§Code/General; §Python).C++ Specific (if C++ files changed)
CONTRIBUTING.md§C++).CONTRIBUTING.md§C++).new/delete; RAII / smart pointers / existing allocators are used.scripts/format.py.csrc/models/llama_legacy/.Python Specific (if Python files changed)
CONTRIBUTING.md§Python).CONTRIBUTING.md§Python).scripts/format.py.python/infinilm/auto_config.py.Testing
examples/test_infer.py), or specify the reason for skipping.examples/bench.py), or specify the reason for skipping.test/bench/test_benchmark.py), or specify the reason for skipping.python/infinilm/server/inference_server.py+scripts/test_perf.py), or specify the reason for skipping.Build, CI, and Tooling
/retestwas requested.Documentation
README.md,CONTRIBUTING.md, or inline docs updated when behavior, build flags, or developer workflow changed.Security and Safety