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

[big-hip]

Summary

Add opt-in, text-only greedy decoding through the Qwen checkpoint's built-in MTP head. Verify 1–4 candidates using the existing scheduler and paged cache, and commit the matching Conv/GDN checkpoint without replaying the target model after rejection.

• Keep model code in csrc/models/qwen3_5/, shared recurrent logic in qwen3_next, and request orchestration in llm/model_runner/mtp_runner.py.
• Reuse ColumnParallelLinear for vocabulary projection, existing Marlin packing for E4M3 block weights, and PagedCompiler for bounded draft graph capture.
• Support independent state-row capacity, packed verification, cancellation/limits/shutdown, and optional byte-bounded exact-prompt snapshots. Wire the same options through inference, benchmark and HTTP entrypoints.
• Keep three MTP correctness test modules and one existing-model linear-layout regression; no research scripts, weights or profiling artifacts in this diff.

Motivation

The existing Qwen model path drops MTP weights and cannot use the checkpoint's head to reduce target-model decoding calls. Hybrid models also need Conv and GDN state committed at the accepted token boundary: rolling back attention KV alone is insufficient.

Speculation is opt-in; random sampling is explicitly rejected. Ordinary Qwen inference also uses vocabulary-parallel projection, corrected norm loading and the checkpoint-configured recurrent-state precision. These shared changes are reviewed separately from MTP acceptance correctness. The generic linear post-load fix ensures tied weights exist before packing and makes layout processing per-linear and repeatable.

Related runtime issue: InfiniTensor/InfiniCore#1565. Runtime PR: InfiniTensor/InfiniCore#1566.

Type of Change

• feat — new feature / new model
• fix — bug fix
• perf — performance improvement (no behavioral change)
• refactor — code restructuring without behavior change
• test — adding or fixing tests only
• docs — documentation only
• build / ci — build system or CI configuration
• chore — tooling, formatting, or other non-code changes
• Breaking change

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

Ordinary-path follow-up (86208ae1)

• Ordinary short GDN Prefill keeps the existing chunked dispatch; per-token recurrence is selected only for checkpointed verification. Ordinary and MTP graph recapture preserve live KV. Async shutdown reports a five-second timeout while retaining in-flight resources until worker cleanup; callers can retry stop().
• Vocabulary parallelism, opt-in draft graphs and byte-bounded prompt snapshots remain available. FP32 recurrent state remains required where configured; this intentionally changes ordinary hybrid-state memory versus upstream BF16 allocation.
• Final TP1: 65 passed, 1 skipped. Final TP2: 66 passed. Added ordinary vocabulary/dense-projection and graph-recapture checks, in-flight shutdown tests, and ordinary remote KV delayed-release ownership checks in the existing test modules.
• Separate upstream/current tiny Qwen2 and Llama runs: eager/Decode-graph logits and greedy tokens exactly match. Tiny latency varies substantially, including A/A repeats of the identical binary, so this does not certify ordinary-model performance non-regression.
• Real 27B FP8, A6000 TP1 / K=2 eager, three repeats, 32 output tokens: median Decode throughput 35.62 → 36.31 tok/s (4-token prompt), 45.67 → 46.50 tok/s (63-token prompt). All 192 compared output tokens and acceptance counts match. Short-prompt TTFT 56.89 → 58.31 ms; 63-token TTFT 143.48 → 139.93 ms. Resident memory 42404 → 42404 MiB, 500 ms sampled peak 43282 → 43254 MiB. Shared, non-exclusive hardware and short runs: no significance claim; TP2 throughput was not remeasured.

Follow-up conditions, raw measurements and reproduction

Saved test output, commands and measurement provenance

The image is a rendering of saved local test output, not a GitHub CI screenshot. Detailed logs, commands, fixture configuration and data provenance accompany it.

Platform/model	Checks	Result / boundary
NVIDIA, 2 × A6000 48 GiB, tiny Qwen FP8-block / BF16 activations	MTP suite + Qwen2 linear pre-transpose/reprocessing regression	66 passed; K=1/2/4, forced acceptance lengths, TP inputs including strided tensors, vocabulary ties/uneven shards, lifecycle and graph recapture
NVIDIA, A6000 TP1	Core MTP suite	65 passed, 1 TP2-only teardown test skipped
NVIDIA, real Qwen3.8-27B-FP8, TP2 / K=2	Two packed requests; cancel one, admit another; close with active/queued requests	Exact tokens against stored same-TP ordinary-Decode graph references; all KV references and state rows reclaimed; 39/44 candidates accepted in the first packed run
NVIDIA, tiny checkpoint	examples/test_infer.py; examples/bench.py	Both passed with MTP enabled; offline smoke also checks a nondefault 32-token page size
NVIDIA, tiny checkpoint	test/bench/test_benchmark.py	Adapter smoke passed on 2 cached MMLU examples; random fixture scored 0/2, not an accuracy claim
NVIDIA, tiny checkpoint	HTTP server + scripts/test_perf.py	64/64 requests succeeded, 20 clients, server max batch 2; stream chunks are not a rigorous token-throughput counter
CPU-side checks	Scheduling/lifecycle and weight-remapping assertions; 6 additional existing Qwen MoE adaptation checks	Passed; neither full Qwen2 nor hybrid MTP inference was validated on CPU. The Qwen2 pre-transpose regression executes on NVIDIA.
Ascend / MetaX / Moore / other accelerators	Full MTP integration	Not validated: no usable allocated environment in this round. No full-path support claim. @wooway777 please review portability of the shared interfaces.

Fresh NVIDIA source build and project formatting checks passed. Integration used the matching Core runtime, including the separate graph-lifetime fix in InfiniTensor/InfiniCore#1560; it was not tested against the unmodified v0.2.9 release.

Benchmark / Performance Impact

The following are archived measurements before final PR cleanup, not new speed claims from the CLI smoke tests. They document the implemented path's benefit. The final branch was rechecked for real-model correctness/lifecycle; timings were not rebenchmarked during the concurrent source build.

2 × A6000 48 GiB over PCIe (no NVLink), TP2/PP1, dense single-MTP-layer Qwen3.8-27B-FP8, E4M3 128×128 weight blocks, BF16 activations, Marlin; 40 KV pages × 64 tokens; greedy; prefix reuse off. Ordinary baseline uses Decode graphs; K=2 MTP uses eager; both use eager Prompt Prefill. Two repetitions per prompt, including Python/scheduler time.

Prompt / output tokens	Ordinary Decode graph	K=2 MTP eager	Change
Chinese 63 / 64	33.91 tok/s	49.57 tok/s	+46.2%
Code 127 / 64	33.48 tok/s	53.29 tok/s	+59.2%
Summary 1023 / 55 (EOS)	30.03 tok/s	48.61 tok/s	+61.9%

Generation rate is sum(output_tokens - 1) / sum(wall_time - TTFT), excluding Prefill and loading. All 366 output tokens in those six runs matched their same-TP baseline. These small samples are not a production capacity estimate or a comparison against vLLM/A100.

TP1 two-request draft batching improved 58.18 → 62.87 tok/s (+8.06%). TP2 draft batching did not reliably improve wall time and remains disabled; target verification is still batched.

Fresh TP2 service check: steady process memory 24,150 / 24,132 MiB; sampled whole-process peak 24,258 MiB per card. nvidia-smi sampling every 100 ms can miss short peaks. This is eager K=2; no new full-27B graph recapture memory claim. State and snapshot budgets are separate from allocator/process memory.

Notes for Reviewers

The review scope is text-only greedy MTP on NVIDIA TP1/TP2, including the documented opt-in draft graphs and bounded exact-prompt snapshots. The ordinary-path follow-up is included in 86208ae1: ordinary short GDN Prefill retains chunked dispatch, graph recapture preserves live KV, and asynchronous shutdown retains in-flight resources on timeout. Vocabulary-parallel projection, corrected weight loading and checkpoint-configured FP32 recurrent state intentionally also affect ordinary Qwen inference; their behavior and memory implications are documented above. Existing-model numerical and lifecycle checks pass, but the validation does not establish unchanged performance for every model, PP/quantization configuration or vendor backend. Additional Qwen MoE checks cover configuration/remapping only, not full MoE inference.

This PR is ready for review within that scope. Merge still requires maintainer review, upstream CI with the matching Core runtime, and resolution of the dependency/overlap order below.

• Base is InfiniLM-v0.2.9c, as requested for this contribution, rather than main.
• Depend on issue/1565 fix(nvidia): complete runtime support for Qwen MTP InfiniCore#1566 for FP8/BOOL casts, communicator device preservation and head-256 Prefill. Graph mode additionally depends on fix(graph): preserve recurrent graph inputs and allocation lifetimes InfiniCore#1560 (currently a separate Draft PR); eager MTP does not require that graph-lifetime change.
• Only generic linear loading/capture-state protection overlaps feat(mamba2): integrate model loading and recurrent inference #575. No Mamba model, checkpoint converter or scan implementation is included; InfiniCore#1562 is not needed. Resolve this small shared overlap according to merge order.
• Review checkpoint ownership first: target verification writes distinct token destinations for both Conv and GDN; the scheduler commits the accepted prefix and reclaims scratch rows exactly once. Attention KV rollback follows the same accepted length.
• Supported scope: greedy text, dense, one shared MTP head, PP1, NVIDIA TP1/TP2; BF16 / aligned E4M3 block loading. K>1 or batching uses eager. Graph mode requires K=1 and max batch=1 and captures ordinary Decode / short draft shapes; target verification remains eager.
• Snapshot prefix reuse is opt-in, exact full-prompt and TP1 only. It copies into request-owned state, has its own byte budget and invalidates after reset/reload. It does not imply partial-prefix, distributed-state-transfer or SLRU support.
• No new attention algorithm/compiler or revival of experimental general Prefill graphs. Random sampling, MoE, multimodal MTP, PP, multi-layer heads and vendor-wide MTP remain out of scope.

CI / ChatOps

Fork push CI (format) and Ruff passed at the current head 86208ae1: https://github.com/big-hip/InfiniLM/actions/runs/35451703600 and https://github.com/big-hip/InfiniLM/actions/runs/35451703133. The hardware ci job was skipped in the fork push workflow; local GPU results are reported separately above.

Please run /retest with the matching Core dependency after assembling the runtime changes. The manual hardware workflow normally checks out InfiniTensor/InfiniCore by branch; the new runtime branch currently lives on the contributor fork. A formatting-only green check is not hardware CI coverage. Fork PR workflows may also require maintainer approval.

---

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; one feature commit and one focused compatibility fix, 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.

At the current head 86208ae1, upstream CI and Ruff report action_required; maintainer approval is required before their jobs run. Fork checks are reported separately above.