[ROCm] Support for the remote dp and tp size for vllm's moriio kv connector

[itej89]

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

This PR enables the support for VLLM MORI IO Connectors interface to communicate TP and DP sizes to the peers in PD Disaggregation Mode.

Related Feature PR on vLLM: vllm-project/vllm#48989

Details:
The MORI IO Connector uses tp size for the remote port offset calculation. To provide a future-proof clean interface we are adding the support for additional registration parameters in the router.

Related vLLM MORI IO PORT offset calculation: Port offset calculation
Related vLLM MORI IO Connector Interface: Interface with TP, DP size data

Test Plan

Verification steps on MI355X-AINIC
STEP 1: Build ROCm VLLM Image

docker build --no-cache  -f ./docker/Dockerfile.rocm_base -t  <IMG_VLLM_ROCM_BASE>  . 

DOCKER_BUILDKIT=1 docker build --no-cache  --target final --progress=plain -t  <FINAL_VLLM_IMAGE_NAME>  -f ./docker/Dockerfile.rocm  --build-
arg BASE_IMAGE=<IMG_VLLM_ROCM_BASE>   .

STEP 1: Build VLLM Router Image

docker build --no-cache  -f Dockerfile.router -t  <VLLM_ROUTER_IMAGE_NAME> . 

STEP 3: Please use below commands to start Prefill and Decode servers

## PREFILL
docker run --rm -i \
    --user "$(id -u):$(id -g)" \
    --device /dev/dri \
    --device /dev/kfd \
    --device /dev/infiniband \
    --network host \
    --ipc host \
    --group-add "$(getent group video | cut -d: -f3)" \
    --group-add "$(getent group render | cut -d: -f3)" \
    --cap-add SYS_PTRACE \
    --cap-add IPC_LOCK \
    --security-opt seccomp=unconfined \
    --shm-size 64G \
    --ulimit nofile=1048576:1048576 \
    --ulimit memlock=-1:-1 \
    -v <DATA_ROOT>:<DATA_ROOT> \
    -v <MORI_DIR>:/tmp/mori \
    -e HOME=<HOME_DIR> \
    -e USER="$(id -un)" \
    -e AITER_JIT_DIR=<JIT_CACHE_DIR> \
    -e FLYDSL_RUNTIME_CACHE_DIR=<FLYDSL_CACHE_DIR> \
    -e VLLM_ROCM_USE_AITER=1 \
    -e TRITON_CACHE_DIR=/tmp/triton_cache \
    -e VLLM_CACHE_ROOT=/tmp/vllm_cache \
    -e MORI_RDMA_DEVICES=<RDMA_DEVICE_LIST> \
    -e MORI_IB_GID_INDEX=1 \
    -e MORI_SHMEM_HEAP_SIZE=16G \
    -e MORI_GPU_ARCHS=<GPU_ARCH> \
    -e NCCL_IB_GID_INDEX=1 \
    --entrypoint bash \
    --name prefill \
    <CONTAINER_IMAGE> \
    -c 'vllm serve <MODEL_PATH> \
        -tp 8 \
        --port <HTTP_PORT> \
        --max-model-len 65536 \
        --gpu-memory-utilization 0.85 \
        --enforce-eager \
        --kv-cache-dtype fp8 \
        --kv-transfer-config "{
            \"kv_connector\": \"MoRIIOConnector\",
            \"kv_role\": \"kv_producer\",
            \"kv_connector_extra_config\": {
                \"proxy_ip\": \"<PROXY_IP>\",
                \"proxy_ping_port\": \"<PROXY_PING_PORT>\",
                \"http_port\": \"<HTTP_PORT>\",
                \"handshake_port\": \"<HANDSHAKE_PORT>\",
                \"notify_port\": \"<NOTIFY_PORT>\"
            }
        }"' \
    2>&1 | tee <LOG_FILE> >/dev/null
    
## DECODE
    docker run --rm -i \
    --user "$(id -u):$(id -g)" \
    --device /dev/dri \
    --device /dev/kfd \
    --device /dev/infiniband \
    --network host \
    --ipc host \
    --group-add "$(getent group video | cut -d: -f3)" \
    --group-add "$(getent group render | cut -d: -f3)" \
    --cap-add SYS_PTRACE \
    --cap-add IPC_LOCK \
    --security-opt seccomp=unconfined \
    --shm-size 64G \
    --ulimit nofile=1048576:1048576 \
    --ulimit memlock=-1:-1 \
    -v <DATA_ROOT>:<DATA_ROOT> \
    -v <MORI_DIR>:/tmp/mori \
    -e HOME=<HOME_DIR> \
    -e USER="$(id -un)" \
    -e AITER_JIT_DIR=<AITER_JIT_DIR> \
    -e FLYDSL_RUNTIME_CACHE_DIR=<FLYDSL_CACHE_DIR> \
    -e VLLM_ROCM_USE_AITER=1 \
    -e TRITON_CACHE_DIR=/tmp/triton_cache \
    -e VLLM_CACHE_ROOT=/tmp/vllm_cache \
    -e MORI_RDMA_DEVICES=<RDMA_DEVICE_LIST> \
    -e MORI_IB_GID_INDEX=1 \
    -e MORI_SHMEM_HEAP_SIZE=16G \
    -e MORI_GPU_ARCHS=<GPU_ARCH> \
    -e NCCL_IB_GID_INDEX=1 \
    --entrypoint bash \
    --name decode \
    <CONTAINER_IMAGE> \
    -c 'vllm serve <MODEL_PATH> \
        -tp 8 \
        --port <HTTP_PORT> \
        --max-model-len 65536 \
        --gpu-memory-utilization 0.85 \
        --enforce-eager \
        --kv-cache-dtype fp8 \
        --kv-transfer-config "{
            \"kv_connector\": \"MoRIIOConnector\",
            \"kv_role\": \"kv_consumer\",
            \"kv_connector_extra_config\": {
                \"proxy_ip\": \"<PROXY_IP>\",
                \"proxy_ping_port\": \"<PROXY_PING_PORT>\",
                \"http_port\": \"<HTTP_PORT>\",
                \"handshake_port\": \"<HANDSHAKE_PORT>\",
                \"notify_port\": \"<NOTIFY_PORT>\"
            }
        }"' \
    2>&1 | tee <LOG_FILE> >/dev/null &
`

STEP 4: Please run the following to start the vLLM Router

docker run --rm \
    --network host \
    --name proxy \
   <VLLM_ROUTER_IMAGE_NAME> \
    vllm-router \
    --vllm-pd-disaggregation \
    --kv-connector moriio \
    --vllm-discovery-address "0.0.0.0:36367" \
    --host 0.0.0.0 \
    --port 30000  2>&1 | tee <LOG_FILE> >/dev/null &

STEP 5: Please run the following to perform the accuracy test

## Exec to prefill node container
docker exec -it prefill bash

## lm-eval command
python3 -m lm_eval \
        --model local-completions \
        --tasks gsm8k \
        --model_args 'model=/data/models2/DeepSeek-V4-Pro,base_url=http://165.245.133.211:30000/v1/completions,num_concurrent=64,max_retries=3,tokenized_requests=False,trust_remote_code=True'

Test Result

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9492	±	0.006
		strict-match	5	exact_match	↑	0.9492	±	0.006

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: ba5165595b

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Use registered remote DP size in MoRI-IO params

In MoRI-IO WRITE mode, when the selected decode instance registers a dp_size that differs from the router's intra_node_data_parallel_size (for example the router is left at the default 1), this lookup discards the registered DP size and the request below still sends remote_dp_size from the router config. That makes the prefill side handshake/notify the wrong number of decode DP ranks; the decode-leg lookup of the selected prefill instance has the same issue, so the registered dp_size should be used for remote_dp_size in both directions.

Useful? React with 👍 / 👎.

[itej89]

Additional Proposals for the future work:

1. Because of the MORI IO interface flexibility, there are number of parameters moving around between registered nodes and the Router. My proposal for a cleaner design would be to have a something like MoriIODispatchContext struct to pass around the parameters.

 struct MoriIODispatchContext {                                  
      transfer_id:     Option<String>,
      decode_base:     Option<String>,
      remote_tp_size:  usize,                                                                                            
      prefill_dp_rank: Option<usize>,                                                                                    
      decode_dp_rank:  Option<usize>,                                                                                    
  } 

2. If my understanding is correct, the Router's policies are considered for the registered URLs. Since DP Ranks aren't being registered, I've added the RR as a back-up (taken from MORIIO reference proxy implementation). We might need to enable DP rank registration in the Router to fully utilize its own policies in EP mode.

Request to provide corrections or suggestion on my understanding.