Skip to content

[NV] dsr1 fp4 b200 trt agg mtp update #642

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
camiloamoreno merged 5 commits into from
Feb 6, 2026
+67 −39
Merged

[NV] dsr1 fp4 b200 trt agg mtp update #642

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
444c823
update nvidia master and benchmark options
camiloamoreno Feb 5, 2026
9670ce5
bugfix
camiloamoreno Feb 5, 2026
f0fd116
perf changelog
camiloamoreno Feb 5, 2026
f5f523c
Merge branch 'main' into nv/dsr1-fp4-b200-trt-agg-mtp-260204
camiloamoreno Feb 5, 2026
04fd054
Merge branch 'main' into nv/dsr1-fp4-b200-trt-agg-mtp-260204
cquil11 Feb 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 21 additions & 30 deletions .github/configs/nvidia-master.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
46 changes: 37 additions & 9 deletions benchmarks/dsr1_fp4_b200_trt_mtp.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
MAX_BATCH_SIZE=$CONC
MTP=3

if [[ "$DP_ATTENTION" == "true" ]]; then
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'"
Expand All @@ -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: $MAX_BATCH_SIZE
enable_attention_dp: $DP_ATTENTION
print_iter_log: true
kv_cache_config:
Expand All @@ -64,13 +67,38 @@ attention_dp_config:
EOF
fi

if [[ "$DP_ATTENTION" == "true" ]]; then
MAX_BATCH_SIZE=$((CONC/TP))
else
MAX_BATCH_SIZE=$CONC
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 == 32 || $CONC == 64 ]]; then
PIECEWISE_CUDA_GRAPHS="true"
elif [[ $CONC == 128 && $DP_ATTENTION == "false" ]]; then
PIECEWISE_CUDA_GRAPHS="true"
fi
elif [[ "$ISL" == "1024" && "$OSL" == "8192" ]]; then
if [[ $CONC == 64 ]]; then
PIECEWISE_CUDA_GRAPHS="true"
fi
fi

MAX_NUM_TOKENS=$(( ((MTP+1)*MAX_BATCH_SIZE+ISL+64+63)/64*64 ))
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



set -x
# Launch TRT-LLM server
Expand Down
9 changes: 9 additions & 0 deletions perf-changelog.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -418,3 +418,12 @@
- "Configurations: TP=8, EP=1, concurrency 4-64 for 1k1k, 1k8k, and 8k1k sequence lengths"
pr-link: https://github.com/InferenceMAX/InferenceMAX/pull/626

- 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"
pr-link: https://github.com/InferenceMAX/InferenceMAX/pull/642


Loading