@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d19110204c
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 11:00 UTC

[P/D] NIXL Integration (#17751)

Signed-off-by: ApostaC <yihua98@uchicago.edu> Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com> Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com> Signed-off-by: Robert Shaw <rshaw@neuralmagic.com> Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Nick Hill <nhill@redhat.com> Signed-off-by: Brent Salisbury <bsalisbu@redhat.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: ApostaC <yihua98@uchicago.edu> Co-authored-by: Robert Shaw <rshaw@neuralmagic.com> Co-authored-by: mgoin <mgoin64@gmail.com> Co-authored-by: Nick Hill <nhill@redhat.com> Co-authored-by: Tyler Michael Smith <tysmith@redhat.com> Co-authored-by: Brent Salisbury <bsalisbu@redhat.com>

Robert Shaw committed May 12, 2025 at 12:46 UTC d19110204c03e9b77ed957fc70c1262ff370f5e2
34 files changed +2724 -109
.buildkite/test-pipeline.yaml
+1
@@ -214,6 +214,7 @@ steps:
214 - pytest -v -s v1/worker
215 - pytest -v -s v1/structured_output
216 - pytest -v -s v1/spec_decode
217 + - pytest -v -s v1/kv_connector/unit
218 - pytest -v -s v1/test_serial_utils.py
219 - pytest -v -s v1/test_stats.py
220 - pytest -v -s v1/test_utils.py
tests/v1/core/test_scheduler.py
+3 -3
@@ -870,7 +870,7 @@ def test_kv_connector_basic():
870 NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE * 2
871 scheduler.connector.get_num_new_matched_tokens = Mock(name="method")
872 scheduler.connector.get_num_new_matched_tokens.return_value = (
873 - NUM_MATCHED_NEW_TOKENS)
873 + NUM_MATCHED_NEW_TOKENS, False)
874
875 ######################################################
876 # FIRST SET OF REQUESTS - External Hit Only
@@ -981,7 +981,7 @@ def test_kv_connector_unable_to_allocate():
981 NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE * 2
982 scheduler.connector.get_num_new_matched_tokens = Mock(name="method")
983 scheduler.connector.get_num_new_matched_tokens.return_value = (
984 - NUM_MATCHED_NEW_TOKENS)
984 + NUM_MATCHED_NEW_TOKENS, False)
985
986 # Create two requests. The second request will not be able to
987 # allocate slots because it will not have enough blocks.
@@ -1060,7 +1060,7 @@ def test_kv_connector_handles_preemption():
1060 NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE
1061 scheduler.connector.get_num_new_matched_tokens = Mock(name="method")
1062 scheduler.connector.get_num_new_matched_tokens.return_value = (
1063 - NUM_MATCHED_NEW_TOKENS)
1063 + NUM_MATCHED_NEW_TOKENS, False)
1064
1065 # Create two requests.
1066 # Both can be scheduled at first, but the second request
tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh new
+171
@@ -0,0 +1,171 @@
1 +#!/bin/bash
2 +set -xe
3 +
4 +# Models to run
5 +MODELS=(
6 + "Qwen/Qwen3-0.6B"
7 +)
8 +
9 +# Number of prefill and decode instances to create
10 +NUM_PREFILL_INSTANCES=${NUM_PREFILL_INSTANCES:-1} # Default to 1
11 +NUM_DECODE_INSTANCES=${NUM_DECODE_INSTANCES:-2} # Default to 2
12 +
13 +# Find the git repository root directory
14 +GIT_ROOT=$(git rev-parse --show-toplevel)
15 +
16 +# Trap the SIGINT signal (triggered by Ctrl+C)
17 +trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT
18 +
19 +# Waits for vLLM to start.
20 +wait_for_server() {
21 + local port=$1
22 + timeout 1200 bash -c "
23 + until curl -s localhost:${port}/v1/completions > /dev/null; do
24 + sleep 1
25 + done" && return 0 || return 1
26 +}
27 +
28 +# Function to clean up previous instances
29 +cleanup_instances() {
30 + echo "Cleaning up any running vLLM instances..."
31 + pkill -f "vllm serve" || true
32 + sleep 2
33 +}
34 +
35 +# Handle to get model-specific arguments for deepseek
36 +get_model_args() {
37 + local model_name=$1
38 + local extra_args=""
39 +
40 + if [[ "$model_name" == "deepseek-ai/deepseek-vl2-tiny" ]]; then
41 + extra_args="--hf_overrides '{\"architectures\": [\"DeepseekVLV2ForCausalLM\"]}' --trust-remote-code"
42 + fi
43 +
44 + echo "$extra_args"
45 +}
46 +
47 +
48 +# Function to run tests for a specific model
49 +run_tests_for_model() {
50 + local model_name=$1
51 + echo "================================"
52 + echo "Testing model: $model_name"
53 + echo "================================"
54 +
55 + # Get model-specific arguments
56 + local model_args=$(get_model_args "$model_name")
57 +
58 + # Arrays to store all hosts and ports
59 + PREFILL_HOSTS=()
60 + PREFILL_PORTS=()
61 + DECODE_HOSTS=()
62 + DECODE_PORTS=()
63 +
64 + # Start prefill instances
65 + for i in $(seq 0 $((NUM_PREFILL_INSTANCES-1))); do
66 + # Calculate GPU ID - we'll distribute across available GPUs
67 + GPU_ID=$((i % $(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)))
68 + # Calculate port number (base port + instance number)
69 + PORT=$((8100 + i))
70 + # Calculate side channel port
71 + SIDE_CHANNEL_PORT=$((5559 + i))
72 +
73 + echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT"
74 +
75 + # Build the command with or without model-specific args
76 + BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT vllm serve $model_name \
77 + --port $PORT \
78 + --enforce-eager \
79 + --disable-log-requests \
80 + --gpu-memory-utilization 0.2 \
81 + --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\"}'"
82 +
83 + if [ -n "$model_args" ]; then
84 + FULL_CMD="$BASE_CMD $model_args"
85 + else
86 + FULL_CMD="$BASE_CMD"
87 + fi
88 +
89 + eval "$FULL_CMD &"
90 +
91 + # Store host and port for proxy configuration
92 + PREFILL_HOSTS+=("localhost")
93 + PREFILL_PORTS+=($PORT)
94 + done
95 +
96 + # Start decode instances
97 + for i in $(seq 0 $((NUM_DECODE_INSTANCES-1))); do
98 + # Calculate GPU ID - we'll distribute across available GPUs, starting from after prefill GPUs
99 + GPU_ID=$(((i + NUM_PREFILL_INSTANCES) % $(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)))
100 + # Calculate port number (base port + instance number)
101 + PORT=$((8200 + i))
102 + # Calculate side channel port
103 + SIDE_CHANNEL_PORT=$((5659 + i))
104 +
105 + echo "Starting decode instance $i on GPU $GPU_ID, port $PORT"
106 +
107 + # Build the command with or without model-specific args
108 + BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT vllm serve $model_name \
109 + --port $PORT \
110 + --enforce-eager \
111 + --disable-log-requests \
112 + --gpu-memory-utilization 0.2 \
113 + --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\"}'"
114 +
115 + if [ -n "$model_args" ]; then
116 + FULL_CMD="$BASE_CMD $model_args"
117 + else
118 + FULL_CMD="$BASE_CMD"
119 + fi
120 +
121 + eval "$FULL_CMD &"
122 +
123 + # Store host and port for proxy configuration
124 + DECODE_HOSTS+=("localhost")
125 + DECODE_PORTS+=($PORT)
126 + done
127 +
128 + # Wait for all instances to start
129 + for PORT in "${PREFILL_PORTS[@]}"; do
130 + echo "Waiting for prefill instance on port $PORT to start..."
131 + wait_for_server $PORT
132 + done
133 +
134 + for PORT in "${DECODE_PORTS[@]}"; do
135 + echo "Waiting for decode instance on port $PORT to start..."
136 + wait_for_server $PORT
137 + done
138 +
139 + # Build the command for the proxy server with all the hosts and ports
140 + PROXY_CMD="python ${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py --port 8192"
141 +
142 + # Add all prefill hosts and ports
143 + PROXY_CMD+=" --prefiller-hosts ${PREFILL_HOSTS[@]}"
144 + PROXY_CMD+=" --prefiller-ports ${PREFILL_PORTS[@]}"
145 +
146 + # Add all decode hosts and ports
147 + PROXY_CMD+=" --decoder-hosts ${DECODE_HOSTS[@]}"
148 + PROXY_CMD+=" --decoder-ports ${DECODE_PORTS[@]}"
149 +
150 + # Start the proxy server
151 + echo "Starting proxy server with command: $PROXY_CMD"
152 + $PROXY_CMD &
153 +
154 + # Wait for the proxy to start
155 + sleep 5
156 +
157 + # Run lm eval for this model
158 + echo "Running tests for $model_name"
159 + TEST_MODEL=$model_name python -m pytest -s -x ${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/test_accuracy.py
160 +
161 + # Clean up before running next model
162 + cleanup_instances
163 + sleep 3
164 +}
165 +
166 +# Run tests for each model
167 +for model in "${MODELS[@]}"; do
168 + run_tests_for_model "$model"
169 +done
170 +
171 +echo "All tests completed!"
tests/v1/kv_connector/nixl_integration/run_edge_case_test.sh new
+123
@@ -0,0 +1,123 @@
1 +#!/bin/bash
2 +set -xe
3 +
4 +# Models to run
5 +MODELS=(
6 + "Qwen/Qwen3-0.6B"
7 +)
8 +
9 +# Find the git repository root directory
10 +GIT_ROOT=$(git rev-parse --show-toplevel)
11 +
12 +# Trap the SIGINT signal (triggered by Ctrl+C)
13 +trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT
14 +
15 +# Waits for vLLM to start.
16 +wait_for_server() {
17 + local port=$1
18 + timeout 1200 bash -c "
19 + until curl -s localhost:${port}/v1/completions > /dev/null; do
20 + sleep 1
21 + done" && return 0 || return 1
22 +}
23 +
24 +# Function to clean up previous instances
25 +cleanup_instances() {
26 + echo "Cleaning up any running vLLM instances..."
27 + pkill -f "vllm serve" || true
28 + sleep 2
29 +}
30 +
31 +# Handle to get model-specific arguments for deepseek
32 +get_model_args() {
33 + local model_name=$1
34 + local extra_args=""
35 +
36 + if [[ "$model_name" == "deepseek-ai/deepseek-vl2-tiny" ]]; then
37 + extra_args="--hf_overrides '{\"architectures\": [\"DeepseekVLV2ForCausalLM\"]}' --trust-remote-code"
38 + fi
39 +
40 + echo "$extra_args"
41 +}
42 +
43 +
44 +# Function to run tests for a specific model
45 +run_tests_for_model() {
46 + local model_name=$1
47 + echo "================================"
48 + echo "Testing model: $model_name"
49 + echo "================================"
50 +
51 + # Get model-specific arguments
52 + local model_args=$(get_model_args "$model_name")
53 +
54 + # Start prefill instance
55 + PREFILL_PORT=8001
56 +
57 + BASE_CMD="CUDA_VISIBLE_DEVICES=0 VLLM_NIXL_SIDE_CHANNEL_PORT=5559 vllm serve $model_name \
58 + --port $PREFILL_PORT \
59 + --enforce-eager \
60 + --disable-log-requests \
61 + --gpu-memory-utilization 0.2 \
62 + --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\"}'"
63 +
64 + if [ -n "$model_args" ]; then
65 + FULL_CMD="$BASE_CMD $model_args"
66 + else
67 + FULL_CMD="$BASE_CMD"
68 + fi
69 +
70 + eval "$FULL_CMD &"
71 +
72 + # Start decode instance
73 + DECODE_PORT=8002
74 +
75 + # Build the command with or without model-specific args
76 + BASE_CMD="CUDA_VISIBLE_DEVICES=1 VLLM_NIXL_SIDE_CHANNEL_PORT=6000 vllm serve $model_name \
77 + --port $DECODE_PORT \
78 + --enforce-eager \
79 + --disable-log-requests \
80 + --gpu-memory-utilization 0.2 \
81 + --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\"}'"
82 +
83 + if [ -n "$model_args" ]; then
84 + FULL_CMD="$BASE_CMD $model_args"
85 + else
86 + FULL_CMD="$BASE_CMD"
87 + fi
88 +
89 + eval "$FULL_CMD &"
90 +
91 + # Wait for all instances to start
92 + echo "Waiting for prefill instance on port $PORT to start..."
93 + wait_for_server $PREFILL_PORT
94 + echo "Waiting for decode instance on port $PORT to start..."
95 + wait_for_server $DECODE_PORT
96 +
97 + # Build the command for the proxy server with all the hosts and ports
98 + PROXY_PORT=8192
99 + PROXY_CMD="python ${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py --port $PROXY_PORT"
100 + PROXY_CMD+=" --prefiller-ports ${PREFILL_PORT}"
101 + PROXY_CMD+=" --decoder-ports ${DECODE_PORT}"
102 + # Start the proxy server
103 + echo "Starting proxy server with command: $PROXY_CMD"
104 + $PROXY_CMD &
105 +
106 + # Wait for the proxy to start
107 + sleep 5
108 +
109 + # Run lm eval for this model
110 + echo "Running tests for $model_name"
111 + PREFILL_PORT=$PREFILL_PORT DECODE_PORT=$DECODE_PORT PROXY_PORT=$PROXY_PORT python -m pytest -s -v ${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/test_edge_cases.py
112 +
113 + # Clean up before running next model
114 + cleanup_instances
115 + sleep 3
116 +}
117 +
118 +# Run tests for each model
119 +for model in "${MODELS[@]}"; do
120 + run_tests_for_model "$model"
121 +done
122 +
123 +echo "All tests completed!"
tests/v1/kv_connector/nixl_integration/test_accuracy.py new
+60
@@ -0,0 +1,60 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import os
3 +
4 +import lm_eval
5 +import openai
6 +
7 +BASE_URL = "http://localhost:8192/v1"
8 +NUM_CONCURRENT = 100
9 +TASK = "gsm8k"
10 +FILTER = "exact_match,strict-match"
11 +RTOL = 0.03
12 +
13 +# Model-specific expected values
14 +EXPECTED_VALUES = {
15 + "Qwen/Qwen3-0.6B": 0.41,
16 +}
17 +
18 +SIMPLE_PROMPT = "The best part about working on vLLM is that I got to meet so many people across various different organizations like UCB, Google, and Meta which means", # noqa: E501
19 +
20 +# Get model name from environment variable
21 +MODEL_NAME = os.environ.get("TEST_MODEL", "Qwen/Qwen3-0.6B")
22 +
23 +
24 +def run_simple_prompt():
25 + client = openai.OpenAI(api_key="EMPTY", base_url=BASE_URL)
26 + completion = client.completions.create(model=MODEL_NAME,
27 + prompt=SIMPLE_PROMPT)
28 +
29 + print("-" * 50)
30 + print(f"Completion results for {MODEL_NAME}:")
31 + print(completion)
32 + print("-" * 50)
33 +
34 +
35 +def test_accuracy():
36 + """Run the end to end accuracy test."""
37 + run_simple_prompt()
38 +
39 + model_args = (f"model={MODEL_NAME},"
40 + f"base_url={BASE_URL}/completions,"
41 + f"num_concurrent={NUM_CONCURRENT},tokenized_requests=False")
42 +
43 + results = lm_eval.simple_evaluate(
44 + model="local-completions",
45 + model_args=model_args,
46 + tasks=TASK,
47 + )
48 +
49 + measured_value = results["results"][TASK][FILTER]
50 + expected_value = EXPECTED_VALUES.get(MODEL_NAME)
51 +
52 + if expected_value is None:
53 + print(f"Warning: No expected value found for {MODEL_NAME}. "
54 + "Skipping accuracy check.")
55 + print(f"Measured value: {measured_value}")
56 + return
57 +
58 + assert (measured_value - RTOL < expected_value
59 + and measured_value + RTOL > expected_value
60 + ), f"Expected: {expected_value} | Measured: {measured_value}"
tests/v1/kv_connector/nixl_integration/test_edge_cases.py new
+77
@@ -0,0 +1,77 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import os
3 +
4 +import openai
5 +
6 +PREFILL_PORT = os.getenv("PREFILL_PORT", None)
7 +DECODE_PORT = os.getenv("DECODE_PORT", None)
8 +PROXY_PORT = os.getenv("PROXY_PORT", None)
9 +
10 +if PREFILL_PORT is None or DECODE_PORT is None or PROXY_PORT is None:
11 + raise ValueError(
12 + "Please set the PREFILL_PORT, DECODE_PORT, and PROXY_PORT.")
13 +
14 +LONG_PROMPT = "Red Hat is the best company in the world to work for because it works on open source software, which means that all the contributions are delivered to the community. As a result, when working on projects like vLLM we are able to meet many amazing people from various organizations like AMD, Google, NVIDIA, " # noqa: E501
15 +PROMPT = "Red Hat is the best company in the world to work for because it works on open source software, which means that all the contributions are delivered to the community. As a result," # noqa: E501
16 +SHORT_PROMPT = "Red Hat is "
17 +
18 +
19 +def test_edge_cases():
20 + # Set the OpenAI API key and base URL
21 + decode_client = openai.OpenAI(
22 + api_key="MY_KEY",
23 + base_url=f"http://localhost:{DECODE_PORT}/v1",
24 + )
25 + prefill_client = openai.OpenAI(
26 + api_key="MY_KEY",
27 + base_url=f"http://localhost:{PREFILL_PORT}/v1",
28 + )
29 + proxy_client = openai.OpenAI(
30 + api_key="MY_KEY",
31 + base_url=f"http://localhost:{PROXY_PORT}/v1",
32 + )
33 +
34 + # Get the list of models
35 + models = decode_client.models.list()
36 + MODEL = models.data[0].id
37 +
38 + # (1) Check that we can handle a very short prompt,
39 + # less than the length of the block size.
40 + completion = proxy_client.completions.create(model=MODEL,
41 + prompt=SHORT_PROMPT,
42 + temperature=0)
43 + proxy_response = completion.choices[0].text
44 + completion = prefill_client.completions.create(model=MODEL,
45 + prompt=SHORT_PROMPT,
46 + temperature=0)
47 + prefill_response = completion.choices[0].text
48 + print(f"SMALL PROMPT: {proxy_response=}")
49 + assert proxy_response == prefill_response
50 +
51 + # (2) Check that we can handle a full prefix cache
52 + # hit on the D worker but not on the P worker.
53 + # (2a): prime the D worker.
54 + completion = decode_client.completions.create(model=MODEL,
55 + prompt=PROMPT,
56 + temperature=0)
57 + decode_response = completion.choices[0].text
58 + # (2b): send via the P/D setup
59 + completion = proxy_client.completions.create(model=MODEL,
60 + prompt=PROMPT,
61 + temperature=0)
62 + proxy_response = completion.choices[0].text
63 + print(f"FULL CACHE HIT: {proxy_response=}")
64 + assert proxy_response == decode_response
65 +
66 + # (3) Check that we can handle a partial prefix cache
67 + # hit on the D worker.
68 + completion = proxy_client.completions.create(model=MODEL,
69 + prompt=LONG_PROMPT,
70 + temperature=0)
71 + proxy_response = completion.choices[0].text
72 + completion = prefill_client.completions.create(model=MODEL,
73 + prompt=LONG_PROMPT,
74 + temperature=0)
75 + prefill_response = completion.choices[0].text
76 + print(f"PARTIAL CACHE HIT: {proxy_response=}")
77 + assert proxy_response == prefill_response
tests/v1/kv_connector/nixl_integration/toy_proxy_server.py new
+260
@@ -0,0 +1,260 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +import argparse
4 +import itertools
5 +import os
6 +import uuid
7 +from contextlib import asynccontextmanager
8 +
9 +import httpx
10 +from fastapi import FastAPI, Request
11 +from fastapi.responses import StreamingResponse
12 +
13 +from vllm.logger import init_logger
14 +
15 +logger = init_logger(__name__)
16 +
17 +
18 +@asynccontextmanager
19 +async def lifespan(app: FastAPI):
20 + """
21 + Lifespan context manager to handle startup and shutdown events.
22 + """
23 + # Startup: Initialize client pools for prefiller and decoder services
24 + app.state.prefill_clients = []
25 + app.state.decode_clients = []
26 +
27 + # Create prefill clients
28 + for i, (host, port) in enumerate(global_args.prefiller_instances):
29 + prefiller_base_url = f'http://{host}:{port}/v1'
30 + app.state.prefill_clients.append({
31 + 'client':
32 + httpx.AsyncClient(timeout=None, base_url=prefiller_base_url),
33 + 'host':
34 + host,
35 + 'port':
36 + port,
37 + 'id':
38 + i
39 + })
40 +
41 + # Create decode clients
42 + for i, (host, port) in enumerate(global_args.decoder_instances):
43 + decoder_base_url = f'http://{host}:{port}/v1'
44 + app.state.decode_clients.append({
45 + 'client':
46 + httpx.AsyncClient(timeout=None, base_url=decoder_base_url),
47 + 'host':
48 + host,
49 + 'port':
50 + port,
51 + 'id':
52 + i
53 + })
54 +
55 + # Initialize round-robin iterators
56 + app.state.prefill_iterator = itertools.cycle(
57 + range(len(app.state.prefill_clients)))
58 + app.state.decode_iterator = itertools.cycle(
59 + range(len(app.state.decode_clients)))
60 +
61 + print(f"Initialized {len(app.state.prefill_clients)} prefill clients "
62 + f"and {len(app.state.decode_clients)} decode clients.")
63 +
64 + yield
65 +
66 + # Shutdown: Close all clients
67 + for client_info in app.state.prefill_clients:
68 + await client_info['client'].aclose()
69 +
70 + for client_info in app.state.decode_clients:
71 + await client_info['client'].aclose()
72 +
73 +
74 +# Update FastAPI app initialization to use lifespan
75 +app = FastAPI(lifespan=lifespan)
76 +
77 +
78 +def parse_args():
79 + parser = argparse.ArgumentParser()
80 +
81 + parser.add_argument("--port", type=int, default=8000)
82 + parser.add_argument("--host", type=str, default="localhost")
83 +
84 + # For prefiller instances
85 + parser.add_argument("--prefiller-hosts",
86 + "--prefiller-host",
87 + type=str,
88 + nargs="+",
89 + default=["localhost"])
90 + parser.add_argument("--prefiller-ports",
91 + "--prefiller-port",
92 + type=int,
93 + nargs="+",
94 + default=[8100])
95 +
96 + # For decoder instances
97 + parser.add_argument("--decoder-hosts",
98 + "--decoder-host",
99 + type=str,
100 + nargs="+",
101 + default=["localhost"])
102 + parser.add_argument("--decoder-ports",
103 + "--decoder-port",
104 + type=int,
105 + nargs="+",
106 + default=[8200])
107 +
108 + args = parser.parse_args()
109 +
110 + # Validate and pair hosts with ports
111 + if len(args.prefiller_hosts) != len(args.prefiller_ports):
112 + raise ValueError(
113 + "Number of prefiller hosts must match number of prefiller ports")
114 +
115 + if len(args.decoder_hosts) != len(args.decoder_ports):
116 + raise ValueError(
117 + "Number of decoder hosts must match number of decoder ports")
118 +
119 + # Create tuples of (host, port) for each service type
120 + args.prefiller_instances = list(
121 + zip(args.prefiller_hosts, args.prefiller_ports))
122 + args.decoder_instances = list(zip(args.decoder_hosts, args.decoder_ports))
123 +
124 + return args
125 +
126 +
127 +def get_next_client(app, service_type: str):
128 + """
129 + Get the next client in round-robin fashion.
130 +
131 + Args:
132 + app: The FastAPI app instance
133 + service_type: Either 'prefill' or 'decode'
134 +
135 + Returns:
136 + The next client to use
137 + """
138 + if service_type == 'prefill':
139 + client_idx = next(app.state.prefill_iterator)
140 + return app.state.prefill_clients[client_idx]
141 + elif service_type == 'decode':
142 + client_idx = next(app.state.decode_iterator)
143 + return app.state.decode_clients[client_idx]
144 + else:
145 + raise ValueError(f"Unknown service type: {service_type}")
146 +
147 +
148 +async def send_request_to_service(client_info: dict, endpoint: str,
149 + req_data: dict, request_id: str):
150 + """
151 + Send a request to a service using a client from the pool.
152 + """
153 + req_data = req_data.copy()
154 + req_data['kv_transfer_params'] = {
155 + "do_remote_decode": True,
156 + "do_remote_prefill": False,
157 + "remote_engine_id": None,
158 + "remote_block_ids": None,
159 + "remote_host": None,
160 + "remote_port": None
161 + }
162 + req_data["stream"] = False
163 + req_data["max_tokens"] = 1
164 + if "stream_options" in req_data:
165 + del req_data["stream_options"]
166 + headers = {
167 + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
168 + "X-Request-Id": request_id
169 + }
170 +
171 + response = await client_info['client'].post(endpoint,
172 + json=req_data,
173 + headers=headers)
174 + response.raise_for_status()
175 +
176 + return response
177 +
178 +
179 +async def stream_service_response(client_info: dict, endpoint: str,
180 + req_data: dict, request_id: str):
181 + """
182 + Asynchronously stream response from a service using a client from the pool.
183 + """
184 + headers = {
185 + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
186 + "X-Request-Id": request_id
187 + }
188 +
189 + async with client_info['client'].stream("POST",
190 + endpoint,
191 + json=req_data,
192 + headers=headers) as response:
193 + response.raise_for_status()
194 + async for chunk in response.aiter_bytes():
195 + yield chunk
196 +
197 +
198 +@app.post("/v1/completions")
199 +async def handle_completions(request: Request):
200 + try:
201 + req_data = await request.json()
202 + request_id = str(uuid.uuid4())
203 +
204 + # Get the next prefill client in round-robin fashion
205 + prefill_client_info = get_next_client(request.app, 'prefill')
206 +
207 + # Send request to prefill service
208 + response = await send_request_to_service(prefill_client_info,
209 + "/completions", req_data,
210 + request_id)
211 +
212 + # Extract the needed fields
213 + response_json = response.json()
214 + kv_transfer_params = response_json.get('kv_transfer_params', {})
215 + if kv_transfer_params:
216 + req_data["kv_transfer_params"] = kv_transfer_params
217 +
218 + # Get the next decode client in round-robin fashion
219 + decode_client_info = get_next_client(request.app, 'decode')
220 +
221 + logger.debug("Using %s %s", prefill_client_info, decode_client_info)
222 +
223 + # Stream response from decode service
224 + async def generate_stream():
225 + async for chunk in stream_service_response(decode_client_info,
226 + "/completions",
227 + req_data,
228 + request_id=request_id):
229 + yield chunk
230 +
231 + return StreamingResponse(generate_stream(),
232 + media_type="application/json")
233 +
234 + except Exception as e:
235 + import sys
236 + import traceback
237 + exc_info = sys.exc_info()
238 + print("Error occurred in disagg prefill proxy server"
239 + " - completions endpoint")
240 + print(e)
241 + print("".join(traceback.format_exception(*exc_info)))
242 + raise
243 +
244 +
245 +@app.get("/healthcheck")
246 +async def healthcheck():
247 + """Simple endpoint to check if the server is running."""
248 + return {
249 + "status": "ok",
250 + "prefill_instances": len(app.state.prefill_clients),
251 + "decode_instances": len(app.state.decode_clients)
252 + }
253 +
254 +
255 +if __name__ == '__main__':
256 + global global_args
257 + global_args = parse_args()
258 +
259 + import uvicorn
260 + uvicorn.run(app, host=global_args.host, port=global_args.port)
tests/v1/kv_connector/unit/__init__.py
tests/v1/kv_connector/unit/test_nixl_connector.py new
+73
@@ -0,0 +1,73 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
4 + NixlConnectorMetadata)
5 +
6 +from .utils import create_request, create_scheduler, create_vllm_config
7 +
8 +
9 +def test_basic_inferface():
10 + """Unit test for basic NixlConnector interface functionality."""
11 +
12 + vllm_config = create_vllm_config()
13 + scheduler = create_scheduler(vllm_config)
14 +
15 + # 2 Full Blocks and 1 Half Block.
16 + BLOCK_SIZE = vllm_config.cache_config.block_size
17 + NUM_EXTERNAL_FULL_BLOCKS = 2
18 + NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
19 +
20 + request = create_request(request_id=1,
21 + num_tokens=NUM_TOKENS,
22 + do_remote_prefill=True)
23 + request_id = request.request_id
24 +
25 + scheduler.add_request(request)
26 +
27 + # Remote Prefill, triggers NixlConnectorMetdata.
28 + scheduler_output = scheduler.schedule()
29 + kv_connector_metadata = scheduler_output.kv_connector_metadata
30 + assert kv_connector_metadata is not None
31 + assert isinstance(kv_connector_metadata, NixlConnectorMetadata)
32 +
33 + assert len(kv_connector_metadata.requests) == 1
34 + assert request_id in kv_connector_metadata.requests
35 + req_meta = kv_connector_metadata.requests[request_id]
36 +
37 + for block_id, block in zip(
38 + req_meta.local_block_ids, scheduler.kv_cache_manager.
39 + single_type_manager.req_to_blocks[request_id]):
40 + assert block_id == block.block_id
41 +
42 +
43 +def test_prompt_less_than_block_size():
44 + """
45 + Test that we can handle case where prompt is < block.
46 +
47 + In this case, the P worker will send empty remote_block_ids.
48 + The D worker should not schedule an async read in this case,
49 + since there is nothing to pull.
50 + """
51 + vllm_config = create_vllm_config()
52 + scheduler = create_scheduler(vllm_config)
53 +
54 + # Half of a block.
55 + BLOCK_SIZE = vllm_config.cache_config.block_size
56 + NUM_TOKENS = int(BLOCK_SIZE * 0.5)
57 +
58 + # Request will have 0 remote blocks.
59 + request = create_request(request_id=1,
60 + num_tokens=NUM_TOKENS,
61 + do_remote_prefill=True,
62 + num_remote_blocks=0)
63 + scheduler.add_request(request)
64 + scheduler_output = scheduler.schedule()
65 +
66 + # This request should not have to read async.
67 + kv_connector_metadata = scheduler_output.kv_connector_metadata
68 + assert kv_connector_metadata is not None
69 + assert isinstance(kv_connector_metadata, NixlConnectorMetadata)
70 + assert len(kv_connector_metadata.requests) == 0
71 +
72 + # This request should be scheduled regularly.
73 + assert len(scheduler_output.scheduled_new_reqs) == 1
tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py new
+181
@@ -0,0 +1,181 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import copy
3 +
4 +from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
5 +from vllm.v1.request import FinishReason, RequestStatus
6 +
7 +from .utils import (assert_scheduler_empty, create_model_runner_output,
8 + create_request, create_scheduler, create_vllm_config)
9 +
10 +
11 +def test_basic_lifecycle():
12 + """Test lifecycle of a Remote Decode request."""
13 +
14 + vllm_config = create_vllm_config()
15 + scheduler = create_scheduler(vllm_config)
16 +
17 + # 2 Full Blocks and 1 Half Block.
18 + BLOCK_SIZE = vllm_config.cache_config.block_size
19 + NUM_EXTERNAL_FULL_BLOCKS = 2
20 + NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
21 +
22 + request = create_request(request_id=1,
23 + max_tokens=1,
24 + num_tokens=NUM_TOKENS,
25 + do_remote_decode=True)
26 +
27 + scheduler.add_request(request)
28 + request_id = request.request_id
29 +
30 + # STEP (1): Prefill.
31 + # (1a): schedule()
32 + scheduler_output = scheduler.schedule()
33 + assert len(scheduler.running) == 1
34 + assert len(scheduler_output.scheduled_new_reqs) == 1
35 +
36 + # (1b): execute_model()
37 + model_runner_output = create_model_runner_output(reqs=[request])
38 +
39 + # (1c): update_from_output()
40 + engine_core_outputs = scheduler.update_from_output(scheduler_output,
41 + model_runner_output)
42 +
43 + # Ensure the request is finished after 1 tokens.
44 + assert request.is_finished()
45 + assert request.status == RequestStatus.FINISHED_LENGTH_CAPPED
46 + output = engine_core_outputs.outputs[0]
47 + assert output.finish_reason == FinishReason.LENGTH
48 + assert output.kv_transfer_params is not None
49 +
50 + # Request freed in Scheduler and in Persistent Batch ...
51 + assert request_id in scheduler.finished_req_ids
52 + assert len(scheduler.running) == 0
53 + assert len(scheduler.waiting) == 0
54 +
55 + # ... but blocks should not be freed.
56 + blocks = scheduler.kv_cache_manager.single_type_manager.req_to_blocks[
57 + request_id]
58 + for block in blocks:
59 + assert block.ref_cnt == 1
60 +
61 + # STEP (2): Send Finished to PB.
62 + # (2a): schedule() - pass finished request to PB.
63 + scheduler_output = scheduler.schedule()
64 + assert len(scheduler.running) == 0
65 + assert len(scheduler_output.finished_req_ids) == 1
66 + assert request_id in scheduler_output.finished_req_ids
67 + assert len(scheduler_output.scheduled_new_reqs) == 0
68 + assert len(scheduler_output.scheduled_cached_reqs) == 0
69 + assert len(scheduler.finished_req_ids) == 0
70 +
71 + # (2b): execute_model()
72 + model_runner_output = EMPTY_MODEL_RUNNER_OUTPUT
73 +
74 + # (2c): update_from_output()
75 + scheduler.update_from_output(scheduler_output, model_runner_output)
76 +
77 + # STEP (3): Finished sending.
78 + # (3a): schedule() - pass finished request to PB.
79 + scheduler_output = scheduler.schedule()
80 + assert len(scheduler.running) == 0
81 + assert len(scheduler_output.finished_req_ids) == 0
82 + assert len(scheduler_output.scheduled_new_reqs) == 0
83 + assert len(scheduler_output.scheduled_cached_reqs) == 0
84 + assert len(scheduler.finished_req_ids) == 0
85 +
86 + # (3b): execute_model()
87 + model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
88 + model_runner_output.finished_sending = [request_id]
89 +
90 + # (3c): update_from_output()
91 + scheduler.update_from_output(scheduler_output, model_runner_output)
92 +
93 + # Confirm we do not have any memory leaks after req lifecycle.
94 + assert_scheduler_empty(scheduler)
95 +
96 +
97 +def test_short_prompt_lifecycle():
98 + """Test lifecycle of a Remote Decode request with short prompt."""
99 +
100 + vllm_config = create_vllm_config()
101 + scheduler = create_scheduler(vllm_config)
102 +
103 + # Not enough tokens for full block.
104 + NUM_TOKENS = vllm_config.cache_config.block_size // 2
105 + request = create_request(request_id=1,
106 + max_tokens=1,
107 + num_tokens=NUM_TOKENS,
108 + do_remote_decode=True)
109 +
110 + scheduler.add_request(request)
111 +
112 + # STEP (1): Prefill.
113 + # (1a): schedule()
114 + scheduler_output = scheduler.schedule()
115 + assert len(scheduler.running) == 1
116 + assert len(scheduler_output.scheduled_new_reqs) == 1
117 +
118 + # (1b): execute_model()
119 + model_runner_output = create_model_runner_output(reqs=[request])
120 +
121 + # (1c): update_from_output()
122 + # Since tokens < block_size, there will be no kv xfer.
123 + # So this should be cleaned up immediately.
124 + _ = scheduler.update_from_output(scheduler_output, model_runner_output)
125 +
126 + # Confirm we do not have any memory leaks after req lifecycle.
127 + # We need one more call to schedule() to clear data for persistent batch.
128 + _ = scheduler.schedule()
129 + assert_scheduler_empty(scheduler)
130 +
131 +
132 +def test_prefix_cache_lifecycle():
133 + """Test that remote decode params still works with a prefix cache hit."""
134 +
135 + vllm_config = create_vllm_config()
136 + scheduler = create_scheduler(vllm_config)
137 +
138 + # Prime the KVCache.
139 + BLOCK_SIZE = vllm_config.cache_config.block_size
140 + NUM_EXTERNAL_FULL_BLOCKS = 3
141 + NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
142 +
143 + request_normal = create_request(request_id=1, num_tokens=NUM_TOKENS)
144 +
145 + scheduler.add_request(request_normal)
146 + scheduler_output = scheduler.schedule()
147 + model_runner_output = create_model_runner_output(reqs=[request_normal],
148 + use_eos=True)
149 + scheduler.update_from_output(scheduler_output, model_runner_output)
150 + scheduler.schedule()
151 + scheduler.update_from_output(scheduler_output, EMPTY_MODEL_RUNNER_OUTPUT)
152 +
153 + #####################
154 + # Actual Test: confirm we send all blocks.
155 +
156 + # Step (1): Send the KV Transfer.
157 + NUM_EXTERNAL_FULL_BLOCKS -= 1
158 + NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
159 +
160 + request_remote = create_request(request_id=1,
161 + num_tokens=NUM_TOKENS,
162 + do_remote_decode=True)
163 +
164 + scheduler.add_request(request_remote)
165 + scheduler_output = scheduler.schedule()
166 + model_runner_output = create_model_runner_output(reqs=[request_remote])
167 + eco = scheduler.update_from_output(scheduler_output, model_runner_output)
168 + kv_transfer_params = eco.outputs[0].kv_transfer_params
169 +
170 + # Ensure we send all block ids, even if there is a cache hit.
171 + assert (len(
172 + kv_transfer_params["remote_block_ids"]) == NUM_EXTERNAL_FULL_BLOCKS)
173 +
174 + # STEP (2): Ensure it is freed.
175 + scheduler_output = scheduler.schedule()
176 + scheduler.schedule()
177 + model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
178 + model_runner_output.finished_sending = [request_remote.request_id]
179 + scheduler.update_from_output(scheduler_output, model_runner_output)
180 + _ = scheduler.schedule()
181 + assert_scheduler_empty(scheduler)
tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py new
+342
@@ -0,0 +1,342 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import copy
3 +
4 +from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
5 +from vllm.v1.request import FinishReason, RequestStatus
6 +
7 +from .utils import (assert_scheduler_empty, create_model_runner_output,
8 + create_request, create_scheduler, create_vllm_config)
9 +
10 +
11 +def test_basic_lifecycle():
12 + """Test lifecycle of a remote prefill."""
13 +
14 + vllm_config = create_vllm_config()
15 + scheduler = create_scheduler(vllm_config)
16 +
17 + # 2 Full Blocks and 1 Half Block.
18 + BLOCK_SIZE = vllm_config.cache_config.block_size
19 + NUM_EXTERNAL_FULL_BLOCKS = 2
20 + NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
21 + START_FREE_BLOCK_QUEUE_SIZE = (
22 + scheduler.kv_cache_manager.block_pool.free_block_queue.num_free_blocks)
23 +
24 + request = create_request(request_id=1,
25 + num_tokens=NUM_TOKENS,
26 + do_remote_prefill=True)
27 +
28 + scheduler.add_request(request)
29 + request_id = request.request_id
30 +
31 + # STEP (1):
32 + # (1a): schedule()
33 + scheduler_output = scheduler.schedule()
34 +
35 + # Nothing running and empty scheduler output.
36 + assert len(scheduler.running) == 0
37 + assert len(scheduler_output.scheduled_new_reqs) == 0
38 + assert len(scheduler_output.scheduled_cached_reqs) == 0
39 + assert len(scheduler_output.num_scheduled_tokens) == 0
40 + assert scheduler_output.total_num_scheduled_tokens == 0
41 +
42 + # Req waiting for KVs with no computed/scheduled toks ...
43 + assert len(scheduler.waiting) == 1
44 + assert request in scheduler.waiting
45 + assert (request.status == RequestStatus.WAITING_FOR_REMOTE_KVS)
46 + assert (request.num_computed_tokens == 0)
47 +
48 + # ... but should have (uncached) blocks allocated to it.
49 + block_pool = scheduler.kv_cache_manager.block_pool
50 + assert (block_pool.free_block_queue.num_free_blocks
51 + < START_FREE_BLOCK_QUEUE_SIZE)
52 + assert len(block_pool.cached_block_hash_to_block) == 0
53 + blocks = scheduler.kv_cache_manager.single_type_manager.req_to_blocks[
54 + request_id]
55 + for block in blocks:
56 + assert block._block_hash is None
57 +
58 + # (1b): forward()
59 + model_runner_output = EMPTY_MODEL_RUNNER_OUTPUT
60 +
61 + # (1c): update_from_output()
62 + engine_core_outputs = scheduler.update_from_output(scheduler_output,
63 + model_runner_output)
64 + assert len(engine_core_outputs.outputs) == 0
65 +
66 + # STEP (2):
67 + # (2a): schedule(): nothing happens!
68 + scheduler_output = scheduler.schedule()
69 + assert len(scheduler.waiting) == 1
70 + assert len(scheduler.running) == 0
71 +
72 + # (2b): forward(): request finishes recv.
73 + model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
74 + model_runner_output.finished_recving = [request_id]
75 +
76 + # (2c): update_from_output():
77 + engine_core_outputs = scheduler.update_from_output(scheduler_output,
78 + model_runner_output)
79 + assert len(scheduler.waiting) == 1
80 + assert (request_id in scheduler.finished_recving_kv_req_ids)
81 +
82 + # STEP (3):
83 + # (3a): schedule(): this should actually schedule.
84 + scheduler_output = scheduler.schedule()
85 + assert len(scheduler.running) == 1
86 +
87 + # Confirm the block are actually allocated.
88 + num_hashed_blocks = 0
89 + blocks = scheduler.kv_cache_manager.single_type_manager.req_to_blocks[
90 + request_id]
91 + for block in blocks:
92 + assert block.ref_cnt == 1
93 + num_hashed_blocks += (1 if block._block_hash is not None else 0)
94 + assert num_hashed_blocks == NUM_EXTERNAL_FULL_BLOCKS
95 +
96 + # Confirm the rest of the prompt is scheduled in this step.
97 + scheduled_req = scheduler_output.scheduled_new_reqs[0]
98 + num_scheduled_tokens = scheduler_output.num_scheduled_tokens[request_id]
99 + num_computed_tokens = scheduled_req.num_computed_tokens
100 + total_prompt_tokens = len(scheduled_req.prompt_token_ids)
101 + assert (num_scheduled_tokens == total_prompt_tokens - num_computed_tokens)
102 +
103 + # (3b): execute_model()
104 + model_runner_output = create_model_runner_output([request])
105 + # (3c): update_from_output()
106 + scheduler.update_from_output(scheduler_output, model_runner_output)
107 +
108 + # Step (4): Hit EOS.
109 + scheduler_output = scheduler.schedule()
110 + model_runner_output = create_model_runner_output([request], use_eos=True)
111 + engine_core_outputs = scheduler.update_from_output(scheduler_output,
112 + model_runner_output)
113 + scheduler.schedule()
114 +
115 + outputs = engine_core_outputs.outputs
116 + assert len(outputs) == 1
117 + output = outputs[0]
118 + assert output.finish_reason == FinishReason.STOP
119 + assert_scheduler_empty(scheduler)
120 +
121 +
122 +def test_interleaved_lifecycle():
123 + """Test Remote Prefills Work Well With Other Requests."""
124 +
125 + vllm_config = create_vllm_config()
126 + scheduler = create_scheduler(vllm_config)
127 +
128 + # 2 Full Blocks and 1 Half Block.
129 + BLOCK_SIZE = vllm_config.cache_config.block_size
130 + NUM_EXTERNAL_FULL_BLOCKS = 2
131 + NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
132 +
133 + request_remote = create_request(request_id=1,
134 + num_tokens=NUM_TOKENS,
135 + do_remote_prefill=True)
136 + request_local_a = create_request(
137 + request_id=2,
138 + num_tokens=NUM_TOKENS,
139 + )
140 + request_local_b = create_request(
141 + request_id=3,
142 + num_tokens=NUM_TOKENS,
143 + )
144 +
145 + # STEP 1: Regular request is running.
146 + scheduler.add_request(request_local_a)
147 + scheduler_output = scheduler.schedule()
148 + assert len(scheduler.running) == 1
149 +
150 + model_runner_output = create_model_runner_output([request_local_a])
151 + scheduler.update_from_output(scheduler_output, model_runner_output)
152 +
153 + # STEP 2: Add a local and remote request.
154 + scheduler.add_request(request_local_b)
155 + scheduler.add_request(request_remote)
156 + scheduler_output = scheduler.schedule()
157 + assert len(scheduler.running) == 2
158 + assert len(scheduler.waiting) == 1
159 + assert len(scheduler_output.scheduled_new_reqs) == 1
160 + assert len(scheduler_output.scheduled_cached_reqs) == 1
161 +
162 + model_runner_output = create_model_runner_output(
163 + [request_local_a, request_local_b])
164 + scheduler.update_from_output(scheduler_output, model_runner_output)
165 +
166 + # STEP 3: continue running, KVs not arrived yet.
167 + scheduler_output = scheduler.schedule()
168 + assert len(scheduler.running) == 2
169 + assert len(scheduler.waiting) == 1
170 + assert len(scheduler_output.scheduled_new_reqs) == 0
171 + assert len(scheduler_output.scheduled_cached_reqs) == 2
172 +
173 + model_runner_output = create_model_runner_output(
174 + reqs=[request_local_a, request_local_b])
175 + scheduler.update_from_output(scheduler_output, model_runner_output)
176 + assert len(scheduler.running) == 2
177 + assert len(scheduler.waiting) == 1
178 + assert len(scheduler_output.scheduled_new_reqs) == 0
179 + assert len(scheduler_output.scheduled_cached_reqs) == 2
180 +
181 + # STEP 4: KVs arrive.
182 + scheduler_output = scheduler.schedule()
183 + assert len(scheduler.running) == 2
184 + assert len(scheduler.waiting) == 1
185 + assert len(scheduler_output.scheduled_new_reqs) == 0
186 + assert len(scheduler_output.scheduled_cached_reqs) == 2
187 +
188 + model_runner_output = create_model_runner_output(
189 + [request_local_a, request_local_b],
190 + finished_recving=[request_remote.request_id])
191 + scheduler.update_from_output(scheduler_output, model_runner_output)
192 +
193 + # STEP 5: RECVed KVs are sent to ModelRunner.
194 + scheduler_output = scheduler.schedule()
195 + assert len(scheduler.running) == 3
196 + assert len(scheduler.waiting) == 0
197 + assert len(scheduler_output.scheduled_new_reqs) == 1
198 + assert len(scheduler_output.scheduled_cached_reqs) == 2
199 +
200 + model_runner_output = create_model_runner_output(
201 + [request_local_a, request_local_b, request_remote])
202 + scheduler.update_from_output(scheduler_output, model_runner_output)
203 +
204 + # STEP 6: Hit EOS and free.
205 + scheduler_output = scheduler.schedule()
206 + model_runner_output = create_model_runner_output(
207 + [request_local_a, request_local_b, request_remote],
208 + use_eos=True,
209 + )
210 + scheduler.update_from_output(scheduler_output, model_runner_output)
211 + scheduler.schedule()
212 + assert_scheduler_empty(scheduler)
213 +
214 +
215 +def test_no_spurious_prefix_caching():
216 + """
217 + With P/D, blocks can be allocated but uncomputed for
218 + multiple engine steps. This test confirms that we do
219 + not accidentally have cache hits against uncomputed
220 + blocks.
221 + """
222 +
223 + vllm_config = create_vllm_config()
224 + scheduler = create_scheduler(vllm_config)
225 +
226 + vllm_config = create_vllm_config()
227 + scheduler = create_scheduler(vllm_config)
228 +
229 + # 2 and a half full external blocks.
230 + BLOCK_SIZE = vllm_config.cache_config.block_size
231 + NUM_EXTERNAL_FULL_BLOCKS = 2
232 + NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
233 +
234 + # Both of these requests have prompts like [1,1,1,1,1, ...]
235 + request_remote = create_request(
236 + request_id=1,
237 + num_tokens=NUM_TOKENS,
238 + do_remote_prefill=True,
239 + use_all_1s_for_prompt_tokens=True,
240 + )
241 +
242 + request_local = create_request(
243 + request_id=2,
244 + num_tokens=NUM_TOKENS,
245 + do_remote_prefill=False,
246 + use_all_1s_for_prompt_tokens=True,
247 + )
248 +
249 + # Schedule the remote prefill request. This should not
250 + # cause any blocks to be cached.
251 + scheduler.add_request(request_remote)
252 + scheduler_output = scheduler.schedule()
253 + scheduler.update_from_output(scheduler_output, EMPTY_MODEL_RUNNER_OUTPUT)
254 + assert len(scheduler.waiting) == 1
255 +
256 + # Schedule the local prefill request. This should
257 + # cause blocks to be cached, but separately from
258 + scheduler.add_request(request_local)
259 + scheduler_output = scheduler.schedule()
260 + assert len(scheduler.running) == 1
261 + assert len(scheduler.waiting) == 1
262 +
263 + local_blocks = scheduler.kv_cache_manager.single_type_manager.req_to_blocks[
264 + request_local.request_id]
265 + remote_blocks = scheduler.kv_cache_manager.single_type_manager.req_to_blocks[ # noqa: E501
266 + request_remote.request_id]
267 +
268 + # Local should have cached blocks (but not all due to preallocate).
269 + num_hashed_blocks = 0
270 + for block in local_blocks:
271 + assert block.ref_cnt == 1
272 + num_hashed_blocks += (1 if block._block_hash is not None else 0)
273 + assert num_hashed_blocks > 0
274 +
275 + # Remote blocks should not be cached.
276 + for block in remote_blocks:
277 + assert block.ref_cnt == 1
278 + assert block._block_hash is None
279 +
280 +
281 +def test_full_block_prompt():
282 + """Test that we handle a prompt that is the full block size."""
283 +
284 + vllm_config = create_vllm_config()
285 + scheduler = create_scheduler(vllm_config)
286 +
287 + # 2 Full Blocks and 1 Half Block.
288 + BLOCK_SIZE = vllm_config.cache_config.block_size
289 + NUM_EXTERNAL_FULL_BLOCKS = 2
290 + NUM_TOKENS = int(BLOCK_SIZE * NUM_EXTERNAL_FULL_BLOCKS)
291 +
292 + request = create_request(request_id=1,
293 + num_tokens=NUM_TOKENS,
294 + do_remote_prefill=True)
295 +
296 + scheduler.add_request(request)
297 + request_id = request.request_id
298 +
299 + # STEP (1): Initialize a recv.
300 + scheduler_output = scheduler.schedule()
301 + # All blocks should be allocated.
302 + num_blocks = len(scheduler.kv_cache_manager.single_type_manager.
303 + req_to_blocks[request_id])
304 + assert num_blocks == NUM_EXTERNAL_FULL_BLOCKS
305 + model_runner_output = EMPTY_MODEL_RUNNER_OUTPUT
306 + scheduler.update_from_output(scheduler_output, model_runner_output)
307 +
308 + # # STEP (2): Recv.
309 + scheduler_output = scheduler.schedule()
310 + model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
311 + model_runner_output.finished_recving = [request_id]
312 + scheduler.update_from_output(scheduler_output, model_runner_output)
313 + assert len(scheduler.waiting) == 1
314 + assert (request_id in scheduler.finished_recving_kv_req_ids)
315 +
316 + # # STEP (3): Run as usual.
317 + scheduler_output = scheduler.schedule()
318 +
319 + # We need to recompute the final token of the prompt to generate
320 + # the first new token, so we should not have a new block.
321 + num_blocks = len(scheduler.kv_cache_manager.single_type_manager.
322 + req_to_blocks[request_id])
323 + assert num_blocks == NUM_EXTERNAL_FULL_BLOCKS
324 + assert (scheduler_output.scheduled_new_reqs[0].num_computed_tokens ==
325 + NUM_TOKENS - 1)
326 + assert (scheduler_output.num_scheduled_tokens[request_id] == 1)
327 +
328 + model_runner_output = create_model_runner_output([request])
329 + scheduler.update_from_output(scheduler_output, model_runner_output)
330 +
331 + # # Step (4): Hit EOS.
332 + scheduler_output = scheduler.schedule()
333 + model_runner_output = create_model_runner_output([request], use_eos=True)
334 + engine_core_outputs = scheduler.update_from_output(scheduler_output,
335 + model_runner_output)
336 + scheduler.schedule()
337 +
338 + outputs = engine_core_outputs.outputs
339 + assert len(outputs) == 1
340 + output = outputs[0]
341 + assert output.finish_reason == FinishReason.STOP
342 + assert_scheduler_empty(scheduler)
tests/v1/kv_connector/unit/utils.py new
+190
@@ -0,0 +1,190 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import Optional
3 +
4 +import torch
5 +
6 +from vllm import SamplingParams
7 +from vllm.config import (CacheConfig, DeviceConfig, KVTransferConfig,
8 + ModelConfig, SchedulerConfig, VllmConfig)
9 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
10 + NixlKVTransferParams)
11 +from vllm.v1.core.sched.scheduler import Scheduler
12 +from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
13 + KVCacheGroupSpec)
14 +from vllm.v1.outputs import ModelRunnerOutput
15 +from vllm.v1.request import Request
16 +from vllm.v1.structured_output import StructuredOutputManager
17 +
18 +EOS_TOKEN_ID = 50256
19 +
20 +
21 +def assert_scheduler_empty(scheduler: Scheduler):
22 + """Confirm the scheduler is "empty" - i.e. no leaks."""
23 + # Scheduler Metadata.
24 + assert len(scheduler.requests) == 0
25 + assert len(scheduler.waiting) == 0
26 + assert len(scheduler.running) == 0
27 + assert len(scheduler.finished_req_ids) == 0
28 + assert len(scheduler.finished_recving_kv_req_ids) == 0
29 + assert len(scheduler._cached_reqs_data) == 0
30 +
31 + # EncoderCacheManager.
32 + assert len(scheduler.encoder_cache_manager.freed) == 0
33 + assert len(scheduler.encoder_cache_manager.cached) == 0
34 +
35 + # KVCache Manager.
36 + assert len(
37 + scheduler.kv_cache_manager.single_type_manager.req_to_blocks) == 0
38 + assert len(scheduler.kv_cache_manager.req_to_block_hashes) == 0
39 + assert len(
40 + scheduler.kv_cache_manager.single_type_manager.num_cached_block) == 0
41 + num_free_blocks = (
42 + scheduler.kv_cache_manager.block_pool.free_block_queue.num_free_blocks)
43 + assert num_free_blocks == (
44 + scheduler.kv_cache_manager.block_pool.num_gpu_blocks - 1)
45 +
46 + # NOTE(rob): just the ref count on blocks will be 0. The hash
47 + # value, etc will remain since we lazily evict for prefix cache.
48 + for block in scheduler.kv_cache_manager.block_pool.blocks:
49 + assert block.ref_cnt == 0
50 +
51 +
52 +def create_vllm_config(
53 + model: str = "facebook/opt-125m",
54 + max_num_seqs: int = 16,
55 + max_num_batched_tokens: int = 64,
56 + block_size: int = 16,
57 +) -> VllmConfig:
58 + """Initialize VllmConfig For Testing."""
59 + scheduler_config = SchedulerConfig(
60 + max_num_seqs=max_num_seqs,
61 + max_num_batched_tokens=max_num_batched_tokens,
62 + max_model_len=max_num_batched_tokens,
63 + )
64 + model_config = ModelConfig(
65 + model=model,
66 + task="auto",
67 + tokenizer=model,
68 + tokenizer_mode="auto",
69 + trust_remote_code=True,
70 + dtype="float16",
71 + seed=42,
72 + )
73 + # Cache config, optionally force APC
74 + cache_config = CacheConfig(
75 + block_size=block_size,
76 + gpu_memory_utilization=0.9,
77 + swap_space=0,
78 + cache_dtype="auto",
79 + enable_prefix_caching=True,
80 + )
81 + kv_transfer_config = KVTransferConfig(
82 + kv_connector="NixlConnector",
83 + kv_role="kv_both",
84 + )
85 + return VllmConfig(scheduler_config=scheduler_config,
86 + model_config=model_config,
87 + cache_config=cache_config,
88 + kv_transfer_config=kv_transfer_config,
89 + device_config=DeviceConfig("cpu"))
90 +
91 +
92 +def create_scheduler(
93 + vllm_config: VllmConfig,
94 + num_blocks: int = 10000,
95 +) -> Scheduler:
96 + """Initialize Scheduler For Testing."""
97 + block_size = vllm_config.cache_config.block_size
98 + kv_cache_config = KVCacheConfig(
99 + num_blocks=num_blocks, # A large number of blocks to hold all requests
100 + tensors={},
101 + kv_cache_groups=[
102 + KVCacheGroupSpec(['layer'],
103 + FullAttentionSpec(block_size, 1, 1, torch.float32,
104 + False))
105 + ],
106 + )
107 + vllm_config.cache_config.num_gpu_blocks = num_blocks
108 + return Scheduler(
109 + vllm_config=vllm_config,
110 + kv_cache_config=kv_cache_config,
111 + log_stats=True,
112 + structured_output_manager=StructuredOutputManager(vllm_config),
113 + )
114 +
115 +
116 +def create_request(
117 + request_id: int,
118 + num_tokens: int = 10,
119 + max_tokens: int = 16,
120 + do_remote_decode: bool = False,
121 + do_remote_prefill: bool = False,
122 + use_all_1s_for_prompt_tokens: bool = False,
123 + num_remote_blocks: int = 3,
124 +) -> Request:
125 + """Make dummy request for testing."""
126 +
127 + if do_remote_decode:
128 + assert not do_remote_prefill
129 + kv_transfer_params = NixlKVTransferParams(do_remote_prefill=False,
130 + do_remote_decode=True)
131 + elif do_remote_prefill:
132 + kv_transfer_params = NixlKVTransferParams(
133 + do_remote_prefill=True,
134 + do_remote_decode=False,
135 + remote_engine_id="my-engine-id",
136 + remote_block_ids=list(range(num_remote_blocks)),
137 + remote_host="my-host",
138 + remote_port=1234)
139 + else:
140 + kv_transfer_params = None
141 +
142 + max_tokens = 1 if do_remote_decode else max_tokens
143 + sampling_params = SamplingParams(max_tokens=max_tokens)
144 +
145 + if use_all_1s_for_prompt_tokens:
146 + prompt_token_ids = [1] * num_tokens
147 + else:
148 + prompt_token_ids = [i * request_id for i in range(num_tokens)]
149 +
150 + req = Request(
151 + request_id=f"id-{request_id}",
152 + prompt_token_ids=prompt_token_ids,
153 + sampling_params=sampling_params,
154 + multi_modal_inputs=None,
155 + multi_modal_placeholders=None,
156 + multi_modal_hashes=None,
157 + eos_token_id=EOS_TOKEN_ID,
158 + arrival_time=0,
159 + )
160 + req.kv_transfer_params = kv_transfer_params
161 + return req
162 +
163 +
164 +def create_model_runner_output(
165 + reqs: list[Request],
166 + finished_sending: Optional[list[str]] = None,
167 + finished_recving: Optional[list[str]] = None,
168 + use_eos: bool = False,
169 +) -> ModelRunnerOutput:
170 + """Make dummy model runner output for testing."""
171 +
172 + # Make request data.
173 + req_ids = [req.request_id for req in reqs]
174 + req_id_to_index = {req_id: idx for idx, req_id in enumerate(req_ids)}
175 +
176 + # Make sampled tokens.
177 + sampled_token = EOS_TOKEN_ID if use_eos else 0
178 + sampled_token_ids = [[sampled_token] for _ in req_ids]
179 +
180 + # Make output data structure.
181 + return ModelRunnerOutput(
182 + req_ids=req_ids,
183 + req_id_to_index=req_id_to_index,
184 + sampled_token_ids=sampled_token_ids,
185 + spec_token_ids=None,
186 + logprobs=None,
187 + prompt_logprobs_dict={},
188 + finished_sending=finished_sending,
189 + finished_recving=finished_recving,
190 + )
vllm/config.py
+5 -1
@@ -8,6 +8,7 @@ import inspect
8 import json
9 import re
10 import textwrap
11 +import uuid
12 import warnings
13 from collections import Counter
14 from contextlib import contextmanager
@@ -3438,6 +3439,9 @@ class KVTransferConfig:
3439 """The KV connector for vLLM to transmit KV caches between vLLM instances.
3440 """
3441
3442 + engine_id: str = str(uuid.uuid4())
3443 + """The engine id for KV transfers."""
3444 +
3445 kv_buffer_device: Optional[str] = "cuda"
3446 """The device used by kv connector to buffer the KV cache.
3447 Currently only support 'cuda'."""
@@ -3448,7 +3452,7 @@ class KVTransferConfig:
3452
3453 kv_role: Optional[KVRole] = None
3454 """Whether this vLLM instance produces, consumes KV cache, or both. Choices
3451 - are 'kv_producer', 'kv_consumer', and 'both'."""
3455 + are 'kv_producer', 'kv_consumer', and 'kv_both'."""
3456
3457 kv_rank: Optional[int] = None
3458 """The rank of this vLLM instance in the KV cache transfer. Typical value:
vllm/distributed/kv_transfer/kv_connector/factory.py
+5
@@ -105,3 +105,8 @@ KVConnectorFactory.register_connector(
105 "LMCacheConnectorV1",
106 "vllm.distributed.kv_transfer.kv_connector.v1.lmcache_connector",
107 "LMCacheConnectorV1")
108 +
109 +KVConnectorFactory.register_connector(
110 + "NixlConnector",
111 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector",
112 + "NixlConnector")
vllm/distributed/kv_transfer/kv_connector/v1/__init__.py
+2 -5
@@ -1,8 +1,5 @@
1 # SPDX-License-Identifier: Apache-2.0
2 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
3 - KVConnectorBase_V1, KVConnectorRole)
3 + KVConnectorBase_V1, KVConnectorRole, KVTransferParams)
4
5 -__all__ = [
6 - "KVConnectorRole",
7 - "KVConnectorBase_V1",
8 -]
5 +__all__ = ["KVConnectorRole", "KVConnectorBase_V1", "KVTransferParams"]
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+82 -7
@@ -23,7 +23,7 @@ The class provides the following primitives:
23 import enum
24 from abc import ABC, abstractmethod
25 from dataclasses import dataclass
26 -from typing import TYPE_CHECKING
26 +from typing import TYPE_CHECKING, Any, Optional
27
28 import torch
29
@@ -34,6 +34,7 @@ if TYPE_CHECKING:
34 from vllm.attention.backends.abstract import AttentionMetadata
35 from vllm.config import VllmConfig
36 from vllm.forward_context import ForwardContext
37 + from vllm.v1.core.kv_cache_manager import KVCacheBlocks
38 from vllm.v1.request import Request
39
40 logger = init_logger(__name__)
@@ -47,12 +48,34 @@ class KVConnectorRole(enum.Enum):
48 WORKER = 1
49
50
51 +class KVTransferParams:
52 + """
53 + Abstract KVTransferParams used to send KVTransfer
54 + parameters between instances of vLLM.
55 +
56 + Specific instances of KVConnector customize this
57 + method for serializing / deserializing msgs sent
58 + via the HTTP protocol.
59 + """
60 +
61 + @staticmethod
62 + def from_raw_dict(
63 + raw_dict: Optional[dict[str,
64 + Any]]) -> Optional["KVTransferParams"]:
65 + return None
66 +
67 +
68 @dataclass
69 class KVConnectorMetadata:
70 + """
71 + Abstract Metadata used to communicate between the
72 + Scheduler KVConnector and Worker KVConnector.
73 + """
74 pass
75
76
77 class KVConnectorBase_V1(ABC):
78 + _KVTransferParams = KVTransferParams
79
80 def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
81 logger.warning(
@@ -66,6 +89,10 @@ class KVConnectorBase_V1(ABC):
89 def role(self) -> KVConnectorRole:
90 return self._role
91
92 + # ==============================
93 + # Worker-side methods
94 + # ==============================
95 +
96 def bind_connector_metadata(
97 self, connector_metadata: KVConnectorMetadata) -> None:
98 """Set the connector metadata from the scheduler.
@@ -97,9 +124,15 @@ class KVConnectorBase_V1(ABC):
124 """
125 return self._connector_metadata
126
100 - # ==============================
101 - # Worker-side methods
102 - # ==============================
127 + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
128 + """
129 + Initialize with the KV caches. Useful for pre-registering the
130 + KV Caches in the KVConnector (e.g. for NIXL).
131 +
132 + Args: kv_caches:
133 + dictionary of layer names, kv cache
134 + """
135 + return
136
137 @abstractmethod
138 def start_load_kv(self, forward_context: "ForwardContext",
@@ -162,15 +195,37 @@ class KVConnectorBase_V1(ABC):
195 """
196 pass
197
198 + def get_finished(
199 + self, finished_req_ids: set[str]
200 + ) -> tuple[Optional[set[str]], Optional[set[str]]]:
201 + """
202 + Notifies worker-side connector ids of requests that have
203 + finished generating tokens.
204 +
205 + Returns:
206 + ids of requests that have finished asynchronous (recving, sending).
207 + The finished saves/sends req ids must belong to a set provided in a
208 + call to this method (this call or a prior one).
209 + """
210 + return None, None
211 +
212 # ==============================
213 # Scheduler-side methods
214 # ==============================
215 +
216 + def set_kv_transfer_params(self, request: "Request"):
217 + """Parse raw KV Transfer params."""
218 + assert request.kv_transfer_params is None
219 + kv_transfer_params = self._KVTransferParams.from_raw_dict(
220 + request.raw_kv_transfer_params)
221 + request.kv_transfer_params = kv_transfer_params
222 +
223 @abstractmethod
224 def get_num_new_matched_tokens(
225 self,
226 request: "Request",
227 num_computed_tokens: int,
173 - ) -> int:
228 + ) -> tuple[int, bool]:
229 """
230 Get number of new tokens that can be loaded from the
231 external KV cache beyond the num_computed_tokens.
@@ -181,13 +236,16 @@ class KVConnectorBase_V1(ABC):
236 computed tokens for this request
237
238 Returns:
184 - the number of tokens that can be loaded from the
185 - external KV cache beyond what is already computed.
239 + * the number of tokens that can be loaded from the
240 + external KV cache beyond what is already computed.
241 + * true if external KV cache tokens will be loaded
242 + asynchronously (between scheduler steps).
243 """
244 pass
245
246 @abstractmethod
247 def update_state_after_alloc(self, request: "Request",
248 + blocks: "KVCacheBlocks",
249 num_external_tokens: int):
250 """
251 Update KVConnector state after block allocation.
@@ -207,3 +265,20 @@ class KVConnectorBase_V1(ABC):
265 scheduler_output (SchedulerOutput): the scheduler output object.
266 """
267 pass
268 +
269 + def request_finished(
270 + self,
271 + request: "Request",
272 + block_ids: list[int],
273 + ) -> tuple[bool, Optional[dict[str, Any]]]:
274 + """
275 + Called when a request has finished, before its blocks are freed.
276 +
277 + Returns:
278 + True if the request is being saved/sent asynchronously and blocks
279 + should not be freed until the request_id is returned from
280 + get_finished().
281 + Optional KVTransferParams to be included in the request outputs
282 + returned by the engine.
283 + """
284 + return False, None
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py
+4 -2
@@ -13,6 +13,7 @@ from vllm.v1.core.sched.output import SchedulerOutput
13 if TYPE_CHECKING:
14 from vllm.attention.backends.abstract import AttentionMetadata
15 from vllm.forward_context import ForwardContext
16 + from vllm.v1.core.kv_cache_manager import KVCacheBlocks
17 from vllm.v1.request import Request
18
19 logger = init_logger(__name__)
@@ -92,7 +93,7 @@ class LMCacheConnectorV1(KVConnectorBase_V1):
93 self,
94 request: "Request",
95 num_computed_tokens: int,
95 - ) -> int:
96 + ) -> tuple[int, bool]:
97 """
98 Get number of new tokens that can be loaded from the
99 external KV cache beyond the num_computed_tokens.
@@ -107,9 +108,10 @@ class LMCacheConnectorV1(KVConnectorBase_V1):
108 external KV cache beyond what is already computed.
109 """
110 return self._lmcache_engine.get_num_new_matched_tokens(
110 - request, num_computed_tokens)
111 + request, num_computed_tokens), False
112
113 def update_state_after_alloc(self, request: "Request",
114 + blocks: "KVCacheBlocks",
115 num_external_tokens: int):
116 """
117 Update KVConnector state after block allocation.
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py new
+805
@@ -0,0 +1,805 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import contextlib
3 +import math
4 +import threading
5 +import time
6 +import uuid
7 +from collections import defaultdict
8 +from dataclasses import dataclass
9 +from typing import TYPE_CHECKING, Any, Iterator
10 +
11 +import msgspec
12 +import torch
13 +import zmq
14 +from typing_extensions import Optional
15 +
16 +from vllm import envs
17 +from vllm.config import VllmConfig
18 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
19 + KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole, KVTransferParams)
20 +from vllm.distributed.parallel_state import (
21 + get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size,
22 + get_tp_group)
23 +from vllm.logger import init_logger
24 +from vllm.utils import round_down
25 +from vllm.v1.core.sched.output import SchedulerOutput
26 +from vllm.v1.request import RequestStatus
27 +
28 +if TYPE_CHECKING:
29 + from vllm.attention.backends.abstract import AttentionMetadata
30 + from vllm.forward_context import ForwardContext
31 + from vllm.v1.core.kv_cache_manager import KVCacheBlocks
32 + from vllm.v1.request import Request
33 +
34 +GET_META_MSG = b"get_meta_msg"
35 +
36 +logger = init_logger(__name__)
37 +
38 +# Lazy import nixl_wrapper to avoid loading nixl_bindings if nixl is not used
39 +try:
40 + from nixl._api import nixl_agent as NixlWrapper
41 + logger.info("NIXL is available")
42 +except ImportError:
43 + logger.warning("NIXL is not available")
44 + NixlWrapper = None
45 +
46 +
47 +@dataclass
48 +class NixlKVTransferParams(KVTransferParams):
49 +
50 + def __init__(
51 + self,
52 + do_remote_prefill: bool,
53 + do_remote_decode: bool,
54 + remote_block_ids: Optional[list[int]] = None,
55 + remote_host: Optional[str] = None,
56 + remote_port: Optional[int] = None,
57 + remote_engine_id: Optional[str] = None,
58 + ):
59 + self.do_remote_prefill = do_remote_prefill
60 + self.do_remote_decode = do_remote_decode
61 + self.remote_block_ids = remote_block_ids
62 + self.remote_host = remote_host
63 + self.remote_port = remote_port
64 + self.remote_engine_id = remote_engine_id
65 +
66 + @staticmethod
67 + def from_raw_dict(
68 + raw_dict: Optional[dict[str,
69 + Any]]) -> Optional["NixlKVTransferParams"]:
70 +
71 + # If no raw transfer params passed, return None.
72 + if raw_dict is None:
73 + return None
74 +
75 + # Validate the request is formatted properly.
76 + if (("do_remote_prefill" not in raw_dict)
77 + or ("do_remote_decode" not in raw_dict)
78 + or ("remote_block_ids" not in raw_dict)
79 + or ("remote_host" not in raw_dict)
80 + or ("remote_port" not in raw_dict)
81 + or ("remote_engine_id" not in raw_dict)):
82 + logger.warning(
83 + "Got invalid KVTransferParams: %s. This "
84 + "request will not utilize KVTransfer", raw_dict)
85 + return None
86 +
87 + return NixlKVTransferParams(
88 + do_remote_prefill=raw_dict["do_remote_prefill"],
89 + do_remote_decode=raw_dict["do_remote_decode"],
90 + remote_block_ids=raw_dict["remote_block_ids"],
91 + remote_host=raw_dict["remote_host"],
92 + remote_port=raw_dict["remote_port"],
93 + remote_engine_id=raw_dict["remote_engine_id"],
94 + )
95 +
96 +
97 +class NixlAgentMetadata(
98 + msgspec.Struct,
99 + omit_defaults=True, # type: ignore[call-arg]
100 + # required for @cached_property.
101 + dict=True):
102 + engine_id: str
103 + agent_metadata: bytes
104 + kv_caches_base_addr: list[int]
105 + num_blocks: int
106 +
107 +
108 +@dataclass
109 +class ReqMeta:
110 + local_block_ids: list[int]
111 + remote_block_ids: list[int]
112 + remote_host: str
113 + remote_port: int
114 + remote_engine_id: str
115 +
116 +
117 +class NixlConnectorMetadata(KVConnectorMetadata):
118 +
119 + def __init__(self):
120 + self.requests: dict[str, ReqMeta] = {}
121 +
122 + def add_new_req(
123 + self,
124 + request_id: str,
125 + local_block_ids: list[int],
126 + kv_transfer_params: NixlKVTransferParams,
127 + ):
128 + assert request_id not in self.requests
129 + assert kv_transfer_params.remote_block_ids is not None
130 + assert kv_transfer_params.remote_engine_id is not None
131 + assert kv_transfer_params.remote_host is not None
132 + assert kv_transfer_params.remote_port is not None
133 +
134 + self.requests[request_id] = ReqMeta(
135 + local_block_ids=local_block_ids,
136 + remote_block_ids=kv_transfer_params.remote_block_ids,
137 + remote_engine_id=kv_transfer_params.remote_engine_id,
138 + remote_host=kv_transfer_params.remote_host,
139 + remote_port=kv_transfer_params.remote_port,
140 + )
141 +
142 +
143 +class NixlConnector(KVConnectorBase_V1):
144 + _KVTransferParams: type[NixlKVTransferParams] = NixlKVTransferParams
145 +
146 + def __init__(self, vllm_config: VllmConfig, role: KVConnectorRole):
147 + assert vllm_config.kv_transfer_config is not None
148 + self.engine_id = vllm_config.kv_transfer_config.engine_id
149 +
150 + if role == KVConnectorRole.SCHEDULER:
151 + self.connector_scheduler : Optional[NixlConnectorScheduler] = \
152 + NixlConnectorScheduler(vllm_config, str(self.engine_id))
153 + self.connector_worker: Optional[NixlConnectorWorker] = None
154 + elif role == KVConnectorRole.WORKER:
155 + self.connector_scheduler = None
156 + self.connector_worker = NixlConnectorWorker(str(self.engine_id))
157 +
158 + ############################################################
159 + # Scheduler Side Methods
160 + ############################################################
161 +
162 + def get_num_new_matched_tokens(
163 + self, request: "Request",
164 + num_computed_tokens: int) -> tuple[int, bool]:
165 + assert self.connector_scheduler is not None
166 + return self.connector_scheduler.get_num_new_matched_tokens(
167 + request, num_computed_tokens)
168 +
169 + def update_state_after_alloc(self, request: "Request",
170 + blocks: "KVCacheBlocks",
171 + num_external_tokens: int):
172 + assert self.connector_scheduler is not None
173 + return self.connector_scheduler.update_state_after_alloc(
174 + request, blocks, num_external_tokens)
175 +
176 + def build_connector_meta(
177 + self,
178 + scheduler_output: SchedulerOutput,
179 + ) -> KVConnectorMetadata:
180 + assert self.connector_scheduler is not None
181 + return self.connector_scheduler.build_connector_meta(scheduler_output)
182 +
183 + def request_finished(
184 + self,
185 + request: "Request",
186 + block_ids: list[int],
187 + ) -> tuple[bool, Optional[dict[str, Any]]]:
188 + assert self.connector_scheduler is not None
189 + return self.connector_scheduler.request_finished(request, block_ids)
190 +
191 + ############################################################
192 + # Worker Side Methods
193 + ############################################################
194 + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
195 + assert self.connector_worker is not None
196 + self.connector_worker.register_kv_caches(kv_caches)
197 +
198 + def get_finished(self,
199 + finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
200 + """Get the finished recving and sending requests."""
201 + assert self.connector_worker is not None
202 + return self.connector_worker.get_finished()
203 +
204 + def start_load_kv(self, forward_context: "ForwardContext",
205 + **kwargs) -> None:
206 + assert self.connector_worker is not None
207 + assert isinstance(self._connector_metadata, NixlConnectorMetadata)
208 + self.connector_worker.start_load_kv(self._connector_metadata)
209 +
210 + def wait_for_layer_load(self, layer_name: str) -> None:
211 + """NixlConnector does not do layerwise saving."""
212 + pass
213 +
214 + def save_kv_layer(self, layer_name: str, kv_layer: torch.Tensor,
215 + attn_metadata: "AttentionMetadata", **kwargs) -> None:
216 + """NixlConnector does not save explicitly."""
217 + pass
218 +
219 + def wait_for_save(self):
220 + """NixlConnector does not save explicitly."""
221 + pass
222 +
223 +
224 +class NixlConnectorScheduler:
225 + """Implementation of Scheduler side methods"""
226 +
227 + def __init__(self, vllm_config: VllmConfig, engine_id: str):
228 + self.vllm_config = vllm_config
229 + self.block_size = vllm_config.cache_config.block_size
230 + self.engine_id = engine_id
231 + logger.info("Initializing NIXL Scheduler %s", engine_id)
232 +
233 + # Requests that need to start recv.
234 + # New requests are added by update_state_after_alloc in
235 + # the scheduler. Used to make metadata passed to Worker.
236 + self._reqs_need_recv: dict[str, tuple[Request, list[int]]] = {}
237 +
238 + def get_num_new_matched_tokens(
239 + self, request: "Request",
240 + num_computed_tokens: int) -> tuple[int, bool]:
241 + """
242 + For remote prefill, pull all prompt blocks from remote
243 + asynchronously relative to engine execution.
244 +
245 + Args:
246 + request (Request): the request object.
247 + num_computed_tokens (int): the number of locally
248 + computed tokens for this request
249 + Returns:
250 + * the number of tokens that can be loaded from the
251 + external KV cache beyond what is already computed.
252 + * true if the external KV cache tokens will be loaded
253 + asynchronously (between scheduler steps).
254 + """
255 +
256 + # No KVTransfer for this request.
257 + if request.kv_transfer_params is None:
258 + return 0, False
259 + assert isinstance(request.kv_transfer_params, NixlKVTransferParams)
260 +
261 + # Remote prefill: get all prompt blocks from remote.
262 + if request.kv_transfer_params.do_remote_prefill:
263 + assert num_computed_tokens % self.block_size == 0
264 + rounded_num_prompt_tokens = round_down(
265 + len(request.prompt_token_ids), self.block_size)
266 + count = max(rounded_num_prompt_tokens - num_computed_tokens, 0)
267 + return count, count > 0
268 +
269 + return 0, False
270 +
271 + def update_state_after_alloc(self, request: "Request",
272 + blocks: "KVCacheBlocks",
273 + num_external_tokens: int):
274 + if request.kv_transfer_params is None:
275 + return
276 +
277 + assert isinstance(request.kv_transfer_params, NixlKVTransferParams)
278 + if request.kv_transfer_params.do_remote_prefill:
279 + # NOTE(rob): if prompt < block_size, no remote blocks
280 + # since the remote only sends fully computed blocks, so
281 + # skip recving for this request. num_external_tokens
282 + # should be 0 if there are no remote blocks.
283 + if request.kv_transfer_params.remote_block_ids:
284 + # Get unhashed blocks to pull from remote.
285 + self._reqs_need_recv[request.request_id] = (
286 + request, blocks.get_unhashed_block_ids())
287 + else:
288 + assert num_external_tokens == 0
289 + # Only trigger 1 KV transfer per request.
290 + request.kv_transfer_params.do_remote_prefill = False
291 +
292 + def build_connector_meta(
293 + self,
294 + scheduler_output: SchedulerOutput,
295 + ) -> KVConnectorMetadata:
296 + meta = NixlConnectorMetadata()
297 +
298 + # Loop through scheduled reqs and convert to ReqMeta.
299 + for req_id, (req, block_ids) in self._reqs_need_recv.items():
300 + assert isinstance(req.kv_transfer_params, NixlKVTransferParams)
301 + meta.add_new_req(
302 + request_id=req_id,
303 + local_block_ids=block_ids,
304 + kv_transfer_params=req.kv_transfer_params,
305 + )
306 +
307 + # Clear the list once workers start the transfers
308 + self._reqs_need_recv.clear()
309 +
310 + return meta
311 +
312 + def request_finished(
313 + self,
314 + request: "Request",
315 + block_ids: list[int],
316 + ) -> tuple[bool, Optional[dict[str, Any]]]:
317 + """
318 + Once a request is finished, determine whether request blocks
319 + should be freed now or will be sent asynchronously and freed later.
320 + """
321 +
322 + if request.kv_transfer_params is None:
323 + return False, None
324 + assert isinstance(request.kv_transfer_params, NixlKVTransferParams)
325 +
326 + if ((not request.kv_transfer_params.do_remote_decode)
327 + or (request.status != RequestStatus.FINISHED_LENGTH_CAPPED)):
328 + return False, None
329 +
330 + # Get computed blocks.
331 + all_full = request.num_computed_tokens % self.block_size == 0
332 + computed_block_ids = (block_ids if all_full else block_ids[:-1])
333 +
334 + # If prompt < block_size, no xfer so free blocks immediately.
335 + delay_free_blocks = len(computed_block_ids) > 0
336 +
337 + return delay_free_blocks, NixlKVTransferParams(
338 + do_remote_prefill=True,
339 + do_remote_decode=False,
340 + remote_block_ids=computed_block_ids,
341 + remote_engine_id=self.engine_id,
342 + remote_host=envs.VLLM_NIXL_SIDE_CHANNEL_HOST,
343 + remote_port=envs.VLLM_NIXL_SIDE_CHANNEL_PORT,
344 + ).__dict__
345 +
346 +
347 +class NixlConnectorWorker:
348 + """Implementation of Worker side methods"""
349 +
350 + def __init__(self, engine_id: str):
351 + if NixlWrapper is None:
352 + logger.error("NIXL is not available")
353 + raise RuntimeError("NIXL is not available")
354 + logger.info("Initializing NIXL wrapper")
355 + logger.info("Initializing NIXL worker %s", engine_id)
356 +
357 + # Agent.
358 + self.nixl_wrapper = NixlWrapper(str(uuid.uuid4()), None)
359 + # Map of engine_id -> agent_name.
360 + self._remote_agents: dict[str, str] = {}
361 +
362 + # Metadata.
363 + self.engine_id = engine_id
364 + self.rank = get_tensor_model_parallel_rank()
365 + self.world_size = get_tensor_model_parallel_world_size()
366 + self.tp_group = get_tp_group()
367 +
368 + # KV Caches and nixl tracking data.
369 + self.kv_caches: dict[str, torch.Tensor] = {}
370 +
371 + # Map of engine_id -> kv_caches_base_addr
372 + self.kv_caches_base_addr: dict[str, list[int]] = {}
373 +
374 + # Number of NIXL regions. Currently one region per cache
375 + # (so 1 per layer for MLA, otherwise 2 per layer)
376 + self.num_regions = 0
377 +
378 + # nixl_prepped_dlist_handle (int).
379 + self.src_xfer_side_handle: int = 0
380 + # Map of engine_id -> nixl_prepped_dlist_handle (int)].
381 + self.dst_xfer_side_handles: dict[str, int] = {}
382 +
383 + # Map of engine_id -> num_blocks.
384 + self.dst_num_blocks: dict[str, int] = {}
385 + self._registered_descs: list[Any] = []
386 +
387 + # In progress transfers.
388 + # [req_id -> list[handle]]
389 + self._recving_transfers: defaultdict[str, list[Any]] = defaultdict(
390 + list[Any])
391 +
392 + # Complete transfer tracker. Used by the rank 0 to track finished
393 + # transactions on ranks 1 to N-1.
394 + # [req_id -> count]
395 + self._done_recving_count: defaultdict[str,
396 + int] = defaultdict(lambda: 0)
397 + self._done_sending_count: defaultdict[str,
398 + int] = defaultdict(lambda: 0)
399 +
400 + # Background thread for establishing new connections.
401 + self._nixl_handshake_listener_t: Optional[threading.Thread] = None
402 +
403 + @staticmethod
404 + def _nixl_handshake_listener(metadata: NixlAgentMetadata,
405 + ready_event: threading.Event, rank: int):
406 + """Background thread for getting new NIXL handshakes."""
407 + # NOTE(rob): this is a simple implementation. We will move
408 + # to a better approach like an ETCD server in the future.
409 +
410 + # NOTE(rob): to support heterogeneous TP, we will have to
411 + # move this into the scheduler rather than worker, since
412 + # each rank needs the metadata of all other ranks (whereas
413 + # in this setup, each rank only gets one other rank's meta.
414 +
415 + encoder = msgspec.msgpack.Encoder()
416 + encoded_data = encoder.encode(metadata)
417 + size_in_bytes = len(encoded_data)
418 + logger.debug("Size of encoded NixlAgentMetadata: %s bytes",
419 + str(size_in_bytes))
420 +
421 + # Listen for new requests for metadata.
422 + host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST
423 + # NOTE(rob): we need each rank to have a unique port. This
424 + # hack to keeps us moving. We will switch when moving to etcd
425 + # or where we have a single ZMQ socket in the scheduler.
426 + port = envs.VLLM_NIXL_SIDE_CHANNEL_PORT + rank
427 + path = f"tcp://{host}:{port}"
428 + logger.debug("Starting listening on path: %s", path)
429 + with zmq_ctx(zmq.ROUTER, path) as sock:
430 + ready_event.set()
431 + while True:
432 + identity, _, msg = sock.recv_multipart()
433 + if msg != GET_META_MSG:
434 + logger.warning(
435 + "Connection listener got unexpected message %s", msg)
436 + sock.send_multipart((identity, b"", encoded_data))
437 +
438 + def _nixl_handshake(self, host: str, port: int):
439 + """Do a NIXL handshake with a remote instance."""
440 +
441 + start_time = time.perf_counter()
442 + # NOTE(rob): we need each rank to have a unique port. This is
443 + # a hack to keep us moving. We will switch when moving to etcd
444 + # or where we have a single ZMQ socket in the scheduler.
445 + path = f"tcp://{host}:{port + self.rank}"
446 + logger.debug("Querying metadata on path: %s", path)
447 + with zmq_ctx(zmq.REQ, path) as sock:
448 + # Send query for the request.
449 + sock.send(GET_META_MSG)
450 + metadata_bytes = sock.recv()
451 + decoder = msgspec.msgpack.Decoder(NixlAgentMetadata)
452 + metadata = decoder.decode(metadata_bytes)
453 + got_metadata_time = time.perf_counter()
454 +
455 + # Register Remote agent.
456 + self.add_remote_agent(metadata)
457 + setup_agent_time = time.perf_counter()
458 +
459 + logger.debug("NIXL handshake: get metadata took: %s",
460 + got_metadata_time - start_time)
461 + logger.debug("NIXL handshake: add agent took: %s",
462 + setup_agent_time - got_metadata_time)
463 +
464 + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
465 + """Register the KV Cache data in nixl."""
466 +
467 + _, first_kv_cache = next(iter(kv_caches.items()))
468 + kv_elem_size = first_kv_cache.element_size()
469 +
470 + # TODO(tms): Find a more robust way to detect and handle MLA
471 + use_mla = len(first_kv_cache.shape) == 3
472 + if use_mla:
473 + # MLA case.
474 + self.num_blocks = first_kv_cache.shape[0]
475 + block_rank = 2 # [block_size, latent_dim]
476 + block_shape = first_kv_cache.shape[-block_rank:]
477 + else:
478 + # [2 (k and v), num_blocks, ...]
479 + self.num_blocks = first_kv_cache.shape[1]
480 + block_rank = 3 # [block_size, kv_heads, head_dim]
481 + block_shape = first_kv_cache.shape[-block_rank:]
482 +
483 + # TODO(tms): self.block_len needs to be per-layer for sliding window,
484 + # hybrid attn, etc
485 + self.block_len = kv_elem_size * math.prod(block_shape)
486 +
487 + logger.debug("Registering KV_Caches. use_mla: %s, shape %s", use_mla,
488 + first_kv_cache.shape)
489 + logger.debug("num_blocks: %s, block_shape: %s", self.num_blocks,
490 + block_shape)
491 + logger.debug("Per layer kv cache size: %s", first_kv_cache.shape)
492 + self.dst_num_blocks[self.engine_id] = self.num_blocks
493 + self.kv_caches = kv_caches
494 + kv_caches_base_addr = []
495 + caches_data = []
496 +
497 + # Note(tms): I modified this from the original region setup code.
498 + # K and V are now in different regions. Advantage is that we can
499 + # elegantly support MLA and any cases where the K and V tensors
500 + # are non-contiguous (it's not locally guaranteed that they will be)
501 + # Disadvantage is that the encoded NixlAgentMetadata is now larger
502 + # (roughly 8KB vs 5KB).
503 + for cache_or_caches in kv_caches.values():
504 + # Normalize to always be a list of caches
505 + cache_list = [cache_or_caches] if use_mla else cache_or_caches
506 + for cache in cache_list:
507 + base_addr = cache.data_ptr()
508 + region_len = self.num_blocks * self.block_len
509 + caches_data.append((base_addr, region_len, self.rank, ""))
510 + kv_caches_base_addr.append(base_addr)
511 + self.kv_caches_base_addr[self.engine_id] = kv_caches_base_addr
512 + self.num_regions = len(caches_data)
513 +
514 + descs = self.nixl_wrapper.get_reg_descs(caches_data, "VRAM")
515 + logger.debug("Registering descs: %s", caches_data)
516 + self.nixl_wrapper.register_memory(descs)
517 + logger.debug("Done registering descs")
518 +
519 + self._registered_descs.append(descs)
520 +
521 + # After KV Caches registered, listen for new connections.
522 + metadata = NixlAgentMetadata(
523 + engine_id=self.engine_id,
524 + agent_metadata=self.nixl_wrapper.get_agent_metadata(),
525 + kv_caches_base_addr=self.kv_caches_base_addr[self.engine_id],
526 + num_blocks=self.num_blocks,
527 + )
528 + ready_event = threading.Event()
529 + self._nixl_handshake_listener_t = threading.Thread(
530 + target=self._nixl_handshake_listener,
531 + args=(metadata, ready_event, self.rank),
532 + daemon=True,
533 + name="nixl_handshake_listener")
534 + self._nixl_handshake_listener_t.start()
535 + ready_event.wait()
536 +
537 + def add_remote_agent(self, nixl_agent_meta: NixlAgentMetadata):
538 + engine_id = nixl_agent_meta.engine_id
539 + if engine_id in self._remote_agents:
540 + return
541 +
542 + self._remote_agents[engine_id] = self.nixl_wrapper.add_remote_agent(
543 + nixl_agent_meta.agent_metadata)
544 + self.kv_caches_base_addr[
545 + engine_id] = nixl_agent_meta.kv_caches_base_addr
546 +
547 + # Create src descs and xfer side handles.
548 + blocks_data = []
549 + for base_addr in self.kv_caches_base_addr[self.engine_id]:
550 + for block_id in range(self.num_blocks):
551 + block_offset = block_id * self.block_len
552 + # (addr, len, device id)
553 + blocks_data.append(
554 + (base_addr + block_offset, self.block_len, self.rank))
555 + logger.debug("Created %s blocks for src engine %s and rank %s",
556 + len(blocks_data), self.engine_id, self.rank)
557 +
558 + # Register with NIXL.
559 + descs = self.nixl_wrapper.get_xfer_descs(blocks_data, "VRAM")
560 + self.src_xfer_side_handle = self.nixl_wrapper.prep_xfer_dlist(
561 + "NIXL_INIT_AGENT", descs)
562 +
563 + # Create dst descs and xfer side handles.
564 + self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks
565 + blocks_data = []
566 + for base_addr in self.kv_caches_base_addr[engine_id]:
567 + for block_id in range(nixl_agent_meta.num_blocks):
568 + block_offset = block_id * self.block_len
569 + # (addr, len, device id)
570 + blocks_data.append(
571 + (base_addr + block_offset, self.block_len, self.rank))
572 + logger.debug("Created %s blocks for dst engine %s and rank %s",
573 + len(blocks_data), engine_id, self.rank)
574 +
575 + # Register with NIXL.
576 + descs = self.nixl_wrapper.get_xfer_descs(blocks_data, "VRAM")
577 + self.dst_xfer_side_handles[
578 + engine_id] = self.nixl_wrapper.prep_xfer_dlist(
579 + self._remote_agents[engine_id], descs)
580 +
581 + def get_finished(self) -> tuple[set[str], set[str]]:
582 + """
583 + Get requests that are done sending or recving.
584 +
585 + In TP>1 setup, each rank exchanges KVs with its counterpart
586 + ranks independently. get_finished() runs in a worker creates
587 + the done_sending and done_recving sets that are sent to the
588 + scheduler via ModelRunnerOutput by Rank 0. To ensure trnxs
589 + are done before adding to finished, Ranks 1 to N-1 communicate
590 + to Rank 0 once their transaction is done + Rank 0 returns
591 + finished sets to Scheduler only once all ranks are done.
592 + """
593 + done_sending = self._get_new_notifs()
594 + done_recving = self._pop_done_transfers(self._recving_transfers)
595 + if len(done_sending) > 0 or len(done_recving) > 0:
596 + logger.debug(
597 + "Rank %s, get_finished: %s requests done sending "
598 + "and %s requests done recving", self.rank, len(done_sending),
599 + len(done_recving))
600 +
601 + if self.world_size == 1:
602 + return done_sending, done_recving
603 +
604 + # Rank 0: get finished from all other ranks.
605 + if self.rank == 0:
606 + for req_id in done_sending:
607 + self._done_sending_count[req_id] += 1
608 + for req_id in done_recving:
609 + self._done_recving_count[req_id] += 1
610 +
611 + # Keep track of how many other ranks have finished.
612 + other_ranks_finished_ids: list[str] = []
613 + for i in range(1, self.world_size):
614 + other_ranks_finished_ids.extend(
615 + self.tp_group.recv_object(src=i))
616 + for req_id in other_ranks_finished_ids:
617 + if (req_id in self._done_recving_count
618 + or req_id in self._recving_transfers):
619 + self._done_recving_count[req_id] += 1
620 + else:
621 + self._done_sending_count[req_id] += 1
622 +
623 + # Return ids that finished on all ranks to the scheduler.
624 + all_done_recving: set[str] = set()
625 + for req_id in list(self._done_recving_count.keys()):
626 + if self._done_recving_count[req_id] == self.world_size:
627 + del self._done_recving_count[req_id]
628 + all_done_recving.add(req_id)
629 +
630 + all_done_sending: set[str] = set()
631 + for req_id in list(self._done_sending_count.keys()):
632 + if self._done_sending_count[req_id] == self.world_size:
633 + del self._done_sending_count[req_id]
634 + all_done_sending.add(req_id)
635 +
636 + return all_done_sending, all_done_recving
637 +
638 + # Ranks 1 to N-1: send finished ids to Rank 0.
639 + else:
640 + finished_req_ids = list(done_recving.union(done_sending))
641 + self.tp_group.send_object(finished_req_ids, dst=0)
642 +
643 + # Unused as only Rank 0 results are sent to scheduler.
644 + return done_sending, done_recving
645 +
646 + def _get_new_notifs(self) -> set[str]:
647 + """Get req_ids which got a remote xfer message."""
648 +
649 + notified_req_ids: set[str] = set()
650 + for req_ids in self.nixl_wrapper.get_new_notifs().values():
651 + for req_id in req_ids:
652 + assert req_id not in notified_req_ids
653 + notified_req_ids.add(req_id.decode("utf-8"))
654 + return notified_req_ids
655 +
656 + def _pop_done_transfers(self, transfers: dict[str, list[int]]) -> set[str]:
657 + """
658 + Pop completed xfers by checking for DONE state.
659 + Args:
660 + transfers: dict of req_id -> list[running_xfer]
661 + Returns:
662 + set of req_ids that have all done xfers
663 + """
664 + done_req_ids: set[str] = set()
665 + for req_id, handles in list(transfers.items()):
666 + running_reqs = []
667 + for handle in handles:
668 + xfer_state = self.nixl_wrapper.check_xfer_state(handle)
669 + if xfer_state == "DONE":
670 + # TODO ptarasiewicz: why abort is throwing errors?
671 + # self.nixl_wrapper.release_xfer_handle(handle)
672 + continue
673 + if xfer_state == "PROC":
674 + running_reqs.append(handle)
675 + else:
676 + raise RuntimeError("Transfer failed with state %s",
677 + xfer_state)
678 + if len(running_reqs) == 0:
679 + done_req_ids.add(req_id)
680 + del transfers[req_id]
681 + else:
682 + transfers[req_id] = running_reqs
683 + return done_req_ids
684 +
685 + def start_load_kv(self, metadata: NixlConnectorMetadata):
686 + """
687 + Start loading by triggering non-blocking nixl_xfer.
688 + We check for these trnxs to complete in each step().
689 + """
690 + for req_id, meta in metadata.requests.items():
691 + logger.debug(
692 + "start_load_kv for request %s from remote engine %s. "
693 + "Num local_block_ids: %s. Num remote_block_ids: %s. ", req_id,
694 + meta.remote_engine_id, len(meta.local_block_ids),
695 + len(meta.remote_block_ids))
696 + self._read_blocks(
697 + request_id=req_id,
698 + dst_engine_id=meta.remote_engine_id,
699 + local_block_ids=meta.local_block_ids,
700 + remote_block_ids=meta.remote_block_ids,
701 + remote_host=meta.remote_host,
702 + remote_port=meta.remote_port,
703 + )
704 +
705 + def _read_blocks(
706 + self,
707 + local_block_ids: list[int],
708 + remote_block_ids: list[int],
709 + remote_host: str,
710 + remote_port: int,
711 + dst_engine_id: str,
712 + request_id: str,
713 + ):
714 + # NOTE(rob): this takes ~2s. We need to get this off the hotpath.
715 + if dst_engine_id not in self._remote_agents:
716 + self._nixl_handshake(remote_host, remote_port)
717 +
718 + # NOTE(rob): having the staging blocks be on the READER side is
719 + # not going to work well (since we will have to call rearrange tensors).
720 + # after we detect the txn is complete (which means we cannot make the
721 + # read trxn async easily). If we want to make "READ" happen cleanly,
722 + # then we will need to have the staging blocks on the remote side.
723 +
724 + # NOTE(rob): according to nvidia the staging blocks are used to
725 + # saturate IB with heterogeneous TP sizes. We should remove the staging
726 + # blocks until we are ready.
727 +
728 + # Full prefix cache hit: do not need to read remote blocks,
729 + # just notify P worker that we have the blocks we need.
730 + num_local_blocks = len(local_block_ids)
731 + if num_local_blocks == 0:
732 + self.nixl_wrapper.send_notif(dst_engine_id,
733 + notif_msg=request_id.encode("utf-8"))
734 + return
735 +
736 + # Partial prefix cache hit: just read uncomputed blocks.
737 + num_remote_blocks = len(remote_block_ids)
738 + assert num_local_blocks <= num_remote_blocks
739 + if num_local_blocks < num_remote_blocks:
740 + remote_block_ids = remote_block_ids[-num_local_blocks:]
741 +
742 + # Get side handles.
743 + local_xfer_side_handle = self.src_xfer_side_handle
744 + remote_xfer_side_handle = self.dst_xfer_side_handles[dst_engine_id]
745 +
746 + # Get descs ids.
747 + remote_block_descs_ids = self._get_block_descs_ids(
748 + dst_engine_id, remote_block_ids)
749 + local_block_descs_ids = self._get_block_descs_ids(
750 + self.engine_id, local_block_ids)
751 + assert len(local_block_descs_ids) == len(remote_block_descs_ids)
752 +
753 + # Prepare transfer with Nixl.
754 + handle = self.nixl_wrapper.make_prepped_xfer(
755 + "READ",
756 + local_xfer_side_handle,
757 + local_block_descs_ids,
758 + remote_xfer_side_handle,
759 + remote_block_descs_ids,
760 + notif_msg=request_id.encode("utf-8"),
761 + )
762 +
763 + # Begin async xfer.
764 + self.nixl_wrapper.transfer(handle)
765 +
766 + # Use handle to check completion in future step().
767 + self._recving_transfers[request_id].append(handle)
768 +
769 + def _get_block_descs_ids(self, engine_id: str,
770 + block_ids: list[int]) -> list[int]:
771 + """Get the descs ids for a set of block ids."""
772 +
773 + # range(1) for MLA, range(2) otherwise.
774 + region_ids = range(self.num_regions)
775 + num_blocks = self.dst_num_blocks[engine_id]
776 +
777 + # Compute the desc ids for each block.
778 + descs_ids: list[int] = []
779 + for reg_id in region_ids:
780 + for block_id in block_ids:
781 + descs_ids.append(reg_id * num_blocks + block_id)
782 + return descs_ids
783 +
784 +
785 +@contextlib.contextmanager
786 +def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]:
787 + """Context manager for a ZMQ socket"""
788 +
789 + ctx: Optional[zmq.Context] = None
790 + try:
791 + ctx = zmq.Context() # type: ignore[attr-defined]
792 +
793 + if socket_type == zmq.ROUTER:
794 + socket = ctx.socket(zmq.ROUTER)
795 + socket.bind(addr)
796 + elif socket_type == zmq.REQ:
797 + socket = ctx.socket(zmq.REQ)
798 + socket.connect(addr)
799 + else:
800 + raise ValueError(f"Unexpected socket type: {socket_type}")
801 +
802 + yield socket
803 + finally:
804 + if ctx is not None:
805 + ctx.destroy(linger=0)
vllm/distributed/kv_transfer/kv_connector/v1/shared_storage_connector.py
+6 -6
@@ -17,6 +17,7 @@ from vllm.v1.core.sched.output import SchedulerOutput
17 if TYPE_CHECKING:
18 from vllm.attention.backends.abstract import AttentionMetadata
19 from vllm.forward_context import ForwardContext
20 + from vllm.v1.core.kv_cache_manager import KVCacheBlocks
21 from vllm.v1.request import Request
22
23 logger = init_logger(__name__)
@@ -132,8 +133,7 @@ class SharedStorageConnector(KVConnectorBase_V1):
133 dst_kv_cache_layer.reshape(dst_kv_cache_layer_shape)
134
135 # Get the metadata
135 - metadata: KVConnectorMetadata = \
136 - self._get_connector_metadata()
136 + metadata: KVConnectorMetadata = self._get_connector_metadata()
137 assert isinstance(metadata, SharedStorageConnectorMetadata)
138
139 if metadata is None:
@@ -225,7 +225,7 @@ class SharedStorageConnector(KVConnectorBase_V1):
225 self,
226 request: "Request",
227 num_computed_tokens: int,
228 - ) -> int:
228 + ) -> tuple[int, bool]:
229 """
230 Get number of new tokens that can be loaded from the
231 external KV cache beyond the num_computed_tokens.
@@ -239,7 +239,6 @@ class SharedStorageConnector(KVConnectorBase_V1):
239 the number of tokens that can be loaded from the
240 external KV cache beyond what is already computed.
241 """
242 -
242 # NOTE: in this debug implementation, we assume that the prompt is
243 # cached_prompt + newly_generated_single_token
244 # Therefore, we use prompt_token_ids[:-1] to determine the folder name
@@ -248,7 +247,7 @@ class SharedStorageConnector(KVConnectorBase_V1):
247 # with the block granularity. And it expects the returned blocks and
248 # num_computed_tokens to also be aligned with the block granularity.
249 if not self._found_match_for_request(request):
251 - return 0
250 + return 0, False
251
252 logger.info("External Cache Hit!")
253
@@ -257,9 +256,10 @@ class SharedStorageConnector(KVConnectorBase_V1):
256 num_tokens_to_check = align_to_block_size(
257 len(request.prompt_token_ids) - 1, self._block_size)
258
260 - return num_tokens_to_check - num_computed_tokens
259 + return num_tokens_to_check - num_computed_tokens, False
260
261 def update_state_after_alloc(self, request: "Request",
262 + blocks: "KVCacheBlocks",
263 num_external_tokens: int):
264 """
265 Update KVConnector state after block allocation.
vllm/entrypoints/openai/protocol.py
+17 -2
@@ -403,6 +403,9 @@ class ChatCompletionRequest(OpenAIBaseModel):
403 "access by 3rd parties, and long enough to be "
404 "unpredictable (e.g., 43 characters base64-encoded, corresponding "
405 "to 256 bit). Not supported by vLLM engine V0."))
406 + kv_transfer_params: Optional[dict[str, Any]] = Field(
407 + default=None,
408 + description="KVTransfer parameters used for disaggregated serving.")
409
410 # doc: end-chat-completion-extra-params
411
@@ -540,7 +543,9 @@ class ChatCompletionRequest(OpenAIBaseModel):
543 output_kind=RequestOutputKind.DELTA if self.stream \
544 else RequestOutputKind.FINAL_ONLY,
545 guided_decoding=guided_decoding,
543 - logit_bias=self.logit_bias)
546 + logit_bias=self.logit_bias,
547 + extra_args=({"kv_transfer_params": self.kv_transfer_params}
548 + if self.kv_transfer_params else None))
549
550 def _get_guided_json_from_tool(
551 self) -> Optional[Union[str, dict, BaseModel]]:
@@ -848,6 +853,10 @@ class CompletionRequest(OpenAIBaseModel):
853 " as strings of the form 'token_id:{token_id}' so that tokens "
854 "that are not JSON-encodable can be identified."))
855
856 + kv_transfer_params: Optional[dict[str, Any]] = Field(
857 + default=None,
858 + description="KVTransfer parameters used for disaggregated serving.")
859 +
860 # doc: end-completion-extra-params
861
862 # Default sampling parameters for completion requests
@@ -973,7 +982,9 @@ class CompletionRequest(OpenAIBaseModel):
982 else RequestOutputKind.FINAL_ONLY,
983 guided_decoding=guided_decoding,
984 logit_bias=self.logit_bias,
976 - allowed_token_ids=self.allowed_token_ids)
985 + allowed_token_ids=self.allowed_token_ids,
986 + extra_args=({"kv_transfer_params": self.kv_transfer_params}
987 + if self.kv_transfer_params else None))
988
989 @model_validator(mode="before")
990 @classmethod
@@ -1223,6 +1234,8 @@ class CompletionResponse(OpenAIBaseModel):
1234 model: str
1235 choices: list[CompletionResponseChoice]
1236 usage: UsageInfo
1237 + kv_transfer_params: Optional[dict[str, Any]] = Field(
1238 + default=None, description="KVTransfer parameters.")
1239
1240
1241 class CompletionResponseStreamChoice(OpenAIBaseModel):
@@ -1412,6 +1425,8 @@ class ChatCompletionResponse(OpenAIBaseModel):
1425 choices: list[ChatCompletionResponseChoice]
1426 usage: UsageInfo
1427 prompt_logprobs: Optional[list[Optional[dict[int, Logprob]]]] = None
1428 + kv_transfer_params: Optional[dict[str, Any]] = Field(
1429 + default=None, description="KVTransfer parameters.")
1430
1431
1432 class DeltaMessage(OpenAIBaseModel):
vllm/entrypoints/openai/serving_chat.py
+1
@@ -1086,6 +1086,7 @@ class OpenAIServingChat(OpenAIServing):
1086 choices=choices,
1087 usage=usage,
1088 prompt_logprobs=clamp_prompt_logprobs(final_res.prompt_logprobs),
1089 + kv_transfer_params=final_res.kv_transfer_params,
1090 )
1091
1092 return response
vllm/entrypoints/openai/serving_completion.py
+1 -1
@@ -482,7 +482,7 @@ class OpenAIServingCompletion(OpenAIServing):
482 model=model_name,
483 choices=choices,
484 usage=usage,
485 - )
485 + kv_transfer_params=final_res_batch[0].kv_transfer_params)
486
487 def _create_completion_logprobs(
488 self,
vllm/envs.py
+10
@@ -112,6 +112,8 @@ if TYPE_CHECKING:
112 VLLM_XGRAMMAR_CACHE_MB: int = 0
113 VLLM_MSGPACK_ZERO_COPY_THRESHOLD: int = 256
114 VLLM_ALLOW_INSECURE_SERIALIZATION: bool = False
115 + VLLM_NIXL_SIDE_CHANNEL_HOST: str = "localhost"
116 + VLLM_NIXL_SIDE_CHANNEL_PORT: int = 5557
117
118
119 def get_default_cache_root():
@@ -747,6 +749,14 @@ environment_variables: dict[str, Callable[[], Any]] = {
749 # insecure method and it is needed for some reason.
750 "VLLM_ALLOW_INSECURE_SERIALIZATION":
751 lambda: bool(int(os.getenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "0"))),
752 +
753 + # IP address used for NIXL handshake between remote agents.
754 + "VLLM_NIXL_SIDE_CHANNEL_HOST":
755 + lambda: os.getenv("VLLM_NIXL_SIDE_CHANNEL_HOST", "localhost"),
756 +
757 + # Port used for NIXL handshake between remote agents.
758 + "VLLM_NIXL_SIDE_CHANNEL_PORT":
759 + lambda: int(os.getenv("VLLM_NIXL_SIDE_CHANNEL_PORT", "5557")),
760 }
761
762 # end-env-vars-definition
vllm/forward_context.py
-21
@@ -11,10 +11,6 @@ import torch.distributed as dist
11
12 import vllm.envs as envs
13 from vllm.config import VllmConfig
14 -from vllm.distributed.kv_transfer import (get_kv_transfer_group,
15 - has_kv_transfer_group,
16 - is_v1_kv_transfer_group)
17 -from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1
14 from vllm.logger import init_logger
15
16 if TYPE_CHECKING:
@@ -106,16 +102,6 @@ def set_forward_context(attn_metadata: Any,
102 attn_metadata=attn_metadata,
103 dp_metadata=dp_metadata)
104
109 - # KVConnector: trigger (possibly async) load before forward.
110 - # Each attn layer will block until the reading is complete.
111 - trigger_kv_transfer = (attn_metadata is not None
112 - and has_kv_transfer_group()
113 - and is_v1_kv_transfer_group())
114 - if trigger_kv_transfer:
115 - kv_connector = get_kv_transfer_group()
116 - assert isinstance(kv_connector, KVConnectorBase_V1)
117 - kv_connector.start_load_kv(_forward_context)
118 -
105 try:
106 yield
107 finally:
@@ -152,11 +138,4 @@ def set_forward_context(attn_metadata: Any,
138 "(batchsize, count, median_time(ms)): %s"),
139 forward_stats)
140
155 - # KVConnector: each attn layer triggers (possibly async) save.
156 - # Ensure all those operations complete before forward() is done.
157 - if trigger_kv_transfer:
158 - kv_connector = get_kv_transfer_group()
159 - assert isinstance(kv_connector, KVConnectorBase_V1)
160 - kv_connector.wait_for_save()
161 -
141 _forward_context = prev_context
vllm/outputs.py
+5 -1
@@ -4,7 +4,7 @@ import time
4 from collections.abc import MutableSequence
5 from collections.abc import Sequence as GenericSequence
6 from dataclasses import dataclass
7 -from typing import Generic, Optional, Union
7 +from typing import Any, Generic, Optional, Union
8
9 import torch
10 from typing_extensions import TypeVar, deprecated
@@ -103,6 +103,7 @@ class RequestOutput:
103 encoder_prompt_token_ids: The token IDs of the encoder prompt.
104 None if decoder-only.
105 num_cached_tokens: The number of tokens with prefix cache hit.
106 + kv_transfer_params: The params for remote K/V transfer.
107 """
108
109 def __init__(
@@ -120,6 +121,7 @@ class RequestOutput:
121 num_cached_tokens: Optional[int] = None,
122 *,
123 multi_modal_placeholders: Optional[MultiModalPlaceholderDict] = None,
124 + kv_transfer_params: Optional[dict[str, Any]] = None,
125 ) -> None:
126 self.request_id = request_id
127 self.prompt = prompt
@@ -133,11 +135,13 @@ class RequestOutput:
135 self.encoder_prompt = encoder_prompt
136 self.encoder_prompt_token_ids = encoder_prompt_token_ids
137 self.num_cached_tokens = num_cached_tokens
138 + self.kv_transfer_params = kv_transfer_params
139
140 def add(self, next_output: "RequestOutput", aggregate: bool) -> None:
141 """Merge subsequent RequestOutput into this one"""
142
143 self.finished |= next_output.finished
144 + self.kv_transfer_params = next_output.kv_transfer_params
145
146 for next_completion in next_output.outputs:
147 for i, completion in enumerate(self.outputs):
vllm/v1/core/kv_cache_manager.py
+32 -1
@@ -36,6 +36,12 @@ class KVCacheBlocks:
36 """Converts the KVCacheBlocks instance to a list of block IDs."""
37 return [block.block_id for block in self.blocks]
38
39 + def get_unhashed_block_ids(self) -> list[int]:
40 + """Get block_ids of unhashed blocks from KVCacheBlocks instance."""
41 + return [
42 + block.block_id for block in self.blocks if block.block_hash is None
43 + ]
44 +
45
46 class KVCacheManager:
47
@@ -116,6 +122,12 @@ class KVCacheManager:
122 - The number of computed tokens.
123 """
124
125 + # Request already has blocks from async load via KVConnector.
126 + num_existing_blocks = len(
127 + self.single_type_manager.req_to_blocks[request.request_id])
128 + if num_existing_blocks > 0:
129 + return KVCacheBlocks.create_empty(), request.num_computed_tokens
130 +
131 # Prefix caching is disabled or
132 # When the request requires prompt logprobs, we skip prefix caching.
133 if (not self.enable_caching
@@ -173,6 +185,7 @@ class KVCacheManager:
185 num_new_tokens: int,
186 new_computed_blocks: Optional[KVCacheBlocks] = None,
187 num_lookahead_tokens: int = 0,
188 + delay_cache_blocks: bool = False,
189 ) -> Optional[KVCacheBlocks]:
190 """Add slots for a request with new tokens to append.
191
@@ -186,6 +199,9 @@ class KVCacheManager:
199 num_lookahead_tokens: The number of speculative tokens to allocate.
200 This is used by spec decode proposers with kv-cache such
201 as eagle.
202 + delay_cache_blocks: Whether to skip caching the blocks. This is
203 + used by P/D when allocating blocks used in a KV transfer
204 + which will complete in a future step.
205
206 Blocks layout:
207 ```
@@ -255,7 +271,9 @@ class KVCacheManager:
271 new_blocks = self.single_type_manager.allocate_new_blocks(
272 request.request_id, num_tokens_need_slot)
273
258 - if not self.enable_caching:
274 + # P/D: delay caching blocks if we have to recv from
275 + # remote. Update state for locally cached blocks.
276 + if not self.enable_caching or delay_cache_blocks:
277 return KVCacheBlocks(new_blocks)
278
279 # Speculated tokens might be rejected in the future, so we does
@@ -350,3 +368,16 @@ class KVCacheManager:
368 A list of KV cache events.
369 """
370 return self.block_pool.take_events()
371 +
372 + def get_block_ids(self, request_id: str) -> list[int]:
373 + """Get the block ids of a request."""
374 + assert request_id in self.single_type_manager.req_to_blocks
375 + return [
376 + block.block_id
377 + for block in self.single_type_manager.req_to_blocks[request_id]
378 + ]
379 +
380 + def get_num_blocks(self, request_id: str):
381 + """Get the number of blocks."""
382 + assert request_id in self.single_type_manager.req_to_blocks
383 + return len(self.single_type_manager.req_to_blocks[request_id])
vllm/v1/core/sched/interface.py
+4
@@ -4,6 +4,7 @@ from collections.abc import Iterable
4 from typing import TYPE_CHECKING, Optional, Union
5
6 if TYPE_CHECKING:
7 + from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1
8 from vllm.v1.core.sched.output import SchedulerOutput
9 from vllm.v1.engine import EngineCoreOutputs
10 from vllm.v1.metrics.stats import SchedulerStats
@@ -137,3 +138,6 @@ class SchedulerInterface(ABC):
138 def shutdown(self) -> None:
139 """Shutdown the scheduler."""
140 raise NotImplementedError
141 +
142 + def get_kv_connector(self) -> Optional["KVConnectorBase_V1"]:
143 + return None
vllm/v1/core/sched/scheduler.py
+152 -36
@@ -5,13 +5,15 @@ from __future__ import annotations
5 import time
6 from collections import defaultdict, deque
7 from collections.abc import Iterable
8 -from typing import Optional, Union
8 +from typing import Any, Optional, Union
9
10 from vllm.config import VllmConfig
11 from vllm.distributed.kv_events import EventPublisherFactory, KVEventBatch
12 from vllm.distributed.kv_transfer.kv_connector.factory import (
13 KVConnectorFactory)
14 -from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
14 +from vllm.distributed.kv_transfer.kv_connector.v1 import (KVConnectorBase_V1,
15 + KVConnectorRole,
16 + KVTransferParams)
17 from vllm.logger import init_logger
18 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
19 from vllm.v1.core.encoder_cache_manager import (EncoderCacheManager,
@@ -96,6 +98,9 @@ class Scheduler(SchedulerInterface):
98 # This is flushed at the end of each scheduling step.
99 self.finished_req_ids: set[str] = set()
100
101 + # P/D: requests in process of recving KV transfers
102 + self.finished_recving_kv_req_ids: set[str] = set()
103 +
104 # OPTIMIZATION: Cache the CachedRequestData objects to avoid creating
105 # them at each scheduling step.
106 # Request id -> deque of CachedRequestData
@@ -307,6 +312,16 @@ class Scheduler(SchedulerInterface):
312
313 request = self.waiting[0]
314
315 + # P/D: skip request if still waiting for remote kvs.
316 + if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
317 + is_ready = self._update_waiting_for_remote_kv(request)
318 + if is_ready:
319 + request.status = RequestStatus.WAITING
320 + else:
321 + self.waiting.popleft()
322 + skipped_waiting_requests.appendleft(request)
323 + continue
324 +
325 # Skip request if the structured output request is still waiting
326 # for FSM compilation.
327 if request.status == RequestStatus.WAITING_FOR_FSM:
@@ -330,49 +345,55 @@ class Scheduler(SchedulerInterface):
345 continue
346
347 # Get already-cached tokens.
333 - computed_blocks, num_computed_tokens = \
348 + new_computed_blocks, num_computed_tokens = \
349 self.kv_cache_manager.get_computed_blocks(
350 request)
351
352 # Get externally-cached tokens if using a KVConnector.
338 - num_external_tokens = (
339 - 0 if self.connector is None else
353 + num_external_tokens, load_kv_async = (
354 + (0, False) if self.connector is None else
355 self.connector.get_num_new_matched_tokens(
356 request, num_computed_tokens))
357
358 # Total computed tokens (local + external).
359 num_computed_tokens += num_external_tokens
360
361 + encoder_inputs_to_schedule = None
362 + new_encoder_budget = encoder_budget
363 +
364 + # P/D: loading remote KV, do not allocate for new work.
365 + if load_kv_async:
366 + num_new_tokens = 0
367 # Number of tokens to be scheduled.
347 - # We use `request.num_tokens` instead of
348 - # `request.num_prompt_tokens` to consider the resumed requests,
349 - # which have output tokens.
350 - num_new_tokens = request.num_tokens - num_computed_tokens
351 - if (0 < self.scheduler_config.long_prefill_token_threshold <
352 - num_new_tokens):
353 - num_new_tokens = (
354 - self.scheduler_config.long_prefill_token_threshold)
355 - num_new_tokens = min(num_new_tokens, token_budget)
356 - assert num_new_tokens > 0
357 -
358 - # Schedule encoder inputs.
359 - if request.has_encoder_inputs:
360 - (encoder_inputs_to_schedule, num_new_tokens,
361 - new_encoder_budget) = self._try_schedule_encoder_inputs(
362 - request, num_computed_tokens, num_new_tokens,
363 - encoder_budget)
364 - if num_new_tokens == 0:
365 - # The request cannot be scheduled.
366 - break
368 else:
368 - encoder_inputs_to_schedule = None
369 - new_encoder_budget = encoder_budget
369 + # We use `request.num_tokens` instead of
370 + # `request.num_prompt_tokens` to consider the resumed
371 + # requests, which have output tokens.
372 + num_new_tokens = request.num_tokens - num_computed_tokens
373 + if (0 < self.scheduler_config.long_prefill_token_threshold
374 + < num_new_tokens):
375 + num_new_tokens = (
376 + self.scheduler_config.long_prefill_token_threshold)
377 + num_new_tokens = min(num_new_tokens, token_budget)
378 + assert num_new_tokens > 0
379 +
380 + # Schedule encoder inputs.
381 + if request.has_encoder_inputs:
382 + (encoder_inputs_to_schedule, num_new_tokens,
383 + new_encoder_budget
384 + ) = self._try_schedule_encoder_inputs(
385 + request, num_computed_tokens, num_new_tokens,
386 + encoder_budget)
387 + if num_new_tokens == 0:
388 + # The request cannot be scheduled.
389 + break
390
391 new_blocks = self.kv_cache_manager.allocate_slots(
392 request,
393 num_new_tokens + num_external_tokens,
374 - computed_blocks,
394 + new_computed_blocks,
395 num_lookahead_tokens=self.num_lookahead_tokens,
396 + delay_cache_blocks=load_kv_async,
397 )
398 if new_blocks is None:
399 # The request cannot be scheduled.
@@ -384,10 +405,18 @@ class Scheduler(SchedulerInterface):
405 if self.connector is not None:
406 self.connector.update_state_after_alloc(
407 request,
408 + new_computed_blocks + new_blocks,
409 num_external_tokens,
410 )
411
412 self.waiting.popleft()
413 + if load_kv_async:
414 + # If loading async, allocate memory and put request
415 + # into the WAITING_FOR_REMOTE_KV state.
416 + skipped_waiting_requests.appendleft(request)
417 + request.status = RequestStatus.WAITING_FOR_REMOTE_KVS
418 + continue
419 +
420 if request.use_structured_output:
421 structured_output_request_ids[
422 request.request_id] = req_index
@@ -407,7 +436,7 @@ class Scheduler(SchedulerInterface):
436 if self.lora_config and request.lora_request:
437 scheduled_loras.add(request.lora_request.lora_int_id)
438 req_to_new_block_ids[request.request_id] = (
410 - computed_blocks + new_blocks).get_block_ids()
439 + self.kv_cache_manager.get_block_ids(request.request_id))
440 num_scheduled_tokens[request.request_id] = num_new_tokens
441 token_budget -= num_new_tokens
442 request.status = RequestStatus.RUNNING
@@ -698,6 +727,7 @@ class Scheduler(SchedulerInterface):
727 stopped = False
728 new_logprobs = None
729 new_token_ids = generated_token_ids
730 + kv_transfer_params = None
731
732 # Append generated tokens and check for stop. Note that if
733 # a request is still being prefilled, we expect the model runner
@@ -709,7 +739,7 @@ class Scheduler(SchedulerInterface):
739 # This must be called before we make the EngineCoreOutput.
740 stopped = check_stop(request, self.max_model_len)
741 if stopped:
712 - self._free_request(request)
742 + kv_transfer_params = self._free_request(request)
743 del new_token_ids[num_new:] # Trim new tokens if needed.
744 break
745
@@ -739,7 +769,8 @@ class Scheduler(SchedulerInterface):
769
770 # Get prompt logprobs for this request.
771 prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
742 - if new_token_ids:
772 + if new_token_ids or kv_transfer_params:
773 +
774 # Add EngineCoreOutput for this Request.
775 outputs.append(
776 EngineCoreOutput(
@@ -749,7 +780,10 @@ class Scheduler(SchedulerInterface):
780 new_logprobs=new_logprobs,
781 new_prompt_logprobs_tensors=prompt_logprobs_tensors,
782 stop_reason=request.stop_reason,
752 - events=request.take_events()))
783 + events=request.take_events(),
784 + kv_transfer_params=kv_transfer_params,
785 + ))
786 +
787 else:
788 # Invariant: EngineCore returns no partial prefill outputs.
789 assert not prompt_logprobs_tensors
@@ -757,6 +791,9 @@ class Scheduler(SchedulerInterface):
791 if not stopped:
792 new_running.append(request)
793
794 + # P/D: update state for finished KV Transfers.
795 + self._update_from_kv_xfer_finished(model_runner_output)
796 +
797 # Return the cached request data to the queue so they can be reused.
798 for req_data in scheduler_output.scheduled_cached_reqs:
799 # NOTE(rob): since we free stopped reqs above, adding stopped reqs
@@ -811,15 +848,27 @@ class Scheduler(SchedulerInterface):
848 request.status = finished_status
849 self._free_request(request)
850
814 - def _free_request(self, request: Request) -> None:
851 + def _free_request(self, request: Request) -> Optional[dict[str, Any]]:
852 +
853 assert request.is_finished()
816 - self.kv_cache_manager.free(request)
817 - self.kv_cache_manager.free_block_hashes(request)
854 +
855 + delay_free_blocks, kv_xfer_params = self._connector_finished(request)
856 self.encoder_cache_manager.free(request)
857 self._cached_reqs_data.pop(request.request_id, None)
820 - del self.requests[request.request_id]
858 self.finished_req_ids.add(request.request_id)
859
860 + if not delay_free_blocks:
861 + self._free_blocks(request)
862 +
863 + return kv_xfer_params
864 +
865 + def _free_blocks(self, request: Request):
866 + assert request.is_finished()
867 + assert request.request_id not in self._cached_reqs_data
868 + self.kv_cache_manager.free(request)
869 + self.kv_cache_manager.free_block_hashes(request)
870 + del self.requests[request.request_id]
871 +
872 def get_num_unfinished_requests(self) -> int:
873 return len(self.waiting) + len(self.running)
874
@@ -863,3 +912,70 @@ class Scheduler(SchedulerInterface):
912 def shutdown(self) -> None:
913 if self.kv_event_publisher:
914 self.kv_event_publisher.shutdown()
915 +
916 + ########################################################################
917 + # P/D Related Methods
918 + ########################################################################
919 +
920 + def get_kv_connector(self) -> Optional[KVConnectorBase_V1]:
921 + return self.connector
922 +
923 + def _connector_finished(
924 + self, request: Request) -> tuple[bool, Optional[KVTransferParams]]:
925 + """Invoke the KV connector request_finished() method if applicable."""
926 + if self.connector is None:
927 + return False, None
928 + block_ids = self.kv_cache_manager.get_block_ids(request.request_id)
929 + return self.connector.request_finished(request, block_ids)
930 +
931 + def _update_waiting_for_remote_kv(self, request: Request) -> bool:
932 + """
933 + P/D: check if the request_id is finished_recving.
934 +
935 + The finished_recving_kv_req_ids list is populated
936 + on the previous steps()'s update_from_output based
937 + on the worker side connector.
938 +
939 + When the kv transfer is ready, we cache the blocks
940 + and the request state will be moved back to WAITING from
941 + WAITING_FOR_REMOTE_KV.
942 + """
943 + if request.request_id not in self.finished_recving_kv_req_ids:
944 + return False
945 +
946 + # Now that the blocks are ready, actually cache them.
947 + block_ids = self.kv_cache_manager.get_block_ids(request.request_id)
948 + num_computed_tokens = len(block_ids) * self.block_size
949 + if num_computed_tokens == request.num_tokens:
950 + num_computed_tokens -= 1
951 + self.kv_cache_manager.single_type_manager.cache_blocks(
952 + request,
953 + self.kv_cache_manager.req_to_block_hashes[request.request_id],
954 + num_computed_tokens,
955 + )
956 +
957 + # Update the request state for scheduling.
958 + request.num_computed_tokens = num_computed_tokens
959 +
960 + # Return that we are ready.
961 + self.finished_recving_kv_req_ids.remove(request.request_id)
962 + return True
963 +
964 + def _update_from_kv_xfer_finished(self,
965 + model_runner_output: ModelRunnerOutput):
966 + """
967 + P/D: update the scheduler state based on the output.
968 +
969 + The Worker side connectors add finished_recving and
970 + finished_sending reqs to the output.
971 + * if finished_sending: free the blocks
972 + # if finished_recving: add to state so we can
973 + scheduler the request during the next step.
974 + """
975 + # P/D: update recv and send status from last step.
976 + for req_id in (model_runner_output.finished_recving or ()):
977 + logger.debug("Finished recving KV transfer for request %s", req_id)
978 + self.finished_recving_kv_req_ids.add(req_id)
979 + for req_id in (model_runner_output.finished_sending or ()):
980 + logger.debug("Finished sending KV transfer for request %s", req_id)
981 + self._free_blocks(self.requests[req_id])
vllm/v1/engine/__init__.py
+1
@@ -105,6 +105,7 @@ class EngineCoreOutput(
105 finish_reason: Optional[FinishReason] = None
106 stop_reason: Union[int, str, None] = None
107 events: Optional[list[EngineCoreEvent]] = None
108 + kv_transfer_params: Optional[dict[str, Any]] = None
109
110 @property
111 def finished(self) -> bool:
vllm/v1/engine/core.py
+9
@@ -182,6 +182,15 @@ class EngineCore:
182 # Start grammar compilation asynchronously
183 self.structured_output_manager.grammar_init(req)
184
185 + if req.raw_kv_transfer_params is not None:
186 + if (kv_connector := self.scheduler.get_kv_connector()):
187 + # Parse raw KV transfer params via connector.
188 + kv_connector.set_kv_transfer_params(req)
189 + else:
190 + logger.warning(
191 + "Got KVTransferParams, but no KVConnector found. "
192 + "Disabling KVTransfer for this request.")
193 +
194 self.scheduler.add_request(req)
195
196 def abort_requests(self, request_ids: list[str]):
vllm/v1/engine/output_processor.py
+9 -3
@@ -3,7 +3,7 @@
3 import asyncio
4 from collections.abc import Iterable
5 from dataclasses import dataclass
6 -from typing import Optional, Union
6 +from typing import Any, Optional, Union
7
8 from vllm.outputs import CompletionOutput, RequestOutput
9 from vllm.sampling_params import RequestOutputKind
@@ -146,6 +146,7 @@ class RequestState:
146 new_token_ids: list[int],
147 finish_reason: Optional[FinishReason],
148 stop_reason: Union[int, str, None],
149 + kv_transfer_params: Optional[dict[str, Any]] = None,
150 ) -> Optional[RequestOutput]:
151
152 finished = finish_reason is not None
@@ -167,13 +168,15 @@ class RequestState:
168 if not outputs:
169 return None
170
170 - return self._new_request_output(request_id, outputs, finished)
171 + return self._new_request_output(request_id, outputs, finished,
172 + kv_transfer_params)
173
174 def _new_request_output(
175 self,
176 request_id: str,
177 outputs: list[CompletionOutput],
178 finished: bool,
179 + kv_transfer_params: Optional[dict[str, Any]] = None,
180 ) -> RequestOutput:
181
182 if self.output_kind == RequestOutputKind.DELTA:
@@ -189,6 +192,7 @@ class RequestState:
192 prompt_logprobs=prompt_logprobs,
193 outputs=outputs,
194 finished=finished,
195 + kv_transfer_params=kv_transfer_params,
196 )
197
198 def _new_completion_output(
@@ -335,6 +339,7 @@ class OutputProcessor:
339 new_token_ids = engine_core_output.new_token_ids
340 finish_reason = engine_core_output.finish_reason
341 stop_reason = engine_core_output.stop_reason
342 + kv_transfer_params = engine_core_output.kv_transfer_params
343
344 req_state.is_prefilling = False
345
@@ -350,7 +355,8 @@ class OutputProcessor:
355
356 # 4) Create and handle RequestOutput objects.
357 if request_output := req_state.make_request_output(
353 - new_token_ids, finish_reason, stop_reason):
358 + new_token_ids, finish_reason, stop_reason,
359 + kv_transfer_params):
360 if req_state.queue is not None:
361 # AsyncLLM: put into queue for handling by generate().
362 req_state.queue.put(request_output)
vllm/v1/outputs.py
+13 -9
@@ -100,12 +100,16 @@ class ModelRunnerOutput:
100 # [prompt_len]
101 prompt_logprobs_dict: dict[str, Optional[LogprobsTensors]]
102
103 -
104 -EMPTY_MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
105 - req_ids=[],
106 - req_id_to_index={},
107 - sampled_token_ids=[],
108 - spec_token_ids=None,
109 - logprobs=None,
110 - prompt_logprobs_dict={},
111 -)
103 + # [req_ids]
104 + finished_sending: Optional[set[str]] = None
105 + finished_recving: Optional[set[str]] = None
106 +
107 +
108 +EMPTY_MODEL_RUNNER_OUTPUT = ModelRunnerOutput(req_ids=[],
109 + req_id_to_index={},
110 + sampled_token_ids=[],
111 + spec_token_ids=None,
112 + logprobs=None,
113 + prompt_logprobs_dict={},
114 + finished_sending=None,
115 + finished_recving=None)
vllm/v1/request.py
+12 -1
@@ -1,8 +1,9 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 import enum
4 -from typing import TYPE_CHECKING, Optional, Union
4 +from typing import TYPE_CHECKING, Any, Optional, Union
5
6 +from vllm.distributed.kv_transfer.kv_connector.v1 import KVTransferParams
7 from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
8 from vllm.sampling_params import SamplingParams
9 from vllm.utils import is_list_of
@@ -61,6 +62,15 @@ class Request:
62 self.num_encoder_inputs = len(self.mm_inputs)
63 self.has_encoder_inputs = self.num_encoder_inputs > 0
64
65 + # P/D: KV transfer parameters (raw and parsed).
66 + raw_params = (None if sampling_params.extra_args is None
67 + else sampling_params.extra_args.get(
68 + "kv_transfer_params", None))
69 + self.raw_kv_transfer_params: Optional[dict[str, Any]] = raw_params
70 + # Each connector parses the raw dictionary and sets this
71 + # attr the first time that the request is processed.
72 + self.kv_transfer_params: Optional[KVTransferParams] = None
73 +
74 # Sanity check
75 assert len(self.mm_inputs) == len(self.mm_positions)
76 if self.mm_hashes:
@@ -150,6 +160,7 @@ class RequestStatus(enum.IntEnum):
160 """Status of a request."""
161 WAITING = enum.auto()
162 WAITING_FOR_FSM = enum.auto()
163 + WAITING_FOR_REMOTE_KVS = enum.auto()
164 RUNNING = enum.auto()
165 PREEMPTED = enum.auto()
166 # Note: anything after PREEMPTED will be considered
vllm/v1/worker/gpu_model_runner.py
+68 -10
@@ -1,5 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 +import copy
4 import gc
5 import time
6 import weakref
@@ -17,8 +18,9 @@ from vllm.config import (CompilationLevel, VllmConfig,
18 get_layers_from_vllm_config)
19 from vllm.distributed.kv_transfer import (get_kv_transfer_group,
20 has_kv_transfer_group)
21 +from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1
22 from vllm.distributed.parallel_state import get_pp_group, graph_capture
21 -from vllm.forward_context import set_forward_context
23 +from vllm.forward_context import get_forward_context, set_forward_context
24 from vllm.logger import init_logger
25 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
26 from vllm.model_executor.model_loader import get_model
@@ -1065,15 +1067,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1067 scheduler_output: "SchedulerOutput",
1068 intermediate_tensors: Optional[IntermediateTensors] = None,
1069 ) -> Union[ModelRunnerOutput, IntermediateTensors]:
1068 - # Update KVConnector with the KVConnector metadata forward().
1069 - if has_kv_transfer_group():
1070 - get_kv_transfer_group().bind_connector_metadata(
1071 - scheduler_output.kv_connector_metadata)
1070
1071 self._update_states(scheduler_output)
1072 if not scheduler_output.total_num_scheduled_tokens:
1075 - # Return empty ModelRunnerOutput if there's no work to do.
1076 - return EMPTY_MODEL_RUNNER_OUTPUT
1073 + if not has_kv_transfer_group():
1074 + # Return empty ModelRunnerOutput if there's no work to do.
1075 + return EMPTY_MODEL_RUNNER_OUTPUT
1076 +
1077 + return self.kv_connector_no_forward(scheduler_output)
1078
1079 # Prepare the decoder inputs.
1080 attn_metadata, logits_indices, spec_decode_metadata = (
@@ -1150,17 +1151,23 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1151 with set_forward_context(attn_metadata,
1152 self.vllm_config,
1153 num_tokens=num_input_tokens):
1153 - output = self.model(
1154 + self.maybe_setup_kv_connector(scheduler_output)
1155 +
1156 + model_output = self.model(
1157 input_ids=input_ids,
1158 positions=positions,
1159 intermediate_tensors=intermediate_tensors,
1160 inputs_embeds=inputs_embeds,
1161 )
1162
1163 + self.maybe_wait_for_kv_save()
1164 + finished_sending, finished_recving = (
1165 + self.get_finished_kv_transfers(scheduler_output))
1166 +
1167 if self.use_aux_hidden_state_outputs:
1161 - hidden_states, aux_hidden_states = output
1168 + hidden_states, aux_hidden_states = model_output
1169 else:
1163 - hidden_states = output
1170 + hidden_states = model_output
1171
1172 if not get_pp_group().is_last_rank:
1173 # For mid-pipeline stages, return the hidden states.
@@ -1341,8 +1348,56 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1348 spec_token_ids=spec_token_ids,
1349 logprobs=logprobs_lists,
1350 prompt_logprobs_dict=prompt_logprobs_dict,
1351 + finished_sending=finished_sending,
1352 + finished_recving=finished_recving,
1353 )
1354
1355 + def kv_connector_no_forward(
1356 + self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
1357 + # KV send/recv even if no work to do.
1358 + with set_forward_context(None, self.vllm_config):
1359 + self.maybe_setup_kv_connector(scheduler_output)
1360 + finished_sending, finished_recving = (
1361 + self.get_finished_kv_transfers(scheduler_output))
1362 +
1363 + if not finished_sending and not finished_recving:
1364 + return EMPTY_MODEL_RUNNER_OUTPUT
1365 +
1366 + output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)
1367 + output.finished_sending = finished_sending
1368 + output.finished_recving = finished_recving
1369 + return output
1370 +
1371 + @staticmethod
1372 + def maybe_setup_kv_connector(scheduler_output: "SchedulerOutput"):
1373 + # Update KVConnector with the KVConnector metadata forward().
1374 + if has_kv_transfer_group():
1375 + kv_connector = get_kv_transfer_group()
1376 + assert isinstance(kv_connector, KVConnectorBase_V1)
1377 + assert scheduler_output.kv_connector_metadata is not None
1378 + kv_connector.bind_connector_metadata(
1379 + scheduler_output.kv_connector_metadata)
1380 +
1381 + # Background KV cache transfers happen here.
1382 + # These transfers are designed to be async and the requests
1383 + # involved may be disjoint from the running requests.
1384 + # Do this here to save a collective_rpc.
1385 + kv_connector.start_load_kv(get_forward_context())
1386 +
1387 + @staticmethod
1388 + def maybe_wait_for_kv_save() -> None:
1389 + if has_kv_transfer_group():
1390 + get_kv_transfer_group().wait_for_save()
1391 +
1392 + @staticmethod
1393 + def get_finished_kv_transfers(
1394 + scheduler_output: "SchedulerOutput",
1395 + ) -> tuple[Optional[set[str]], Optional[set[str]]]:
1396 + if has_kv_transfer_group():
1397 + return get_kv_transfer_group().get_finished(
1398 + scheduler_output.finished_req_ids)
1399 + return None, None
1400 +
1401 def generate_draft_token_ids(
1402 self,
1403 sampled_token_ids: list[list[int]],
@@ -1813,6 +1868,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1868 self.vllm_config.compilation_config.static_forward_context,
1869 self.kv_caches)
1870
1871 + if has_kv_transfer_group():
1872 + get_kv_transfer_group().register_kv_caches(kv_caches)
1873 +
1874 self.attn_metadata_builder = self.attn_backend.get_builder_cls()(
1875 weakref.proxy(self),
1876 kv_cache_config.kv_cache_groups[0].kv_cache_spec,