single node nvtx test dsr1

Author: hbarclay

benchmarks/benchmark_lib.sh

@@ -9,6 +9,14 @@ export PYTHONDONTWRITEBYTECODE=1
 export PYTHONPYCACHEPREFIX="${PYTHONPYCACHEPREFIX:-/tmp/inferencex-pycache}"
 mkdir -p "$PYTHONPYCACHEPREFIX" 2>/dev/null || true
 
+# When profiling, enable layerwise NVTX markers for per-layer annotations
+# (module names, input shapes, weight shapes) in torch profiler traces.
+# Single-node scripts should append $SGLANG_PROFILE_ARGS to their launch command.
+SGLANG_PROFILE_ARGS=""
+if [[ "${PROFILE:-}" == "1" ]]; then
+    SGLANG_PROFILE_ARGS="--enable-layerwise-nvtx-marker --disable-cuda-graph"
+fi
+
 # --------------------------------
 # GPU monitoring helpers
 # --------------------------------

benchmarks/single_node/dsr1_fp4_b200.sh

@@ -45,7 +45,7 @@ PYTHONNOUSERSITE=1 python3 -m sglang.launch_server --model-path $MODEL --host 0.
 --cuda-graph-max-bs 256 --max-running-requests 256 --mem-fraction-static 0.85 --kv-cache-dtype fp8_e4m3 \
 --chunked-prefill-size 16384 \
 --ep-size $EP_SIZE --quantization modelopt_fp4 --enable-flashinfer-allreduce-fusion --scheduler-recv-interval $SCHEDULER_RECV_INTERVAL \
---enable-symm-mem --disable-radix-cache --attention-backend trtllm_mla --moe-runner-backend flashinfer_trtllm --stream-interval 10 $EVAL_CONTEXT_ARGS > $SERVER_LOG 2>&1 &
+--enable-symm-mem --disable-radix-cache --attention-backend trtllm_mla --moe-runner-backend flashinfer_trtllm --stream-interval 10 $EVAL_CONTEXT_ARGS $SGLANG_PROFILE_ARGS > $SERVER_LOG 2>&1 &
 
 SERVER_PID=$!
 

benchmarks/single_node/dsr1_fp4_mi355x.sh

@@ -49,7 +49,7 @@ python3 -m sglang.launch_server --model-path=$MODEL --trust-remote-code \
 --max-prefill-tokens=$PREFILL_SIZE \
 --cuda-graph-max-bs=128 \
 --attention-backend aiter \
---kv-cache-dtype fp8_e4m3 $EVAL_CONTEXT_ARGS > $SERVER_LOG 2>&1 &
+--kv-cache-dtype fp8_e4m3 $EVAL_CONTEXT_ARGS $SGLANG_PROFILE_ARGS > $SERVER_LOG 2>&1 &
 
 SERVER_PID=$!
 

benchmarks/single_node/dsr1_fp8_b200.sh

@@ -77,7 +77,7 @@ PYTHONNOUSERSITE=1 python3 -m sglang.launch_server --model-path=$MODEL --host=0.
 --cuda-graph-max-bs $CUDA_GRAPH_MAX_BATCH_SIZE --max-running-requests $MAX_RUNNING_REQUESTS \
 --mem-fraction-static $MEM_FRAC_STATIC --kv-cache-dtype fp8_e4m3 --chunked-prefill-size $CHUNKED_PREFILL_SIZE --max-prefill-tokens $MAX_PREFILL_TOKENS \
 --enable-flashinfer-allreduce-fusion --scheduler-recv-interval $SCHEDULER_RECV_INTERVAL --disable-radix-cache \
---attention-backend trtllm_mla --stream-interval 30 --ep-size $EP_SIZE --moe-runner-backend flashinfer_trtllm --quantization fp8 $EVAL_CONTEXT_ARGS > $SERVER_LOG 2>&1 &
+--attention-backend trtllm_mla --stream-interval 30 --ep-size $EP_SIZE --moe-runner-backend flashinfer_trtllm --quantization fp8 $EVAL_CONTEXT_ARGS $SGLANG_PROFILE_ARGS > $SERVER_LOG 2>&1 &
 
 SERVER_PID=$!
 

benchmarks/single_node/dsr1_fp8_b200_mtp.sh

@@ -90,7 +90,7 @@ PYTHONNOUSERSITE=1 python3 -m sglang.launch_server \
     --speculative-num-steps $SPECULATIVE_NUM_STEPS \
     --speculative-num-draft-tokens $SPECULATIVE_DRAFT_TOKENS \
     --speculative-eagle-topk $SPECULATIVE_EAGLE_TOPK \
-    $EVAL_CONTEXT_ARGS > $SERVER_LOG 2>&1 &
+    $EVAL_CONTEXT_ARGS $SGLANG_PROFILE_ARGS > $SERVER_LOG 2>&1 &
 
 SERVER_PID=$!
 

benchmarks/single_node/dsr1_fp8_h200.sh

@@ -41,7 +41,7 @@ if [[ $ISL -eq 1024 && $OSL -eq 1024 ]]; then
     --chunked-prefill-size 32768 --max-prefill-tokens 32768 --mem-fraction-static 0.82 \
     --attention-backend flashinfer --stream-interval 10 \
     --decode-log-interval 1 \
-    $EVAL_CONTEXT_ARGS > $SERVER_LOG 2>&1 &
+    $EVAL_CONTEXT_ARGS $SGLANG_PROFILE_ARGS > $SERVER_LOG 2>&1 &
 else
     PYTHONNOUSERSITE=1 python3 -m sglang.launch_server --model-path $MODEL \
     --host 0.0.0.0 --port $PORT --trust-remote-code \
@@ -50,7 +50,7 @@ else
     --chunked-prefill-size 32768 --max-prefill-tokens 32768 --mem-fraction-static 0.82 \
     --attention-backend flashinfer --stream-interval 10 \
     --decode-log-interval 1 \
-    $EVAL_CONTEXT_ARGS > $SERVER_LOG 2>&1 &
+    $EVAL_CONTEXT_ARGS $SGLANG_PROFILE_ARGS > $SERVER_LOG 2>&1 &
 fi
 
 SERVER_PID=$!

benchmarks/single_node/dsr1_fp8_mi300x.sh

@@ -55,7 +55,7 @@ python3 -m sglang.launch_server \
 --max-prefill-tokens=131072 \
 --kv-cache-dtype fp8_e4m3 \
 --attention-backend aiter \
---disable-radix-cache $EVAL_CONTEXT_ARGS > $SERVER_LOG 2>&1 &
+--disable-radix-cache $EVAL_CONTEXT_ARGS $SGLANG_PROFILE_ARGS > $SERVER_LOG 2>&1 &
 
 SERVER_PID=$!
 

benchmarks/single_node/dsr1_fp8_mi325x.sh

@@ -47,7 +47,7 @@ python3 -m sglang.launch_server \
 --kv-cache-dtype fp8_e4m3 \
 --attention-backend aiter \
 --disable-radix-cache \
-$EVAL_CONTEXT_ARGS > $SERVER_LOG 2>&1 &
+$EVAL_CONTEXT_ARGS $SGLANG_PROFILE_ARGS > $SERVER_LOG 2>&1 &
 
 SERVER_PID=$!
 

benchmarks/single_node/dsr1_fp8_mi355x.sh

@@ -47,7 +47,7 @@ python3 -m sglang.launch_server \
     --num-continuous-decode-steps 4 \
     --max-prefill-tokens 196608 \
     --kv-cache-dtype fp8_e4m3 \
-    --cuda-graph-max-bs "$CONC" $EVAL_CONTEXT_ARGS > $SERVER_LOG 2>&1 &
+    --cuda-graph-max-bs "$CONC" $EVAL_CONTEXT_ARGS $SGLANG_PROFILE_ARGS > $SERVER_LOG 2>&1 &
 
 SERVER_PID=$!