Skip to content

feat(qwen): add paged greedy MTP with FP8 block weights - #584

Open
big-hip wants to merge 7 commits into
InfiniTensor:InfiniLM-v0.2.9cfrom
big-hip:feat/qwen-greedy-mtp
Open

big-hip wants to merge 7 commits into
InfiniTensor:InfiniLM-v0.2.9cfrom
big-hip:feat/qwen-greedy-mtp

Conversation

@big-hip

@big-hip big-hip commented Sep 19, 2026

Copy link
Copy Markdown

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.

  • Reuse the existing scheduler, paged KV allocator, linear operators and loader for FP8-block/Marlin weights, vocabulary-parallel projection, packed verification, state-row budgeting and request cleanup.
  • Scope: one shared MTP layer, dense text, PP1, NVIDIA TP1/TP2, eager MTP. Ordinary inference retains Decode graphs. MTP requires --disable-prefix-caching; no separate prompt snapshot cache or MTP graph capture is included.
  • Model code stays in csrc/models/qwen3_5/, shared recurrence in qwen3_next, and orchestration in llm/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.
  • Breaking change.

Test Results of Involved Models on Supported Platforms (Please attach screenshots)

Current paired heads: InfiniLM 31ef4290 + InfiniCore a3ac4df4, 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.

Rendered current-head local regression logs, not GitHub CI

Historical model/entrypoint evidence:

Platform / check Verified revision and scope
A6000 TP1/TP2 Archived greedy equivalence, batching, EOS/limits, cancellation, active/queued close and complete KV/state reclamation
2×5090, real 27B, TP2 78d19f74: ordinary graph vs K2, 3 prompts × 3 repeats × 64 outputs, exact token equality
5090 dynamic batching 78d19f74: 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 probe
Public entrypoints / tiny fixture Single inference, offline benchmark, evaluation adapter and HTTP smoke passed historically; 64/64 HTTP requests. The evaluation score is not a model-accuracy claim
CPU / other accelerators CPU control-flow/configuration/remapping checks only. Full MTP has not been validated on Ascend/MetaX/Moore; the tested Marlin backend is NVIDIA-specific. @wooway777 please review portability

Source/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, PyTorch 2.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.

Prompt tokens Ordinary Decode graph K2 eager Gain
Chinese 63 63.35 tok/s 110.44 tok/s +74.3%
Code 127 62.40 tok/s 118.27 tok/s +89.5%
Summary 1023 52.00 tok/s 89.67 tok/s +72.4%

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

  • Base: 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.
  • Textual merge checks show overlapping changes with feat(engine): integrate prefix eviction and chunked parallel execution #573 (chunk input flags, engine/configuration/service wiring) and feat(mamba2): integrate model loading and recurrent inference #575 (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.
  • Review checkpoint ownership, per-request causal/page-table metadata and accepted-length KV rollback first. Short verification reuses Decode Attention without copying KV payloads.
  • TP1 batches draft continuation; TP2 uses separate draft calls because end-to-end batching benefit was inconsistent. Random sampling, multimodal/MoE/PP, multiple MTP layers and remote recurrent-state transfer remain outside scope.
  • Prompt snapshots had value for repeated full prompts but introduced a separate cache/ownership mechanism; they are deferred. K1 draft graphs were removed after paired measurements found no consistent benefit. Their earlier implementations remain in commit history. Neither is required for the retained MTP decode gain.

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

Every contributor must verify every item below before requesting
review. Tick each box only after the check has actually been performed —
do not tick speculatively. If an item truly does not apply, replace the
checkbox with N/A and briefly explain why in an inline comment.

Title, Branch, and Commits

  • PR title follows Conventional Commits (e.g. feat(nvidia): …, fix(cuda/gemm): …).
  • Branch name follows <type>/xxx-yyyy-zzzz where <type> matches the PR title's Conventional Commits type and words are joined with hyphens (see CONTRIBUTING.md §Branches).
  • Each commit message follows Conventional Commits.
  • Small PR is a single squashable commit; or, for a large PR, every commit is meaningful, well-formed, and independently reviewable (see CONTRIBUTING.md §Pull Requests).
  • No stray merge commits; a feature commit, focused correctness/compatibility fixes and scope/test cleanups, based on target InfiniLM-v0.2.9c (270feb3e).
  • No fixup! / squash! / wip commits remain.
  • N/A Existing legacy issue format: this new branch uses Conventional Commits.

Scope and Design

  • Changes are minimal — nothing unrelated to the stated motivation was added (CONTRIBUTING.md §Code/General).
  • No dead code, commented-out blocks, debug prints, printf/std::cout/print(...) left behind, or TODO without an owner and issue link.
  • No unrelated formatting churn that would obscure the diff.
  • Public API changes (if any) are intentional, documented, and reflected in affected callers/tests.

General Code Hygiene (applies to all languages)

  • The code is self-explanatory; comments were added only where the why is non-obvious (CONTRIBUTING.md §Code/General).
  • Every modified or added file ends with a single trailing newline (CONTRIBUTING.md §Code/General).
  • No trailing whitespace, tab/space mixing, or stray BOMs.
  • Identifiers in comments and error messages are wrapped in backticks (e.g. the `seqlens_k` tensor) (CONTRIBUTING.md §Code/General).
  • All comments and error messages are in English (CONTRIBUTING.md §Code/General).
  • Comments and error messages are complete sentences — capitalized first letter, terminal punctuation — unless the language/framework convention says otherwise (CONTRIBUTING.md §Code/General; §Python).

C++ Specific (if C++ files changed)

  • Code follows the Google C++ Style Guide strictly.
  • Error and warning message wording follows the LLVM Coding Standards (CONTRIBUTING.md §C++).
  • Constructor initializer list order matches member declaration order (CONTRIBUTING.md §C++).
  • No raw new/delete; RAII / smart pointers / existing allocators are used.
  • Changed files are formatted by scripts/format.py.
  • No changes/reference to csrc/models/llama_legacy/.

Python Specific (if Python files changed)

  • Code is PEP 8 compliant.
  • Comments are complete English sentences, starting with a capital letter and ending with punctuation; Markdown backticks are used for code references (CONTRIBUTING.md §Python).
  • Docstrings (if any) follow PEP 257 (CONTRIBUTING.md §Python).
  • Changed files are formatted by scripts/format.py.
  • No changes/reference to python/infinilm/auto_config.py.

Testing

  • For any platform that could not be tested, an explicit reason is given in the table and a reviewer with access has been tagged.
  • Passed single request test (examples/test_infer.py), or specify the reason for skipping.
  • Passed offline performance test (examples/bench.py), or specify the reason for skipping.
  • Passed sanity test (test/bench/test_benchmark.py), or specify the reason for skipping.
  • Passed service test (python/infinilm/server/inference_server.py + scripts/test_perf.py), or specify the reason for skipping.

Build, CI, and Tooling

  • The project builds cleanly from a fresh directory on at least one affected platform.
  • CI has been triggered manually (Actions → CI on this branch), or /retest was requested.

Documentation

  • README.md, CONTRIBUTING.md, or inline docs updated when behavior, build flags, or developer workflow changed.
  • N/A Breaking change: the new mode is opt-in; unsupported MTP combinations fail explicitly.

Security and Safety

  • No secrets, access tokens, internal URLs, customer data, or personal hardware identifiers have been committed.
  • Third-party code is license-compatible and attributed.
  • No unsafe pointer arithmetic, uninitialized reads, or missing bounds checks were introduced.

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.
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
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.
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant