This repository is a vLLM-based research fork for the paper:
NIKA: Optimal KV Cache Transfer for Minimizing the Latency of Disaggregated LLM Inference
This README focuses on one goal: how to start the disaggregated inference system (prefill + decode + proxy) in Docker.
- Prefill server: GPU 0, port 8100
- Decode server: GPU 1, port 8200
- Proxy server: CPU container, port 8000
- Benchmark/client container: optional traffic generation
Proxy route:
- Send a short prefill request (
max_tokens=1) to prefill server. - Send the original request to decode server.
- Stream decode response back to client.
- Linux host with NVIDIA drivers and Docker GPU runtime.
- At least 2 GPUs (recommended for prefill/decode split).
- Hugging Face access for
meta-llama/Llama-3.2-1B-Instruct. - Local clone of this repository.
docker pull docker.io/erictsai90/myvllm:5090
git clone https://github.com/erictsaii/vllm-research.git ~/vllm-research
cd ~/vllm-researchdocker network create vllm-pd-netdocker run --rm -it \
--name vllm-proxy \
--network vllm-pd-net \
-v ~/vllm-research:/workspace/vllm \
-e OPENAI_API_KEY=EMPTY \
erictsai90/myvllm:5090 \
bashdocker run --rm -it \
--name vllm-prefill \
--network vllm-pd-net \
--device nvidia.com/gpu=0 \
--ipc=host \
--cap-add NET_ADMIN \
-p 8100:8100 \
-v ~/.cache/huggingface:/root/.cache/huggingface \
-v ~/vllm-research:/workspace/vllm \
-e NCCL_SOCKET_IFNAME=eth0 \
erictsai90/myvllm:5090 \
bashdocker run --rm -it \
--name vllm-decode \
--network vllm-pd-net \
--device nvidia.com/gpu=1 \
--ipc=host \
-p 8200:8200 \
-v ~/.cache/huggingface:/root/.cache/huggingface \
-v ~/vllm-research:/workspace/vllm \
-e NCCL_SOCKET_IFNAME=eth0 \
erictsai90/myvllm:5090 \
bashdocker run --rm -it \
--name vllm-benchmark \
--network vllm-pd-net \
-v ~/vllm-research:/workspace/vllm \
erictsai90/myvllm:5090 \
bashcd /workspace/vllm/research/scripts
export VLLM_HOST_IP=$(hostname -I | awk '{print $1}')
bash run_prefill.shUse the same KV producer IP as prefill.
cd /workspace/vllm/research/scripts
export VLLM_HOST_IP=vllm-prefill
bash run_decode.shcd /workspace/vllm
python research/disagg_prefill_proxy_server.pyFrom host machine (or benchmark container), send traffic to port 8000:
curl -X POST http://localhost:8000/v1/completions \
-H 'Content-Type: application/json' \
-d '{
"model": "meta-llama/Llama-3.2-1B-Instruct",
"prompt": "Explain NIKA disaggregated inference in one paragraph.",
"max_tokens": 64,
"temperature": 0.0,
"stream": false
}'In vllm-benchmark container:
cd /workspace/vllm/benchmarks/disagg_benchmarks
bash ~/vllm-research/research/scripts/performance_benchmark.sh