[None][feat] Support DeepSeek-V4 in layer_wise_benchmarks

[ruodil]

DeepSeek-V4's multi-head hyper-connection (mHC) decoder layers do not fit the generic single-layer harness in tools/layer_wise_benchmarks/runner.py:

• create_kv_cache_manager: the is_mla branch omitted vocab_size, which DeepseekV4CacheManager.init requires (keyword-only, no default).
• create_run_pack: V4 layers take the initial residual as hc_state shaped [num_tokens, hc_mult, hidden_size] and their MoE routing needs input_ids; both are now synthesized when the config exposes hc_mult.
• run_pack(check=True): guarded the NaN/Inf/zero check with isinstance since mHC layers return an HCState, not a Tensor.

Non-V4 models are unaffected (all V4 paths gate on hc_mult / MLA class).

Usage for V4: --moe-backend DEEPGEMM and a --layer-indices slice that starts at layer 0 and covers both compress ratios (e.g. 0,1,2,3).

Dev Engineer Review

• Added DeepSeek-V4 support in runner.py by:
  • Passing vocab_size to the MLA KV-cache manager.
  • Synthesizing hc_state and input_ids for mHC layers.
  • Restricting NaN/Inf/zero checks to tensor outputs.
  • Gating V4-specific behavior to preserve non-V4 model behavior.
• V4 benchmarks require --moe-backend DEEPGEMM and layer indices beginning at 0 and covering both compression ratios, e.g. 0,1,2,3.
• No public API or configuration-file changes were identified.

QA Engineer Review

No test changes.

Description

Test Coverage

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.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[coderabbitai]

Walkthrough

The layer-wise benchmark runner now prepares DeepSeek-V4 hyper-connection inputs, conditionally validates tensor outputs, and passes vocabulary size when constructing MLA KV-cache managers.

Changes

Benchmark runner behavior

Layer / File(s)	Summary
DeepSeek-V4 input preparation tensorrt_llm/tools/layer_wise_benchmarks/runner.py	Benchmark pack creation reshapes hidden states for configured hc_mult values and adds input_ids for MoE routing.
Validation and MLA cache wiring tensorrt_llm/tools/layer_wise_benchmarks/runner.py	Output checks now require tensor-valued hidden states, and MLA KV-cache manager construction receives config.vocab_size.

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

Suggested reviewers: dc3671

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title matches the PR’s main change and follows the required [None][feat] format.
Description check	✅ Passed	The description clearly explains the issue and solution and is mostly complete despite lacking a dedicated test coverage section.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

🧹 Nitpick comments (1)

tensorrt_llm/tools/layer_wise_benchmarks/runner.py (1)

664-678: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add a regression test for the V4 input contract.

Cover hc_mult=4 and a non-V4 configuration, asserting the expanded hidden-state shape and generated input_ids dtype/range. DeepSeek-V4 relies on expanded mHC residual streams and token-ID-based hash routing. (nvidia.github.io)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/tools/layer_wise_benchmarks/runner.py` around lines 664 - 678,
Add a regression test covering the input-preparation logic around the hc_mult
handling, using hc_mult=4 and a non-V4 configuration. Assert that V4 hidden
states expand to [num_tokens, 4, hidden_size], generated input_ids use
torch.int32 and remain within [0, vocab_size), while the non-V4 path preserves
the original hidden-state shape and does not add input_ids.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@tensorrt_llm/tools/layer_wise_benchmarks/runner.py`:
- Around line 664-678: Add a regression test covering the input-preparation
logic around the hc_mult handling, using hc_mult=4 and a non-V4 configuration.
Assert that V4 hidden states expand to [num_tokens, 4, hidden_size], generated
input_ids use torch.int32 and remain within [0, vocab_size), while the non-V4
path preserves the original hidden-state shape and does not add input_ids.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: a39b3908-7632-403f-a51a-3f1331349703

📥 Commits

Reviewing files that changed from the base of the PR and between 0a401b4 and 263b526.

📒 Files selected for processing (1)

• tensorrt_llm/tools/layer_wise_benchmarks/runner.py

[ruodil]

/bot run

[tensorrt-cicd]

PR_Github #61203 [ run ] triggered by Bot. Commit: 263b526 Link to invocation

[tensorrt-cicd]

PR_Github #61203 [ run ] completed with state FAILURE. Commit: 263b526
/LLM/main/L0_MergeRequest_PR pipeline #49447 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