feat: vllm prefill router (#3155)

Signed-off-by: PeaBrane <yanrpei@gmail.com>
Signed-off-by: Kyle H <kylhuang@nvidia.com>

Author: PeaBrane

benchmarks/router/README.md

@@ -66,6 +66,24 @@ First, start the vLLM worker engines in a terminal.
     --tensor-parallel-size 2
 ```
 
+#### Prefill Workers
+
+You can also launch separate decode and prefill workers for disaggregated serving. This allows you to dedicate specific GPUs to prefill (prompt processing) and decode (token generation) tasks:
+
+```bash
+# Launch 4 decode workers (GPUs 0-3)
+./run_engines.sh \
+    --num-workers 4 \
+    --model-path deepseek-ai/DeepSeek-R1-Distill-Llama-8B
+
+# Launch 4 prefill workers (GPUs 4-7)
+./run_engines.sh \
+    --prefills \
+    --num-workers 4 \
+    --base-gpu-offset 4 \
+    --model-path deepseek-ai/DeepSeek-R1-Distill-Llama-8B
+```
+
 #### Alternative: Launch vLLM Mock Workers
 
 We also supports running lightweight mock engines that simulate vLLM behavior without performing actual model inference. Mocker engines are useful for testing router logic and performance without GPU requirements. Use the `--mockers` flag to run mocker engines instead of real vLLM workers.
@@ -106,6 +124,27 @@ python -m dynamo.frontend --help
 
 For detailed explanations of router arguments (especially KV cache routing parameters), see the [KV Cache Routing documentation](../../docs/architecture/kv_cache_routing.md).
 
+#### Launching a Prefill Router (Optional)
+
+If you're using disaggregated serving with separate prefill and decode workers, you should also launch a prefill router. The prefill router handles routing prefill requests to dedicated prefill workers. When using a prefill router, it's recommended to start the frontend (decode router) with `--kv-overlap-score-weight 0` for pure load balancing (as prefix-aware routing is now handled by the prefill router):
+
+```bash
+# Start the decode router with pure load balancing
+python -m dynamo.frontend \
+    --router-mode kv \
+    --kv-cache-block-size 64 \
+    --router-reset-states \
+    --http-port 8000 \
+    --kv-overlap-score-weight 0
+
+# In another terminal, start the prefill router (currently only supports vLLM)
+python -m dynamo.vllm_prefill_router \
+    --namespace dynamo \
+    --block-size 64
+```
+
+The prefill router will automatically coordinate with the decode router to handle request routing between prefill and decode workers.
+
 **Note**: If you're unsure whether your backend engines correctly emit KV events for certain models (e.g., hybrid models like gpt-oss or nemotron nano 2), use the `--no-kv-events` flag to disable KV event tracking and use approximate KV indexing instead:
 
 ```bash

benchmarks/router/ping.sh

@@ -3,8 +3,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-# Get port from first argument, default to 8080 if not provided
-PORT=${1:-8080}
+# Get port from first argument, default to 8000 if not provided
+PORT=${1:-8000}
 
 curl -X POST http://localhost:${PORT}/v1/chat/completions \
     -H "Content-Type: application/json" \

benchmarks/router/prefix_ratio_benchmark.py

@@ -309,7 +309,7 @@ def main():
         "--url",
         type=str,
         nargs="+",  # Accept multiple URLs
-        default=["http://localhost:8080"],
+        default=["http://localhost:8000"],
         # default=["http://localhost:8090", "http://localhost:8090"],
         help="Server URL(s). Can specify multiple URLs for parallel benchmarking",
     )

benchmarks/router/real_data_benchmark.py

@@ -118,7 +118,7 @@ def main():
     parser.add_argument(
         "--url",
         type=str,
-        default="http://localhost:8080",
+        default="http://localhost:8000",
         help="Server URL",
     )
     parser.add_argument(

benchmarks/router/run_engines.sh

@@ -8,6 +8,8 @@ NUM_WORKERS=8
 MODEL_PATH="deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
 TENSOR_PARALLEL_SIZE=1
 USE_MOCKERS=false
+USE_PREFILLS=false
+BASE_GPU_OFFSET=0
 EXTRA_ARGS=()
 
 # Parse arguments
@@ -29,6 +31,14 @@ while [[ $# -gt 0 ]]; do
             USE_MOCKERS=true
             shift
             ;;
+        --prefills)
+            USE_PREFILLS=true
+            shift
+            ;;
+        --base-gpu-offset)
+            BASE_GPU_OFFSET="$2"
+            shift 2
+            ;;
         --)
             shift
             EXTRA_ARGS+=("$@")
@@ -71,14 +81,22 @@ if ! [[ "$TENSOR_PARALLEL_SIZE" =~ ^[0-9]+$ ]] || [ "$TENSOR_PARALLEL_SIZE" -lt
     exit 1
 fi
 
+if ! [[ "$BASE_GPU_OFFSET" =~ ^[0-9]+$ ]]; then
+    echo "Error: BASE_GPU_OFFSET must be a non-negative integer"
+    exit 1
+fi
+
 # Calculate total GPUs needed
 TOTAL_GPUS_NEEDED=$((NUM_WORKERS * TENSOR_PARALLEL_SIZE))
+LAST_GPU=$((BASE_GPU_OFFSET + TOTAL_GPUS_NEEDED - 1))
 echo "Configuration:"
 echo "  Engine Type: $([ "$USE_MOCKERS" = true ] && echo "Mocker" || echo "vLLM")"
+echo "  Worker Type: $([ "$USE_PREFILLS" = true ] && echo "Prefill" || echo "Decode")"
 echo "  Workers: $NUM_WORKERS"
 echo "  Model: $MODEL_PATH"
 echo "  Tensor Parallel Size: $TENSOR_PARALLEL_SIZE"
 echo "  Total GPUs needed: $TOTAL_GPUS_NEEDED"
+echo "  GPU Range: $BASE_GPU_OFFSET-$LAST_GPU"
 echo "  Engine args: ${EXTRA_ARGS[*]}"
 echo ""
 
@@ -93,14 +111,15 @@ cleanup() {
 
 trap cleanup SIGINT SIGTERM
 
-echo "Starting $NUM_WORKERS workers..."
+WORKER_TYPE=$([ "$USE_PREFILLS" = true ] && echo "prefill" || echo "decode")
+echo "Starting $NUM_WORKERS $WORKER_TYPE workers..."
 
 for i in $(seq 1 $NUM_WORKERS); do
     {
-        echo "[Worker-$i] Starting..."
+        echo "[${WORKER_TYPE^} Worker-$i] Starting..."
 
-        # Calculate GPU indices for this worker
-        START_GPU=$(( (i - 1) * TENSOR_PARALLEL_SIZE ))
+        # Calculate GPU indices for this worker (with base offset)
+        START_GPU=$(( BASE_GPU_OFFSET + (i - 1) * TENSOR_PARALLEL_SIZE ))
         END_GPU=$(( START_GPU + TENSOR_PARALLEL_SIZE - 1 ))
 
         # Build CUDA_VISIBLE_DEVICES string
@@ -124,17 +143,22 @@ for i in $(seq 1 $NUM_WORKERS); do
                 --endpoint dyn://test.mocker.generate \
                 "${EXTRA_ARGS[@]}"
         else
-            echo "[Worker-$i] Using GPUs: $GPU_DEVICES"
+            echo "[${WORKER_TYPE^} Worker-$i] Using GPUs: $GPU_DEVICES"
             # Run vLLM engine with PYTHONHASHSEED=0 for deterministic event IDs in KV-aware routing
+            VLLM_ARGS=()
+            VLLM_ARGS+=("--model" "$MODEL_PATH")
+            VLLM_ARGS+=("--tensor-parallel-size" "$TENSOR_PARALLEL_SIZE")
+            if [ "$USE_PREFILLS" = true ]; then
+                VLLM_ARGS+=("--is-prefill-worker")
+            fi
+            VLLM_ARGS+=("${EXTRA_ARGS[@]}")
+
             exec env PYTHONHASHSEED=0 CUDA_VISIBLE_DEVICES=$GPU_DEVICES python -m dynamo.vllm \
-                --model "$MODEL_PATH" \
-                --endpoint dyn://test.vllm.generate \
-                --tensor-parallel-size $TENSOR_PARALLEL_SIZE \
-                "${EXTRA_ARGS[@]}"
+                "${VLLM_ARGS[@]}"
         fi
     } &
     PIDS+=($!)
-    echo "Started worker $i (PID: $!)"
+    echo "Started $WORKER_TYPE worker $i (PID: $!)"
 done
 
 echo "All workers started. Press Ctrl+C to stop."

components/backends/vllm/launch/agg_router.sh

@@ -4,11 +4,29 @@
 set -e
 trap 'echo Cleaning up...; kill 0' EXIT
 
-# run ingress
-python -m dynamo.frontend --router-mode kv --http-port=8000 &
+# Set deterministic hash for KV event IDs
+export PYTHONHASHSEED=0
+
+# Common configuration
+MODEL="Qwen/Qwen3-0.6B"
+BLOCK_SIZE=64
+
+# run frontend + KV router
+python -m dynamo.frontend \
+    --router-mode kv \
+    --http-port 8000 \
+    --router-reset-states &
 
 # run workers
 # --enforce-eager is added for quick deployment. for production use, need to remove this flag
-CUDA_VISIBLE_DEVICES=0 python3 -m dynamo.vllm --model Qwen/Qwen3-0.6B --enforce-eager --connector none &
+CUDA_VISIBLE_DEVICES=0 python3 -m dynamo.vllm \
+    --model $MODEL \
+    --block-size $BLOCK_SIZE \
+    --enforce-eager \
+    --connector none &
 
-CUDA_VISIBLE_DEVICES=1 python3 -m dynamo.vllm --model Qwen/Qwen3-0.6B --enforce-eager --connector none
+CUDA_VISIBLE_DEVICES=1 python3 -m dynamo.vllm \
+    --model $MODEL \
+    --block-size $BLOCK_SIZE \
+    --enforce-eager \
+    --connector none

components/backends/vllm/launch/disagg_router.sh

@@ -2,19 +2,48 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 set -e
-
 trap 'echo Cleaning up...; kill 0' EXIT
 
-# run ingress
-python -m dynamo.frontend --router-mode kv --http-port=8000 &
+# Set deterministic hash for KV event IDs
+export PYTHONHASHSEED=0
+
+# Common configuration
+MODEL="Qwen/Qwen3-0.6B"
+BLOCK_SIZE=64
+
+# run decode router with kv-overlap-score-weight 0 for pure load balancing
+python -m dynamo.frontend \
+    --router-mode kv \
+    --http-port 8000 \
+    --kv-overlap-score-weight 0 \
+    --router-reset-states &
 
-# routing will happen between the two decode workers
+# run prefill router service
+python -m dynamo.vllm_prefill_router \
+    --namespace dynamo \
+    --block-size $BLOCK_SIZE &
+
+# two decode workers
 # --enforce-eager is added for quick deployment. for production use, need to remove this flag
-CUDA_VISIBLE_DEVICES=0 python3 -m dynamo.vllm --model Qwen/Qwen3-0.6B --enforce-eager &
+CUDA_VISIBLE_DEVICES=0 python3 -m dynamo.vllm \
+    --model $MODEL \
+    --block-size $BLOCK_SIZE \
+    --enforce-eager &
 
-CUDA_VISIBLE_DEVICES=1 python3 -m dynamo.vllm --model Qwen/Qwen3-0.6B --enforce-eager &
+CUDA_VISIBLE_DEVICES=1 python3 -m dynamo.vllm \
+    --model $MODEL \
+    --block-size $BLOCK_SIZE \
+    --enforce-eager &
 
+# two prefill workers
 CUDA_VISIBLE_DEVICES=2 python3 -m dynamo.vllm \
-    --model Qwen/Qwen3-0.6B \
+    --model $MODEL \
+    --block-size $BLOCK_SIZE \
+    --enforce-eager \
+    --is-prefill-worker &
+
+CUDA_VISIBLE_DEVICES=3 python3 -m dynamo.vllm \
+    --model $MODEL \
+    --block-size $BLOCK_SIZE \
     --enforce-eager \
     --is-prefill-worker

components/backends/vllm/src/dynamo/vllm/handlers.py

@@ -94,9 +94,13 @@ def __init__(
         engine,
         default_sampling_params,
         prefill_worker_client=None,
+        prefill_router_client=None,
+        prefill_router_free_client=None,
     ):
         super().__init__(runtime, component, engine, default_sampling_params)
         self.prefill_worker_client = prefill_worker_client
+        self.prefill_router_client = prefill_router_client
+        self.prefill_router_free_client = prefill_router_free_client
         self.can_prefill = 0
         self._prefill_check_task = None
 
@@ -143,7 +147,11 @@ async def generate(self, request, context):
             if value is not None and hasattr(sampling_params, key):
                 setattr(sampling_params, key, value)
 
-        # TODO Change to prefill queue
+        # TODO: Change to prefill queue
+        # TODO: (PeaBrane) eventually, do not use a router_client and a free_client directly.
+        # This is least intrusive for now, but quite error prone. Should consider (major) refactoring
+        # TODO: (PeaBrane) longer term, decode workers should not handle prefill routing at all.
+        # Prefill routing logic should be integrated directly into the frontend service potentially.
         if self.can_prefill:
             # Create a copy for prefill with specific modifications
             prefill_sampling_params = deepcopy(sampling_params)
@@ -162,12 +170,37 @@ async def generate(self, request, context):
                 "request_id": request_id,
             }
 
+            used_prefill_router = False
             try:
-                prefill_response = await anext(
-                    await self.prefill_worker_client.round_robin(
-                        prefill_request, context=context
+                prefill_worker_id = None
+                if (
+                    self.prefill_router_client is not None
+                    and self.prefill_router_client.instance_ids()
+                ):
+                    used_prefill_router = True
+                    best_worker_response = await anext(
+                        await self.prefill_router_client.generate(
+                            {
+                                "token_ids": request["token_ids"],
+                                "request_id": request_id,
+                            }
+                        )
                     )
-                )
+                    prefill_worker_id = best_worker_response.data().get("worker_id")
+
+                if prefill_worker_id is not None:
+                    prefill_response = await anext(
+                        await self.prefill_worker_client.direct(
+                            prefill_request, prefill_worker_id, context=context
+                        )
+                    )
+                else:
+                    prefill_response = await anext(
+                        await self.prefill_worker_client.round_robin(
+                            prefill_request, context=context
+                        )
+                    )
+
             except Exception as e:
                 # TODO: Cancellation does not propagate until the first token is received
                 if context.is_stopped() or context.is_killed():
@@ -176,6 +209,15 @@ async def generate(self, request, context):
                     return
                 raise e
 
+            finally:
+                if used_prefill_router:
+                    await anext(
+                        await self.prefill_router_free_client.generate(
+                            {"request_id": request_id}
+                        )
+                    )
+                    logger.debug(f"Freed router state for request {request_id}")
+
             prefill_response = MyRequestOutput.model_validate_json(
                 prefill_response.data()
             )