@@ -52,8 +52,10 @@ MODEL_PATH="${MODEL_PATH:-${MODEL_DIR}/${MODEL_NAME}}"
5252source $WS_PATH /env.sh
5353
5454host_ip=$( ip route get 1.1.1.1 2> /dev/null | awk ' /src/ {print $7}' )
55- # RDMA IP for Nixl KV transfer (prefer 192.168.x.x subnet if available)
56- rdma_ip=$( hostname -I | tr ' ' ' \n' | grep ' ^192\.168\.' | head -1)
55+ # RDMA IP for Nixl KV transfer: prefer Thor data-plane subnet 172.29.x.x, then 192.168.x.x,
56+ # then fall back to the default-gateway interface IP.
57+ rdma_ip=$( hostname -I | tr ' ' ' \n' | grep ' ^172\.29\.' | head -1)
58+ rdma_ip=" ${rdma_ip:- $(hostname -I | tr ' ' ' \n' | grep ' ^192\.168\.' | head -1)} "
5759rdma_ip=" ${rdma_ip:- $host_ip } "
5860host_name=$( hostname)
5961
@@ -217,6 +219,33 @@ echo "Decode node IPs: ${DECODE_ARGS}"
217219# MoRI-IO proxy ZMQ registration port (must match vllm-router --vllm-discovery-address)
218220PROXY_PING_PORT=" ${PROXY_PING_PORT:- 36367} "
219221
222+ # MoRIIO connector extra config: optional caller overrides merged with the
223+ # standard proxy endpoints. Prefer only proven keys; unknown keys are silently
224+ # ignored by the connector.
225+ MORI_KV_EXTRA_CONFIG_JSON=" ${MORI_KV_EXTRA_CONFIG_JSON:- {\" read_mode\" : true,\" allow_full_cudagraph\" : true} }"
226+
227+ KV_TRANSFER_CONFIG_BASE=$( python3 -c "
228+ import json, os
229+ extra = json.loads(os.environ.get('MORI_KV_EXTRA_CONFIG_JSON', '{}'))
230+ print(json.dumps(extra, separators=(',', ':')))
231+ " )
232+
233+ # Prefill nodes are kv_producer; decode nodes are kv_consumer. The proxy IP,
234+ # ping port, and HTTP port are always nested inside kv_connector_extra_config.
235+ build_kv_config () {
236+ local role=" $1 "
237+ python3 -c "
238+ import json
239+ extra = json.loads('''${KV_TRANSFER_CONFIG_BASE} ''')
240+ extra.update({'proxy_ip': '${NODE0_ADDR} ', 'proxy_ping_port': '${PROXY_PING_PORT} ', 'http_port': '${SERVER_PORT} '})
241+ cfg = {'kv_connector': 'MoRIIOConnector', 'kv_role': '${role} ', 'kv_connector_extra_config': extra}
242+ print(json.dumps(cfg, separators=(',', ':')))
243+ "
244+ }
245+
246+ KV_TRANSFER_CONFIG_PRODUCER=$( build_kv_config kv_producer)
247+ KV_TRANSFER_CONFIG_CONSUMER=$( build_kv_config kv_consumer)
248+
220249# vLLM runtime environment (static vars moved to env.sh; these depend on per-node state)
221250setup_vllm_env () {
222251 export VLLM_NIXL_SIDE_CHANNEL_HOST=${rdma_ip}
@@ -256,7 +285,7 @@ if [ "$NODE_RANK" -eq 0 ]; then
256285 --served-model-name ${SERVED_MODEL} \
257286 --port $SERVER_PORT \
258287 --trust-remote-code \
259- --kv-transfer-config '{ \" kv_connector \" : \" MoRIIOConnector \" , \" kv_role \" : \" kv_producer \" , \" kv_connector_extra_config \" : { \" proxy_ip \" : \" ${NODE0_ADDR} \" , \" proxy_ping_port \" : \" ${PROXY_PING_PORT} \" , \" http_port \" : \" ${SERVER_PORT} \" }}' \
288+ --kv-transfer-config " ${KV_TRANSFER_CONFIG_PRODUCER} " \
260289 ${PREFILL_SERVER_CONFIG} "
261290
262291 if [[ " $DRY_RUN " -eq 1 ]]; then
@@ -301,6 +330,12 @@ if [ "$NODE_RANK" -eq 0 ]; then
301330 echo " Benchmarking on ${host_name} :${host_ip} "
302331 cd $WS_PATH
303332
333+ if [[ " ${DECODE_MTP_SIZE:- 0} " -gt 0 ]]; then
334+ export IS_MTP=true
335+ else
336+ export IS_MTP=false
337+ fi
338+
304339 export ROUTER_PORT=$ROUTER_PORT
305340 BENCH_CMD=" bash $WS_PATH /bench.sh ${xP} ${yD} $(( GPUS_PER_NODE* xP)) $(( GPUS_PER_NODE* yD)) \
306341 $MODEL_DIR $MODEL_NAME /run_logs/slurm_job-${SLURM_JOB_ID} ${BENCH_INPUT_LEN} \
@@ -410,7 +445,6 @@ if [ "$NODE_RANK" -eq 0 ]; then
410445 echo " ERROR: eval failed; exiting node-0 with rc=1"
411446 exit 1
412447 fi
413-
414448elif [ " $NODE_RANK " -gt 0 ] && [ " $NODE_RANK " -lt " $xP " ]; then
415449 echo " ${host_name} :${host_ip} is Additional Prefill Node (Model: ${MODEL_NAME} )"
416450 echo " Using prefill config: $PREFILL_SERVER_CONFIG "
@@ -422,7 +456,7 @@ elif [ "$NODE_RANK" -gt 0 ] && [ "$NODE_RANK" -lt "$xP" ]; then
422456 --served-model-name ${SERVED_MODEL} \
423457 --port $SERVER_PORT \
424458 --trust-remote-code \
425- --kv-transfer-config '{ \" kv_connector \" : \" MoRIIOConnector \" , \" kv_role \" : \" kv_producer \" , \" kv_connector_extra_config \" : { \" proxy_ip \" : \" ${NODE0_ADDR} \" , \" proxy_ping_port \" : \" ${PROXY_PING_PORT} \" , \" http_port \" : \" ${SERVER_PORT} \" }}' \
459+ --kv-transfer-config " ${KV_TRANSFER_CONFIG_PRODUCER} " \
426460 ${PREFILL_SERVER_CONFIG} "
427461
428462 if [[ " $DRY_RUN " -eq 1 ]]; then
478512 --served-model-name ${SERVED_MODEL} \
479513 --port $SERVER_PORT \
480514 --trust-remote-code \
481- --kv-transfer-config '{ \" kv_connector \" : \" MoRIIOConnector \" , \" kv_role \" : \" kv_consumer \" , \" kv_connector_extra_config \" : { \" proxy_ip \" : \" ${NODE0_ADDR} \" , \" proxy_ping_port \" : \" ${PROXY_PING_PORT} \" , \" http_port \" : \" ${SERVER_PORT} \" }}' \
515+ --kv-transfer-config " ${KV_TRANSFER_CONFIG_CONSUMER} " \
482516 ${DECODE_SERVER_CONFIG} "
483517
484518 if [[ " $DRY_RUN " -eq 1 ]]; then
0 commit comments