feat: Add vars to multi-node trtllm slurm scripts to support xP yD deployments (#2429)

components/backends/trtllm/multinode/multinode-examples.md

@@ -186,6 +186,10 @@ deployment across 8 nodes:
 ./srun_disaggregated.sh
 ```
 
+> [!Tip]
+> To launch multiple replicas of the configured prefill/decode workers, you can set
+> NUM_PREFILL_WORKERS and NUM_DECODE_WORKERS respectively (default: 1).
+
 ## Understanding the Output
 
 1. The `srun_aggregated.sh` launches two `srun` jobs. The first launches

components/backends/trtllm/multinode/srun_disaggregated.sh

@@ -16,9 +16,11 @@ MOUNTS="${MOUNTS:-${DEFAULT_MOUNT}}"
 NUM_GPUS_PER_NODE=${NUM_GPUS_PER_NODE:-4}
 
 NUM_PREFILL_NODES=${NUM_PREFILL_NODES:-4}
+NUM_PREFILL_WORKERS=${NUM_PREFILL_WORKERS:-1}
 PREFILL_ENGINE_CONFIG="${PREFILL_ENGINE_CONFIG:-/mnt/engine_configs/deepseek_r1/wide_ep/wide_ep_prefill.yaml}"
 
 NUM_DECODE_NODES=${NUM_DECODE_NODES:-4}
+NUM_DECODE_WORKERS=${NUM_DECODE_WORKERS:-1}
 DECODE_ENGINE_CONFIG="${DECODE_ENGINE_CONFIG:-/mnt/engine_configs/deepseek_r1/wide_ep/wide_ep_decode.yaml}"
 
 DISAGGREGATION_STRATEGY=${DISAGGREGATION_STRATEGY:-"decode_first"}
@@ -59,38 +61,42 @@ srun \
 # NOTE: Output streamed to stdout for ease of understanding the example, but
 # in practice you would probably set `srun --output ... --error ...` to pipe
 # the stdout/stderr to files.
-echo "Launching multi-node prefill worker in background."
-DISAGGREGATION_MODE=prefill \
-ENGINE_CONFIG=${PREFILL_ENGINE_CONFIG} \
-srun \
-  --mpi pmix \
-  --oversubscribe \
-  --container-image "${IMAGE}" \
-  --container-mounts "${MOUNTS}" \
-  --container-env ETCD_ENDPOINTS,NATS_SERVER,HEAD_NODE_IP,HEAD_NODE,DISAGGREGATION_MODE,DISAGGREGATION_STRATEGY,ENGINE_CONFIG \
-  --verbose \
-  --label \
-  -A "${ACCOUNT}" \
-  -J "${ACCOUNT}-dynamo.trtllm" \
-  --nodes "${NUM_PREFILL_NODES}" \
-  --ntasks-per-node "${NUM_GPUS_PER_NODE}" \
-  --jobid "${SLURM_JOB_ID}" \
-  /mnt/multinode/start_trtllm_worker.sh &
+for ((i=1; i<=${NUM_PREFILL_WORKERS}; i++)); do
+  echo "Launching multi-node prefill worker in background."
+  DISAGGREGATION_MODE=prefill \
+  ENGINE_CONFIG=${PREFILL_ENGINE_CONFIG} \
+  srun \
+    --mpi pmix \
+    --oversubscribe \
+    --container-image "${IMAGE}" \
+    --container-mounts "${MOUNTS}" \
+    --container-env ETCD_ENDPOINTS,NATS_SERVER,HEAD_NODE_IP,HEAD_NODE,DISAGGREGATION_MODE,DISAGGREGATION_STRATEGY,ENGINE_CONFIG \
+    --verbose \
+    --label \
+    -A "${ACCOUNT}" \
+    -J "${ACCOUNT}-dynamo.trtllm" \
+    --nodes "${NUM_PREFILL_NODES}" \
+    --ntasks-per-node "${NUM_GPUS_PER_NODE}" \
+    --jobid "${SLURM_JOB_ID}" \
+    /mnt/multinode/start_trtllm_worker.sh &
+done
 
-echo "Launching multi-node decode worker in background."
-DISAGGREGATION_MODE=decode \
-ENGINE_CONFIG=${DECODE_ENGINE_CONFIG} \
-srun \
-  --mpi pmix \
-  --oversubscribe \
-  --container-image "${IMAGE}" \
-  --container-mounts "${MOUNTS}" \
-  --container-env ETCD_ENDPOINTS,NATS_SERVER,HEAD_NODE_IP,HEAD_NODE,DISAGGREGATION_MODE,DISAGGREGATION_STRATEGY,ENGINE_CONFIG \
-  --verbose \
-  --label \
-  -A "${ACCOUNT}" \
-  -J "${ACCOUNT}-dynamo.trtllm" \
-  --nodes "${NUM_DECODE_NODES}" \
-  --ntasks-per-node "${NUM_GPUS_PER_NODE}" \
-  --jobid "${SLURM_JOB_ID}" \
-  /mnt/multinode/start_trtllm_worker.sh &
+for ((i=1; i<=${NUM_DECODE_WORKERS}; i++)); do
+  echo "Launching multi-node decode worker in background."
+  DISAGGREGATION_MODE=decode \
+  ENGINE_CONFIG=${DECODE_ENGINE_CONFIG} \
+  srun \
+    --mpi pmix \
+    --oversubscribe \
+    --container-image "${IMAGE}" \
+    --container-mounts "${MOUNTS}" \
+    --container-env ETCD_ENDPOINTS,NATS_SERVER,HEAD_NODE_IP,HEAD_NODE,DISAGGREGATION_MODE,DISAGGREGATION_STRATEGY,ENGINE_CONFIG \
+    --verbose \
+    --label \
+    -A "${ACCOUNT}" \
+    -J "${ACCOUNT}-dynamo.trtllm" \
+    --nodes "${NUM_DECODE_NODES}" \
+    --ntasks-per-node "${NUM_GPUS_PER_NODE}" \
+    --jobid "${SLURM_JOB_ID}" \
\ No newline at end of file
+    /mnt/multinode/start_trtllm_worker.sh &
+done
\ No newline at end of file