[Feature] Enable cache aware DP routing for standalone ATOMesh

[simondanielsson]

Motivation

Enable cache aware DP routing for data parallel attention using standalone ATOMesh router.

Note: #1505 added logic to wire up the data-parallel-rank field from the ATOMesh router to ATOM, however

• it only targeted /v1/completions, not /v1/chat/completions
• Streaming paths not supported
• Standalone atommesh mesh-only mode not supported
• /server_info endpoint used by the atommesh router doesn't exist in ATOM

Technical Details

Running DPA with ATOM on real workloads requires prefix cache hit rate similar to TP, which requires a cache aware router like atommesh. Atommesh selects which DP rank to route to and injects data-parallel-rank into the request body accordingly. However, ATOM currently ignores this field and defaults to round-robin DP routing.

This PR makes ATOM account for this field injected by atommesh. We also expose the endpoint /server_info which atommesh uses to identify the number of DP replicas.

Test Plan

Run ATOM with DPA + Atommesh router and check if prefix cache hits are as expected (i.e, better than round robin). Tested on DSv3 on MI350.

1. Run ATOM with DPA in this branch

docker run \
  --rm \
  --name p-replay-dsv3 \
  --init --network host --ipc host --privileged \
  --cap-add SYS_PTRACE --security-opt seccomp=unconfined \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  --shm-size 256G \
  --group-add video \
  --device /dev/kfd --device /dev/dri --device /dev/infiniband \
  -v /sys:/sys \
  -v "/home/pedaniel@amd.com/.cache/huggingface/:/root/.cache/huggingface" \
  -e HF_HOME=/root/.cache/huggingface \
  -e HF_TOKEN=$HF_TOKEN \
  -e HF_HUB_OFFLINE=1 \
  -e HF_HUB_ENABLE_HF_TRANSFER=0 \
  -e NCCL_MIN_NCHANNELS=112 \
  -e SAFETENSORS_FAST_GPU=1 \
  --entrypoint python3 \
  atom-dpr:nightly-patched \
  -m atom.entrypoints.openai_server --model deepseek-ai/DeepSeek-V3-0324 \
    --port 8000 \
    -tp 8 \
    --enable-dp-attention \
    --gpu-memory-utilization 0.8 \
    --max-model-len 16384

2. Run atommesh in cache-aware mode

docker run \
  --rm \
  --name atomesh-router \
  --init --network host --ipc host --privileged \
  --cap-add SYS_PTRACE --security-opt seccomp=unconfined \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  --shm-size 256G \
  --group-add video \
  --device /dev/kfd --device /dev/dri --device /dev/infiniband \
  -e USE_ATOMESH_ENTRYPOINTS=1 \
  --entrypoint python3 \
  atommesh:nightly \
  -m atom.entrypoints.openai_server mesh-only \
    --host 0.0.0.0 \
    --port 30000 \
    --worker-urls "http://localhost:8000" \
    --dp-aware \
    --policy cache_aware

Expand to see dockerfile for building atommesh

# docker build -f Dockerfile.atommesh -t atommesh:nightly .
ARG BASE_IMAGE=rocm/atom-dev:nightly_202606161823
FROM ${BASE_IMAGE}

# Location of the editable ATOM checkout inside the base image.
ARG ATOM_SRC=/app/ATOM

# Build atomesh only if libmesh.so is not already in the base image.
# rustup is used so we get a current stable toolchain (edition 2021 crate).
# No protoc/OpenSSL needed: build.rs does no codegen and TLS is rustls.
RUN set -eux; \
    if [ ! -f "${ATOM_SRC}/atom/mesh/target/release/libmesh.so" ]; then \
        echo "libmesh.so not found in base image — building atomesh from source"; \
        if ! command -v cargo >/dev/null 2>&1; then \
            curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \
                | sh -s -- -y --default-toolchain stable --profile minimal; \
        fi; \
        . "$HOME/.cargo/env" 2>/dev/null || export PATH="$HOME/.cargo/bin:${PATH}"; \
        cd "${ATOM_SRC}"; \
        ATOM_MESH_BUILD=1 python -m pip install -e . ; \
        cd "${ATOM_SRC}/atom/mesh"; \
        strip target/release/atomesh || true; \
        cp target/release/atomesh /usr/local/bin/atomesh; \
    else \
        echo "libmesh.so already present in base image — skipping build"; \
        [ -x /usr/local/bin/atomesh ] || cp "${ATOM_SRC}/atom/mesh/target/release/atomesh" /usr/local/bin/atomesh; \
    fi; \
    /usr/local/bin/atomesh --version

# Standalone mode: Python owns the ATOM engine + tokenizer; Rust Atomesh serves
# HTTP and routing. Selected via USE_ATOMESH_ENTRYPOINTS=1 (no `mesh-only`).
ENV USE_ATOMESH_ENTRYPOINTS=1 \
    AITER_LOG_LEVEL=WARNING \
    SAFETENSORS_FAST_GPU=1

WORKDIR ${ATOM_SRC}

# All engine flags (model, -tp, --enable-dp-attention, etc.) are passed at
# `docker run` time as CMD args; --host/--port configure the mesh HTTP server.
ENTRYPOINT ["python", "-m", "atom.entrypoints.openai_server"]
CMD ["--help"]

Run a basic prefix repetition benchmark, like the one shipped by vLLM. This has 5 unique prefixes each of length 512 tokens, each running 100/5=20 times. So we would expect about (20-1)*5*512=48k cached tokens. The (20-1) comes from the fact that the first request of each prefix does not hit the cache, but the subsequent 19 requests of the same prefix does hit it.

ATOM logs the cache hit rate every 100 requests so we can read it off there.

docker run --rm -it     --entrypoint bash     --network host     --ipc host     vllm/vllm-openai-rocm:latest     -c 'vllm bench serve \
          --backend openai \
          --model deepseek-ai/DeepSeek-V3-0324 \
          --dataset-name prefix_repetition \
          --num-prompts 100 \
          --prefix-repetition-prefix-len 512 \
          --prefix-repetition-suffix-len 128 \
          --prefix-repetition-num-prefixes 5 --port 30000 \
          --prefix-repetition-output-len 128 --max-concurrency 1'

Test Result

Checking the logs, as expected we get 48k cached tokens.

# We see lots of these logs, showing the shared prefix is indeed cached
[atom 10:41:42] Scheduled prefill batch: 1 reqs, 128 new tokens (cached: [512], new: [128]), req_ids: (96,)  
...
# and finally we see 48128 cache hits
[atom 10:41:46] [Cache Stats Interval] Reqs: 100, Cached/Total tokens: 48128/63365, Hit rate: 75.95%   
[atom 10:41:46] [Cache Stats         ] Reqs: 100, Cached/Total tokens: 48128/63365, Hit rate: 75.95%

A round robin routing would yield a much lower hit rate.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[ZhangLirong-amd]

hi, thanks for your support! And we also have a similar pr #1505, @JiaoliangYu will discuss and work with you

[simondanielsson]

@ZhangLirong-amd @JiaoliangYu This PR is still needed to support cache aware routing using

• standalone ATOMesh mode (python3 -m atom.entrypoints.openai_server mesh-only)
• the /v1/chat/completions endpoint

Would appreciate a review, thanks 🙏