Fixes for gb300

Signed-off-by: jthomson04 <[email protected]>

Author: jthomson04

components/backends/trtllm/performance_sweeps/benchmark_disagg.slurm

@@ -5,8 +5,7 @@ MULTI_ROUND="${MULTI_ROUND:-8}"
 
 # set MOUNT_DIR
 MOUNT_DIR="${MOUNT_DIR:-${PWD}}"
-CONTAINER_NAME=disaggr-test
-
+CONTAINER_NAME=disaggr-test-$(date +%s)-$$
 
 STREAMING=true
 CTX_GPU_FRAC=0.85
@@ -35,6 +34,7 @@ image=${21}
 isl=${22}
 osl=${23}
 benchmark_kind=${24}
+ntasks_per_node=${25}
 
 CACHE_TRANSCEIVER_MAX_NUM_TOKENS=${CACHE_TRANSCEIVER_MAX_NUM_TOKENS:-$((${isl} + ${osl} + 512))}
 
@@ -179,7 +179,9 @@ for ((i=1; i<=DECODE_COUNT; i++)); do
       --nodes ${num_gen_nodes} \
       --ntasks $gen_tp_size \
       --oversubscribe \
+      --gpus-per-node $ntasks_per_node \
       --overlap \
+      -e UCX_NET_DEVICES,TRTLLM_UCX_INTERFACE \
       bash ${SCRIPTS_DIR}/scripts/start_disagg_worker.sh ${full_logdir}/decode_config.yaml ${ctx_gpus} ${served_model_name} ${model_path} 'decode' $gen_enable_attention_dp &> ${full_logdir}/output_decode_worker_${i}.log &
   echo "$!" >> "$PID_FILE"
 done
@@ -203,8 +205,10 @@ for ((i=1; i<=PREFILL_COUNT; i++)); do
         --mpi=pmix --overlap -w ${nodes[node_idx]} \
         --oversubscribe \
         --overlap \
-        --ntasks $(( ctx_tp_size < 4 ? ctx_tp_size : 4 )) \
+        --ntasks $(( ctx_tp_size < ntasks_per_node ? ctx_tp_size : ntasks_per_node )) \
+        --gpus-per-node $ntasks_per_node \
         --nodes 1 \
+        -e UCX_NET_DEVICES,TRTLLM_UCX_INTERFACE \
         bash ${SCRIPTS_DIR}/scripts/start_disagg_worker.sh ${full_logdir}/prefill_config.yaml ${ctx_gpus} ${served_model_name} ${model_path} 'prefill' $ctx_enable_attention_dp &> ${full_logdir}/output_prefill_worker_${i}.log &
   prefill_pids+=($!)
   echo "$!" >> "$PID_FILE"

components/backends/trtllm/performance_sweeps/scripts/start_disagg_worker.sh

@@ -55,10 +55,6 @@ if [[ "${model_path,,}" != *r1* ]]; then
     export OVERRIDE_QUANT_ALGO=W4A8_MXFP4_MXFP8
 fi
 
-# NOTE: Set (or unset) these depending on what cluster you're using
-export TRTLLM_UCX_INTERFACE=enP6p9s0np0
-export UCX_NET_DEVICES=mlx5_0:1,mlx5_1:1,mlx5_3:1,mlx5_4:1,enP6p9s0np0
-
 trtllm-llmapi-launch python3 -m dynamo.trtllm \
     --model-path ${model_path} \
     --served-model-name ${model_name} \

components/backends/trtllm/performance_sweeps/submit_disagg.sh

@@ -95,11 +95,11 @@ run_single() {
     # TODO: expose kind to the command line
     local kind="dynamo_disagg"
 
-    gen_nodes=$(((gen_tp_size + 3)/4 * gen_num))
+    gen_nodes=$(((gen_tp_size + NTASKS_PER_NODE - 1)/NTASKS_PER_NODE * gen_num))
     total_nodes=$((ctx_num + gen_nodes))
-    total_tasks=$((total_nodes * 4))
+    total_tasks=$((total_nodes * NTASKS_PER_NODE))
     set -x
-    sbatch --nodes=${total_nodes} --ntasks=${total_tasks} --ntasks-per-node=${NTASKS_PER_NODE} --segment=${total_nodes} ${slurm_args} benchmark_disagg.slurm ${ctx_num} ${ctx_tp_size} ${ctx_ep_size} ${ctx_enable_attention_dp} 30 20000 ${gen_num} ${gen_tp_size} ${gen_ep_size} ${gen_batch_size} ${gen_max_num_tokens} ${gen_enable_attention_dp} ${gen_gpu_memory_fraction} ${gen_eplb_num_slots} ${gen_mtp_size} "${gen_concurrency_list}" ${gen_nodes} ${kind} ${MODEL_PATH} ${SERVED_MODEL_NAME} ${IMAGE} ${ISL} ${OSL} ${BENCHMARK_KIND}
+    sbatch --nodes=${total_nodes} --gpus-per-node ${NTASKS_PER_NODE} --ntasks=${total_tasks} --ntasks-per-node=${NTASKS_PER_NODE} --segment=${total_nodes} ${slurm_args} benchmark_disagg.slurm ${ctx_num} ${ctx_tp_size} ${ctx_ep_size} ${ctx_enable_attention_dp} 30 20000 ${gen_num} ${gen_tp_size} ${gen_ep_size} ${gen_batch_size} ${gen_max_num_tokens} ${gen_enable_attention_dp} ${gen_gpu_memory_fraction} ${gen_eplb_num_slots} ${gen_mtp_size} "${gen_concurrency_list}" ${gen_nodes} ${kind} ${MODEL_PATH} ${SERVED_MODEL_NAME} ${IMAGE} ${ISL} ${OSL} ${BENCHMARK_KIND} ${NTASKS_PER_NODE}
     set +x
 }