[Tabrizian]

[None][perf] MTP draft-step glue trim (TRTLLM_MTP_TAIL_TRIM, default on; "0" = original path)

Description

In the one-model MTP-Eagle draft loop a sizeable share of the per-step draft tail is host/launch glue rather than model compute: a single-CTA torch.max reduce over each rank's vocab shard followed by cast/offset/stack kernels, an NCCL allgather of the per-rank (index, value) pairs, index-gather and int32->int64 copy kernels for hidden states / position ids that from draft step 1 on select arange(batch), DeepGEMM MQA-logits schedule rebuilds (schedule launches + KV-length copies + GVR row order) before draft passes that never read them, and a duplicated on_update_kv_lens() at the step boundary under the overlap scheduler. This PR removes that glue without changing any numeric result.

What changed:

• Draft argmax pack (tensorrt_llm/_torch/speculative/mtp_tail_trim.py, new): one Triton launch computes each rank's (global argmax index, max value) pair and writes it straight into the exchange row, replacing torch.max(dim=-1) + int32 offset + two float casts + stack. Same tie-break (lowest index) and NaN rule (NaN wins, lowest NaN index) as at::native GreaterOrNan; bf16 is widened bit-exactly in-kernel and a NaN max is rounded back to the same 0x7FFF bf16 payload torch produces.
• Dedicated MNNVL mailbox (same file): the pairs are exchanged through a one-shot MNNVL SUM allreduce over a dedicated Lamport workspace (each rank fills its own slot of a zero row, so the fp32 sum is exact) instead of the NCCL allgather. The workspace is private because the one-shot kernel clears the previous call's dirty Lamport buffer with the current launch's threads: on the model's shared workspace a one-row exchange would spend its few threads clearing the payload left by the preceding model allreduce. The mailbox packs into a persistent staging buffer and is built with the model workspace's own helpers, which this PR factors out of distributed/ops.py as pure extractions (_build_mnnvl_workspace: communicator split, multicast buffer, converged construction, protocol init; _launch_mnnvl_allreduce: the kernel launch MNNVLAllReduce.forward now also calls), so it works under both the MPI and the Ray (process-group) orchestrators and a failed build hands its communicator back the same way the model workspace does. NCCL allgather remains the path off an MNNVL fabric, under non-pure-TP worlds (PP/CP/attention-DP), above the one-shot payload, and when the exchange is first reached under CUDA-graph capture; the construction result is converged across the group so every rank takes the same path.
• Draft loop slices (eagle3.py): from draft step 1 on gather_ids is batch_indices_cuda[:batch_size] == arange(batch_size), so one row_sel per step turns the hidden-state and position-id row gathers into prefix slices; step 0 selects the hidden-state rows once for both the LM head input and the next step; the last step skips the selections nothing consumes. With the gate off every site keeps its own hidden_states[gather_ids] gather, launch-for-launch as before.
• DSA schedule skip (dsa/metadata.py, dsa/indexer.py, sparse/params.py, eagle3.py): update_for_spec_dec before an index-shared draft step (the indexer takes the reuse-TopK path) and after the last draft step no longer rebuilds the DeepGEMM MQA-logits schedule, the 2D/expanded indexer KV lengths or the GVR row order; slot mappings / indptrs are still refreshed. The metadata marks the schedule stale and the indexer raises if a TopK pass ever reads a stale schedule. The draft loop reaches this through a small runtime-checkable Protocol (MTPIndexerScheduleTrimMetadata: set_skip_indexer_sched_metadata, mtp_next_step_reuses_topk), so it stays decoupled from the DSA backend; whether the following steps reuse TopK is computed once after draft step 0.
• Step boundary (model_engine.py, attention/backends/interface.py): with overlap scheduling + spec decode, _preprocess_inputs calls on_update_kv_lens() twice per step and the first call's work is fully overwritten by the second after the kv_lens correction when the hook is a pure function of the current seq/kv lengths. Metadata classes opt in through an explicit kv_lens_hook_idempotent class attribute (True on the audited DSA metadata, False on the base class and on the DeepSeek-V4 metadata, whose override has not been audited); the first call is skipped for those only.

Gate: TRTLLM_MTP_TAIL_TRIM (unset / 1 = trimmed path, 0 = original path, read once and cached so the choice is fixed before CUDA-graph capture). The default is on; TRTLLM_MTP_TAIL_TRIM=0 gives a same-image A/B. Existing kernels and the non-MTP Eagle3 path are untouched.

Numerics: bit-identical. The pack kernel reproduces _get_local_max_and_combined bit for bit (ties -> lowest index, NaN wins with the lowest NaN index, bf16 -inf / NaN payloads preserved); the mailbox's fp32 sum of one non-zero slot per rank is exact, the only difference being that a NaN max returns with the canonical NaN payload, which the consuming argmax treats identically. The slice and schedule-skip trims feed every draft step the same tensors as before.

Measured (GB300, NVFP4 MoE model, TP8, bs1, MTP depth 5, low-latency disaggregated decode point, decode step ~9 ms; trim on vs TRTLLM_MTP_TAIL_TRIM=0, same image):

run	device step p50	tok/s/user p50 / p10
gen-only microbenchmark, off	9.05 ms	389.7 / 376.8
gen-only microbenchmark, on	8.71 ms (-0.34 ms)	406.6 / 395.6
paired 900 s serving run, off	9.05 ms	395.5 / 370.0
paired 900 s serving run, on	8.74 ms (-0.31 ms)	409.9 / 386.4

Every percentile moved in favour of the trimmed arm; 0 request errors in each arm; identical draft acceptance between arms.

Measurements were taken on a 1.3.0rc28-based branch; this PR is the change rebuilt on main's MNNVL workspace helpers with the draft-loop hook expressed against the MTPIndexShareMetadata / MTPIndexerScheduleTrimMetadata protocols.

Follow-ups (not in this PR): checkpoint_prepare / checkpoint_restore do not yet cover the mailbox's dedicated workspace (the model workspace registry is keyed by mapping only); the ADP + LM-head-TP draft-argmax branch still takes the original _get_local_max_and_combined + allgather path.

Test Coverage

• tests/unittest/_torch/speculative/test_mtp_tail_trim.py (new, GPU):
  • test_local_argmax_pack_matches_torch: pack kernel vs _get_local_max_and_combined bitwise over bf16 / fp16 / fp32, vocab shards 77 .. 151936, rows 1 .. 9, with all-equal rows, ties, all -inf, NaN and +inf columns, at two rank offsets, in both the [rows, 2] and the mailbox [rows, 128] slot layouts.
  • test_local_argmax_pack_strided_rows_and_cuda_graph: non-contiguous row stride and CUDA-graph capture + replay on fresh inputs.
  • test_draft_loop_trim_is_bit_identical: a mocked MTP-Eagle draft loop with the trim on vs off for draft lengths 1 / 2 / 5 must produce the same draft tokens, feed every draft step the same input_ids / position_ids / hidden_states, and leave the same kv_lens after each step (runs on CPU as well).
  • test_mailbox_matches_nccl_allgather (2 GPUs, mpi_pool_executor): dedicated mailbox vs NCCL allgather bitwise and through the real _get_draft_tokens_from_gathered, across row counts including the max_rows + 1 NCCL fallback, once directly after a model-sized MNNVL allreduce on the shared workspace (the dedicated-workspace regression check), and under CUDA-graph replay of model allreduce + pack + mailbox on fresh logits.
• tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py::test_indexer_decode_mtp_topk_reuse: new stale_sched parameter runs the index-shared reuse steps with the MQA-logits schedule marked stale and poisoned (-1 fill), asserts the reuse copies the stash 1:1, and asserts a step that computes TopK on the stale schedule raises.

PR Checklist

• PR description clearly explains what and why.
• Follows the TRT-LLM coding guidelines to the best of my knowledge.
• Test cases are provided for the new code path.
• No API changes; no new dependencies; no documentation change needed.
• Please check this after reviewing the above items as appropriate for this PR.

[Tabrizian]

/bot run

[tensorrt-cicd]

PR_Github #75635 [ run ] triggered by Bot. Commit: 2897ca6 Link to invocation

[Tabrizian]

/bot run

[Tabrizian]

/bot run

[tensorrt-cicd]

PR_Github #75643 [ run ] triggered by Bot. Commit: 4780f93 Link to invocation

[tensorrt-cicd]

PR_Github #75644 [ run ] triggered by Bot. Commit: 4780f93 Link to invocation

[tensorrt-cicd]

PR_Github #75643 [ run ] completed with state ABORTED. Commit: 4780f93

Link to invocation

[tensorrt-cicd]

PR_Github #75635 [ run ] completed with state ABORTED. Commit: 2897ca6

Link to invocation

[tensorrt-cicd]

PR_Github #75644 [ run ] completed with state SUCCESS. Commit: 4780f93
/LLM/main/L0_MergeRequest_PR pipeline #62354 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[Tabrizian]

/bot run

[Tabrizian]

/bot run

[tensorrt-cicd]

PR_Github #75674 [ run ] triggered by Bot. Commit: 325ec90 Link to invocation

[tensorrt-cicd]

PR_Github #75675 [ run ] triggered by Bot. Commit: 325ec90 Link to invocation

[tensorrt-cicd]

PR_Github #75674 [ run ] completed with state ABORTED. Commit: 325ec90

Link to invocation

[Tabrizian]

/bot run

[tensorrt-cicd]

PR_Github #75684 [ run ] triggered by Bot. Commit: 325ec90 Link to invocation

[tensorrt-cicd]

PR_Github #75675 [ run ] completed with state ABORTED. Commit: 325ec90

Link to invocation

[tensorrt-cicd]

PR_Github #75684 [ run ] completed with state SUCCESS. Commit: 325ec90
/LLM/main/L0_MergeRequest_PR pipeline #62388 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[Tabrizian]

/bot run

[tensorrt-cicd]

PR_Github #75711 [ run ] triggered by Bot. Commit: 325ec90 Link to invocation

[tensorrt-cicd]

PR_Github #75711 [ run ] completed with state SUCCESS. Commit: 325ec90
/LLM/main/L0_MergeRequest_PR pipeline #62414 completed with status: 'UNSTABLE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Either:

• Wait for the PR to be fully approved — the label is added automatically once approval is complete. Having unresolved open comments is fine, or
• If needed, ask a member of NVIDIA/trt-llm-ci-approvers to add the label manually.
  Then re-trigger CI with the same bot command (no rebase needed).

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[Tabrizian]

Superseded: split into three self-contained PRs for review, each based independently on main:

• [None][perf] MTP draft loop: exchange the TP argmax pairs with one Triton pack + an AllReduce SUM #19787 MTP draft loop: exchange the TP argmax pairs with one Triton pack + an AllReduce SUM (the exchange now goes through the generic AllReduce module so it follows the configured allreduce strategy, with a tagged MNNVL workspace instead of a private one)
• [None][perf] MTP draft loop: select rows with a prefix slice from draft step 1 #19785 MTP draft loop: prefix slices instead of gathers from draft step 1
• [None][perf] DSA: skip the indexer schedule rebuild for index-sharing draft steps and the duplicated kv-lens hook #19786 DSA + MTP: skip the dead indexer schedule rebuilds before index-shared draft steps and the duplicated kv-lens hook

Closing this one.