[NVIDIA] [WIP] MTP for dsr1 TRT without chat template

[ankursingh-nv]

Note

Introduces MTP-enabled TRT configs and end-to-end launch support for DeepSeek-R1 on NVIDIA single-node.

• Adds dsr1-fp4-b200-trt-mtp, dsr1-fp8-b200-trt-mtp, dsr1-fp8-h200-trt-mtp with spec-decoding: mtp, updated TP/EP/DP-attn search spaces, and notes for MTP sizing
• New benchmark scripts: benchmarks/dsr1_fp4_b200_trt_mtp_slurm.sh, dsr1_fp8_b200_trt_mtp_slurm.sh, dsr1_fp8_h200_trt_mtp_slurm.sh configuring MOE backend, DP attention, and MTP server options
• Runner scripts (runners/launch_*) detect SPEC_DECODING=mtp and select _mtp benchmark scripts via suffix
• benchmarks/benchmark_lib.sh: add optional --use-chat-template flag and refactor to build the benchmark command array
• Update perf-changelog.yaml to record new MTP TRT single-node configurations

^{Written by Cursor Bugbot for commit 5f34e19. This will update automatically on new commits. Configure here.}

[gemini-code-assist]

Summary of Changes

Hello @ankursingh-nv, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the benchmarking capabilities for DeepSeek-R1 models by integrating Multi-Turn Prediction (MTP) support within the TensorRT-LLM framework. It introduces new configuration profiles and specialized SLURM scripts for B200 and H200 GPUs, enabling performance evaluation across different precisions and operational parameters. The changes also refine the general benchmarking utility to accommodate diverse model requirements, such as optional chat template usage.

Highlights

• Multi-Turn Prediction (MTP) Benchmarking: New configurations and dedicated benchmark scripts have been added to support Multi-Turn Prediction (MTP) for DeepSeek-R1 (dsr1) models using TensorRT-LLM.
• Expanded GPU Support: MTP benchmarks are now integrated for both B200 and H200 GPUs, covering FP4 and FP8 precisions, with specific optimizations for different sequence lengths and concurrency levels.
• Dynamic MoE Backend and MTP Layer Configuration: The new benchmark scripts dynamically determine the MoE backend (CUTLASS or DEEPGEMM) and the number of MTP layers (1 or 3) based on whether Data Parallel Attention (DP_ATTENTION) is enabled.
• Benchmark Utility Enhancement: The core run_benchmark_serving utility has been updated to include an optional --use-chat-template flag, providing greater flexibility for various benchmarking scenarios.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces MTP (speculative decoding) support for DeepSeek-R1 TRT on NVIDIA hardware. It adds new configurations, benchmark scripts, and updates runner scripts accordingly. The changes are generally well-structured, but there are a few areas for improvement in the new benchmark scripts to enhance robustness and prevent potential race conditions. I've also noted a change in the perf-changelog.yaml that appears to be unrelated to this PR's scope.

I am having trouble creating individual review comments. Click here to see my feedback.

benchmarks/dsr1_fp4_b200_trt_mtp_slurm.sh (35)

The script uses a hardcoded filename dsr1-fp4-mtp.yml for the temporary configuration. If multiple instances of this script run concurrently in the same directory, they could overwrite each other's configuration files, leading to race conditions and incorrect benchmark results. It's safer to use mktemp to create a unique temporary file.

Consider also adding a trap at the beginning of the script to ensure the temporary file is cleaned up on exit:
trap 'rm -f "$EXTRA_CONFIG_FILE"' EXIT

EXTRA_CONFIG_FILE=$(mktemp /tmp/dsr1-fp4-mtp.XXXXXX.yml)

benchmarks/dsr1_fp8_b200_trt_mtp_slurm.sh (35)

The script uses a hardcoded filename dsr1-fp8-mtp.yml. This filename is also used in benchmarks/dsr1_fp8_h200_trt_mtp_slurm.sh. If these scripts were ever to run in the same working directory, they would conflict. To prevent race conditions and ensure correctness, a unique temporary file should be created using mktemp.

Consider also adding a trap at the beginning of the script to ensure the temporary file is cleaned up on exit:
trap 'rm -f "$EXTRA_CONFIG_FILE"' EXIT

EXTRA_CONFIG_FILE=$(mktemp /tmp/dsr1-fp8-mtp.XXXXXX.yml)

benchmarks/dsr1_fp8_h200_trt_mtp_slurm.sh (35)

The script uses a hardcoded filename dsr1-fp8-mtp.yml. This filename is also used in benchmarks/dsr1_fp8_b200_trt_mtp_slurm.sh. If these scripts were ever to run in the same working directory, they would conflict. To prevent race conditions and ensure correctness, a unique temporary file should be created using mktemp.

Consider also adding a trap at the beginning of the script to ensure the temporary file is cleaned up on exit:
trap 'rm -f "$EXTRA_CONFIG_FILE"' EXIT

EXTRA_CONFIG_FILE=$(mktemp /tmp/dsr1-fp8-mtp.XXXXXX.yml)

benchmarks/dsr1_fp4_b200_trt_mtp_slurm.sh (46)

There is a trailing space on this line. While it may not affect YAML parsing for a boolean value, it's best to remove it for cleanliness and to avoid potential issues with stricter parsers.

    enable_block_reuse: false

benchmarks/dsr1_fp4_b200_trt_mtp_slurm.sh (74-83)

The mpirun command is missing the PYTHONNOUSERSITE=1 prefix, which is present in benchmarks/dsr1_fp8_h200_trt_mtp_slurm.sh. This variable prevents user-site packages from being used, which can help avoid environment-related issues and improve reproducibility. For consistency and robustness, it should be added here as well.

PYTHONNOUSERSITE=1 mpirun -n 1 --oversubscribe --allow-run-as-root \
    trtllm-serve $MODEL --port=$PORT \
    --trust_remote_code \
    --backend=pytorch \
    --max_batch_size=$MAX_BATCH_SIZE \
    --max_seq_len=$MAX_MODEL_LEN \
    --max_num_tokens=$MAX_NUM_TOKENS \
    --tp_size=$TP --ep_size=$EP_SIZE \
    --extra_llm_api_options=$EXTRA_CONFIG_FILE \
    > $SERVER_LOG 2>&1 &

benchmarks/dsr1_fp8_b200_trt_mtp_slurm.sh (46)

There is a trailing space on this line. While it may not affect YAML parsing for a boolean value, it's best to remove it for cleanliness and to avoid potential issues with stricter parsers.

    enable_block_reuse: false

benchmarks/dsr1_fp8_b200_trt_mtp_slurm.sh (74-83)

The mpirun command is missing the PYTHONNOUSERSITE=1 prefix, which is present in benchmarks/dsr1_fp8_h200_trt_mtp_slurm.sh. This variable prevents user-site packages from being used, which can help avoid environment-related issues and improve reproducibility. For consistency and robustness, it should be added here as well.

PYTHONNOUSERSITE=1 mpirun -n 1 --oversubscribe --allow-run-as-root \
    trtllm-serve $MODEL --port=$PORT \
    --trust_remote_code \
    --backend=pytorch \
    --max_batch_size=$MAX_BATCH_SIZE \
    --max_seq_len=$MAX_MODEL_LEN \
    --max_num_tokens=$MAX_NUM_TOKENS \
    --tp_size=$TP --ep_size=$EP_SIZE \
    --extra_llm_api_options=$EXTRA_CONFIG_FILE \
    > $SERVER_LOG 2>&1 &

benchmarks/dsr1_fp8_h200_trt_mtp_slurm.sh (51)

There is a trailing space on this line. While it may not affect YAML parsing for a boolean value, it's best to remove it for cleanliness and to avoid potential issues with stricter parsers.

    enable_block_reuse: false

perf-changelog.yaml (160-164)

This addition to the changelog appears to be unrelated to the current pull request, which is focused on adding MTP support for DSR1 TRT. The changelog entry refers to updating an SGLang image for MI355x and references PR #395. It's generally better to keep pull requests focused on a single logical change. Please consider moving this to a separate PR.