Run chained MTP drafting in the batch Engine

[Tianlei Wu (tianleiwu)]

Summary

• run the auxiliary MTP decoder transactionally beside the target decoder
• support per-request draft lengths and compact active requests between head stages
• keep chained draft IDs on device and perform one final host materialization
• expose cumulative Engine speculative telemetry through C, C++, and Python
• defer drafts for a selected transaction that mixes prefill and speculative rows, then resume MTP on decode-only transactions

Stack

PR 6/6. Depends on tlwu/20260823/gdn-paged-mtp-config.

Correctness fix

Packed recurrent operators choose one execution plan for the whole ragged batch. A long prefill can select chunked GatedDeltaNet, which publishes final state but not the intermediate recurrent checkpoints required to reject drafts. The Engine previously treated those checkpoint outputs as valid, so a rejected request could publish an uninitialized checkpoint and diverge on its next target call.

After paged selection, the composite cache planner now clears draft counts only when the selected transaction contains both prefill and draft rows. Decode requests still advance by one token, prefills continue, and K=3 resumes when the selected transaction is decode-only.

Validation

• engine_unit_tests: 323/323 passed
• focused CompositeDefersDraftsWhilePrefillSharesTheStep regression passed
• clang-format dry run and git diff --check passed
• post-fix layer tracing is exact for 2,048/2,048 rows in transaction 1 and 1,803/1,803 rows in transaction 2
• batch-1 token parity: 438/438; 49.02 to 91.93 tok/s (1.88x)
• sustained concurrency-16 token parity: 2,004/2,004; 288.55 to 316.00 tok/s (+9.5%)
• seven-shard MMLU-Pro: true plain 664/800 (83.00%), post-fix K=3 666/800 (83.25%), zero unextracted
• seven-shard GPQA: true plain 147/198 (74.24%), post-fix K=3 151/198 (76.26%), zero unextracted
• weighted accepted/evaluated draft rate: 84.98% MMLU-Pro and 84.73% GPQA

Release disposition

K=3 MTP is quality-safe for continuous batching with this safeguard. The prior batch-1-only disposition and numerical-shape explanation are superseded by the recurrent-checkpoint rollback root cause above. Keep the stack in draft pending CI.

[Tianlei Wu (tianleiwu)]

Stack: #2454 → #2455 → #2456 → #2457 → #2458 → 6/6

[Tianlei Wu (tianleiwu)]

H200 RC4 evidence (corrected true-plain control)

Correctness: engine_unit_tests 322/322 passed; clang-format and git diff --check passed. Focused coordinator coverage includes mixed K, acceptance outcomes, rollback, max-length clamping, cache reuse, and one final draft readback.

Controlled batch-1 throughput (1,024 generated tokens, K=3, medians): paged Engine versus dense RC2 eager is +2.96%, -5.08%, -3.63%, +1.97%, and -0.73% at 512, 2K, 8K, 32K, and 256K total context, respectively (geomean -0.95%). Versus pre-captured RC2 graphs the geomean is -6.33%; hybrid fixed-state Engine graph capture is not available yet.

Quality gate (thinking, greedy, 8,192-token cap, historical seven-shard geometry):

arm	MMLU-Pro	GPQA
true plain primary	664/800 (83.00%), 0 unextracted	147/198 (74.24%), 0 unextracted
Engine MTP K=1 diagnostic	638/800 (79.75%), 0 unextracted	150/198 (75.76%), 3 unextracted
Engine MTP K=3 original	594/800 (74.25%), 19 unextracted	139/198 (70.20%), 10 unextracted
prior RC3 paged	666/800 (83.25%)	149/198 (75.25%)

The earlier result labeled plain was not plain: omitting request-level speculative options left model-configured MTP enabled. The corrected control proves the checkpoint-enabled target graph preserves historical quality when run without MTP.

Disposition: the primary model configuration omits model.mtp. K=3 is retained in a sibling latency artifact with max_batch_size=1, where the 64-prompt diagnostic scores 54/64, equal to true plain. Keep the stack draft pending CI; do not add a generic batched-MTP ban.

[Tianlei Wu (tianleiwu)]

Quality follow-up (corrected)

The result previously labeled plain still had model-configured MTP enabled. A true no-MTP control scores 54/64 on the focused diagnostic and 664/800 on MMLU-Pro, matching RC2/RC3 quality. K=1, K=2, and K=3 score 53/64, 51/64, and 48/64 at concurrency 16; K=3 returns to 54/64 at concurrency 1.

Transactional state/checkpoint publication, verification-row selection, per-request bonus sampling, FP4/FP8 GEMV dispatch, paged attention selection, and isolated packed GatedDeltaNet state/checkpoints were tested and do not explain the divergence. The remaining mechanism is full-graph batch-shape numerical sensitivity during speculative target verification: mathematically valid target rows can follow different greedy trajectories in this quantized model.

Release fix: use the true-plain configuration as the batch primary and a separate K=3, max_batch_size=1 configuration for latency. The Engine implementation remains unchanged.

[Tianlei Wu (tianleiwu)]

Rollback fix and corrected quality disposition

This supersedes the earlier quality follow-ups in this thread. The continuous-batch regression was not inherent quantized batch-shape sensitivity: mixed prefill/speculative transactions could select chunked GatedDeltaNet, which does not populate the intermediate recurrent checkpoints needed for draft rejection. Commit c2ee5a37b2 defers drafts for only those mixed transactions and resumes K=3 on decode-only steps.

Post-fix validation: 323/323 Engine tests; exact 438/438 batch-1 and 2,004/2,004 sustained-c16 tokens versus plain; MMLU-Pro 666/800 versus plain 664/800; GPQA 151/198 versus plain 147/198; zero unextracted answers in both. Batch-1 decode reaches 91.93 versus 49.02 tok/s (1.88x), and sustained c16 reaches 316.00 versus 288.55 tok/s (+9.5%). K=3 is therefore quality-safe for continuous batching with the safeguard.

[Tianlei Wu (tianleiwu)]

Superseded by #2495, which rebuilds the MTP batch-engine runtime on compact fixed-state replay and composite transactions without dense recurrent checkpoints or an ONNX Runtime dependency bump.