@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6066284914
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 14, 23:45 UTC

[P/D] Support CPU Transfer in NixlConnector (#18293)

Signed-off-by: Juncheng Gu <juncgu@gmail.com> Signed-off-by: Richard Liu <ricliu@google.com> Co-authored-by: Richard Liu <39319471+richardsliu@users.noreply.github.com> Co-authored-by: Richard Liu <ricliu@google.com>

Juncheng Gu committed Jul 24, 2025 at 09:58 UTC 6066284914c7984b37dc8bed69ac29bf21f8c49e
12 files changed +893 -110
requirements/tpu.txt
+1
@@ -10,6 +10,7 @@ jinja2>=3.1.6
10 ray[default]
11 ray[data]
12 setuptools==78.1.0
13 +nixl==0.3.0
14
15 # Install torch_xla
16 --pre
tests/v1/kv_connector/nixl_integration/run_tpu_disagg_accuracy_test.sh new
+162
@@ -0,0 +1,162 @@
1 +#!/bin/bash
2 +set -xe
3 +
4 +# Hosts / ports
5 +PREFILL_HOST=${PREFILL_HOST:-"localhost"}
6 +PREFILL_PORT=${PREFILL_PORT:-8100}
7 +PREFILL_NIXL_SIDE_PORT=${PREFILL_NIXL_SIDE_PORT:-5577}
8 +DECODE_HOST=${DECODE_HOST:-"localhost"}
9 +DECODE_PORT=${DECODE_PORT:-8200}
10 +PROXY_HOST=${PROXY_HOST:-"localhost"}
11 +PROXY_PORT=${PROXY_PORT:-8192}
12 +BASELINE_HOST=${BASELINE_HOST:-"localhost"}
13 +BASELINE_PORT=${BASELINE_PORT:-9290}
14 +
15 +
16 +# Model to run.
17 +MODEL_NAME=${MODEL_NAME:-"meta-llama/Llama-3.2-3B-Instruct"}
18 +MAX_MODEL_LEN=${MAX_MODEL_LEN:-1024}
19 +BLOCK_SIZE=${BLOCK_SIZE:-32}
20 +
21 +
22 +# execution env
23 +GIT_ROOT=$(git rev-parse --show-toplevel)
24 +EXP_ROOT="${GIT_ROOT}/tests/v1/kv_connector/nixl_integration"
25 +CONDA_PATH=${CONDA_PATH:-"/home/${USER}/anaconda3"}
26 +CONDA_ENV_NAME=${CONDA_ENV_NAME:-"nixl"}
27 +
28 +OUTPUT_FILE=${OUTPUT_FILE:-"${EXP_ROOT}/.tpu_accuracy_test_outputs.txt"}
29 +
30 +# Trap the SIGINT signal (triggered by Ctrl+C)
31 +trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT
32 +
33 +
34 +# Waits for vLLM server to start.
35 +wait_for_server() {
36 + local host=$1
37 + local port=$2
38 + timeout 1200 bash -c "
39 + until curl -s ${host}:${port}/v1/completions > /dev/null; do
40 + sleep 1
41 + done" && return 0 || return 1
42 +}
43 +
44 +# Cleanup function
45 +cleanup() {
46 + echo "Caught Ctrl+C, cleaning up..."
47 + # Cleanup commands
48 + pgrep python | xargs kill -9 || true
49 + # pkill -f python || true
50 + echo "Cleanup complete. Exiting."
51 +}
52 +
53 +launch_baseline() {
54 + BASELINE_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME};
55 + VLLM_LOGGING_LEVEL=DEBUG \
56 + VLLM_USE_V1=1 \
57 + PJRT_DEVICE=TPU \
58 + VLLM_WORKER_MULTIPROC_METHOD=spawn \
59 + VLLM_ENABLE_V1_MULTIPROCESSING=0 vllm serve $MODEL_NAME \
60 + --host ${BASELINE_HOST} \
61 + --port ${BASELINE_PORT} \
62 + --max-model-len ${MAX_MODEL_LEN}\
63 + --seed 42 \
64 + --block-size ${BLOCK_SIZE} \
65 + --gpu-memory-utilization 0.5 \
66 + --disable-log-requests \
67 + --enforce-eager"
68 + echo ${BASELINE_BASE_CMD}
69 + ssh -tt ${BASELINE_HOST} "${BASELINE_BASE_CMD}" &
70 +}
71 +
72 +launch_pd() {
73 + PREFILL_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME};
74 + UCX_TLS=tcp \
75 + VLLM_MULTIPROC_EXECUTE_MODEL_TIMEOUT_S=200 \
76 + VLLM_LOGGING_LEVEL=DEBUG \
77 + VLLM_USE_V1=1 \
78 + VLLM_NIXL_SIDE_CHANNEL_HOST=${PREFILL_HOST} \
79 + VLLM_NIXL_SIDE_CHANNEL_PORT=${PREFILL_NIXL_SIDE_PORT} \
80 + PJRT_DEVICE=TPU \
81 + VLLM_WORKER_MULTIPROC_METHOD=spawn \
82 + VLLM_ENABLE_V1_MULTIPROCESSING=0 vllm serve $MODEL_NAME \
83 + --host ${PREFILL_HOST} \
84 + --port ${PREFILL_PORT} \
85 + --max-model-len ${MAX_MODEL_LEN}\
86 + --seed 42 \
87 + --block-size ${BLOCK_SIZE} \
88 + --enforce-eager \
89 + --gpu-memory-utilization 0.5 \
90 + --disable-log-requests \
91 + --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"cpu\"}'"
92 +
93 +
94 + DECODE_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME};
95 + UCX_TLS=tcp \
96 + VLLM_MULTIPROC_EXECUTE_MODEL_TIMEOUT_S=200 \
97 + VLLM_LOGGING_LEVEL=DEBUG \
98 + VLLM_USE_V1=1 \
99 + PJRT_DEVICE=TPU \
100 + VLLM_WORKER_MULTIPROC_METHOD=spawn \
101 + VLLM_ENABLE_V1_MULTIPROCESSING=0 vllm serve $MODEL_NAME \
102 + --host ${DECODE_HOST} \
103 + --port ${DECODE_PORT} \
104 + --max-model-len ${MAX_MODEL_LEN}\
105 + --seed 42 \
106 + --block-size ${BLOCK_SIZE} \
107 + --enforce-eager \
108 + --gpu-memory-utilization 0.5 \
109 + --disable-log-requests \
110 + --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"cpu\"}'"
111 +
112 + echo ${PREFILL_BASE_CMD}
113 + echo ${DECODE_BASE_CMD}
114 + sleep 2
115 +
116 + # execute on hosts
117 + ssh -tt ${PREFILL_HOST} "${PREFILL_BASE_CMD}" &
118 + ssh -tt ${DECODE_HOST} "${DECODE_BASE_CMD}" &
119 + sleep 1
120 + wait_for_server ${PREFILL_HOST} ${PREFILL_PORT}
121 + sleep 1
122 + wait_for_server ${DECODE_HOST} ${DECODE_PORT}
123 + sleep 1
124 +}
125 +
126 +launch_pd_proxy(){
127 + PROXY_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME};
128 + python3 ${EXP_ROOT}/toy_proxy_server.py \
129 + --prefiller-host ${PREFILL_HOST} --prefiller-port ${PREFILL_PORT} \
130 + --decoder-host ${DECODE_HOST} --decoder-port ${DECODE_PORT} \
131 + --host=${PROXY_HOST} --port ${PROXY_PORT}"
132 + echo ${PROXY_BASE_CMD}
133 + ssh -tt ${PROXY_HOST} "${PROXY_BASE_CMD}" &
134 +}
135 +
136 +run_tests(){
137 + local service_url=$1
138 + local mode=$2
139 + python3 ${EXP_ROOT}/test_disagg_accuracy.py --service_url=${service_url} --model_name=${MODEL_NAME} --mode=${mode} --file_name=${OUTPUT_FILE}
140 +}
141 +
142 +
143 +# run non-disagg. baseline & save outputs
144 +launch_baseline
145 +sleep 2
146 +wait_for_server ${BASELINE_HOST} ${BASELINE_PORT}
147 +run_tests "http://${BASELINE_HOST}:${BASELINE_PORT}" "baseline"
148 +cleanup
149 +sleep 10
150 +
151 +
152 +# run disagg. & do exact-match with the outputs from baseline
153 +launch_pd
154 +launch_pd_proxy
155 +sleep 10
156 +run_tests "http://${PROXY_HOST}:${PROXY_PORT}" "disagg"
157 +echo "-----P/D success----"
158 +
159 +rm ${OUTPUT_FILE}
160 +cleanup
161 +
162 +exit 0
\ No newline at end of file
tests/v1/kv_connector/nixl_integration/run_tpu_edge_case_test.sh new
+128
@@ -0,0 +1,128 @@
1 +#!/bin/bash
2 +set -xe
3 +
4 +# Hosts / ports
5 +PREFILL_HOST=${PREFILL_HOST:-"localhost"}
6 +PREFILL_PORT=${PREFILL_PORT:-8100}
7 +PREFILL_NIXL_SIDE_PORT=${PREFILL_NIXL_SIDE_PORT:-5577}
8 +DECODE_HOST=${DECODE_HOST:-"localhost"}
9 +DECODE_PORT=${DECODE_PORT:-8200}
10 +PROXY_HOST=${PROXY_HOST:-"localhost"}
11 +PROXY_PORT=${PROXY_PORT:-8192}
12 +BASELINE_HOST=${BASELINE_HOST:-"localhost"}
13 +BASELINE_PORT=${BASELINE_PORT:-9290}
14 +
15 +
16 +# Model to run.
17 +MODEL_NAME=${MODEL_NAME:-"meta-llama/Llama-3.2-3B-Instruct"}
18 +MAX_MODEL_LEN=${MAX_MODEL_LEN:-1024}
19 +BLOCK_SIZE=${BLOCK_SIZE:-32}
20 +
21 +
22 +# execution env
23 +GIT_ROOT=$(git rev-parse --show-toplevel)
24 +EXP_ROOT="${GIT_ROOT}/tests/v1/kv_connector/nixl_integration"
25 +CONDA_PATH=${CONDA_PATH:-"/home/${USER}/anaconda3"}
26 +CONDA_ENV_NAME=${CONDA_ENV_NAME:-"nixl"}
27 +
28 +OUTPUT_FILE=${OUTPUT_FILE:-"${EXP_ROOT}/.tpu_accuracy_test_outputs.txt"}
29 +
30 +# Trap the SIGINT signal (triggered by Ctrl+C)
31 +trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT
32 +
33 +# Waits for vLLM server to start.
34 +wait_for_server() {
35 + local host=$1
36 + local port=$2
37 + timeout 1200 bash -c "
38 + until curl -s ${host}:${port}/v1/completions > /dev/null; do
39 + sleep 1
40 + done" && return 0 || return 1
41 +}
42 +
43 +# Cleanup function
44 +cleanup() {
45 + echo "Caught Ctrl+C, cleaning up..."
46 + # Cleanup commands
47 + pgrep python | xargs kill -9 || true
48 + # pkill -f python || true
49 + echo "Cleanup complete. Exiting."
50 +}
51 +
52 +
53 +launch_pd() {
54 + PREFILL_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME};
55 + UCX_TLS=tcp \
56 + VLLM_MULTIPROC_EXECUTE_MODEL_TIMEOUT_S=200 \
57 + VLLM_LOGGING_LEVEL=DEBUG \
58 + VLLM_USE_V1=1 \
59 + VLLM_NIXL_SIDE_CHANNEL_HOST=${PREFILL_HOST} \
60 + VLLM_NIXL_SIDE_CHANNEL_PORT=${PREFILL_NIXL_SIDE_PORT} \
61 + PJRT_DEVICE=TPU \
62 + VLLM_WORKER_MULTIPROC_METHOD=spawn \
63 + VLLM_ENABLE_V1_MULTIPROCESSING=0 vllm serve $MODEL_NAME \
64 + --host ${PREFILL_HOST} \
65 + --port ${PREFILL_PORT} \
66 + --max-model-len ${MAX_MODEL_LEN}\
67 + --seed 42 \
68 + --block-size ${BLOCK_SIZE} \
69 + --enforce-eager \
70 + --gpu-memory-utilization 0.5 \
71 + --disable-log-requests \
72 + --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"cpu\"}'"
73 +
74 +
75 + DECODE_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME};
76 + UCX_TLS=tcp \
77 + VLLM_MULTIPROC_EXECUTE_MODEL_TIMEOUT_S=200 \
78 + VLLM_LOGGING_LEVEL=DEBUG \
79 + VLLM_USE_V1=1 \
80 + PJRT_DEVICE=TPU \
81 + VLLM_WORKER_MULTIPROC_METHOD=spawn \
82 + VLLM_ENABLE_V1_MULTIPROCESSING=0 vllm serve $MODEL_NAME \
83 + --host ${DECODE_HOST} \
84 + --port ${DECODE_PORT} \
85 + --max-model-len ${MAX_MODEL_LEN}\
86 + --seed 42 \
87 + --block-size ${BLOCK_SIZE} \
88 + --enforce-eager \
89 + --gpu-memory-utilization 0.5 \
90 + --disable-log-requests \
91 + --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"cpu\"}'"
92 +
93 + echo ${PREFILL_BASE_CMD}
94 + echo ${DECODE_BASE_CMD}
95 + sleep 2
96 +
97 + # execute on hosts
98 + ssh -tt ${PREFILL_HOST} "${PREFILL_BASE_CMD}" &
99 + ssh -tt ${DECODE_HOST} "${DECODE_BASE_CMD}" &
100 + sleep 1
101 + wait_for_server ${PREFILL_HOST} ${PREFILL_PORT}
102 + sleep 1
103 + wait_for_server ${DECODE_HOST} ${DECODE_PORT}
104 + sleep 1
105 +}
106 +
107 +launch_pd_proxy(){
108 + PROXY_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME};
109 + python3 ${EXP_ROOT}/toy_proxy_server.py \
110 + --prefiller-host ${PREFILL_HOST} --prefiller-port ${PREFILL_PORT} \
111 + --decoder-host ${DECODE_HOST} --decoder-port ${DECODE_PORT} \
112 + --host=${PROXY_HOST} --port ${PROXY_PORT}"
113 + echo ${PROXY_BASE_CMD}
114 + ssh -tt ${PROXY_HOST} "${PROXY_BASE_CMD}" &
115 +}
116 +
117 +
118 +# run disagg. & do exact-match with the outputs from baseline
119 +launch_pd
120 +launch_pd_proxy
121 +sleep 10
122 +
123 +PREFILL_HOST=${PREFILL_HOST} \
124 +PREFILL_PORT=${PREFILL_PORT} \
125 +DECODE_HOST=${DECODE_HOST} \
126 +DECODE_PORT=${DECODE_PORT} \
127 +PROXY_HOST=${PROXY_HOST} \
128 +PROXY_PORT=${PROXY_PORT} python -m pytest -s -v ${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/test_edge_cases.py
\ No newline at end of file
tests/v1/kv_connector/nixl_integration/test_disagg_accuracy.py new
+162
@@ -0,0 +1,162 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import argparse
4 +import json
5 +import os
6 +import time
7 +
8 +import openai
9 +import requests
10 +
11 +MAX_OUTPUT_LEN = 30
12 +
13 +SAMPLE_PROMPTS = (
14 + "Red Hat is the best company in the world to work for because it works on "
15 + "open source software, which means that all the contributions are "
16 + "delivered to the community. As a result, when working on projects like "
17 + "vLLM we are able to meet many amazing people from various organizations "
18 + "like AMD, Google, NVIDIA, ",
19 + "We hold these truths to be self-evident, that all men are created equal, "
20 + "that they are endowed by their Creator with certain unalienable Rights, "
21 + "that among these are Life, Liberty and the pursuit of Happiness.--That "
22 + "to secure these rights, Governments are instituted among Men, deriving "
23 + "their just powers from the consent of the governed, ",
24 +)
25 +
26 +
27 +def check_vllm_server(url: str, timeout=5, retries=3) -> bool:
28 + """
29 + Checks if the vLLM server is ready by sending a GET request to the
30 + /health endpoint.
31 +
32 + Args:
33 + url (str): The base URL of the vLLM server.
34 + timeout (int): Timeout in seconds for the request.
35 + retries (int): Number of retries if the server is not ready.
36 +
37 + Returns:
38 + bool: True if the server is ready, False otherwise.
39 + """
40 + for attempt in range(retries):
41 + try:
42 + response = requests.get(url, timeout=timeout)
43 + if response.status_code == 200:
44 + return True
45 + else:
46 + print(f"Attempt {attempt + 1}: Server returned status code "
47 + "{response.status_code}")
48 + except requests.exceptions.RequestException as e:
49 + print(f"Attempt {attempt + 1}: Error connecting to server: {e}")
50 + time.sleep(1) # Wait before retrying
51 + return False
52 +
53 +
54 +def run_simple_prompt(base_url: str, model_name: str,
55 + input_prompt: str) -> str:
56 + client = openai.OpenAI(api_key="EMPTY", base_url=base_url)
57 + completion = client.completions.create(model=model_name,
58 + prompt=input_prompt,
59 + max_tokens=MAX_OUTPUT_LEN,
60 + temperature=0.0,
61 + seed=42)
62 +
63 + # print("-" * 50)
64 + # print(f"Completion results for {model_name}:")
65 + # print(completion)
66 + # print("-" * 50)
67 + return completion.choices[0].text
68 +
69 +
70 +def main():
71 + """
72 + This script demonstrates how to accept two optional string arguments
73 + ("service_url" and "file_name") from the command line, each with a
74 + default value of an empty string, using the argparse module.
75 + """
76 + parser = argparse.ArgumentParser(description="vLLM client script")
77 +
78 + parser.add_argument(
79 + "--service_url", # Name of the first argument
80 + type=str,
81 + required=True,
82 + help="The vLLM service URL.")
83 +
84 + parser.add_argument(
85 + "--model_name", # Name of the first argument
86 + type=str,
87 + required=True,
88 + help="model_name",
89 + )
90 +
91 + parser.add_argument(
92 + "--mode", # Name of the second argument
93 + type=str,
94 + default="baseline",
95 + help="mode: baseline==non-disagg, or disagg",
96 + )
97 +
98 + parser.add_argument(
99 + "--file_name", # Name of the second argument
100 + type=str,
101 + default=".vllm_output.txt",
102 + help="the file that saves the output tokens ",
103 + )
104 +
105 + args = parser.parse_args()
106 +
107 + for arg in vars(args):
108 + print(f"{arg}: {getattr(args, arg)}")
109 +
110 + if args.mode == "baseline":
111 + # non-disagg
112 + health_check_url = f"{args.service_url}/health"
113 + else:
114 + # disagg proxy
115 + health_check_url = f"{args.service_url}/healthcheck"
116 + if not os.path.exists(args.file_name):
117 + raise ValueError(
118 + f"In disagg mode, the output file {args.file_name} from "
119 + "non-disagg. baseline does not exist.")
120 +
121 + service_url = f"{args.service_url}/v1"
122 +
123 + if not check_vllm_server(health_check_url):
124 + raise RuntimeError(
125 + f"vllm server: {args.service_url} is not ready yet!")
126 +
127 + output_strs = dict()
128 + for prompt in SAMPLE_PROMPTS:
129 + output_str = run_simple_prompt(base_url=service_url,
130 + model_name=args.model_name,
131 + input_prompt=prompt)
132 + print(f"Prompt: {prompt}, output: {output_str}")
133 + output_strs[prompt] = output_str
134 +
135 + if args.mode == "baseline":
136 + # baseline: save outputs
137 + try:
138 + with open(args.file_name, 'w') as json_file:
139 + json.dump(output_strs, json_file, indent=4)
140 + except OSError as e:
141 + print(f"Error writing to file: {e}")
142 + raise
143 + else:
144 + # disagg. verify outputs
145 + baseline_outputs = None
146 + try:
147 + with open(args.file_name) as json_file:
148 + baseline_outputs = json.load(json_file)
149 + except OSError as e:
150 + print(f"Error writing to file: {e}")
151 + raise
152 + assert isinstance(baseline_outputs, dict)
153 + assert len(baseline_outputs) == len(output_strs)
154 + for prompt, output in baseline_outputs.items():
155 + assert prompt in output_strs, f"{prompt} not included"
156 + assert output == output_strs[prompt], (
157 + f"baseline_output: {output} != PD output: {output_strs[prompt]}"
158 + )
159 +
160 +
161 +if __name__ == "__main__":
162 + main()
tests/v1/kv_connector/nixl_integration/test_edge_cases.py
+6 -3
@@ -4,8 +4,11 @@ import os
4
5 import openai
6
7 +PREFILL_HOST = os.getenv("PREFILL_HOST", "localhost")
8 PREFILL_PORT = os.getenv("PREFILL_PORT", None)
9 +DECODE_HOST = os.getenv("DECODE_HOST", "localhost")
10 DECODE_PORT = os.getenv("DECODE_PORT", None)
11 +PROXY_HOST = os.getenv("PROXY_HOST", "localhost")
12 PROXY_PORT = os.getenv("PROXY_PORT", None)
13
14 if PREFILL_PORT is None or DECODE_PORT is None or PROXY_PORT is None:
@@ -21,15 +24,15 @@ def test_edge_cases():
24 # Set the OpenAI API key and base URL
25 decode_client = openai.OpenAI(
26 api_key="MY_KEY",
24 - base_url=f"http://localhost:{DECODE_PORT}/v1",
27 + base_url=f"http://{DECODE_HOST}:{DECODE_PORT}/v1",
28 )
29 prefill_client = openai.OpenAI(
30 api_key="MY_KEY",
28 - base_url=f"http://localhost:{PREFILL_PORT}/v1",
31 + base_url=f"http://{PREFILL_HOST}:{PREFILL_PORT}/v1",
32 )
33 proxy_client = openai.OpenAI(
34 api_key="MY_KEY",
32 - base_url=f"http://localhost:{PROXY_PORT}/v1",
35 + base_url=f"http://{PROXY_HOST}:{PROXY_PORT}/v1",
36 )
37
38 # Get the list of models
tests/v1/kv_connector/nixl_integration/toy_proxy_server.py
+3 -3
@@ -3,6 +3,7 @@
3
4 import argparse
5 import itertools
6 +import logging
7 import os
8 import uuid
9 from contextlib import asynccontextmanager
@@ -11,9 +12,8 @@ import httpx
12 from fastapi import FastAPI, Request
13 from fastapi.responses import StreamingResponse
14
14 -from vllm.logger import init_logger
15 -
16 -logger = init_logger(__name__)
15 +logger = logging.getLogger(__name__)
16 +logger.setLevel(logging.DEBUG)
17
18
19 @asynccontextmanager
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+14 -1
@@ -32,7 +32,7 @@ The class provides the following primitives:
32
33 import enum
34 from abc import ABC, abstractmethod
35 -from typing import TYPE_CHECKING, Any, Optional
35 +from typing import TYPE_CHECKING, Any, Callable, Literal, Optional
36
37 import torch
38
@@ -46,6 +46,12 @@ if TYPE_CHECKING:
46 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
47 from vllm.v1.request import Request
48
49 +# s_tensor_list, d_tensor_list, s_indices, d_indices, direction
50 +CopyBlocksOp = Callable[[
51 + dict[str, torch.Tensor], dict[
52 + str, torch.Tensor], list[int], list[int], Literal["h2d", "d2h"]
53 +], None]
54 +
55 logger = init_logger(__name__)
56
57
@@ -127,6 +133,13 @@ class KVConnectorBase_V1(ABC):
133 """
134 return
135
136 + def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
137 + """
138 + Set the xPU-specific ops for copying KV between host and device.
139 + Needed when host buffer is used for kv transfer (e.g., in NixlConnector)
140 + """
141 + return
142 +
143 @abstractmethod
144 def start_load_kv(self, forward_context: "ForwardContext",
145 **kwargs) -> None:
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+229 -43
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import contextlib
4 +import logging
5 import math
6 import queue
7 import threading
@@ -20,14 +21,14 @@ from vllm import envs
21 from vllm.attention.selector import backend_name_to_enum, get_attn_backend
22 from vllm.config import VllmConfig
23 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
23 - KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole)
24 + CopyBlocksOp, KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole)
25 from vllm.distributed.parallel_state import (
26 get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size,
27 get_tp_group)
28 from vllm.distributed.utils import divide
29 from vllm.forward_context import ForwardContext
30 from vllm.logger import init_logger
30 -from vllm.platforms import _Backend
31 +from vllm.platforms import _Backend, current_platform
32 from vllm.utils import make_zmq_path, make_zmq_socket, round_down
33 from vllm.v1.core.sched.output import SchedulerOutput
34 from vllm.v1.request import RequestStatus
@@ -40,6 +41,7 @@ if TYPE_CHECKING:
41 Transfer = tuple[int, float] # (xfer_handle, start_time)
42 EngineId = str
43 ReqId = str
44 +
45 GET_META_MSG = b"get_meta_msg"
46
47 logger = init_logger(__name__)
@@ -52,6 +54,13 @@ except ImportError:
54 logger.warning("NIXL is not available")
55 NixlWrapper = None
56
57 +# Supported xPUs and types of kv transfer buffer.
58 +# {xPU: tuple of supported kv buffer types}
59 +_NIXL_SUPPORTED_XPUS = {
60 + "cuda": ("cuda", ),
61 + "tpu": ("cpu", ),
62 +}
63 +
64
65 class NixlAgentMetadata(
66 msgspec.Struct,
@@ -80,6 +89,7 @@ class NixlConnectorMetadata(KVConnectorMetadata):
89
90 def __init__(self):
91 self.reqs_to_recv: dict[ReqId, ReqMeta] = {}
92 + self.reqs_to_save: dict[ReqId, ReqMeta] = {}
93 self.reqs_to_send: dict[ReqId, float] = {}
94
95 def add_new_req(
@@ -87,8 +97,12 @@ class NixlConnectorMetadata(KVConnectorMetadata):
97 request_id: ReqId,
98 local_block_ids: list[int],
99 kv_transfer_params: dict[str, Any],
100 + load_remote_cache: bool = True,
101 + save_to_host: bool = False,
102 ):
91 - self.reqs_to_recv[request_id] = ReqMeta(
103 + # save and load are mutually exclusive
104 + assert load_remote_cache ^ save_to_host
105 + _req = ReqMeta(
106 local_block_ids=local_block_ids,
107 remote_block_ids=kv_transfer_params["remote_block_ids"],
108 remote_engine_id=kv_transfer_params["remote_engine_id"],
@@ -97,6 +111,10 @@ class NixlConnectorMetadata(KVConnectorMetadata):
111 # P workers don't need to receive tp_size from proxy here.
112 tp_size=kv_transfer_params.get("tp_size", 1),
113 )
114 + if save_to_host:
115 + self.reqs_to_save[request_id] = _req
116 + if load_remote_cache:
117 + self.reqs_to_recv[request_id] = _req
118
119
120 class NixlConnector(KVConnectorBase_V1):
@@ -155,6 +173,10 @@ class NixlConnector(KVConnectorBase_V1):
173 assert self.connector_worker is not None
174 self.connector_worker.register_kv_caches(kv_caches)
175
176 + def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
177 + assert self.connector_worker is not None
178 + self.connector_worker.set_host_xfer_buffer_ops(copy_operation)
179 +
180 def get_finished(self,
181 finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
182 """Get the finished recving and sending requests."""
@@ -177,8 +199,11 @@ class NixlConnector(KVConnectorBase_V1):
199 pass
200
201 def wait_for_save(self):
180 - """NixlConnector does not save explicitly."""
181 - pass
202 + assert self.connector_worker is not None
203 + assert isinstance(self._connector_metadata, NixlConnectorMetadata)
204 + if self.connector_worker.use_host_buffer and \
205 + self.connector_worker.copy_blocks:
206 + self.connector_worker.save_kv_to_host(self._connector_metadata)
207
208
209 class NixlConnectorScheduler:
@@ -193,12 +218,15 @@ class NixlConnectorScheduler:
218 envs.VLLM_NIXL_SIDE_CHANNEL_PORT +
219 vllm_config.parallel_config.data_parallel_rank *
220 vllm_config.parallel_config.tensor_parallel_size)
221 + self.use_host_buffer = \
222 + vllm_config.kv_transfer_config.kv_buffer_device == "cpu"
223 logger.info("Initializing NIXL Scheduler %s", engine_id)
224
225 # Requests that need to start recv/send.
226 # New requests are added by update_state_after_alloc in
227 # the scheduler. Used to make metadata passed to Worker.
228 self._reqs_need_recv: dict[ReqId, tuple[Request, list[int]]] = {}
229 + self._reqs_need_save: dict[ReqId, tuple[Request, list[int]]] = {}
230 # Reqs to send and their expiration time
231 self._reqs_need_send: dict[ReqId, float] = {}
232
@@ -248,7 +276,25 @@ class NixlConnectorScheduler:
276 "num_external_tokens=%s, kv_transfer_params=%s",
277 num_external_tokens, params)
278
251 - if params is not None and params.get("do_remote_prefill"):
279 + if not params:
280 + return
281 + if self.use_host_buffer and params.get("do_remote_decode"):
282 + # NOTE: when accelerator is not directly supported by Nixl,
283 + # prefilled blocks need to be saved to host memory before transfer.
284 +
285 + # figure out full computed blocks to save
286 + block_ids = blocks.get_block_ids()[0]
287 + all_full = request.num_tokens % self.block_size == 0
288 + full_block_ids = (block_ids if all_full else block_ids[:-1])
289 + # TODO: skip the blocks that are already in the host xfer buffer.
290 + # Currently, the host xfer buffer block is 1-to-1 mapped to device
291 + # kv blocks, so host blocks won't be flushed as long as its device
292 + # block is not overwritten; and it will be safe to skip saving them
293 + # to host xfer buffer.
294 + if full_block_ids:
295 + self._reqs_need_save[request.request_id] = \
296 + (request, full_block_ids)
297 + elif params.get("do_remote_prefill"):
298 if params.get("remote_block_ids"):
299 if all(p in params for p in ("remote_engine_id", "remote_host",
300 "remote_port")):
@@ -260,6 +306,7 @@ class NixlConnectorScheduler:
306 # Get unhashed blocks to pull from remote.
307 self._reqs_need_recv[request.request_id] = (
308 request, local_block_ids)
309 +
310 else:
311 logger.warning(
312 "Got invalid KVTransferParams: %s. This "
@@ -284,10 +331,21 @@ class NixlConnectorScheduler:
331 kv_transfer_params=req.kv_transfer_params,
332 )
333
287 - # Clear the list once workers start the transfers
288 - self._reqs_need_recv.clear()
334 + for req_id, (req, block_ids) in self._reqs_need_save.items():
335 + assert req.kv_transfer_params is not None
336 + meta.add_new_req(
337 + request_id=req_id,
338 + local_block_ids=block_ids,
339 + kv_transfer_params=req.kv_transfer_params,
340 + load_remote_cache=False,
341 + save_to_host=True,
342 + )
343
344 meta.reqs_to_send = self._reqs_need_send
345 +
346 + # Clear the list once workers start the transfers
347 + self._reqs_need_recv.clear()
348 + self._reqs_need_save.clear()
349 self._reqs_need_send = {}
350
351 return meta
@@ -379,9 +437,36 @@ class NixlConnectorWorker:
437 self.tp_rank = get_tensor_model_parallel_rank()
438 self.world_size = get_tensor_model_parallel_world_size()
439 self.tp_group = get_tp_group()
440 + self.num_blocks = 0
441
442 # KV Caches and nixl tracking data.
384 - self.kv_caches: dict[str, torch.Tensor] = {}
443 + self.device_type = current_platform.device_type
444 + self.kv_buffer_device: str = \
445 + vllm_config.kv_transfer_config.kv_buffer_device
446 + if self.device_type not in _NIXL_SUPPORTED_XPUS:
447 + raise RuntimeError(f"{self.device_type} is not supported.")
448 + elif self.kv_buffer_device not in _NIXL_SUPPORTED_XPUS[
449 + self.device_type]:
450 + raise RuntimeError(
451 + f"{self.device_type} with {self.kv_buffer_device} kv_buffer "
452 + "is not supported.")
453 + self.device_kv_caches: dict[str, torch.Tensor] = {}
454 +
455 + # cpu kv buffer for xfer
456 + # used when xPU memory can not be registered under nixl
457 + self.host_xfer_buffers: dict[str, torch.Tensor] = {}
458 + self.use_host_buffer = self.kv_buffer_device == "cpu"
459 + if self.kv_buffer_device == "cuda":
460 + self.nixl_memory_type = "VRAM"
461 + elif self.kv_buffer_device == "cpu":
462 + self.nixl_memory_type = "DRAM"
463 + else:
464 + raise RuntimeError(
465 + f"{self.device_type} with {self.kv_buffer_device} kv_buffer "
466 + "is not supported.")
467 +
468 + # Note: host xfer buffer ops when use_host_buffer is True
469 + self.copy_blocks: Optional[CopyBlocksOp] = None
470
471 # Map of engine_id -> kv_caches_base_addr. For TP case, each local
472 # rank will still only pull from a single remote TP worker.
@@ -404,6 +489,7 @@ class NixlConnectorWorker:
489
490 # In progress transfers.
491 # [req_id -> list[handle]]
492 + self._recving_metadata: dict[ReqId, ReqMeta] = {}
493 self._recving_transfers = defaultdict[ReqId, list[Transfer]](list)
494 # Track the expiration time of requests that are waiting to be sent.
495 self._reqs_to_send: dict[ReqId, float] = {}
@@ -440,6 +526,7 @@ class NixlConnectorWorker:
526 self.backend_name = backend.get_name()
527 attn_backend = backend_name_to_enum(self.backend_name)
528 self._use_flashinfer = attn_backend == _Backend.FLASHINFER_VLLM_V1
529 + self._use_pallas_v1 = attn_backend == _Backend.PALLAS_VLLM_V1
530 logger.debug("Detected attention backend %s", self.backend_name)
531
532 self._tp_size: dict[EngineId, int] = {self.engine_id: self.world_size}
@@ -529,6 +616,31 @@ class NixlConnectorWorker:
616 # Remote rank -> agent name.
617 return {p_remote_rank: remote_agent_name}
618
619 + def initialize_host_xfer_buffer(
620 + self, kv_caches: dict[str, torch.Tensor]) -> None:
621 + """
622 + Initialize transfer buffer in CPU mem for accelerators
623 + NOT directly supported by NIXL (e.g., tpu)
624 + """
625 + xfer_buffers: dict[str, torch.Tensor] = {}
626 + try:
627 + for layer_name, kv_cache in kv_caches.items():
628 + kv_shape = kv_cache.shape
629 + kv_dtype = kv_cache.dtype
630 + xfer_buffers[layer_name] = torch.empty(kv_shape,
631 + dtype=kv_dtype,
632 + device="cpu")
633 + except MemoryError as e:
634 + logger.error("NIXLConnectorWorker gets %s.", e)
635 + raise
636 +
637 + self.host_xfer_buffers = xfer_buffers
638 +
639 + def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
640 + """Assign copy (d2h, h2d) operations when host buffer is used."""
641 + assert self.use_host_buffer
642 + self.copy_blocks = copy_operation
643 +
644 def _background_nixl_handshake(self, req_id: str,
645 remote_engine_id: EngineId, meta: ReqMeta):
646 # Do NIXL handshake in background and add to _ready_requests when done.
@@ -562,47 +674,76 @@ class NixlConnectorWorker:
674 _, first_kv_cache = next(iter(kv_caches.items()))
675 kv_elem_size = first_kv_cache.element_size()
676
677 + if self.use_host_buffer:
678 + self.initialize_host_xfer_buffer(kv_caches=kv_caches)
679 + assert len(self.host_xfer_buffers) == len(kv_caches), (
680 + f"host_buffer: {len(self.host_xfer_buffers)}, "
681 + f"kv_caches: {len(kv_caches)}")
682 + xfer_buffers = self.host_xfer_buffers
683 + else:
684 + xfer_buffers = kv_caches
685 + assert not self.host_xfer_buffers, (
686 + "host_xfer_buffer should not be initialized when "
687 + f"kv_buffer_device is {self.kv_buffer_device}")
688 +
689 # TODO(tms): Find a more robust way to detect and handle MLA
690 # NOTE (NickLucche) To move blocks efficiently with NIXL, the expected
691 # KV memory layout is HND, as opposed to the default NHD. Note that it
692 # will only affects the strides. For MLA instead, we make require no
693 # such thing and resort to the standard layout.
694 use_mla = len(first_kv_cache.shape) == 3
571 - assert use_mla == self.use_mla
572 -
573 - # TODO (NickLucche) not compatible with hybrid allocator. Enforce check
574 - # once it goes live, as a single kv layout is expected for xfers.
575 - if use_mla:
576 - # MLA case.
695 + if self.device_type == "tpu":
696 + assert not use_mla, f"{self.kv_buffer_device} does not support MLA."
697 + assert self._use_pallas_v1, f"attn backend: {self.backend_name}"
698 + # tpu (v1) kv shape per layer:
699 + # (num_blocks, block_size, num_kv_heads * 2, head_size)
700 self.num_blocks = first_kv_cache.shape[0]
578 - block_rank = 2 # [block_size, latent_dim]
701 + block_rank = 3 # [block_size, kv_heads, head_dim]
702 block_shape = first_kv_cache.shape[-block_rank:]
580 - block_size, kv_latent_dim = block_shape
581 - self.slot_size_bytes = kv_elem_size * kv_latent_dim
582 - else:
583 - # [2 (k and v), num_blocks, ...]
584 - if self._use_flashinfer:
585 - # FlashInfer swaps 2<->num_blocks dimensions.
703 + block_size, n_kv_heads_x_2, head_dim = block_shape
704 + self.slot_size_bytes = kv_elem_size * n_kv_heads_x_2 * head_dim
705 + elif self.device_type == "cuda":
706 + assert use_mla == self.use_mla
707 + # TODO (NickLucche) not compatible with hybrid allocator.
708 + # Enforce check once it goes live, as a single kv layout
709 + # is expected for xfers.
710 + if use_mla:
711 + # MLA case.
712 self.num_blocks = first_kv_cache.shape[0]
587 - block_rank = 4 # [2, block_size, kv_heads, head_dim]
713 + block_rank = 2 # [block_size, latent_dim]
714 + block_shape = first_kv_cache.shape[-block_rank:]
715 + block_size, kv_latent_dim = block_shape
716 + self.slot_size_bytes = kv_elem_size * kv_latent_dim
717 else:
589 - self.num_blocks = first_kv_cache.shape[1]
590 - block_rank = 3 # [block_size, kv_heads, head_dim]
591 - block_shape = first_kv_cache.shape[-block_rank:]
592 - block_size, n_kv_heads, head_dim = block_shape[-3:]
593 - # head size in bytes.
594 - self.slot_size_bytes = kv_elem_size * n_kv_heads * head_dim
595 - assert block_size == self.block_size
718 + # [2 (k and v), num_blocks, ...]
719 + if self._use_flashinfer:
720 + # FlashInfer swaps 2<->num_blocks dimensions.
721 + self.num_blocks = first_kv_cache.shape[0]
722 + block_rank = 4 # [2, block_size, kv_heads, head_dim]
723 + else:
724 + self.num_blocks = first_kv_cache.shape[1]
725 + block_rank = 3 # [block_size, kv_heads, head_dim]
726 + block_shape = first_kv_cache.shape[-block_rank:]
727 + block_size, n_kv_heads, head_dim = block_shape[-3:]
728 + # head size in bytes.
729 + self.slot_size_bytes = kv_elem_size * n_kv_heads * head_dim
730 + assert block_size == self.block_size
731 + else:
732 + raise RuntimeError(
733 + f"{self.device_type} ({self.backend_name}) is not supported.")
734 +
735 # TODO(tms): self.block_len needs to be per-layer for sliding window,
736 # hybrid attn, etc
737 # block size in bytes
738 self.block_len = kv_elem_size * math.prod(block_shape)
739 logger.info(
601 - "Registering KV_Caches: use_mla: %s, num_blocks: %s, "
602 - "block_shape: %s, per_layer_kv_cache_shape: %s", use_mla,
603 - self.num_blocks, block_shape, first_kv_cache.shape)
740 + "Registering KV_Caches. use_mla: %s, kv_buffer_device: %s, "
741 + "use_host_buffer: %s, num_blocks: %s, block_shape: %s, "
742 + "per_layer_kv_cache_shape: %s", use_mla, self.kv_buffer_device,
743 + self.use_host_buffer, self.num_blocks, block_shape,
744 + first_kv_cache.shape)
745 self.dst_num_blocks[self.engine_id] = self.num_blocks
605 - self.kv_caches = kv_caches
746 + self.device_kv_caches = kv_caches
747 kv_caches_base_addr = []
748 caches_data = []
749
@@ -614,19 +755,21 @@ class NixlConnectorWorker:
755 # (roughly 8KB vs 5KB).
756 # Conversely for FlashInfer, K and V are transferred in the same tensor
757 # to better exploit the memory layout (ie num_blocks is the first dim).
617 - for cache_or_caches in kv_caches.values():
758 + for cache_or_caches in xfer_buffers.values():
759 # Normalize to always be a list of caches
619 - cache_list = [cache_or_caches] if use_mla or self._use_flashinfer \
620 - else cache_or_caches
760 + cache_list = [cache_or_caches] if use_mla \
761 + or self._use_pallas_v1 or self._use_flashinfer \
762 + else cache_or_caches
763 for cache in cache_list:
764 base_addr = cache.data_ptr()
765 region_len = self.num_blocks * self.block_len
624 - caches_data.append(
625 - (base_addr, region_len, cache.device.index, ""))
766 + # NOTE: use tp_rank for device_id since multi-node TP
767 + # is rarely used.
768 + caches_data.append((base_addr, region_len, self.tp_rank, ""))
769 kv_caches_base_addr.append(base_addr)
770 self.kv_caches_base_addr[self.engine_id] = kv_caches_base_addr
771 self.num_regions = len(caches_data)
629 - self.num_layers = len(self.kv_caches.keys())
772 + self.num_layers = len(xfer_buffers.keys())
773
774 # TODO(mgoin): remove this once we have hybrid memory allocator
775 # Optimization for models with local attention (Llama 4)
@@ -648,7 +791,8 @@ class NixlConnectorWorker:
791 self.block_window_per_layer)
792 assert len(self.block_window_per_layer) == self.num_layers
793
651 - descs = self.nixl_wrapper.get_reg_descs(caches_data, "VRAM")
794 + descs = self.nixl_wrapper.get_reg_descs(caches_data,
795 + self.nixl_memory_type)
796 logger.debug("Registering descs: %s", caches_data)
797 self.nixl_wrapper.register_memory(descs)
798 logger.debug("Done registering descs")
@@ -666,11 +810,13 @@ class NixlConnectorWorker:
810 block_offset = block_id * self.block_len
811 addr = base_addr + block_offset
812 # (addr, len, device id)
813 + # TODO: does device_id matter to DRAM?
814 blocks_data.append((addr, self.block_len, self.tp_rank))
815 logger.debug("Created %s blocks for src engine %s and rank %s",
816 len(blocks_data), self.engine_id, self.tp_rank)
817
673 - descs = self.nixl_wrapper.get_xfer_descs(blocks_data, "VRAM")
818 + descs = self.nixl_wrapper.get_xfer_descs(blocks_data,
819 + self.nixl_memory_type)
820 # NIXL_INIT_AGENT to be used for preparations of local descs.
821 self.src_xfer_side_handle = self.nixl_wrapper.prep_xfer_dlist(
822 "NIXL_INIT_AGENT", descs)
@@ -755,6 +901,8 @@ class NixlConnectorWorker:
901 tp_ratio = divide(self._tp_size[self.engine_id],
902 self._tp_size[engine_id])
903 assert tp_ratio > 0, "Decode TP cannot be smaller than prefill TP"
904 + assert not self._use_pallas_v1 or tp_ratio == 1, \
905 + "TPU (pallas_v1) DOES NOT support heterogeneous TP yet."
906
907 # Handle tp_size>num_kv_heads: replicate KV cache.
908 total_num_kv_heads = self.model_config.get_total_num_kv_heads()
@@ -813,13 +961,43 @@ class NixlConnectorWorker:
961 self.tp_rank)
962
963 # Register with NIXL.
816 - descs = self.nixl_wrapper.get_xfer_descs(blocks_data, "VRAM")
964 + descs = self.nixl_wrapper.get_xfer_descs(blocks_data,
965 + self.nixl_memory_type)
966 self.dst_xfer_side_handles[
967 engine_id] = self.nixl_wrapper.prep_xfer_dlist(
968 remote_agent_name, descs)
969
970 return remote_agent_name
971
972 + def sync_recved_kv_to_device(self, req_id: str, meta: ReqMeta):
973 + """copy recved kv from host buffer to device."""
974 + assert self.use_host_buffer
975 + assert self.copy_blocks is not None
976 +
977 + local_block_ids = meta.local_block_ids
978 + self.copy_blocks(self.host_xfer_buffers, self.device_kv_caches,
979 + local_block_ids, local_block_ids, "h2d")
980 + if logger.isEnabledFor(logging.DEBUG):
981 + logger.debug(
982 + "synced recved kv of request[%s] to device kv buffer,"
983 + "local_block_ids: %s. ", req_id,
984 + ",".join(map(str, meta.local_block_ids)))
985 +
986 + def save_kv_to_host(self, metadata: NixlConnectorMetadata):
987 + """copy kv from device to host buffer."""
988 + assert self.use_host_buffer
989 + assert self.copy_blocks is not None
990 +
991 + for req_id, meta in metadata.reqs_to_save.items():
992 + if logger.isEnabledFor(logging.DEBUG):
993 + logger.debug(
994 + "save_load_kv for request[%s] to host xfer buffer."
995 + "local_block_ids: %s. ", req_id,
996 + ",".join(map(str, meta.local_block_ids)))
997 + # blocking
998 + self.copy_blocks(self.device_kv_caches, self.host_xfer_buffers,
999 + meta.local_block_ids, meta.local_block_ids, "d2h")
1000 +
1001 def get_finished(self) -> tuple[set[str], set[str]]:
1002 """
1003 Get requests that are done sending or recving on this specific worker.
@@ -834,6 +1012,12 @@ class NixlConnectorWorker:
1012 "and %s requests done recving", self.tp_rank,
1013 len(done_sending), len(done_recving))
1014
1015 + if self.use_host_buffer:
1016 + for req_id in done_recving:
1017 + meta = self._recving_metadata.pop(req_id)
1018 + assert meta, f"{req_id} not found in recving_metadata list"
1019 + self.sync_recved_kv_to_device(req_id, meta)
1020 +
1021 # Handle timeout to avoid stranding blocks on remote.
1022 now = time.perf_counter()
1023 while self._reqs_to_send:
@@ -904,6 +1088,8 @@ class NixlConnectorWorker:
1088 "Num local_block_ids: %s. Num remote_block_ids: %s. ", req_id,
1089 remote_engine_id, len(meta.local_block_ids),
1090 len(meta.remote_block_ids))
1091 + if self.use_host_buffer:
1092 + self._recving_metadata[req_id] = meta
1093 if remote_engine_id not in self._remote_agents:
1094 # Initiate handshake with remote engine to exchange metadata.
1095 with self._handshake_lock:
vllm/v1/worker/gpu_model_runner.py
+6 -52
@@ -1,7 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -import copy
4 import gc
5 import time
6 from contextlib import contextmanager
@@ -23,12 +22,10 @@ from vllm.config import (CompilationLevel, VllmConfig,
22 from vllm.distributed.eplb.eplb_state import EplbState
23 from vllm.distributed.kv_transfer import (get_kv_transfer_group,
24 has_kv_transfer_group)
26 -from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1
25 from vllm.distributed.parallel_state import (
26 get_pp_group, get_tp_group, graph_capture, is_global_first_rank,
27 prepare_communication_buffer_for_model)
30 -from vllm.forward_context import (DPMetadata, get_forward_context,
31 - set_forward_context)
28 +from vllm.forward_context import DPMetadata, set_forward_context
29 from vllm.logger import init_logger
30 from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaBase
31 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
@@ -66,6 +63,8 @@ from vllm.v1.spec_decode.medusa import MedusaProposer
63 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
64 from vllm.v1.spec_decode.ngram_proposer import NgramProposer
65 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
66 +from vllm.v1.worker.kv_connector_model_runner_mixin import (
67 + KVConnectorModelRunnerMixin)
68 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
69
70 from ..sample.logits_processor import LogitsProcessorManager
@@ -88,7 +87,7 @@ else:
87 logger = init_logger(__name__)
88
89
91 -class GPUModelRunner(LoRAModelRunnerMixin):
90 +class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
91
92 def __init__(
93 self,
@@ -1357,7 +1356,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1356 # Return empty ModelRunnerOutput if there's no work to do.
1357 return EMPTY_MODEL_RUNNER_OUTPUT
1358
1360 - return self.kv_connector_no_forward(scheduler_output)
1359 + return self.kv_connector_no_forward(scheduler_output,
1360 + self.vllm_config)
1361
1362 # Prepare the decoder inputs.
1363 (attn_metadata, attention_cuda_graphs, logits_indices,
@@ -1745,52 +1745,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1745 spec_token_ids = draft_token_ids.tolist()
1746 return spec_token_ids
1747
1748 - @staticmethod
1749 - def maybe_setup_kv_connector(scheduler_output: "SchedulerOutput"):
1750 - # Update KVConnector with the KVConnector metadata forward().
1751 - if has_kv_transfer_group():
1752 - kv_connector = get_kv_transfer_group()
1753 - assert isinstance(kv_connector, KVConnectorBase_V1)
1754 - assert scheduler_output.kv_connector_metadata is not None
1755 - kv_connector.bind_connector_metadata(
1756 - scheduler_output.kv_connector_metadata)
1757 -
1758 - # Background KV cache transfers happen here.
1759 - # These transfers are designed to be async and the requests
1760 - # involved may be disjoint from the running requests.
1761 - # Do this here to save a collective_rpc.
1762 - kv_connector.start_load_kv(get_forward_context())
1763 -
1764 - @staticmethod
1765 - def maybe_wait_for_kv_save() -> None:
1766 - if has_kv_transfer_group():
1767 - get_kv_transfer_group().wait_for_save()
1768 -
1769 - @staticmethod
1770 - def get_finished_kv_transfers(
1771 - scheduler_output: "SchedulerOutput",
1772 - ) -> tuple[Optional[set[str]], Optional[set[str]]]:
1773 - if has_kv_transfer_group():
1774 - return get_kv_transfer_group().get_finished(
1775 - scheduler_output.finished_req_ids)
1776 - return None, None
1777 -
1778 - def kv_connector_no_forward(
1779 - self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
1780 - # KV send/recv even if no work to do.
1781 - with set_forward_context(None, self.vllm_config):
1782 - self.maybe_setup_kv_connector(scheduler_output)
1783 - finished_sending, finished_recving = (
1784 - self.get_finished_kv_transfers(scheduler_output))
1785 -
1786 - if not finished_sending and not finished_recving:
1787 - return EMPTY_MODEL_RUNNER_OUTPUT
1788 -
1789 - output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)
1790 - output.finished_sending = finished_sending
1791 - output.finished_recving = finished_recving
1792 - return output
1793 -
1748 def propose_ngram_draft_token_ids(
1749 self,
1750 sampled_token_ids: list[list[int]],
vllm/v1/worker/kv_connector_model_runner_mixin.py new
+70
@@ -0,0 +1,70 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Define KV connector functionality mixin for model runners.
5 +"""
6 +import copy
7 +from typing import TYPE_CHECKING, Optional
8 +
9 +from vllm.config import VllmConfig
10 +from vllm.distributed.kv_transfer import (get_kv_transfer_group,
11 + has_kv_transfer_group)
12 +from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1
13 +from vllm.forward_context import get_forward_context, set_forward_context
14 +from vllm.logger import init_logger
15 +from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, ModelRunnerOutput
16 +
17 +if TYPE_CHECKING:
18 + from vllm.v1.core.sched.output import SchedulerOutput
19 +
20 +logger = init_logger(__name__)
21 +
22 +
23 +# Defined as a kv connector functionality mixin for ModelRunner (GPU, TPU)
24 +class KVConnectorModelRunnerMixin:
25 +
26 + @staticmethod
27 + def maybe_setup_kv_connector(scheduler_output: "SchedulerOutput"):
28 + # Update KVConnector with the KVConnector metadata forward().
29 + if has_kv_transfer_group():
30 + kv_connector = get_kv_transfer_group()
31 + assert isinstance(kv_connector, KVConnectorBase_V1)
32 + assert scheduler_output.kv_connector_metadata is not None
33 + kv_connector.bind_connector_metadata(
34 + scheduler_output.kv_connector_metadata)
35 +
36 + # Background KV cache transfers happen here.
37 + # These transfers are designed to be async and the requests
38 + # involved may be disjoint from the running requests.
39 + # Do this here to save a collective_rpc.
40 + kv_connector.start_load_kv(get_forward_context())
41 +
42 + @staticmethod
43 + def maybe_wait_for_kv_save() -> None:
44 + if has_kv_transfer_group():
45 + get_kv_transfer_group().wait_for_save()
46 +
47 + @staticmethod
48 + def get_finished_kv_transfers(
49 + scheduler_output: "SchedulerOutput",
50 + ) -> tuple[Optional[set[str]], Optional[set[str]]]:
51 + if has_kv_transfer_group():
52 + return get_kv_transfer_group().get_finished(
53 + scheduler_output.finished_req_ids)
54 + return None, None
55 +
56 + def kv_connector_no_forward(self, scheduler_output: "SchedulerOutput",
57 + vllm_config: VllmConfig) -> ModelRunnerOutput:
58 + # KV send/recv even if no work to do.
59 + with set_forward_context(None, vllm_config):
60 + self.maybe_setup_kv_connector(scheduler_output)
61 + finished_sending, finished_recving = (
62 + self.get_finished_kv_transfers(scheduler_output))
63 +
64 + if not finished_sending and not finished_recving:
65 + return EMPTY_MODEL_RUNNER_OUTPUT
66 +
67 + output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)
68 + output.finished_sending = finished_sending
69 + output.finished_recving = finished_recving
70 + return output
vllm/v1/worker/tpu_model_runner.py
+101 -4
@@ -3,7 +3,7 @@
3 import bisect
4 import gc
5 import time
6 -from typing import TYPE_CHECKING, Any, Optional, cast
6 +from typing import TYPE_CHECKING, Any, Literal, Optional, Union, cast
7 from unittest.mock import patch
8
9 import numpy as np
@@ -20,6 +20,8 @@ from vllm.attention.layer import Attention
20 from vllm.compilation.wrapper import TorchCompileWrapperWithCustomDispatcher
21 from vllm.config import (ParallelConfig, VllmConfig,
22 get_layers_from_vllm_config, update_config)
23 +from vllm.distributed.kv_transfer import (get_kv_transfer_group,
24 + has_kv_transfer_group)
25 from vllm.forward_context import set_forward_context
26 from vllm.logger import init_logger
27 from vllm.lora.layers import BaseLayerWithLoRA
@@ -46,6 +48,8 @@ from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsLists,
48 LogprobsTensors, ModelRunnerOutput)
49 from vllm.v1.sample.tpu.metadata import TPUSupportedSamplingMetadata
50 from vllm.v1.sample.tpu.sampler import Sampler as TPUSampler
51 +from vllm.v1.worker.kv_connector_model_runner_mixin import (
52 + KVConnectorModelRunnerMixin)
53 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
54 from vllm.v1.worker.tpu_input_batch import CachedRequestState, InputBatch
55
@@ -97,7 +101,7 @@ MIN_NUM_SEQS = 8
101 # The dummy_run should be comprehensive, ensuring all potential input shapes and
102 # branch predictions are included as subgraph inputs to facilitate
103 # pre-compilation.
100 -class TPUModelRunner(LoRAModelRunnerMixin):
104 +class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
105
106 def __init__(
107 self,
@@ -971,8 +975,12 @@ class TPUModelRunner(LoRAModelRunnerMixin):
975 # Update cached state
976 self._update_states(scheduler_output)
977 if not scheduler_output.total_num_scheduled_tokens:
974 - # Return empty ModelRunnerOutput if there's no work to do.
975 - return EMPTY_MODEL_RUNNER_OUTPUT
978 + if not has_kv_transfer_group():
979 + # Return empty ModelRunnerOutput if there's no work to do.
980 + return EMPTY_MODEL_RUNNER_OUTPUT
981 +
982 + return self.kv_connector_no_forward(scheduler_output,
983 + self.vllm_config)
984
985 if self.is_multimodal_model:
986 # Run the multimodal encoder if any.
@@ -986,6 +994,12 @@ class TPUModelRunner(LoRAModelRunnerMixin):
994 start_index = 0
995 combined_selected_tokens: list[torch.Tensor] = []
996 combined_logprobs: list[LogprobsLists] = []
997 +
998 + # NOTE: setup current batch's metadata for kv connector.
999 + # Currently, only verified with NixlConnector
1000 + with set_forward_context(None, self.vllm_config):
1001 + self.maybe_setup_kv_connector(scheduler_output)
1002 +
1003 while start_index < self.input_batch.num_reqs:
1004 attn_metadata, logits_indices, padded_num_reqs, num_reqs,\
1005 end_index = self._prepare_inputs(scheduler_output, start_index)
@@ -1032,6 +1046,14 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1046
1047 start_index = end_index
1048
1049 + # NOTE: current kv load and save get h2d/d2h copies involved.
1050 + # Those copies are blocking. Once they become async., kv_save
1051 + # should be called right after each single forward pass,
1052 + # instead of the forwards of the entire input batch.
1053 + self.maybe_wait_for_kv_save()
1054 + finished_sending, finished_recving = (
1055 + self.get_finished_kv_transfers(scheduler_output))
1056 +
1057 selected_token_ids = torch.cat(combined_selected_tokens, dim=0)
1058 if tpu_sampling_metadata.logprobs:
1059
@@ -1126,6 +1148,8 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1148 logprobs=logprobs_lists,
1149 prompt_logprobs_dict=prompt_logprobs_dict,
1150 pooler_output=[],
1151 + finished_sending=finished_sending,
1152 + finished_recving=finished_recving,
1153 )
1154
1155 # Check there are no new graphs compiled - all the graphs should be
@@ -1637,6 +1661,10 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1661 for cache in self.kv_caches:
1662 xs.mark_sharding(cache, self.mesh, (None, 'x', None, None))
1663
1664 + if has_kv_transfer_group():
1665 + get_kv_transfer_group().register_kv_caches(kv_caches)
1666 + get_kv_transfer_group().set_host_xfer_buffer_ops(copy_kv_blocks)
1667 +
1668 def reset_dynamo_cache(self):
1669 if self.is_multimodal_model:
1670 compiled_model = self.model.get_language_model().model
@@ -1851,6 +1879,75 @@ def _get_padded_token_len(paddings: list[int], x: int) -> int:
1879 return paddings[index]
1880
1881
1882 +def _make_src_and_dst_indices(
1883 + src_block_ids: list[int],
1884 + dst_block_ids: list[int],
1885 + src_device: Union[torch.device, str],
1886 + dst_device: Union[torch.device, str],
1887 +) -> tuple[torch.Tensor, torch.Tensor]:
1888 + src_indices = torch.tensor(src_block_ids,
1889 + device=src_device,
1890 + dtype=torch.int64)
1891 + dst_indices = torch.tensor(dst_block_ids,
1892 + device=dst_device,
1893 + dtype=torch.int64)
1894 + return src_indices, dst_indices
1895 +
1896 +
1897 +@torch.compile(backend="openxla")
1898 +def _insert_blocks_to_tpu(
1899 + cpu_cache: torch.Tensor,
1900 + tpu_cache: torch.Tensor,
1901 + cpu_block_indices: torch.Tensor,
1902 + tpu_block_indices: torch.Tensor,
1903 +) -> None:
1904 + torch.ops.xla.dynamo_set_buffer_donor_(tpu_cache, True)
1905 + tpu_cache[tpu_block_indices] = cpu_cache[cpu_block_indices].to(
1906 + tpu_cache.device)
1907 +
1908 +
1909 +@torch.compile(backend="openxla")
1910 +def _swap_out_tpu_blocks(
1911 + tpu_cache: torch.Tensor,
1912 + cpu_cache: torch.Tensor,
1913 + tpu_block_indices: torch.Tensor,
1914 + cpu_block_indices: torch.Tensor,
1915 +) -> None:
1916 + """ tpu blocks to cpu blocks"""
1917 + torch.ops.xla.dynamo_set_buffer_donor_(tpu_cache, True)
1918 + cpu_cache[cpu_block_indices] = tpu_cache[tpu_block_indices].cpu()
1919 +
1920 +
1921 +def copy_kv_blocks(
1922 + src_kv_caches: dict[str, torch.Tensor],
1923 + dst_kv_caches: dict[str, torch.Tensor],
1924 + src_block_ids: list[int],
1925 + dst_block_ids: list[int],
1926 + direction: Literal["h2d", "d2h"],
1927 +) -> None:
1928 + """Copy kv blocks between different buffers."""
1929 + if not src_kv_caches or not dst_kv_caches or \
1930 + not src_block_ids or not dst_block_ids or \
1931 + len(src_block_ids) != len(dst_block_ids):
1932 + return
1933 +
1934 + src_device = next(iter(src_kv_caches.values())).device
1935 + dst_device = next(iter(dst_kv_caches.values())).device
1936 +
1937 + src_indices, dst_indices = _make_src_and_dst_indices(
1938 + src_block_ids=src_block_ids,
1939 + dst_block_ids=dst_block_ids,
1940 + src_device=src_device,
1941 + dst_device=dst_device)
1942 +
1943 + _copy_fn = _insert_blocks_to_tpu if direction == "h2d" else \
1944 + _swap_out_tpu_blocks
1945 + for layer_name in src_kv_caches:
1946 + src_tensor = src_kv_caches[layer_name]
1947 + dst_tensor = dst_kv_caches[layer_name]
1948 + _copy_fn(src_tensor, dst_tensor, src_indices, dst_indices)
1949 +
1950 +
1951 def _get_padded_num_kv_cache_update_slices(
1952 num_tokens: int, max_num_reqs: int, page_size: int,
1953 num_slices_per_kv_cache_update_block: int) -> int:
vllm/v1/worker/tpu_worker.py
+11 -4
@@ -12,9 +12,11 @@ import torch_xla.debug.profiler as xp
12 import torch_xla.runtime as xr
13
14 import vllm.envs as envs
15 -from vllm.config import ParallelConfig, VllmConfig
15 +from vllm.config import VllmConfig
16 from vllm.distributed import (ensure_model_parallel_initialized,
17 init_distributed_environment)
18 +from vllm.distributed.kv_transfer import (ensure_kv_transfer_initialized,
19 + has_kv_transfer_group)
20 from vllm.logger import init_logger
21 from vllm.lora.request import LoRARequest
22 from vllm.model_executor import set_random_seed
@@ -118,7 +120,7 @@ class TPUWorker:
120
121 # Initialize the distributed environment.
122 self._init_tpu_worker_distributed_environment(
121 - self.parallel_config, self.rank, self.distributed_init_method,
123 + self.vllm_config, self.rank, self.distributed_init_method,
124 self.local_rank)
125
126 # Device initialization should happen after initializing
@@ -242,7 +244,9 @@ class TPUWorker:
244 scheduler_output: "SchedulerOutput",
245 ) -> Optional[ModelRunnerOutput]:
246 output = self.model_runner.execute_model(scheduler_output)
245 - return output if self.is_driver_worker else None
247 + # every worker's output is needed when kv_transfer_group is setup
248 + return output if self.is_driver_worker or has_kv_transfer_group(
249 + ) else None
250
251 def profile(self, is_start: bool = True):
252 if self.rank < 1:
@@ -294,7 +298,7 @@ class TPUWorker:
298
299 def _init_tpu_worker_distributed_environment(
300 self,
297 - parallel_config: ParallelConfig,
301 + vllm_config: VllmConfig,
302 rank: int,
303 distributed_init_method: Optional[str] = None,
304 local_rank: int = -1,
@@ -306,6 +310,7 @@ class TPUWorker:
310 # the input objects on CPU. The all-reduce and all-gather ops on TPU
311 # are invoked by `xm.all_reduce` and `xm.all_gather` which use their
312 # own context.
313 + parallel_config = vllm_config.parallel_config
314 init_distributed_environment(
315 world_size=parallel_config.world_size,
316 rank=rank,
@@ -317,6 +322,8 @@ class TPUWorker:
322 parallel_config.tensor_parallel_size,
323 parallel_config.pipeline_parallel_size)
324
325 + ensure_kv_transfer_initialized(vllm_config)
326 +
327
328 try:
329 from tpu_commons.worker import TPUWorker as TPUCommonsWorker