[Bugfix][ROCm][Disagg] Fold MoRIIO WRITE notify port with local DP rank (follow-up to #51681)

[avininjamay8]

Purpose

Follow-up to #51681. That PR moves MoRIIO WRITE routing onto each WriteTask so concurrent requests do not race when picking the decode pod IP. After it lands, the worker no longer stashes remote_dp_size_local, but _finalize_if_complete still read it from the worker to fold the write_done notify port with decode_dp_rank % dp_size_local. With DP16 and dp_size_local=8, decode ranks 8–15 on the headless child pod need the per-pod local rank for the port offset (e.g. rank 10 → port for rank 2); otherwise write_done misses the notify socket and decode hangs.

This change carries remote_dp_size_local on RemoteAllocInfo and uses it in _finalize_if_complete so the notify port stays correct after #51681 removes the worker stash.

Changes

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py and moriio_engine.py only:

• Add remote_dp_size_local to RemoteAllocInfo.
• Copy _dp_local into request_info in _execute_write_task.
• Read request_info.remote_dp_size_local in _finalize_if_complete instead of worker.remote_dp_size_local.

No WriteTask routing, engine, scheduler, or harness changes (those are in #51681).

Test Plan

MoRIIO 2P2D wide-EP (DP16, WRITE) disaggregated serving, ROUTER_TYPE=proxy, RUN_AFTER_HEALTH=accuracy, image vllm/vllm-openai-rocm:nightly.

MiniMax-M3-MXFP8 — 2P2D wide-EP DP16

export VLLM_ROCM_USE_AITER=1
export VLLM_ROCM_USE_AITER_MOE=1
export VLLM_ROCM_USE_AITER_RMSNORM=1
export VLLM_DP_DISABLE_BATCH_QUEUE=1
export ALL2ALL_BACKEND=allgather_reducescatter

# PREFILL — master (node0, DP ranks 0–7)
vllm serve /data/models2/MiniMax-M3-MXFP8 \
    --host $PREFILL_IP --port 20005 -tp 1 \
    --data-parallel-size 16 \
    --data-parallel-size-local 8 \
    --data-parallel-address $PREFILL_IP \
    --data-parallel-rpc-port 13345 \
    --enable-expert-parallel \
    --all2all-backend allgather_reducescatter \
    --trust-remote-code \
    --attention-backend TRITON_ATTN \
    --block-size 128 \
    --language-model-only \
    --kv-cache-dtype fp8 \
    --gpu-memory-utilization 0.85 \
    --enforce-eager \
    --no-async-scheduling \
    --no-disable-nccl-for-dp-synchronization \
    --kv-transfer-config '{"kv_connector":"MoRIIOConnector","kv_role":"kv_producer","kv_port":"9711",
      "kv_connector_extra_config":{"proxy_ip":"'$PREFILL_IP'","proxy_port":"10001",
      "proxy_ping_port":"36367","http_port":"20005","local_ping_port":"61555",
      "handshake_port":"6301","notify_port":"61005"}}'

# PREFILL — child (node1, headless, DP ranks 8–15)
vllm serve /data/models2/MiniMax-M3-MXFP8 \
    -tp 1 \
    --data-parallel-size 16 \
    --data-parallel-size-local 8 \
    --data-parallel-address $PREFILL_IP \
    --data-parallel-rpc-port 13345 \
    --data-parallel-start-rank 8 \
    --headless \
    --enable-expert-parallel \
    --all2all-backend allgather_reducescatter \
    --trust-remote-code \
    --attention-backend TRITON_ATTN \
    --block-size 128 \
    --language-model-only \
    --kv-cache-dtype fp8 \
    --gpu-memory-utilization 0.85 \
    --enforce-eager \
    --no-async-scheduling

# DECODE — master (node2, DP ranks 0–7)
vllm serve /data/models2/MiniMax-M3-MXFP8 \
    --host $DECODE_IP --port 20005 -tp 1 \
    --data-parallel-size 16 \
    --data-parallel-size-local 8 \
    --data-parallel-address $DECODE_IP \
    --data-parallel-rpc-port 13345 \
    --enable-expert-parallel \
    --all2all-backend allgather_reducescatter \
    --trust-remote-code \
    --attention-backend TRITON_ATTN \
    --block-size 128 \
    --language-model-only \
    --kv-cache-dtype fp8 \
    --gpu-memory-utilization 0.85 \
    --no-async-scheduling \
    --no-disable-nccl-for-dp-synchronization \
    --kv-transfer-config '{"kv_connector":"MoRIIOConnector","kv_role":"kv_consumer","kv_port":"9711",
      "kv_connector_extra_config":{"proxy_ip":"'$PREFILL_IP'","proxy_port":"10001",
      "proxy_ping_port":"36367","http_port":"20005","local_ping_port":"61555",
      "handshake_port":"6301","notify_port":"61005"}}'

# DECODE — child (node3, headless, DP ranks 8–15)
vllm serve /data/models2/MiniMax-M3-MXFP8 \
    -tp 1 \
    --data-parallel-size 16 \
    --data-parallel-size-local 8 \
    --data-parallel-address $DECODE_IP \
    --data-parallel-rpc-port 13345 \
    --data-parallel-start-rank 8 \
    --headless \
    --enable-expert-parallel \
    --all2all-backend allgather_reducescatter \
    --trust-remote-code \
    --attention-backend TRITON_ATTN \
    --block-size 128 \
    --language-model-only \
    --kv-cache-dtype fp8 \
    --gpu-memory-utilization 0.85 \
    --no-async-scheduling

GSM8K evaluation

lm_eval --model local-completions \
    --tasks gsm8k \
    --model_args "model=/data/models2/MiniMax-M3-MXFP8,base_url=http://127.0.0.1:10001/v1/completions,num_concurrent=64,max_retries=3,tokenized_requests=False,trust_remote_code=True,timeout=7200" \
    --limit 250

Test Result

MiniMax-M3-MXFP8, MoRIIO 2P2D wide-EP DP16 (proxy). GSM8K flexible-extract, threshold 0.80.

Variant	Health	GSM8K exact_match
#51681 only (write_done notify port wrong for headless ranks 8–15)	Fail (decode hang)	—
#51681 + this fix	Pass	0.8125 (PASS)

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[avininjamay8]

This is already accounted for Vincent in his new commit to the PR : #51681.