[None][fix] Announce MTP shapes to attention metadata in layer-wise benchmarks

[dc3671]

Description

seq_len_q > 1 in the layer-wise benchmarks means MTP: each request submits 1 + num_draft tokens. In serving, the executor announces that via update_spec_dec_param() — the only place max_draft_tokens is set on the attention metadata. The harness never called it, so max_draft_tokens stayed at its default of 0 while create_run_pack() submitted batch_size * seq_len_q tokens.

The DSA indexer's 2D context_lens buffer (kv_lens_cuda_2d, shaped (max_num_sequences, 1 + max_draft_tokens)) therefore stayed one column wide. The call site already derives the real width locally (dsa.py: next_n = num_gen_tokens // num_generations, then slices kv_lens_cuda_2d[:num_generations, :next_n]), so slicing 4 columns out of a 1-column buffer yielded 1, and DeepGEMM aborted in csrc/apis/attention.hpp:

DG_HOST_ASSERT(batch_size == __batch_size and next_n == _next_n)

Concretely, batch_size 32 / seq_len_q 4 submitted 128 tokens against 32 rows of KV length and never ran an iteration.

Fix: call update_spec_dec_param() for DSA generation packs with seq_len_q > 1, so the DSA override rebuilds kv_lens_cuda_2d, the expanded MTP buffers and the radix/heuristic scratch at the right widths. Masking stays disabled (is_spec_decoding_enabled=False) — only shapes are announced.

The guard is hasattr(attn_metadata, "kv_lens_cuda_2d") rather than the method name: update_spec_dec_param lives on the base metadata class, so a hasattr check on the method would admit every backend, and the base sets max_total_draft_tokens unconditionally — a value that reaches the attention op cache key and FMHA/XQA kernel selection. kv_lens_cuda_2d exists only on DSA metadata, so non-DSA runs are untouched by construction. On DSA the value now matches what the executor sets for a real MTP step.

Longer term the cleaner shape is for max_draft_tokens to be a declared field on DSAMetadataParams, alongside the seven other construction-time settings __post_init__ already copies before allocating the indexer buffers — then no announcement or rebuild is needed at all. That is a runtime change under different ownership and is deliberately not in this PR.

Test Coverage

No new unit test: the layer-wise benchmarks are a manual profiling harness with no CI target, and reproducing this needs 8 GB200 ranks plus a DeepSeek-V4 checkpoint. The change cannot execute for anything CI exercises (seq_len_q == 1, CTX packs, and any non-DSA backend).

Verified on GB200, 8 ranks / 2 nodes, DeepSeek-V4-Pro, MEGAMOE_DEEPGEMM, kv_cache_dtype fp8:

config	before	after
batch_size 32, seq_len_q 4, seq_len_kv_cache 2049	aborts in attention.hpp	1985.4 us mean
batch_size 32, seq_len_q 1, seq_len_kv_cache 2049	1628 us	1628 us (unchanged)

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

Dev Engineer Review

Runner.create_run_pack now updates speculative-decoding shape metadata only for multi-token GEN runs with DSA metadata. It sets draft-token widths while keeping speculative decoding and tree modes disabled. This addresses DSA buffer sizing for seq_len_q > 1.

The implementation should narrow the guard to DSA metadata. The base metadata class also defines update_spec_dec_param(), so a hasattr check can modify max_total_draft_tokens for non-DSA backends. This value can affect attention configuration, cache keys, and FMHA/XQA kernel selection.

No configuration files or test-list files changed.

QA Engineer Review

No test changes.

[coderabbitai]

Walkthrough

Changes

Generation metadata

Layer / File(s)	Summary
Configure multi-token generation metadata tensorrt_llm/tools/layer_wise_benchmarks/runner.py	When a GEN run has multiple query tokens and kv_lens_cuda_2d metadata, create_run_pack sets draft and total draft lengths to seq_len_q - 1. Speculative decoding and both tree modes remain disabled.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#16925: Updates token-to-request mappings in DSA attention backends while this PR updates DSA speculative-decoding metadata in runner.py.

Suggested reviewers: hyukn, chenfeiz0326

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly identifies the fix for announcing MTP shapes in layer-wise benchmarks and follows the required format.
Description check	✅ Passed	The description includes all required sections, explains the issue and solution, documents test coverage, and records manual verification results.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[BowenFu]

The DSA half checks out: update_spec_dec_param is on the base interface (interface.py:426-437), the argument names match all three implementations, and the DSA override rebuilds kv_lens_cuda_2d (dsa.py:1205-1208), the expanded buffers (:1209-1211), the heuristic scratch (:1212-1222) and the radix aux scratch (:1223-1229) unconditionally — none of that is gated on is_spec_decoding_enabled, so passing False really does get you shapes without masking. Guards are right too: run_type is a create_run_pack() parameter, and CTX / seq_len_q == 1 are untouched.

One thing before this lands: the call is not inert on non-DSA backends. hasattr is always true — the method is on the base metadata class — so a GEN pack with seq_len_q > 1 on TrtllmAttentionMetadata now takes this path too. The allocation block is gated on is_spec_decoding_enabled (trtllm.py:1010), so nothing is allocated, but self.max_total_draft_tokens = seq_len_q - 1 is set unconditionally at trtllm.py:1007, and that value reaches the attention op (fmha/fallback.py:100), is stored at attentionOp.cpp:1236-1238, is part of the op cache key (attentionOp.h:573-575), and is copied into the FMHA/XQA fixed params (attentionOp.cpp:3241-3246).

So an existing TRTLLM-backend layer-wise run at seq_len_q > 1 can now select a different attention configuration than the numbers it was producing yesterday. For a profiling harness that is the one kind of change that matters, and the PR describes itself as inert. Either narrow the guard to the backend that needs it rather than hasattr, or say in the description that non-DSA MTP benchmark numbers move and why that is the more correct baseline.

Not blocking on the missing unit test — the reasoning for skipping it is sound.

[dc3671]

Good catch, and the trace is right — I checked each link: fallback.py:100 passes metadata.max_total_draft_tokens ungated, thop/attentionOp.cpp:1268 turns it into mSpecDecodingTargetMaxGenLen, that member is in the data() cache-key tuple at attentionOp.h:584, and it reaches kernel selection via attentionOp.cpp:3275 → xqaDispatcher.cpp:285,536 → fmhaKernels.h:501-502. The hasattr guard was useless — I had assumed the method was DSA-only; it is on the base metadata class.

Pushed a fix: the guard is now hasattr(attn_metadata, "kv_lens_cuda_2d"), which exists only on DSA metadata, so non-DSA backends never enter the call and max_total_draft_tokens stays None for them.

One amendment to your framing, in the interest of not overclaiming in the other direction: the DSA override calls super().update_spec_dec_param(...) first, so DSA metadata gets that same assignment. Narrowing the guard confines the effect to the backend under test rather than eliminating it. I think that is correct for DSA — the executor sets the same value for a real MTP step, so the benchmark now matches serving instead of diverging from it — but "inert" was wrong and I have removed that claim from the description.

On the deeper point: the harness should not have to announce this at all. The consumer already derives the real width locally (dsa.py: next_n = num_gen_tokens // num_generations, then slices kv_lens_cuda_2d[:num_generations, :next_n]); only the buffer's capacity comes from the separately announced max_draft_tokens. The natural fix is a declared max_draft_tokens field on DSAMetadataParams, beside the seven other construction-time settings __post_init__ already copies before calling create_buffers_for_indexer — then the buffer is born the right width, update_spec_dec_param keeps serving's runtime-change path, and no rebuild is needed. I left that out of this PR since it is a runtime change under trt-llm-torch-attention-devs rather than a harness fix, and noted it in the description. Happy to open it separately if you agree it is the right shape.

[dc3671]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #63775 [ run ] triggered by Bot. Commit: 5bc9c09 Link to invocation

[tensorrt-cicd]

PR_Github #63775 [ run ] completed with state SUCCESS. Commit: 5bc9c09
/LLM/main/L0_MergeRequest_PR pipeline #51724 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

[dc3671]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #63885 [ run ] triggered by Bot. Commit: 5bc9c09 Link to invocation

[tensorrt-cicd]

PR_Github #63885 [ run ] completed with state SUCCESS. Commit: 5bc9c09
/LLM/main/L0_MergeRequest_PR pipeline #51825 completed with status: 'SUCCESS'

CI Report

Link to invocation