From 444c823a3169dc9769a0dae11e3a293520e8406a Mon Sep 17 00:00:00 2001
From: Camilo Moreno Ramirez <camilom@nvidia.com>
Date: Wed, 4 Feb 2026 17:23:19 -0800
Subject: [PATCH 1/3] update nvidia master and benchmark options

---
 .github/configs/nvidia-master.yaml  | 51 ++++++++++++-----------------
 benchmarks/dsr1_fp4_b200_trt_mtp.sh | 42 +++++++++++++++++++-----
 2 files changed, 55 insertions(+), 38 deletions(-)

diff --git a/.github/configs/nvidia-master.yaml b/.github/configs/nvidia-master.yaml
index a5cad52066..935701af72 100644
--- a/.github/configs/nvidia-master.yaml
+++ b/.github/configs/nvidia-master.yaml
@@ -865,7 +865,7 @@ dsr1-fp4-b200-trt:
     - { tp: 8, ep: 8, dp-attn: true, conc-start: 128, conc-end: 256 }
 
 dsr1-fp4-b200-trt-mtp:
-  image: nvcr.io#nvidia/tensorrt-llm/release:1.1.0rc2.post2
+  image: nvcr.io#nvidia/tensorrt-llm/release:1.2.0rc6.post3
   model: nvidia/DeepSeek-R1-0528-FP4-V2
   model-prefix: dsr1
   runner: b200-trt
@@ -876,44 +876,35 @@ dsr1-fp4-b200-trt-mtp:
   - isl: 1024
     osl: 1024
     search-space:
-    # If TP=4:
-    #   If CONC >= 16, then EP=4
-    #   If CONC >= 128, DP_ATTN=true, MOE_BACKEND=CUTLASS, MTP=1
+    # TP=4 configurations
     - { tp: 4, conc-start: 4, conc-end: 8, spec-decoding: mtp }
-    - { tp: 4, ep: 4, conc-start: 16, conc-end: 64, spec-decoding: mtp }
-    - { tp: 4, ep: 4, dp-attn: true, conc-start: 128, conc-end: 256, spec-decoding: mtp }
-    # If TP=8:
-    #   If CONC >= 16, then EP=8
-    #   If CONC >= 64, DP_ATTN=true, MOE_BACKEND=CUTLASS, MTP=1
-    - { tp: 8, conc-start: 4, conc-end: 8, spec-decoding: mtp }
-    - { tp: 8, ep: 8, conc-start: 16, conc-end: 32, spec-decoding: mtp }
-    - { tp: 8, ep: 8, dp-attn: true, conc-start: 64, conc-end: 256, spec-decoding: mtp }
+    - { tp: 4, ep: 4, dp-attn: true, conc-start: 256, conc-end: 256, spec-decoding: mtp }
+    # TP=8 configurations
+    - { tp: 8, conc-start: 4, conc-end: 4, spec-decoding: mtp }
+    - { tp: 8, conc-start: 128, conc-end: 128, spec-decoding: mtp }
+    - { tp: 8, ep: 8, conc-start: 32, conc-end: 128, spec-decoding: mtp }
+    - { tp: 8, ep: 8, dp-attn: true, conc-start: 32, conc-end: 64, spec-decoding: mtp }
   - isl: 1024
     osl: 8192
     search-space:
-    # If TP=4:
-    #   If CONC >= 32, then EP=4
-    #   If CONC >= 128, DP_ATTN=true, MOE_BACKEND=CUTLASS, MTP=1
-    - { tp: 4, conc-start: 4, conc-end: 16, spec-decoding: mtp }
-    - { tp: 4, ep: 4, conc-start: 32, conc-end: 64, spec-decoding: mtp }
-    - { tp: 4, ep: 4, dp-attn: true, conc-start: 128, conc-end: 256, spec-decoding: mtp }
-    # If TP=8:
-    #   If CONC >= 8, then EP=8
-    #   If CONC >= 128, DP_ATTN=true, MOE_BACKEND=CUTLASS, MTP=1
-    - { tp: 8, conc-start: 4, conc-end: 4, spec-decoding: mtp }
-    - { tp: 8, ep: 8, conc-start: 8, conc-end: 64, spec-decoding: mtp }
+      # TP=4 configurations
+    - { tp: 4, conc-start: 16, conc-end: 16, spec-decoding: mtp }
+    - { tp: 4, ep: 4, conc-start: 8, conc-end: 8, spec-decoding: mtp }
+    - { tp: 4, ep: 4, dp-attn: true, conc-start: 256, conc-end: 256, spec-decoding: mtp }
+      # TP=8 configurations
+    - { tp: 8, conc-start: 4, conc-end: 8, spec-decoding: mtp }
+    - { tp: 8, ep: 8, conc-start: 32, conc-end: 64, spec-decoding: mtp }
     - { tp: 8, ep: 8, dp-attn: true, conc-start: 128, conc-end: 256, spec-decoding: mtp }
   - isl: 8192
     osl: 1024
     search-space:
-    # If TP=4:
-    #   If CONC >= 32, then EP=4, DP_ATTN=true, MOE_BACKEND=CUTLASS, MTP=1
+      # TP=4 configurations
     - { tp: 4, conc-start: 4, conc-end: 16, spec-decoding: mtp }
-    - { tp: 4, ep: 4, dp-attn: true, conc-start: 32, conc-end: 256, spec-decoding: mtp }
-    # If TP=8:
-    #   If CONC >= 32, then EP=8, DP_ATTN=true, MOE_BACKEND=CUTLASS, MTP=1
-    - { tp: 8, conc-start: 4, conc-end: 16, spec-decoding: mtp }
-    - { tp: 8, ep: 8, dp-attn: true, conc-start: 32, conc-end: 256, spec-decoding: mtp }
+    - { tp: 4, ep: 4, conc-start: 32, conc-end: 32, spec-decoding: mtp }
+    - { tp: 4, ep: 4, dp-attn: true, conc-start: 256, conc-end: 256, spec-decoding: mtp }
+      # TP=8 configurations
+    - { tp: 8, conc-start: 4, conc-end: 4, spec-decoding: mtp }
+    - { tp: 8, ep: 8, dp-attn: true, conc-start: 64, conc-end: 256, spec-decoding: mtp }
 
 dsr1-fp8-b200-sglang:
   image: lmsysorg/sglang:v0.5.6-cu129-amd64
diff --git a/benchmarks/dsr1_fp4_b200_trt_mtp.sh b/benchmarks/dsr1_fp4_b200_trt_mtp.sh
index 0821bda9da..e674fb4a13 100644
--- a/benchmarks/dsr1_fp4_b200_trt_mtp.sh
+++ b/benchmarks/dsr1_fp4_b200_trt_mtp.sh
@@ -23,12 +23,15 @@ echo "TP: $TP, CONC: $CONC, ISL: $ISL, OSL: $OSL, EP_SIZE: $EP_SIZE, DP_ATTENTIO
 hf download "$MODEL"
 
 # ========= Determine MOE_BACKEND and MTP based on DP_ATTENTION =========
+MOE_BACKEND="TRTLLM"
+PIECEWISE_CUDA_GRAPHS="false"
+CUDA_GRAPH_MAX_BATCH_SIZE=$CONC
+MTP=3
+
 if [[ "$DP_ATTENTION" == "true" ]]; then
+    CUDA_GRAPH_MAX_BATCH_SIZE=$(( CONC < 4 ? CONC : CONC / 4 ))
     MOE_BACKEND="CUTLASS"
     MTP=1
-else
-    MOE_BACKEND="TRTLLM"
-    MTP=3
 fi
 
 echo "MOE_BACKEND='$MOE_BACKEND', MTP='$MTP'"
@@ -40,7 +43,7 @@ EXTRA_CONFIG_FILE="dsr1-fp4-mtp.yml"
 cat > $EXTRA_CONFIG_FILE << EOF
 cuda_graph_config:
     enable_padding: true
-    max_batch_size: 512
+    max_batch_size: $CUDA_GRAPH_MAX_BATCH_SIZE
 enable_attention_dp: $DP_ATTENTION
 print_iter_log: true
 kv_cache_config:
@@ -64,12 +67,35 @@ attention_dp_config:
 EOF
 fi
 
-if [[ "$DP_ATTENTION" == "true" ]]; then
-    MAX_BATCH_SIZE=$((CONC/TP))
-else
-    MAX_BATCH_SIZE=$CONC
+# set of configs using piecewise_cuda_graphs
+if [[ "$ISL" == "1024" && "$OSL" == "1024" ]]; then
+    if [[ $CONC -eq 32 || $CONC -eq 64 ]]; then
+        PIECEWISE_CUDA_GRAPHS="true"
+    elif [[ $CONC -eq 128 && $DP_ATTENTION == "false" ]]; then
+        PIECEWISE_CUDA_GRAPHS="true"
+    fi
+elif [[ "$ISL" == "1024" && "$OSL" == "8192" ]]; then
+    if [[ $CONC -eq 64 ]]; then
+        PIECEWISE_CUDA_GRAPHS="true"
+    fi
 fi
 
+if [[ "$PIECEWISE_CUDA_GRAPHS" == "true" ]]; then
+    # [2^i for i in range(8)] + [i for i in range(256, max_num_tokens, 256)] + [max_num_tokens]
+    capture_tokens=(1 2 4 8 16 32 64 128)
+    capture_tokens+=( $(seq 256 256 $MAX_NUM_TOKENS))
+    if [ $((MAX_NUM_TOKENS%256)) -ne 0 ]; then
+        capture_tokens+=($MAX_NUM_TOKENS)
+    fi
+    CAPTURE_TOKENS_LIST=$(printf "%s, " "${capture_tokens[@]}")
+
+    cat << EOF >> $EXTRA_CONFIG_FILE
+torch_compile_config:
+    capture_num_tokens: [${CAPTURE_TOKENS_LIST%, }]
+    enable_piecewise_cuda_graph: true 
+EOF
+fi  # end of set of configs using piecewise_cuda_graphs
+
 MAX_NUM_TOKENS=$(( ((MTP+1)*MAX_BATCH_SIZE+ISL+64+63)/64*64 ))
 
 set -x

From 9670ce5ec6f1b806fc70e9b9c2ceb34db3f132db Mon Sep 17 00:00:00 2001
From: Camilo Moreno Ramirez <camilom@nvidia.com>
Date: Thu, 5 Feb 2026 13:24:25 -0800
Subject: [PATCH 2/3] bugfix

---
 benchmarks/dsr1_fp4_b200_trt_mtp.sh | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/benchmarks/dsr1_fp4_b200_trt_mtp.sh b/benchmarks/dsr1_fp4_b200_trt_mtp.sh
index e674fb4a13..c91ce9dc1d 100644
--- a/benchmarks/dsr1_fp4_b200_trt_mtp.sh
+++ b/benchmarks/dsr1_fp4_b200_trt_mtp.sh
@@ -25,11 +25,11 @@ hf download "$MODEL"
 # ========= Determine MOE_BACKEND and MTP based on DP_ATTENTION =========
 MOE_BACKEND="TRTLLM"
 PIECEWISE_CUDA_GRAPHS="false"
-CUDA_GRAPH_MAX_BATCH_SIZE=$CONC
+MAX_BATCH_SIZE=$CONC
 MTP=3
 
 if [[ "$DP_ATTENTION" == "true" ]]; then
-    CUDA_GRAPH_MAX_BATCH_SIZE=$(( CONC < 4 ? CONC : CONC / 4 ))
+    MAX_BATCH_SIZE=$(( CONC < 4 ? CONC : CONC / 4 ))
     MOE_BACKEND="CUTLASS"
     MTP=1
 fi
@@ -43,7 +43,7 @@ EXTRA_CONFIG_FILE="dsr1-fp4-mtp.yml"
 cat > $EXTRA_CONFIG_FILE << EOF
 cuda_graph_config:
     enable_padding: true
-    max_batch_size: $CUDA_GRAPH_MAX_BATCH_SIZE
+    max_batch_size: $MAX_BATCH_SIZE
 enable_attention_dp: $DP_ATTENTION
 print_iter_log: true
 kv_cache_config:
@@ -67,15 +67,17 @@ attention_dp_config:
 EOF
 fi
 
+MAX_NUM_TOKENS=$(( ((MTP+1)*MAX_BATCH_SIZE+ISL+64+63)/64*64 ))
+
 # set of configs using piecewise_cuda_graphs
 if [[ "$ISL" == "1024" && "$OSL" == "1024" ]]; then
-    if [[ $CONC -eq 32 || $CONC -eq 64 ]]; then
+    if [[ $CONC == 32 || $CONC == 64 ]]; then
         PIECEWISE_CUDA_GRAPHS="true"
-    elif [[ $CONC -eq 128 && $DP_ATTENTION == "false" ]]; then
+    elif [[ $CONC == 128 && $DP_ATTENTION == "false" ]]; then
         PIECEWISE_CUDA_GRAPHS="true"
     fi
 elif [[ "$ISL" == "1024" && "$OSL" == "8192" ]]; then
-    if [[ $CONC -eq 64 ]]; then
+    if [[ $CONC == 64 ]]; then
         PIECEWISE_CUDA_GRAPHS="true"
     fi
 fi
@@ -96,7 +98,7 @@ torch_compile_config:
 EOF
 fi  # end of set of configs using piecewise_cuda_graphs
 
-MAX_NUM_TOKENS=$(( ((MTP+1)*MAX_BATCH_SIZE+ISL+64+63)/64*64 ))
+
 
 set -x
 # Launch TRT-LLM server

From f0fd1169740cf140f4fa7dc0c47fe93fa10a40e2 Mon Sep 17 00:00:00 2001
From: Camilo Moreno Ramirez <camilom@nvidia.com>
Date: Thu, 5 Feb 2026 15:27:47 -0800
Subject: [PATCH 3/3] perf changelog

---
 perf-changelog.yaml | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/perf-changelog.yaml b/perf-changelog.yaml
index aa8ad57f9b..56fe9331e5 100644
--- a/perf-changelog.yaml
+++ b/perf-changelog.yaml
@@ -361,3 +361,10 @@
     - "8k1k: 14 scenarios (7 MTP, 7 STP) for long context workloads"
     - "Prefill workers: 1-5P, Decode workers: 1-4D, TP/EP: 8/16/32"
   pr-link: https://github.com/InferenceMAX/InferenceMAX/pull/617
+
+- config-keys:
+    - dsr1-fp4-b200-trt-mtp
+  description:
+    - "Upgrade TensorRT-LLM container from release:1.1.0rc2.post2 to release:1.2.0rc6.post3"
+    - "Enable dynamic piecewise CUDA graphs for several conditions"
+    - "Adjust TP8/TP4 search space to reduce overlapping points"
\ No newline at end of file