fix: Add mpi arg to all srun commands (#5948)

examples/basics/multinode/trtllm/srun_aggregated.sh

@@ -41,6 +41,7 @@ fi
 # the stdout/stderr to files.
 echo "Launching frontend services in background."
 srun \
+  --mpi pmix \
   --overlap \
   --container-image "${IMAGE}" \
   --container-mounts "${MOUNTS}" \
@@ -71,4 +72,4 @@ srun \
   --nodes "${NUM_NODES}" \
   --ntasks-per-node "${NUM_GPUS_PER_NODE}" \
   --jobid "${SLURM_JOB_ID}" \
-  /mnt/examples/basics/multinode/trtllm/start_trtllm_worker.sh &
\ No newline at end of file
+  /mnt/examples/basics/multinode/trtllm/start_trtllm_worker.sh &

examples/basics/multinode/trtllm/srun_disaggregated.sh

@@ -44,6 +44,7 @@ fi
 # the stdout/stderr to files.
 echo "Launching frontend services in background."
 srun \
+  --mpi pmix \
   --overlap \
   --container-image "${IMAGE}" \
   --container-mounts "${MOUNTS}" \
@@ -97,4 +98,4 @@ for ((i=1; i<=${NUM_DECODE_WORKERS}; i++)); do
     --ntasks-per-node "${NUM_GPUS_PER_NODE}" \
     --jobid "${SLURM_JOB_ID}" \
     /mnt/examples/basics/multinode/trtllm/start_trtllm_worker.sh &
-done
\ No newline at end of file
+done