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

[P/D][V1] KV Connector API V1 (#15960)

Signed-off-by: ApostaC <yihua98@uchicago.edu> Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com> Signed-off-by: remi <remi@mistral.ai> Co-authored-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Rémi Delacourt <54138269+Flechman@users.noreply.github.com> Co-authored-by: Tyler Michael Smith <tysmith@redhat.com>

Yihua Cheng committed Apr 17, 2025 at 15:22 UTC 3408e471597e7a36ca79fab5fc849f4fb5576df8
24 files changed +1377 -83
examples/offline_inference/disaggregated-prefill-v1/decode_example.py new
+36
@@ -0,0 +1,36 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from vllm import LLM, SamplingParams
4 +from vllm.config import KVTransferConfig
5 +
6 +# Read prompts from output.txt
7 +prompts = []
8 +try:
9 + with open("output.txt") as f:
10 + for line in f:
11 + prompts.append(line.strip())
12 + print(f"Loaded {len(prompts)} prompts from output.txt")
13 +except FileNotFoundError:
14 + print("Error: output.txt file not found")
15 + exit(-1)
16 +
17 +sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
18 +
19 +llm = LLM(
20 + model="meta-llama/Llama-3.2-1B-Instruct",
21 + enforce_eager=True,
22 + gpu_memory_utilization=0.8,
23 + max_num_batched_tokens=64,
24 + max_num_seqs=16,
25 + kv_transfer_config=KVTransferConfig.from_cli(
26 + '{"kv_connector":"SharedStorageConnector","kv_role":"kv_both",'
27 + '"kv_connector_extra_config": {"shared_storage_path": "local_storage"}}'
28 + )) #, max_model_len=2048, max_num_batched_tokens=2048)
29 +
30 +# 1ST generation (prefill instance)
31 +outputs = llm.generate(prompts, sampling_params)
32 +
33 +for output in outputs:
34 + prompt = output.prompt
35 + generated_text = output.outputs[0].text
36 + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
examples/offline_inference/disaggregated-prefill-v1/prefill_example.py new
+43
@@ -0,0 +1,43 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from vllm import LLM, SamplingParams
4 +from vllm.config import KVTransferConfig
5 +
6 +context = "Hi " * 1000
7 +context2 = "Hey " * 500
8 +prompts = [
9 + context + "Hello, my name is",
10 + context + "The capital of France is",
11 + context2 + "Your name is",
12 + context2 + "The capital of China is",
13 +]
14 +
15 +sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1)
16 +
17 +llm = LLM(model="meta-llama/Llama-3.2-1B-Instruct",
18 + enforce_eager=True,
19 + gpu_memory_utilization=0.8,
20 + kv_transfer_config=KVTransferConfig.from_cli(
21 + '{"kv_connector":"SharedStorageConnector","kv_role":"kv_both", '
22 + '"kv_connector_extra_config": '
23 + '{"shared_storage_path": "local_storage"}}')
24 + ) #, max_model_len=2048, max_num_batched_tokens=2048)
25 +
26 +# 1ST generation (prefill instance)
27 +outputs = llm.generate(
28 + prompts,
29 + sampling_params,
30 +)
31 +
32 +new_prompts = []
33 +for output in outputs:
34 + prompt = output.prompt
35 + generated_text = output.outputs[0].text
36 + new_prompts.append(prompt + generated_text)
37 + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
38 +
39 +# Write new_prompts to output.txt
40 +with open("output.txt", "w") as f:
41 + for prompt in new_prompts:
42 + f.write(prompt + "\n")
43 +print(f"Saved {len(new_prompts)} prompts to output.txt")
examples/offline_inference/disaggregated-prefill-v1/run.sh new
+5
@@ -0,0 +1,5 @@
1 +rm -rf local_storage/
2 +rm output.txt
3 +
4 +VLLM_ENABLE_V1_MULTIPROCESSING=0 CUDA_VISIBLE_DEVICES=0 python3 prefill_example.py
5 +VLLM_ENABLE_V1_MULTIPROCESSING=0 CUDA_VISIBLE_DEVICES=0 python3 decode_example.py
tests/v1/core/test_scheduler.py
+406 -9
@@ -1,10 +1,12 @@
1 # SPDX-License-Identifier: Apache-2.0
2 from typing import Optional
3 +from unittest.mock import Mock
4
5 import pytest
6 import torch
7
7 -from vllm.config import CacheConfig, ModelConfig, SchedulerConfig, VllmConfig
8 +from vllm.config import (CacheConfig, KVTransferConfig, ModelConfig,
9 + SchedulerConfig, VllmConfig)
10 from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
11 from vllm.sampling_params import SamplingParams
12 from vllm.v1.core.sched.output import SchedulerOutput
@@ -25,6 +27,9 @@ def create_scheduler(
27 enable_prefix_caching: Optional[bool] = None,
28 long_prefill_token_threshold: int = 0,
29 disable_chunked_mm_input: bool = False,
30 + use_kv_connector: bool = False,
31 + num_blocks: int = 10000,
32 + block_size: int = 16,
33 ) -> Scheduler:
34 '''Create scheduler under test.
35
@@ -60,31 +65,36 @@ def create_scheduler(
65 'enable_prefix_caching': enable_prefix_caching
66 })
67 cache_config = CacheConfig(
63 - block_size=16,
68 + block_size=block_size,
69 gpu_memory_utilization=0.9,
70 swap_space=0,
71 cache_dtype="auto",
72 **kwargs_cache,
73 )
74 + kv_transfer_config = KVTransferConfig(
75 + kv_connector="SharedStorageConnector",
76 + kv_role="kv_both",
77 + kv_connector_extra_config={"shared_storage_path": "local_storage"},
78 + ) if use_kv_connector else None
79 +
80 vllm_config = VllmConfig(
81 scheduler_config=scheduler_config,
82 model_config=model_config,
83 cache_config=cache_config,
84 + kv_transfer_config=kv_transfer_config,
85 )
86 kv_cache_config = KVCacheConfig(
75 - num_blocks=10000, # A large number of blocks to hold all requests
87 + num_blocks=num_blocks, # A large number of blocks to hold all requests
88 tensors={},
89 kv_cache_groups=[
90 KVCacheGroupSpec(['layer'],
79 - FullAttentionSpec(16, 1, 1, torch.float32, False))
91 + FullAttentionSpec(block_size, 1, 1, torch.float32,
92 + False))
93 ],
94 )
82 - cache_config.num_gpu_blocks = 10000
95 + cache_config.num_gpu_blocks = num_blocks
96 return Scheduler(
84 - scheduler_config,
85 - model_config,
86 - cache_config,
87 - lora_config=None,
97 + vllm_config=vllm_config,
98 kv_cache_config=kv_cache_config,
99 log_stats=True,
100 structured_output_manager=StructuredOutputManager(vllm_config),
@@ -761,3 +771,390 @@ def test_schedule_spec_decoding_stats(spec_tokens, output_tokens, expected):
771 stats = scheduler_stats.spec_decoding_stats
772 assert stats.num_draft_tokens == expected[0]
773 assert stats.num_accepted_tokens == expected[1]
774 +
775 +
776 +def _assert_right_scheduler_output(
777 + output: SchedulerOutput,
778 + num_requests: int,
779 + expected_num_scheduled_tokens: int,
780 +):
781 + """Check if SchedulerOutput is correct after remote KV cache hit."""
782 +
783 + # We should inject the kv_connector_metadata.
784 + assert len(output.kv_connector_metadata.requests) == num_requests
785 +
786 + # Only num_tokens - matched_num_new_tokens should be scheduled.
787 + for _, num_scheduled_tokens in output.num_scheduled_tokens.items():
788 + assert num_scheduled_tokens == expected_num_scheduled_tokens
789 +
790 +
791 +def _assert_right_kv_cache_manager(
792 + scheduler: Scheduler,
793 + req_ids: list[str],
794 + num_tokens: int,
795 + block_size: int,
796 + num_requests: int,
797 + num_total_blocks: int,
798 +):
799 + """Check whether KVCacheManager is correct after allocate."""
800 +
801 + # Make sure the request stats are right.
802 + EXPECTED_ACTUAL_BLOCKS = num_tokens // block_size
803 + EXPECTED_TOTAL_BLOCKS = (EXPECTED_ACTUAL_BLOCKS +
804 + scheduler.kv_cache_manager.num_preallocate_blocks)
805 + for req_id in req_ids:
806 + blocks = scheduler.kv_cache_manager.req_to_blocks[req_id]
807 + hashes = scheduler.kv_cache_manager.req_to_block_hashes[req_id]
808 + assert (scheduler.kv_cache_manager.num_cached_block[req_id] ==
809 + EXPECTED_ACTUAL_BLOCKS)
810 + assert len(blocks) == EXPECTED_TOTAL_BLOCKS
811 + assert len(hashes) == EXPECTED_ACTUAL_BLOCKS
812 +
813 + # Make sure we actually touched all the blocks.
814 + BLOCKS_PER_REQ = (num_tokens / block_size +
815 + scheduler.kv_cache_manager.num_preallocate_blocks)
816 + assert (scheduler.kv_cache_manager.block_pool.get_num_free_blocks() ==
817 + num_total_blocks - num_requests * BLOCKS_PER_REQ)
818 +
819 +
820 +def _step_until_done(
821 + scheduler: Scheduler,
822 + output: SchedulerOutput,
823 + model_runner_output: ModelRunnerOutput,
824 +):
825 + """Loop over schedule(), update_from_output() until finished."""
826 +
827 + all_finished = False
828 + _ = scheduler.update_from_output(output, model_runner_output)
829 + while not all_finished:
830 + # Schedule + a few iterations until stopping.
831 + output = scheduler.schedule()
832 + assert len(scheduler.running)
833 + for _, num_scheduled_tokens in output.num_scheduled_tokens.items():
834 + # We should be in the decode phase now.
835 + assert num_scheduled_tokens == 1
836 + assert len(output.kv_connector_metadata.requests) == 0
837 + ecos = scheduler.update_from_output(output, model_runner_output)
838 + all_done = True
839 + for eco in ecos.outputs:
840 + if eco.finish_reason is None:
841 + all_done = False
842 + all_finished = all_done
843 +
844 +
845 +def test_kv_connector_basic():
846 + """
847 + Test whether Scheduler with KVConnector schedules tokens, allocates
848 + memory, and cleans up requests as expected under normal operation.
849 + """
850 +
851 + # Setup Scheduler.
852 + scheduler = create_scheduler(
853 + enable_prefix_caching=True,
854 + use_kv_connector=True,
855 + )
856 + NUM_TOTAL_BLOCKS = (
857 + scheduler.kv_cache_manager.block_pool.get_num_free_blocks())
858 + BLOCK_SIZE = scheduler.cache_config.block_size
859 +
860 + # Mock External Cache Hit.
861 + NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE * 2
862 + scheduler.connector.get_num_new_matched_tokens = Mock(name="method")
863 + scheduler.connector.get_num_new_matched_tokens.return_value = (
864 + NUM_MATCHED_NEW_TOKENS)
865 +
866 + ######################################################
867 + # FIRST SET OF REQUESTS - External Hit Only
868 + NUM_REQUESTS = 2
869 + NUM_TOKENS = NUM_MATCHED_NEW_TOKENS * 2
870 + MAX_TOKENS = 3
871 + requests = create_requests(num_requests=NUM_REQUESTS,
872 + num_tokens=NUM_TOKENS,
873 + max_tokens=MAX_TOKENS)
874 + req_ids = []
875 + req_to_index = {}
876 + for i, request in enumerate(requests):
877 + scheduler.add_request(request)
878 + req_ids.append(request.request_id)
879 + req_to_index[request.request_id] = i
880 +
881 + MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
882 + req_ids=req_ids,
883 + req_id_to_index=req_to_index,
884 + sampled_token_ids=[[1000]] * len(req_ids),
885 + spec_token_ids=None,
886 + logprobs=None,
887 + prompt_logprobs_dict={},
888 + )
889 +
890 + # Ensure ScheduleOutput is correct.
891 + output = scheduler.schedule()
892 + _assert_right_scheduler_output(
893 + output=output,
894 + num_requests=NUM_REQUESTS,
895 + # Just the incremental tokens should be scheduled.
896 + expected_num_scheduled_tokens=NUM_TOKENS - NUM_MATCHED_NEW_TOKENS,
897 + )
898 +
899 + # Ensure KVCacheManager is correct.
900 + _assert_right_kv_cache_manager(scheduler, req_ids, NUM_TOKENS, BLOCK_SIZE,
901 + NUM_REQUESTS, NUM_TOTAL_BLOCKS)
902 +
903 + # Continue Generation until done.
904 + _step_until_done(scheduler, output, MODEL_RUNNER_OUTPUT)
905 + _ = scheduler.schedule()
906 + # Confirm we clean up the memory properly.
907 + assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() \
908 + == NUM_TOTAL_BLOCKS
909 +
910 + ######################################################
911 + # SECOND SET OF REQUESTS - Local And External Hit
912 + NUM_TOKENS_PREFIX = NUM_TOKENS
913 + # We will get a local prefix cache hit for the first
914 + # NUM_TOKENS_PREFIX tokens since they are used above.
915 + NUM_TOKENS = NUM_TOKENS_PREFIX * 2
916 + requests = create_requests(num_requests=NUM_REQUESTS,
917 + num_tokens=NUM_TOKENS,
918 + max_tokens=MAX_TOKENS)
919 + req_ids = []
920 + req_to_index = {}
921 + for i, request in enumerate(requests):
922 + scheduler.add_request(request)
923 + req_ids.append(request.request_id)
924 + req_to_index[request.request_id] = i
925 +
926 + MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
927 + req_ids=req_ids,
928 + req_id_to_index=req_to_index,
929 + sampled_token_ids=[[1000]] * len(req_ids),
930 + spec_token_ids=None,
931 + logprobs=None,
932 + prompt_logprobs_dict={},
933 + )
934 +
935 + # We should get a local cache hit of NUM_TOKENS_PREFIX and
936 + # a remote KV cache hit of NUM_MATCHED_NEW_TOKENS.
937 + output = scheduler.schedule()
938 + _assert_right_scheduler_output(
939 + output=output,
940 + num_requests=NUM_REQUESTS,
941 + # Just the incremental tokens after local + remote cache hit.
942 + expected_num_scheduled_tokens=(NUM_TOKENS - NUM_TOKENS_PREFIX -
943 + NUM_MATCHED_NEW_TOKENS))
944 +
945 + # Ensure KVCacheManager is correct.
946 + _assert_right_kv_cache_manager(scheduler, req_ids, NUM_TOKENS, BLOCK_SIZE,
947 + NUM_REQUESTS, NUM_TOTAL_BLOCKS)
948 +
949 + # Continue Generation until done.
950 + _step_until_done(scheduler, output, MODEL_RUNNER_OUTPUT)
951 + _ = scheduler.schedule()
952 + # Confirm we clean up the memory properly.
953 + assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() \
954 + == NUM_TOTAL_BLOCKS
955 +
956 +
957 +def test_kv_connector_unable_to_allocate():
958 + """
959 + Test whether scheduler with KVConnector is able to handle
960 + unable to allocate (run out of blocks in allocate_slots().
961 + """
962 +
963 + # Setup Scheduler With Mock External Cache Hit.
964 + BLOCK_SIZE = 4
965 + NUM_BLOCKS = 10
966 + scheduler = create_scheduler(
967 + enable_prefix_caching=True,
968 + use_kv_connector=True,
969 + block_size=BLOCK_SIZE,
970 + num_blocks=NUM_BLOCKS,
971 + )
972 + NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE * 2
973 + scheduler.connector.get_num_new_matched_tokens = Mock(name="method")
974 + scheduler.connector.get_num_new_matched_tokens.return_value = (
975 + NUM_MATCHED_NEW_TOKENS)
976 +
977 + # Create two requests. The second request will not be able to
978 + # allocate slots because it will not have enough blocks.
979 + NUM_REQUESTS = 2
980 + NUM_TOKENS = (NUM_BLOCKS // 2 + 1) * BLOCK_SIZE
981 + MAX_TOKENS = 2
982 + requests = create_requests(num_requests=NUM_REQUESTS,
983 + num_tokens=NUM_TOKENS,
984 + max_tokens=MAX_TOKENS)
985 + req_ids = []
986 + req_to_index = {}
987 + for i, request in enumerate(requests):
988 + scheduler.add_request(request)
989 + req_ids.append(request.request_id)
990 + req_to_index[request.request_id] = i
991 +
992 + MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
993 + req_ids=req_ids,
994 + req_id_to_index=req_to_index,
995 + sampled_token_ids=[[1000]] * len(req_ids),
996 + spec_token_ids=None,
997 + logprobs=None,
998 + prompt_logprobs_dict={},
999 + )
1000 +
1001 + # Just one request should be running.
1002 + output = scheduler.schedule()
1003 + _assert_right_scheduler_output(output,
1004 + num_requests=1,
1005 + expected_num_scheduled_tokens=NUM_TOKENS -
1006 + NUM_MATCHED_NEW_TOKENS)
1007 + assert len(scheduler.running) == 1
1008 + assert len(scheduler.waiting) == 1
1009 +
1010 + # All memory should be freed, with one request waiting.
1011 + _step_until_done(scheduler, output, MODEL_RUNNER_OUTPUT)
1012 + assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() \
1013 + == NUM_BLOCKS - 1
1014 + assert len(scheduler.running) == 0
1015 + assert len(scheduler.waiting) == 1
1016 +
1017 + # Just one request should be running.
1018 + output = scheduler.schedule()
1019 + _assert_right_scheduler_output(output,
1020 + num_requests=1,
1021 + expected_num_scheduled_tokens=NUM_TOKENS -
1022 + NUM_MATCHED_NEW_TOKENS)
1023 + assert len(scheduler.running) == 1
1024 + assert len(scheduler.waiting) == 0
1025 +
1026 + # All memory should be freed, with no requests waiting / running.
1027 + _step_until_done(scheduler, output, MODEL_RUNNER_OUTPUT)
1028 + assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() \
1029 + == NUM_BLOCKS - 1
1030 + assert len(scheduler.running) == 0
1031 + assert len(scheduler.waiting) == 0
1032 +
1033 +
1034 +def test_kv_connector_handles_preemption():
1035 + """
1036 + Test whether scheduler with KVConnector is able to handle
1037 + unable to allocate (run out of blocks in allocate_slots().
1038 + """
1039 +
1040 + # Setup Scheduler With Mock External Cache Hit.
1041 + BLOCK_SIZE = 2
1042 + # NOTE: there is 1 null block, so this is 6 blocks.
1043 + NUM_BLOCKS = 7
1044 + scheduler = create_scheduler(
1045 + enable_prefix_caching=True,
1046 + use_kv_connector=True,
1047 + block_size=BLOCK_SIZE,
1048 + num_blocks=NUM_BLOCKS,
1049 + )
1050 + scheduler.kv_cache_manager.num_preallocate_blocks = 0
1051 +
1052 + NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE
1053 + scheduler.connector.get_num_new_matched_tokens = Mock(name="method")
1054 + scheduler.connector.get_num_new_matched_tokens.return_value = (
1055 + NUM_MATCHED_NEW_TOKENS)
1056 +
1057 + # Create two requests.
1058 + # Both can be scheduled at first, but the second request
1059 + # will be preempted and re-scheduled.
1060 + NUM_REQUESTS = 2
1061 + NUM_TOKENS = BLOCK_SIZE * 2 + 1
1062 + MAX_TOKENS = BLOCK_SIZE * 2
1063 + requests = create_requests(num_requests=NUM_REQUESTS,
1064 + num_tokens=NUM_TOKENS,
1065 + max_tokens=MAX_TOKENS)
1066 + req_ids = []
1067 + req_to_index = {}
1068 + for i, request in enumerate(requests):
1069 + scheduler.add_request(request)
1070 + req_ids.append(request.request_id)
1071 + req_to_index[request.request_id] = i
1072 +
1073 + MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
1074 + req_ids=req_ids,
1075 + req_id_to_index=req_to_index,
1076 + sampled_token_ids=[[1000]] * len(req_ids),
1077 + spec_token_ids=None,
1078 + logprobs=None,
1079 + prompt_logprobs_dict={},
1080 + )
1081 +
1082 + # All can be scheduled - 1st token.
1083 + output = scheduler.schedule()
1084 + _assert_right_scheduler_output(
1085 + output,
1086 + # 2 remote kv cache hits.
1087 + num_requests=2,
1088 + expected_num_scheduled_tokens=NUM_TOKENS - NUM_MATCHED_NEW_TOKENS)
1089 + assert len(scheduler.running) == 2
1090 + _ = scheduler.update_from_output(output, MODEL_RUNNER_OUTPUT)
1091 +
1092 + # All can be scheduled - 2nd token.
1093 + output = scheduler.schedule()
1094 + _assert_right_scheduler_output(
1095 + output,
1096 + # no connector_metadata
1097 + num_requests=0,
1098 + expected_num_scheduled_tokens=1)
1099 + assert len(scheduler.running) == 2
1100 + _ = scheduler.update_from_output(output, MODEL_RUNNER_OUTPUT)
1101 +
1102 + # This will generate a new block and cause a preemption - 3rd token.
1103 + output = scheduler.schedule()
1104 + _assert_right_scheduler_output(
1105 + output,
1106 + # no connector_metadata
1107 + num_requests=0,
1108 + expected_num_scheduled_tokens=1)
1109 + assert len(scheduler.running) == 1
1110 + assert len(scheduler.waiting) == 1
1111 + _ = scheduler.update_from_output(output, MODEL_RUNNER_OUTPUT)
1112 + assert len(scheduler.running) == 1
1113 + assert len(scheduler.waiting) == 1
1114 +
1115 + # Only 1 can be scheduled - 4th (and last token).
1116 + output = scheduler.schedule()
1117 + _assert_right_scheduler_output(
1118 + output,
1119 + # no connector_metadata
1120 + num_requests=0,
1121 + expected_num_scheduled_tokens=1)
1122 + assert len(scheduler.waiting) == 1
1123 + assert len(scheduler.running) == 1
1124 + _ = scheduler.update_from_output(output, MODEL_RUNNER_OUTPUT)
1125 + assert len(scheduler.running) == 0
1126 + assert len(scheduler.waiting) == 1
1127 + # All memory should be freed since nothing is running.
1128 + assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() \
1129 + == NUM_BLOCKS - 1
1130 +
1131 + # Restarts the preempted request - generate 3rd token.
1132 + # This will have a local and remote cache hit.
1133 + output = scheduler.schedule()
1134 + _assert_right_scheduler_output(
1135 + output,
1136 + # 1 remote kv_cache hit!
1137 + num_requests=1,
1138 + # Only 1 block was preempted and there is a single
1139 + # remote hit. So only single new token is scheduled.
1140 + expected_num_scheduled_tokens=1,
1141 + )
1142 + assert len(scheduler.running) == 1
1143 + assert len(scheduler.waiting) == 0
1144 + _ = scheduler.update_from_output(output, MODEL_RUNNER_OUTPUT)
1145 + assert len(scheduler.running) == 1
1146 + assert len(scheduler.waiting) == 0
1147 +
1148 + # Only 1 can be scheduled - 4th (and last token).
1149 + output = scheduler.schedule()
1150 + _assert_right_scheduler_output(
1151 + output,
1152 + # no connector_metadata
1153 + num_requests=0,
1154 + expected_num_scheduled_tokens=1)
1155 + assert len(scheduler.running) == 1
1156 + _ = scheduler.update_from_output(output, MODEL_RUNNER_OUTPUT)
1157 + assert len(scheduler.running) == 0
1158 + # All memory should be freed since nothing is running.
1159 + assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() \
1160 + == NUM_BLOCKS - 1
vllm/attention/layer.py
+44 -1
@@ -10,6 +10,9 @@ import vllm.envs as envs
10 from vllm.attention import AttentionType
11 from vllm.attention.selector import backend_name_to_enum, get_attn_backend
12 from vllm.config import CacheConfig, get_current_vllm_config
13 +from vllm.distributed.kv_transfer import (get_kv_transfer_group,
14 + has_kv_transfer_group,
15 + is_v1_kv_transfer_group)
16 from vllm.forward_context import ForwardContext, get_forward_context
17 from vllm.model_executor.layers.linear import UnquantizedLinearMethod
18 from vllm.model_executor.layers.quantization.base_config import (
@@ -329,17 +332,54 @@ class MultiHeadAttention(nn.Module):
332 return out.reshape(bsz, q_len, -1)
333
334
335 +def wait_for_kv_layer_from_connector(layer_name: str):
336 + if not has_kv_transfer_group() or not is_v1_kv_transfer_group():
337 + return
338 +
339 + connector = get_kv_transfer_group()
340 +
341 + forward_context: ForwardContext = get_forward_context()
342 + attn_metadata = forward_context.attn_metadata
343 + if attn_metadata is None:
344 + return
345 +
346 + connector.wait_for_layer_load(layer_name)
347 +
348 +
349 +def maybe_save_kv_layer_to_connector(
350 + layer_name: str,
351 + kv_cache_layer: List[torch.Tensor],
352 +):
353 + if not has_kv_transfer_group() or not is_v1_kv_transfer_group():
354 + return
355 +
356 + connector = get_kv_transfer_group()
357 +
358 + forward_context: ForwardContext = get_forward_context()
359 + attn_metadata = forward_context.attn_metadata
360 + if attn_metadata is None:
361 + return
362 +
363 + connector.save_kv_layer(layer_name, kv_cache_layer, attn_metadata)
364 +
365 +
366 def unified_attention(
367 query: torch.Tensor,
368 key: torch.Tensor,
369 value: torch.Tensor,
370 layer_name: str,
371 ) -> torch.Tensor:
372 + wait_for_kv_layer_from_connector(layer_name)
373 +
374 forward_context: ForwardContext = get_forward_context()
375 attn_metadata = forward_context.attn_metadata
376 self = forward_context.no_compile_layers[layer_name]
377 kv_cache = self.kv_cache[forward_context.virtual_engine]
342 - return self.impl.forward(self, query, key, value, kv_cache, attn_metadata)
378 + output = self.impl.forward(self, query, key, value, kv_cache,
379 + attn_metadata)
380 +
381 + maybe_save_kv_layer_to_connector(layer_name, kv_cache)
382 + return output
383
384
385 def unified_attention_fake(
@@ -367,6 +407,7 @@ def unified_attention_with_output(
407 output: torch.Tensor,
408 layer_name: str,
409 ) -> None:
410 + wait_for_kv_layer_from_connector(layer_name)
411 forward_context: ForwardContext = get_forward_context()
412 attn_metadata = forward_context.attn_metadata
413 self = forward_context.no_compile_layers[layer_name]
@@ -379,6 +420,8 @@ def unified_attention_with_output(
420 attn_metadata,
421 output=output)
422
423 + maybe_save_kv_layer_to_connector(layer_name, kv_cache)
424 +
425
426 def unified_attention_with_output_fake(
427 query: torch.Tensor,
vllm/distributed/kv_transfer/__init__.py
+11
@@ -0,0 +1,11 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from vllm.distributed.kv_transfer.kv_transfer_state import (
4 + ensure_kv_transfer_initialized, get_kv_transfer_group,
5 + has_kv_transfer_group, is_v1_kv_transfer_group)
6 +
7 +__all__ = [
8 + "get_kv_transfer_group", "has_kv_transfer_group",
9 + "is_v1_kv_transfer_group", "ensure_kv_transfer_initialized",
10 + "KVConnectorBaseType"
11 +]
vllm/distributed/kv_transfer/kv_connector/base.py
+4
@@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, List, Tuple, Union
12
13 import torch
14
15 +from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1
16 from vllm.sequence import IntermediateTensors
17
18 if TYPE_CHECKING:
@@ -121,3 +122,6 @@ class KVConnectorBase(ABC):
122 """
123
124 raise NotImplementedError
125 +
126 +
127 +KVConnectorBaseType = Union[KVConnectorBase, KVConnectorBase_V1]
vllm/distributed/kv_transfer/kv_connector/factory.py
+47 -5
@@ -3,14 +3,22 @@
3 import importlib
4 from typing import TYPE_CHECKING, Callable, Dict, Type
5
6 +import vllm.envs as envs
7 +from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBaseType
8 +from vllm.distributed.kv_transfer.kv_connector.v1 import (KVConnectorBase_V1,
9 + KVConnectorRole)
10 +from vllm.logger import init_logger
11 +
12 from .base import KVConnectorBase
13
14 if TYPE_CHECKING:
15 from vllm.config import VllmConfig
16
17 +logger = init_logger(__name__)
18 +
19
20 class KVConnectorFactory:
13 - _registry: Dict[str, Callable[[], Type[KVConnectorBase]]] = {}
21 + _registry: Dict[str, Callable[[], Type[KVConnectorBaseType]]] = {}
22
23 @classmethod
24 def register_connector(cls, name: str, module_path: str,
@@ -19,22 +27,51 @@ class KVConnectorFactory:
27 if name in cls._registry:
28 raise ValueError(f"Connector '{name}' is already registered.")
29
22 - def loader() -> Type[KVConnectorBase]:
30 + def loader() -> Type[KVConnectorBaseType]:
31 module = importlib.import_module(module_path)
32 return getattr(module, class_name)
33
34 cls._registry[name] = loader
35
36 @classmethod
29 - def create_connector(cls, rank: int, local_rank: int,
30 - config: "VllmConfig") -> KVConnectorBase:
37 + def create_connector_v0(cls, rank: int, local_rank: int,
38 + config: "VllmConfig") -> KVConnectorBase:
39 + if envs.VLLM_USE_V1:
40 + raise ValueError("Attempting to initialize a V0 Connector, "
41 + f"but found {envs.VLLM_USE_V1=}")
42 +
43 connector_name = config.kv_transfer_config.kv_connector
44 if connector_name not in cls._registry:
45 raise ValueError(f"Unsupported connector type: {connector_name}")
46
47 connector_cls = cls._registry[connector_name]()
48 + assert issubclass(connector_cls, KVConnectorBase)
49 return connector_cls(rank, local_rank, config)
50
51 + @classmethod
52 + def create_connector_v1(
53 + cls,
54 + config: "VllmConfig",
55 + role: KVConnectorRole,
56 + ) -> KVConnectorBase_V1:
57 + if not envs.VLLM_USE_V1:
58 + raise ValueError("Attempting to initialize a V1 Connector, "
59 + f"but found {envs.VLLM_USE_V1=}")
60 +
61 + connector_name = config.kv_transfer_config.kv_connector
62 + connector_cls = cls._registry[connector_name]()
63 + assert issubclass(connector_cls, KVConnectorBase_V1)
64 + logger.info("Creating v1 connector with name: %s", connector_name)
65 + # NOTE(Kuntai): v1 connector is explicitly separated into two roles.
66 + # Scheduler connector:
67 + # - Co-locate with scheduler process
68 + # - Should only be used inside the Scheduler class
69 + # Worker connector:
70 + # - Co-locate with worker process
71 + # - Should only be used inside the forward context & attention layer
72 + # We build separately to enforce strict separation
73 + return connector_cls(config, role)
74 +
75
76 # Register various connectors here.
77 # The registration should not be done in each individual file, as we want to
@@ -57,4 +94,9 @@ KVConnectorFactory.register_connector(
94 KVConnectorFactory.register_connector(
95 "MooncakeStoreConnector",
96 "vllm.distributed.kv_transfer.kv_connector.mooncake_store_connector",
60 - "MooncakeStoreConnector")
\ No newline at end of file
97 + "MooncakeStoreConnector")
98 +
99 +KVConnectorFactory.register_connector(
100 + "SharedStorageConnector",
101 + "vllm.distributed.kv_transfer.kv_connector.v1.shared_storage_connector",
102 + "SharedStorageConnector")
vllm/distributed/kv_transfer/kv_connector/v1/__init__.py new
+8
@@ -0,0 +1,8 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
3 + KVConnectorBase_V1, KVConnectorRole)
4 +
5 +__all__ = [
6 + "KVConnectorRole",
7 + "KVConnectorBase_V1",
8 +]
vllm/distributed/kv_transfer/kv_connector/v1/base.py new
+209
@@ -0,0 +1,209 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""
3 +KVConnectorBase_V1 Class for Distributed KV Cache & Hidden State
4 +communication in vLLM v1
5 +
6 +The class provides the following primitives:
7 + Scheduler-side: runs in the scheduler, binds metadata, which
8 + is used by the worker-side to load/save KV cache.
9 + get_num_new_matched_tokens() - get number of new tokens
10 + that exist in the remote KV cache
11 + update_state_after_alloc() - update KVConnector state after
12 + temporary buffer alloc by the CacheManager.
13 +
14 + Worker-side: runs in each worker, loads/saves KV cache to/from
15 + the Connector based on the metadata.
16 + start_load_kv() - starts loading all KVs (maybe async)
17 + wait_for_layer_load() - blocks until layer i load is done
18 +
19 + save_kv_layer() - starts saving KV for layer i (maybe async)
20 + wait_for_save() - blocks until all saves are done
21 +"""
22 +
23 +import enum
24 +from abc import ABC, abstractmethod
25 +from dataclasses import dataclass
26 +from typing import TYPE_CHECKING
27 +
28 +import torch
29 +
30 +from vllm.logger import init_logger
31 +from vllm.v1.core.sched.output import SchedulerOutput
32 +
33 +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.request import Request
38 +
39 +logger = init_logger(__name__)
40 +
41 +
42 +class KVConnectorRole(enum.Enum):
43 + # Connector running in the scheduler process
44 + SCHEDULER = 0
45 +
46 + # Connector running in the worker process
47 + WORKER = 1
48 +
49 +
50 +@dataclass
51 +class KVConnectorMetadata:
52 + pass
53 +
54 +
55 +class KVConnectorBase_V1(ABC):
56 +
57 + def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
58 + logger.warning(
59 + "Initializing KVConnectorBase_V1. This API is experimental and "
60 + "subject to change in the future as we iterate the design.")
61 + self._connector_metadata = KVConnectorMetadata()
62 + self._vllm_config = vllm_config
63 + self._role = role
64 +
65 + @property
66 + def role(self) -> KVConnectorRole:
67 + return self._role
68 +
69 + def bind_connector_metadata(
70 + self, connector_metadata: KVConnectorMetadata) -> None:
71 + """Set the connector metadata from the scheduler.
72 +
73 + This function should be called by the model runner every time
74 + before the model execution. The metadata will be used for runtime
75 + KV cache loading and saving.
76 +
77 + Args:
78 + connector_metadata (dict): the connector metadata.
79 + """
80 + self._connector_metadata = connector_metadata
81 +
82 + def clear_connector_metadata(self) -> None:
83 + """Clear the connector metadata.
84 +
85 + This function should be called by the model runner every time
86 + after the model execution.
87 + """
88 + self._connector_metadata = KVConnectorMetadata()
89 +
90 + def _get_connector_metadata(self) -> KVConnectorMetadata:
91 + """Get the connector metadata.
92 +
93 + This function should only be called inside the connector.
94 +
95 + Returns:
96 + ConnectorMetadata: the connector metadata.
97 + """
98 + return self._connector_metadata
99 +
100 + # ==============================
101 + # Worker-side methods
102 + # ==============================
103 +
104 + @abstractmethod
105 + def start_load_kv(self, forward_context: "ForwardContext",
106 + **kwargs) -> None:
107 + """
108 + Start loading the KV cache from the connector to vLLM's paged
109 + KV buffer. This is called from the forward context before the
110 + forward pass to enable async loading during model execution.
111 +
112 + Args:
113 + forward_context (ForwardContext): the forward context.
114 + **kwargs: additional arguments for the load operation
115 +
116 + Note:
117 + The number of elements in kv_caches and layer_names should be
118 + the same.
119 +
120 + """
121 + pass
122 +
123 + @abstractmethod
124 + def wait_for_layer_load(self, layer_name: str) -> None:
125 + """
126 + Block until the KV for a specific layer is loaded into vLLM's
127 + paged buffer. This is called from within attention layer to ensure
128 + async copying from start_load_kv is complete.
129 +
130 + This interface will be useful for layer-by-layer pipelining.
131 +
132 + Args:
133 + layer_name: the name of that layer
134 + """
135 + pass
136 +
137 + @abstractmethod
138 + def save_kv_layer(self, layer_name: str, kv_layer: torch.Tensor,
139 + attn_metadata: "AttentionMetadata", **kwargs) -> None:
140 + """
141 + Start saving a layer of KV cache from vLLM's paged buffer
142 + to the connector. This is called from within attention layer to
143 + enable async copying during execution.
144 +
145 + Args:
146 + layer_name (str): the name of the layer.
147 + kv_layer (torch.Tensor): the paged KV buffer of the current
148 + layer in vLLM.
149 + attn_metadata (AttentionMetadata): the attention metadata.
150 + **kwargs: additional arguments for the save operation.
151 + """
152 + pass
153 +
154 + @abstractmethod
155 + def wait_for_save(self):
156 + """
157 + Block until all the save operations is done. This is called
158 + as the forward context exits to ensure that the async saving
159 + from save_kv_layer is complete before finishing the forward.
160 +
161 + This prevents overwrites of paged KV buffer before saving done.
162 + """
163 + pass
164 +
165 + # ==============================
166 + # Scheduler-side methods
167 + # ==============================
168 + @abstractmethod
169 + def get_num_new_matched_tokens(
170 + self,
171 + request: "Request",
172 + num_computed_tokens: int,
173 + ) -> int:
174 + """
175 + Get number of new tokens that can be loaded from the
176 + external KV cache beyond the num_computed_tokens.
177 +
178 + Args:
179 + request (Request): the request object.
180 + num_computed_tokens (int): the number of locally
181 + computed tokens for this request
182 +
183 + Returns:
184 + the number of tokens that can be loaded from the
185 + external KV cache beyond what is already computed.
186 + """
187 + pass
188 +
189 + @abstractmethod
190 + def update_state_after_alloc(self, request: "Request",
191 + num_external_tokens: int):
192 + """
193 + Update KVConnector state after block allocation.
194 + """
195 + pass
196 +
197 + @abstractmethod
198 + def build_connector_meta(
199 + self, scheduler_output: SchedulerOutput) -> KVConnectorMetadata:
200 + """
201 + Build the connector metadata for this step.
202 +
203 + This function should NOT modify fields in the scheduler_output.
204 + Also, calling this function will reset the state of the connector.
205 +
206 + Args:
207 + scheduler_output (SchedulerOutput): the scheduler output object.
208 + """
209 + pass
vllm/distributed/kv_transfer/kv_connector/v1/shared_storage_connector.py new
+382
@@ -0,0 +1,382 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import hashlib
3 +import os
4 +from dataclasses import dataclass
5 +from typing import TYPE_CHECKING
6 +
7 +import safetensors
8 +import torch
9 +
10 +from vllm.config import VllmConfig
11 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
12 + KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole)
13 +from vllm.logger import init_logger
14 +from vllm.v1.attention.backends.mla.common import MLACommonMetadata
15 +from vllm.v1.core.sched.output import SchedulerOutput
16 +
17 +if TYPE_CHECKING:
18 + from vllm.attention.backends.abstract import AttentionMetadata
19 + from vllm.forward_context import ForwardContext
20 + from vllm.v1.request import Request
21 +
22 +logger = init_logger(__name__)
23 +
24 +
25 +@dataclass
26 +class ReqMeta:
27 + # Request tokens
28 + token_ids: torch.Tensor
29 + # Slot mappings, should have the same length as token_ids
30 + slot_mapping: torch.Tensor
31 + # Is store or load
32 + is_store: bool
33 +
34 + @staticmethod
35 + def make_meta(token_ids: list[int], block_ids: list[int], block_size: int,
36 + is_store: bool) -> "ReqMeta":
37 + valid_num_tokens = align_to_block_size(len(token_ids), block_size)
38 + token_ids_tensor = torch.tensor(token_ids)[:valid_num_tokens]
39 + block_ids_tensor = torch.tensor(block_ids)
40 + num_blocks = block_ids_tensor.shape[0]
41 + block_offsets = torch.arange(0, block_size)
42 + slot_mapping = block_offsets.reshape((1, block_size)) + \
43 + block_ids_tensor.reshape((num_blocks, 1)) * block_size
44 + slot_mapping = slot_mapping.flatten()[:valid_num_tokens]
45 + return ReqMeta(
46 + token_ids=token_ids_tensor,
47 + slot_mapping=slot_mapping,
48 + is_store=is_store,
49 + )
50 +
51 +
52 +@dataclass
53 +class SharedStorageConnectorMetadata(KVConnectorMetadata):
54 + requests: list[ReqMeta]
55 +
56 + def __init__(self):
57 + self.requests = []
58 +
59 + def add_request(
60 + self,
61 + token_ids: list[int],
62 + block_ids: list[int],
63 + block_size: int,
64 + is_store: bool,
65 + ) -> None:
66 + self.requests.append(
67 + ReqMeta.make_meta(token_ids, block_ids, block_size, is_store))
68 +
69 +
70 +class SharedStorageConnector(KVConnectorBase_V1):
71 + # NOTE: This is Simple debug implementation of the KV connector.
72 + # It save / load the KV cache to / from the disk.
73 + # It does extra work which will overwrite the existing prefix-cache in GPU
74 + # - to remove the overhead, need to add some "mask" in the ReqMeta class
75 +
76 + def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
77 + super().__init__(vllm_config=vllm_config, role=role)
78 + self._block_size = vllm_config.cache_config.block_size
79 + self._requests_need_load: dict[str, Request] = {}
80 + transfer_config = vllm_config.kv_transfer_config
81 + self._storage_path = transfer_config.get_from_extra_config(
82 + "shared_storage_path", "/tmp")
83 + logger.info(vllm_config.kv_transfer_config)
84 + logger.info("Shared storage path is %s", self._storage_path)
85 +
86 + def start_load_kv(self, forward_context: "ForwardContext",
87 + **kwargs) -> None:
88 + """Start loading the KV cache from the connector buffer to vLLM's
89 + paged KV buffer.
90 +
91 + Args:
92 + forward_context (ForwardContext): the forward context.
93 + **kwargs: additional arguments for the load operation
94 +
95 + Note:
96 + The number of elements in kv_caches and layer_names should be
97 + the same.
98 + """
99 + attn_metadata = forward_context.attn_metadata
100 +
101 + def inject_kv_into_layer(
102 + dst_kv_cache_layer: torch.Tensor,
103 + src_kv_cache: torch.Tensor,
104 + slot_mapping: torch.Tensor,
105 + ) -> None:
106 + """Inject the KV cache into the layer.
107 +
108 + Args:
109 + dst_kv_cache_layer (torch.Tensor): the destination KV cache
110 + layer. In shape [2, num_pages, page_size, xxx] if not
111 + using MLA, [num_pages, page_size, xxx] otherwise.
112 + src_kv_cache (torch.Tensor): the source KV cache. In shape
113 + [2, num_tokens, xxx] if not using MLA, [num_tokens, xxx]
114 + otherwise.
115 + slot_mapping (torch.Tensor): the slot mapping. In shape
116 + [num_tokens].
117 + """
118 + dst_kv_cache_layer_shape = dst_kv_cache_layer.shape
119 + if isinstance(attn_metadata, MLACommonMetadata):
120 + num_pages = dst_kv_cache_layer_shape[0]
121 + page_size = dst_kv_cache_layer_shape[1]
122 + dst_kv_cache_layer = dst_kv_cache_layer.reshape(
123 + num_pages * page_size, -1)
124 + dst_kv_cache_layer[slot_mapping, ...] = src_kv_cache
125 + dst_kv_cache_layer.reshape(dst_kv_cache_layer_shape)
126 + else:
127 + num_pages = dst_kv_cache_layer_shape[1]
128 + page_size = dst_kv_cache_layer_shape[2]
129 + dst_kv_cache_layer = dst_kv_cache_layer.reshape(
130 + 2, num_pages * page_size, -1)
131 + dst_kv_cache_layer[:, slot_mapping, ...] = src_kv_cache
132 + dst_kv_cache_layer.reshape(dst_kv_cache_layer_shape)
133 +
134 + # Get the metadata
135 + metadata: KVConnectorMetadata = \
136 + self._get_connector_metadata()
137 + assert isinstance(metadata, SharedStorageConnectorMetadata)
138 +
139 + if metadata is None:
140 + logger.warning(
141 + "In connector.start_load_kv, but the connector metadata is None"
142 + )
143 + return
144 +
145 + attn_metadata = forward_context.attn_metadata
146 + if attn_metadata is None:
147 + logger.warning(
148 + "In connector.start_load_kv, but the attn_metadata is None")
149 + return
150 +
151 + # Load the KV for each request each layer
152 + for request in metadata.requests:
153 + if request.is_store:
154 + continue
155 + logger.info("Inject KV cache of %d tokens to the paged memory",
156 + len(request.slot_mapping))
157 + for layer_name in forward_context.no_compile_layers:
158 + attn_layer = forward_context.no_compile_layers[layer_name]
159 + kv_cache_layer = attn_layer.kv_cache[\
160 + forward_context.virtual_engine]
161 +
162 + filename = self._generate_filename_debug(
163 + layer_name, request.token_ids)
164 + kv_cache = safetensors.torch.load_file(
165 + filename)["kv_cache"].cuda()
166 + inject_kv_into_layer(kv_cache_layer, kv_cache,
167 + request.slot_mapping)
168 +
169 + def wait_for_layer_load(self, layer_name: str) -> None:
170 + """Blocking until the KV for a specific layer is loaded into vLLM's
171 + paged buffer.
172 +
173 + This interface will be useful for layer-by-layer pipelining.
174 +
175 + Args:
176 + layer_name: the name of that layer
177 + """
178 + return
179 +
180 + def save_kv_layer(self, layer_name: str, kv_layer: torch.Tensor,
181 + attn_metadata: "AttentionMetadata", **kwargs) -> None:
182 + """Start saving the KV cache of the layer from vLLM's paged buffer
183 + to the connector.
184 +
185 + Args:
186 + layer_name (str): the name of the layer.
187 + kv_layer (torch.Tensor): the paged KV buffer of the current
188 + layer in vLLM.
189 + attn_metadata (AttentionMetadata): the attention metadata.
190 + **kwargs: additional arguments for the save operation.
191 + """
192 +
193 + def extract_kv_from_layer(
194 + layer: torch.Tensor,
195 + slot_mapping: torch.Tensor,
196 + ) -> torch.Tensor:
197 + """Extract the KV cache from the layer.
198 +
199 + Assume the shape of the layer is (2, num_pages, page_size, xxx)
200 + if MLA is not used, and (num_pages, page_size, xxx) otherwise.
201 + """
202 + if isinstance(attn_metadata, MLACommonMetadata):
203 + num_pages, page_size = layer.shape[0], layer.shape[1]
204 + return layer.reshape(num_pages * page_size, -1)[slot_mapping,
205 + ...]
206 + num_pages, page_size = layer.shape[1], layer.shape[2]
207 + return layer.reshape(2, num_pages * page_size, -1)[:, slot_mapping,
208 + ...]
209 +
210 + connector_metadata = self._get_connector_metadata()
211 + assert isinstance(connector_metadata, SharedStorageConnectorMetadata)
212 + for request in connector_metadata.requests:
213 + if request.is_store:
214 + filename = self._generate_filename_debug(
215 + layer_name, request.token_ids)
216 + kv_cache = extract_kv_from_layer(kv_layer,
217 + request.slot_mapping)
218 + tensors = {"kv_cache": kv_cache.detach().cpu()}
219 + safetensors.torch.save_file(tensors, filename)
220 +
221 + def wait_for_save(self):
222 + return
223 +
224 + def get_num_new_matched_tokens(
225 + self,
226 + request: "Request",
227 + num_computed_tokens: int,
228 + ) -> int:
229 + """
230 + Get number of new tokens that can be loaded from the
231 + external KV cache beyond the num_computed_tokens.
232 +
233 + Args:
234 + request (Request): the request object.
235 + num_computed_tokens (int): the number of locally
236 + computed tokens for this request
237 +
238 + Returns:
239 + the number of tokens that can be loaded from the
240 + external KV cache beyond what is already computed.
241 + """
242 +
243 + # NOTE: in this debug implementation, we assume that the prompt is
244 + # cached_prompt + newly_generated_single_token
245 + # Therefore, we use prompt_token_ids[:-1] to determine the folder name
246 +
247 + # NOTE: in current v1 scheduler, the num_computed_tokens is aligned
248 + # with the block granularity. And it expects the returned blocks and
249 + # num_computed_tokens to also be aligned with the block granularity.
250 + if not self._found_match_for_request(request):
251 + return 0
252 +
253 + logger.info("External Cache Hit!")
254 +
255 + # Now, first num_tokens_to_check tokens are hit, we need to prepare
256 + # the metadata for the worker connector to correctly load the KV
257 + num_tokens_to_check = align_to_block_size(
258 + len(request.prompt_token_ids) - 1, self._block_size)
259 +
260 + return num_tokens_to_check - num_computed_tokens
261 +
262 + def update_state_after_alloc(self, request: "Request",
263 + num_external_tokens: int):
264 + """
265 + Update KVConnector state after block allocation.
266 +
267 + If blocks were allocated, add to _requests_need_load,
268 + such that we load the KVs in the next forward pass.
269 + """
270 + if num_external_tokens > 0:
271 + self._requests_need_load[request.request_id] = request
272 +
273 + def build_connector_meta(
274 + self,
275 + scheduler_output: SchedulerOutput,
276 + ) -> KVConnectorMetadata:
277 + """Build the connector metadata for this step.
278 +
279 + This function should NOT modify any fields in the scheduler_output.
280 + Also, calling this function will reset the state of the connector.
281 +
282 + Args:
283 + scheduler_output (SchedulerOutput): the scheduler output object.
284 + """
285 + meta = SharedStorageConnectorMetadata()
286 +
287 + total_need_load = 0
288 + for new_req in scheduler_output.scheduled_new_reqs:
289 + if new_req.req_id in self._requests_need_load:
290 + meta.add_request(token_ids=new_req.prompt_token_ids,
291 + block_ids=new_req.block_ids,
292 + block_size=self._block_size,
293 + is_store=False)
294 + total_need_load += 1
295 + else:
296 + # NOTE: here, we set the store and load being exclusive,
297 + # but a single request can have both store and load.
298 + # NOTE(rob): for this debug implementation, we only cache
299 + # the original prompt tokens.
300 + if not self._found_match_for_request(new_req):
301 + meta.add_request(token_ids=new_req.prompt_token_ids,
302 + block_ids=new_req.block_ids,
303 + block_size=self._block_size,
304 + is_store=True)
305 +
306 + for cached_req in scheduler_output.scheduled_cached_reqs:
307 + # NOTE(rob): here we rely on the resumed requests being
308 + # the first N requests in the list scheduled_cache_reqs.
309 + if not cached_req.resumed_from_preemption:
310 + break
311 + if cached_req.req_id in self._requests_need_load:
312 + # NOTE(rob): cached_req_data does not have the full
313 + # list of token ids (only new tokens). So we look it
314 + # up in the actual request object.
315 + request = self._requests_need_load[cached_req.req_id]
316 + total_tokens = (len(cached_req.new_token_ids) +
317 + cached_req.num_computed_tokens)
318 + token_ids = request.all_token_ids[:total_tokens]
319 +
320 + # NOTE(rob): For resumed req, new_block_ids is all
321 + # of the block_ids for the request.
322 + block_ids = cached_req.new_block_ids
323 +
324 + meta.add_request(token_ids=token_ids,
325 + block_ids=block_ids,
326 + block_size=self._block_size,
327 + is_store=False)
328 + total_need_load += 1
329 +
330 + assert total_need_load == len(self._requests_need_load)
331 + self._requests_need_load.clear()
332 + return meta
333 +
334 + # ==============================
335 + # Helper functions
336 + # ==============================
337 +
338 + def _found_match_for_request(
339 + self,
340 + request: "Request",
341 + ) -> bool:
342 + """Check if the cache is hit for the request.
343 + """
344 + num_tokens_to_check = align_to_block_size(
345 + len(request.prompt_token_ids) - 1, self._block_size)
346 + foldername = self._generate_foldername_debug(torch.tensor(
347 + request.prompt_token_ids)[:num_tokens_to_check],
348 + create_folder=False)
349 + return os.path.exists(foldername)
350 +
351 + def _generate_foldername_debug(
352 + self,
353 + input_ids: torch.Tensor,
354 + create_folder=False,
355 + ) -> str:
356 + """Generate a folder name based on the hash of the bytes of the input
357 + ids.
358 + """
359 + input_ids_bytes = input_ids.numpy().tobytes()
360 + input_ids_hash = hashlib.md5(input_ids_bytes).hexdigest()
361 + foldername = os.path.join(self._storage_path, input_ids_hash)
362 + if create_folder:
363 + os.makedirs(foldername, exist_ok=True)
364 + return foldername
365 +
366 + def _generate_filename_debug(
367 + self,
368 + layer_name: str,
369 + input_ids: torch.Tensor,
370 + ) -> str:
371 + """Generate a file name based on the layer name and the hash
372 + of the bytes of the input ids.
373 + """
374 + foldername = self._generate_foldername_debug(input_ids,
375 + create_folder=True)
376 + return os.path.join(foldername, f"{layer_name}.safetensors")
377 +
378 +
379 +def align_to_block_size(num_tokens: int, block_size) -> int:
380 + """Align the number of tokens to the block size.
381 + """
382 + return (num_tokens - 1) // block_size * block_size
vllm/distributed/kv_transfer/kv_connector_agent.py renamed
+1 -1
@@ -46,7 +46,7 @@ class KVTransferAgent:
46 assert self.config.kv_transfer_config.is_kv_transfer_instance, "KV"\
47 "TransferAgent should only be used when kv_connector is set."
48
49 - self.connector = KVConnectorFactory.create_connector(
49 + self.connector = KVConnectorFactory.create_connector_v0(
50 rank, local_rank, config)
51
52 def send_kv_caches_and_hidden_states(
vllm/distributed/kv_transfer/kv_transfer_state.py new
+70
@@ -0,0 +1,70 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import TYPE_CHECKING, Optional
3 +
4 +from vllm import envs
5 +from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBaseType
6 +from vllm.distributed.kv_transfer.kv_connector.factory import (
7 + KVConnectorFactory)
8 +from vllm.distributed.kv_transfer.kv_connector.v1 import (KVConnectorBase_V1,
9 + KVConnectorRole)
10 +from vllm.distributed.parallel_state import get_world_group
11 +
12 +if TYPE_CHECKING:
13 + from vllm.config import VllmConfig
14 +
15 +_KV_CONNECTOR_AGENT: Optional[KVConnectorBaseType] = None
16 +
17 +
18 +def get_kv_transfer_group() -> KVConnectorBaseType:
19 + assert _KV_CONNECTOR_AGENT is not None, (
20 + "disaggregated KV cache transfer parallel group is not initialized")
21 + return _KV_CONNECTOR_AGENT
22 +
23 +
24 +def has_kv_transfer_group() -> bool:
25 + return _KV_CONNECTOR_AGENT is not None
26 +
27 +
28 +def is_v1_kv_transfer_group(
29 + connector: Optional[KVConnectorBaseType] = None) -> bool:
30 + """Check if the KV connector is the v1 connector.
31 + If the argument is None, it will check the global KV connector
32 +
33 + Args:
34 + connector: The KV connector to check. If None, it will check the
35 + global KV connector.
36 +
37 + Note:
38 + This function will no-longer be needed after the v1 KV connector
39 + becomes the default.
40 + """
41 + if connector is None:
42 + connector = _KV_CONNECTOR_AGENT
43 +
44 + if connector is None:
45 + return False
46 +
47 + return isinstance(connector, KVConnectorBase_V1)
48 +
49 +
50 +def ensure_kv_transfer_initialized(vllm_config: "VllmConfig") -> None:
51 + """
52 + Initialize KV cache transfer parallel group.
53 + """
54 +
55 + global _KV_CONNECTOR_AGENT
56 +
57 + if vllm_config.kv_transfer_config is None:
58 + return
59 +
60 + if (vllm_config.kv_transfer_config.is_kv_transfer_instance
61 + and _KV_CONNECTOR_AGENT is None):
62 + if envs.VLLM_USE_V1:
63 + _KV_CONNECTOR_AGENT = KVConnectorFactory.create_connector_v1(
64 + config=vllm_config, role=KVConnectorRole.WORKER)
65 + else:
66 + _KV_CONNECTOR_AGENT = KVConnectorFactory.create_connector_v0(
67 + rank=get_world_group().rank,
68 + local_rank=get_world_group().local_rank,
69 + config=vllm_config,
70 + )
vllm/distributed/parallel_state.py
+1 -34
@@ -29,15 +29,13 @@ from collections import namedtuple
29 from contextlib import contextmanager, nullcontext
30 from dataclasses import dataclass
31 from multiprocessing import shared_memory
32 -from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple,
33 - Union)
32 +from typing import Any, Callable, Dict, List, Optional, Tuple, Union
33 from unittest.mock import patch
34
35 import torch
36 import torch.distributed
37 from torch.distributed import Backend, ProcessGroup
38
40 -import vllm.distributed.kv_transfer.kv_transfer_agent as kv_transfer
39 import vllm.envs as envs
40 from vllm.distributed.device_communicators.base_device_communicator import (
41 DeviceCommunicatorBase)
@@ -46,9 +44,6 @@ from vllm.logger import init_logger
44 from vllm.utils import (direct_register_custom_op, resolve_obj_by_qualname,
45 supports_custom_op)
46
49 -if TYPE_CHECKING:
50 - from vllm.config import VllmConfig
51 -
47
48 @dataclass
49 class GraphCaptureContext:
@@ -772,14 +767,6 @@ def get_pp_group() -> GroupCoordinator:
767 # kept for backward compatibility
768 get_pipeline_model_parallel_group = get_pp_group
769
775 -_KV_TRANSFER: Optional[kv_transfer.KVTransferAgent] = None
776 -
777 -
778 -def get_kv_transfer_group() -> kv_transfer.KVTransferAgent:
779 - assert _KV_TRANSFER is not None, (
780 - "disaggregated KV cache transfer parallel group is not initialized")
781 - return _KV_TRANSFER
782 -
770
771 @contextmanager
772 def graph_capture(device: torch.device):
@@ -962,26 +949,6 @@ def initialize_model_parallel(
949 _DP.rank_in_group, _PP.rank_in_group, _TP.rank_in_group)
950
951
965 -def ensure_kv_transfer_initialized(vllm_config: "VllmConfig") -> None:
966 - """
967 - Initialize KV cache transfer parallel group.
968 - """
969 -
970 - global _KV_TRANSFER
971 -
972 - if vllm_config.kv_transfer_config is None:
973 - return
974 -
975 - if all([
976 - vllm_config.kv_transfer_config.is_kv_transfer_instance,
977 - _KV_TRANSFER is None
978 - ]):
979 - _KV_TRANSFER = kv_transfer.KVTransferAgent(
980 - rank=get_world_group().rank,
981 - local_rank=get_world_group().local_rank,
982 - config=vllm_config)
983 -
984 -
952 def ensure_model_parallel_initialized(
953 tensor_model_parallel_size: int,
954 pipeline_model_parallel_size: int,
vllm/engine/arg_utils.py
-6
@@ -1527,12 +1527,6 @@ class EngineArgs:
1527 recommend_to_remove=False)
1528 return False
1529
1530 - # No Disaggregated Prefill so far.
1531 - if self.kv_transfer_config != EngineArgs.kv_transfer_config:
1532 - _raise_or_fallback(feature_name="--kv-transfer-config",
1533 - recommend_to_remove=False)
1534 - return False
1535 -
1530 # No FlashInfer or XFormers so far.
1531 V1_BACKENDS = [
1532 "FLASH_ATTN_VLLM_V1", "FLASH_ATTN", "PALLAS", "PALLAS_VLLM_V1",
vllm/forward_context.py
+23
@@ -11,6 +11,10 @@ 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
18 from vllm.logger import init_logger
19
20 if TYPE_CHECKING:
@@ -98,6 +102,17 @@ def set_forward_context(attn_metadata: Any,
102 virtual_engine=virtual_engine,
103 attn_metadata=attn_metadata,
104 dp_metadata=dp_metadata)
105 +
106 + # KVConnector: trigger (possibly async) load before forward.
107 + # Each attn layer will block until the reading is complete.
108 + trigger_kv_transfer = (attn_metadata is not None
109 + and has_kv_transfer_group()
110 + and is_v1_kv_transfer_group())
111 + if trigger_kv_transfer:
112 + kv_connector = get_kv_transfer_group()
113 + assert isinstance(kv_connector, KVConnectorBase_V1)
114 + kv_connector.start_load_kv(_forward_context)
115 +
116 try:
117 yield
118 finally:
@@ -133,4 +148,12 @@ def set_forward_context(attn_metadata: Any,
148 logger.info(("Batchsize forward time stats "
149 "(batchsize, count, median_time(ms)): %s"),
150 forward_stats)
151 +
152 + # KVConnector: each attn layer triggers (possibly async) save.
153 + # Ensure all those operations complete before forward() is done.
154 + if trigger_kv_transfer:
155 + kv_connector = get_kv_transfer_group()
156 + assert isinstance(kv_connector, KVConnectorBase_V1)
157 + kv_connector.wait_for_save()
158 +
159 _forward_context = prev_context
vllm/v1/core/kv_cache_manager.py
+3 -2
@@ -171,8 +171,9 @@ class KVCacheManager:
171
172 Args:
173 request: The request to allocate slots.
174 - num_tokens: The number of tokens to allocate. Note that this does
175 - not include the tokens that have already been computed.
174 + num_tokens: The number of tokens to allocate, including external
175 + tokens. Note that this does not include tokens that have
176 + already been computed locally (i.e. new_computed_blocks).
177 new_computed_blocks: A list of new computed blocks just hitting the
178 prefix caching.
179 num_lookahead_tokens: The number of speculative tokens to allocate.
vllm/v1/core/sched/output.py
+5
@@ -9,6 +9,8 @@ if TYPE_CHECKING:
9 import numpy as np
10 import numpy.typing as npt
11
12 + from vllm.distributed.kv_transfer.kv_connector.v1.base import (
13 + KVConnectorMetadata)
14 from vllm.lora.request import LoRARequest
15 from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
16 from vllm.sampling_params import SamplingParams
@@ -121,3 +123,6 @@ class SchedulerOutput:
123 structured_output_request_ids: dict[str, int]
124 # the bitmask for the whole batch
125 grammar_bitmask: Optional[npt.NDArray[np.int32]]
126 +
127 + # KV Cache Connector metadata.
128 + kv_connector_metadata: Optional[KVConnectorMetadata] = None
vllm/v1/core/sched/scheduler.py
+56 -14
@@ -7,8 +7,10 @@ from collections import deque
7 from collections.abc import Iterable
8 from typing import Optional, Union
9
10 -from vllm.config import (CacheConfig, LoRAConfig, ModelConfig, SchedulerConfig,
11 - SpeculativeConfig)
10 +from vllm.config import VllmConfig
11 +from vllm.distributed.kv_transfer.kv_connector.factory import (
12 + KVConnectorFactory)
13 +from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
14 from vllm.logger import init_logger
15 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
16 from vllm.v1.core.encoder_cache_manager import (EncoderCacheManager,
@@ -34,20 +36,17 @@ class Scheduler(SchedulerInterface):
36
37 def __init__(
38 self,
37 - scheduler_config: SchedulerConfig,
38 - model_config: ModelConfig,
39 - cache_config: CacheConfig,
40 - lora_config: Optional[LoRAConfig],
39 + vllm_config: VllmConfig,
40 kv_cache_config: KVCacheConfig,
41 structured_output_manager: StructuredOutputManager,
43 - speculative_config: SpeculativeConfig = None,
42 mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
43 include_finished_set: bool = False,
44 log_stats: bool = False,
45 ) -> None:
48 - self.scheduler_config = scheduler_config
49 - self.cache_config = cache_config
50 - self.lora_config = lora_config
46 + self.vllm_config = vllm_config
47 + self.scheduler_config = vllm_config.scheduler_config
48 + self.cache_config = vllm_config.cache_config
49 + self.lora_config = vllm_config.lora_config
50 self.kv_cache_config = kv_cache_config
51 self.log_stats = log_stats
52 self.structured_output_manager = structured_output_manager
@@ -64,11 +63,22 @@ class Scheduler(SchedulerInterface):
63 self.scheduler_config.max_num_batched_tokens
64 self.max_model_len = self.scheduler_config.max_model_len
65
66 + # Create KVConnector for the Scheduler. Note that each Worker
67 + # will have a corresponding KVConnector with Role=WORKER.
68 + # KV Connector pushes/pull of remote KVs for P/D and offloading.
69 + self.connector = None
70 + if self.vllm_config.kv_transfer_config is not None:
71 + self.connector = KVConnectorFactory.create_connector_v1(
72 + config=self.vllm_config, role=KVConnectorRole.SCHEDULER)
73 +
74 + num_gpu_blocks = self.cache_config.num_gpu_blocks
75 + assert num_gpu_blocks is not None and num_gpu_blocks > 0
76 +
77 # Create the KV cache manager.
78 self.kv_cache_manager = KVCacheManager(
79 kv_cache_config=kv_cache_config,
80 max_model_len=self.max_model_len,
71 - enable_caching=cache_config.enable_prefix_caching,
81 + enable_caching=self.cache_config.enable_prefix_caching,
82 caching_hash_algo=self.cache_config.prefix_caching_hash_algo,
83 log_stats=self.log_stats)
84 self.block_size = self.cache_config.block_size
@@ -99,8 +109,8 @@ class Scheduler(SchedulerInterface):
109 # This can be changed when we make encoder cache for embedding caching
110 # across requests.
111 encoder_compute_budget, encoder_cache_size = compute_encoder_budget(
102 - model_config=model_config,
103 - scheduler_config=scheduler_config,
112 + model_config=vllm_config.model_config,
113 + scheduler_config=vllm_config.scheduler_config,
114 mm_registry=mm_registry,
115 )
116
@@ -115,6 +125,7 @@ class Scheduler(SchedulerInterface):
125 cache_size=encoder_cache_size)
126
127 self.num_lookahead_tokens = 0
128 + speculative_config = vllm_config.speculative_config
129 if speculative_config and speculative_config.method == "eagle":
130 self.num_lookahead_tokens = \
131 speculative_config.num_speculative_tokens
@@ -304,6 +315,16 @@ class Scheduler(SchedulerInterface):
315 # Get already-cached tokens.
316 computed_blocks, num_computed_tokens = \
317 self.kv_cache_manager.get_computed_blocks(request)
318 +
319 + # Get externally-cached tokens if using a KVConnector.
320 + num_external_tokens = (
321 + 0 if self.connector is None else
322 + self.connector.get_num_new_matched_tokens(
323 + request, num_computed_tokens))
324 +
325 + # Total computed tokens (local + external).
326 + num_computed_tokens += num_external_tokens
327 +
328 # Number of tokens to be scheduled.
329 # We use `request.num_tokens` instead of
330 # `request.num_prompt_tokens` to consider the resumed requests,
@@ -330,11 +351,21 @@ class Scheduler(SchedulerInterface):
351 new_encoder_budget = encoder_budget
352
353 new_blocks = self.kv_cache_manager.allocate_slots(
333 - request, num_new_tokens, computed_blocks)
354 + request, num_new_tokens + num_external_tokens,
355 + computed_blocks)
356 if new_blocks is None:
357 # The request cannot be scheduled.
358 break
359
360 + # KVConnector: update internal state after allocation.
361 + # This information is used to determine if a load is
362 + # needed for this request.
363 + if self.connector is not None:
364 + self.connector.update_state_after_alloc(
365 + request,
366 + num_external_tokens,
367 + )
368 +
369 self.waiting.popleft()
370 if request.use_structured_output:
371 structured_output_request_ids[
@@ -443,6 +474,14 @@ class Scheduler(SchedulerInterface):
474 grammar_bitmask=grammar_bitmask,
475 )
476
477 + # NOTE(Kuntai): this function is designed for multiple purposes:
478 + # 1. Plan the KV cache store
479 + # 2. Wrap up all the KV cache load / save ops into an opaque object
480 + # 3. Clear the internal states of the connector
481 + if self.connector is not None:
482 + meta = self.connector.build_connector_meta(scheduler_output)
483 + scheduler_output.kv_connector_metadata = meta
484 +
485 # Advance the number of computed tokens for the request AFTER
486 # the request is scheduled.
487 # 1. The scheduler_output of the current step has to include the
@@ -508,6 +547,9 @@ class Scheduler(SchedulerInterface):
547 If an encoder input cannot be scheduled due to cache or budget
548 limitations, the method adjusts `num_new_tokens` to schedule only the
549 decoder tokens up to just before the unschedulable encoder input.
550 +
551 + Note that num_computed_tokens includes both locally cached
552 + blocks and externally cached blocks (via KVConnector).
553 """
554 encoder_inputs_to_schedule: list[int] = []
555 mm_positions = request.mm_positions
vllm/v1/engine/core.py
+1 -5
@@ -92,12 +92,8 @@ class EngineCore:
92 vllm_config.scheduler_config.scheduler_cls)
93
94 self.scheduler: SchedulerInterface = Scheduler(
95 - scheduler_config=vllm_config.scheduler_config,
96 - model_config=vllm_config.model_config,
97 - cache_config=vllm_config.cache_config,
98 - lora_config=vllm_config.lora_config,
95 + vllm_config=vllm_config,
96 kv_cache_config=kv_cache_config,
100 - speculative_config=vllm_config.speculative_config,
97 structured_output_manager=self.structured_output_manager,
98 include_finished_set=vllm_config.parallel_config.data_parallel_size
99 > 1,
vllm/v1/worker/gpu_model_runner.py
+11
@@ -13,6 +13,8 @@ import torch.nn as nn
13 from vllm.attention import AttentionType, get_attn_backend
14 from vllm.attention.layer import Attention
15 from vllm.config import CompilationLevel, VllmConfig
16 +from vllm.distributed.kv_transfer import (get_kv_transfer_group,
17 + has_kv_transfer_group)
18 from vllm.distributed.parallel_state import get_pp_group, graph_capture
19 from vllm.forward_context import set_forward_context
20 from vllm.logger import init_logger
@@ -987,6 +989,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
989 scheduler_output: "SchedulerOutput",
990 intermediate_tensors: Optional[IntermediateTensors] = None,
991 ) -> Union[ModelRunnerOutput, torch.Tensor]:
992 + # Update KVConnector with the KVConnector metadata forward().
993 + if has_kv_transfer_group():
994 + get_kv_transfer_group().bind_connector_metadata(
995 + scheduler_output.kv_connector_metadata)
996 +
997 self._update_states(scheduler_output)
998 if not scheduler_output.total_num_scheduled_tokens:
999 # Return empty ModelRunnerOutput if there's no work to do.
@@ -1228,6 +1235,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1235 # in the next step.
1236 del draft_probs
1237
1238 + # Clear KVConnector state after all KVs are generated.
1239 + if has_kv_transfer_group():
1240 + get_kv_transfer_group().clear_connector_metadata()
1241 +
1242 return ModelRunnerOutput(
1243 req_ids=self.input_batch.req_ids,
1244 req_id_to_index=self.input_batch.req_id_to_index,
vllm/v1/worker/gpu_worker.py
+7 -3
@@ -9,11 +9,12 @@ import torch.distributed
9 import torch.nn as nn
10
11 import vllm.envs as envs
12 -from vllm.config import ParallelConfig, VllmConfig
12 +from vllm.config import VllmConfig
13 from vllm.device_allocator.cumem import CuMemAllocator
14 from vllm.distributed import (ensure_model_parallel_initialized,
15 init_distributed_environment,
16 set_custom_all_reduce)
17 +from vllm.distributed.kv_transfer import ensure_kv_transfer_initialized
18 from vllm.distributed.parallel_state import get_pp_group
19 from vllm.logger import init_logger
20 from vllm.lora.request import LoRARequest
@@ -110,7 +111,7 @@ class Worker(WorkerBase):
111 raise RuntimeError(
112 f"Not support device type: {self.device_config.device}")
113 # Initialize the distributed environment.
113 - init_worker_distributed_environment(self.parallel_config, self.rank,
114 + init_worker_distributed_environment(self.vllm_config, self.rank,
115 self.distributed_init_method,
116 self.local_rank)
117 # Set random seed.
@@ -285,12 +286,13 @@ class Worker(WorkerBase):
286
287
288 def init_worker_distributed_environment(
288 - parallel_config: ParallelConfig,
289 + vllm_config: VllmConfig,
290 rank: int,
291 distributed_init_method: Optional[str] = None,
292 local_rank: int = -1,
293 ) -> None:
294 """Initialize the distributed environment."""
295 + parallel_config = vllm_config.parallel_config
296 set_custom_all_reduce(not parallel_config.disable_custom_all_reduce)
297
298 init_distributed_environment(parallel_config.world_size, rank,
@@ -299,6 +301,8 @@ def init_worker_distributed_environment(
301 ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
302 parallel_config.pipeline_parallel_size)
303
304 + ensure_kv_transfer_initialized(vllm_config)
305 +
306
307 def _check_if_gpu_supports_dtype(torch_dtype: torch.dtype):
308 # Check if the GPU supports the dtype.
vllm/worker/model_runner.py
+2 -1
@@ -23,7 +23,8 @@ from vllm.attention.backends.abstract import AttentionState
23 from vllm.attention.backends.utils import CommonAttentionState
24 from vllm.config import CompilationLevel, VllmConfig
25 from vllm.core.scheduler import SchedulerOutputs
26 -from vllm.distributed import get_kv_transfer_group, get_pp_group
26 +from vllm.distributed import get_pp_group
27 +from vllm.distributed.kv_transfer import get_kv_transfer_group
28 from vllm.distributed.parallel_state import (get_tensor_model_parallel_rank,
29 graph_capture)
30 from vllm.forward_context import get_forward_context, set_forward_context
vllm/worker/worker.py
+2 -2
@@ -10,10 +10,10 @@ import torch.distributed
10 import vllm.envs as envs
11 from vllm.config import VllmConfig
12 from vllm.device_allocator.cumem import CuMemAllocator
13 -from vllm.distributed import (ensure_kv_transfer_initialized,
14 - ensure_model_parallel_initialized,
13 +from vllm.distributed import (ensure_model_parallel_initialized,
14 init_distributed_environment,
15 set_custom_all_reduce)
16 +from vllm.distributed.kv_transfer import ensure_kv_transfer_initialized
17 from vllm.logger import init_logger
18 from vllm.lora.request import LoRARequest
19 from vllm.model_executor import set_random_seed