A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 11:00 UTC
Construct `KVTransferConfig` properly from Python instead of using JSON blobs without CLI (#17994)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
May 12, 2025 at 19:25 UTC
72a3f6b898d8397c406debc49593e15aa7cbb4bc
5 files changed
+37
-31
examples/lmcache/disagg_prefill_lmcache_v0.py
+8
-6
@@ -49,9 +49,10 @@ def run_prefill(prefill_done, prompts):
49
50
sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1)
51
52
- ktc = KVTransferConfig.from_cli(
53
- '{"kv_connector":"LMCacheConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2}'
54
- )
52
+ ktc = KVTransferConfig(kv_connector="LMCacheConnector",
53
+ kv_role="kv_producer",
54
+ kv_rank=0,
55
+ kv_parallel_size=2)
56
# Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
57
# memory. Reduce the value if your GPU has less memory.
58
llm = LLM(model="mistralai/Mistral-7B-Instruct-v0.2",
@@ -78,9 +79,10 @@ def run_decode(prefill_done, prompts, timeout=1):
79
80
sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
81
81
- ktc = KVTransferConfig.from_cli(
82
- '{"kv_connector":"LMCacheConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2}'
83
- )
82
+ ktc = KVTransferConfig(kv_connector="LMCacheConnector",
83
+ kv_role="kv_consumer",
84
+ kv_rank=1,
85
+ kv_parallel_size=2)
86
# Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
87
# of memory. Reduce the value if your GPU has less memory.
88
llm = LLM(model="mistralai/Mistral-7B-Instruct-v0.2",
examples/lmcache/kv_cache_sharing_lmcache_v1.py
+4
-4
@@ -49,8 +49,8 @@ def run_store(store_done, prompts):
49
50
sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
51
52
- ktc = KVTransferConfig.from_cli(
53
- '{"kv_connector":"LMCacheConnectorV1", "kv_role":"kv_both"}')
52
+ ktc = KVTransferConfig(kv_connector="LMCacheConnectorV1",
53
+ kv_role="kv_both")
54
# Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
55
# memory. Reduce the value if your GPU has less memory.
56
llm = LLM(model="mistralai/Mistral-7B-Instruct-v0.2",
@@ -76,8 +76,8 @@ def run_retrieve(store_done, prompts, timeout=1):
76
77
sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
78
79
- ktc = KVTransferConfig.from_cli(
80
- '{"kv_connector":"LMCacheConnectorV1", "kv_role":"kv_both"}')
79
+ ktc = KVTransferConfig(kv_connector="LMCacheConnectorV1",
80
+ kv_role="kv_both")
81
# Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
82
# of memory. Reduce the value if your GPU has less memory.
83
llm = LLM(model="mistralai/Mistral-7B-Instruct-v0.2",
examples/offline_inference/disaggregated-prefill-v1/decode_example.py
+11
-10
@@ -16,16 +16,17 @@ except FileNotFoundError:
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)
19
+llm = LLM(model="meta-llama/Llama-3.2-1B-Instruct",
20
+ enforce_eager=True,
21
+ gpu_memory_utilization=0.8,
22
+ max_num_batched_tokens=64,
23
+ max_num_seqs=16,
24
+ kv_transfer_config=KVTransferConfig(
25
+ kv_connector="SharedStorageConnector",
26
+ kv_role="kv_both",
27
+ kv_connector_extra_config={
28
+ "shared_storage_path": "local_storage"
29
+ })) #, max_model_len=2048, max_num_batched_tokens=2048)
30
31
# 1ST generation (prefill instance)
32
outputs = llm.generate(prompts, sampling_params)
examples/offline_inference/disaggregated-prefill-v1/prefill_example.py
+6
-5
@@ -17,11 +17,12 @@ sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1)
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)
20
+ kv_transfer_config=KVTransferConfig(
21
+ kv_connector="SharedStorageConnector",
22
+ kv_role="kv_both",
23
+ kv_connector_extra_config={
24
+ "shared_storage_path": "local_storage"
25
+ })) #, max_model_len=2048, max_num_batched_tokens=2048)
26
27
# 1ST generation (prefill instance)
28
outputs = llm.generate(
examples/offline_inference/disaggregated_prefill.py
+8
-6
@@ -32,9 +32,10 @@ def run_prefill(prefill_done):
32
# This instance is the prefill node (kv_producer, rank 0).
33
# The number of parallel instances for KV cache transfer is set to 2,
34
# as required for PyNcclConnector.
35
- ktc = KVTransferConfig.from_cli(
36
- '{"kv_connector":"PyNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2}'
37
- )
35
+ ktc = KVTransferConfig(kv_connector="PyNcclConnector",
36
+ kv_role="kv_producer",
37
+ kv_rank=0,
38
+ kv_parallel_size=2)
39
40
# Set GPU memory utilization to 0.8 for an A6000 GPU with 40GB
41
# memory. You may need to adjust the value to fit your GPU.
@@ -71,9 +72,10 @@ def run_decode(prefill_done):
72
# This instance is the decode node (kv_consumer, rank 1).
73
# The number of parallel instances for KV cache transfer is set to 2,
74
# as required for PyNcclConnector.
74
- ktc = KVTransferConfig.from_cli(
75
- '{"kv_connector":"PyNcclConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2}'
76
- )
75
+ ktc = KVTransferConfig(kv_connector="PyNcclConnector",
76
+ kv_role="kv_consumer",
77
+ kv_rank=1,
78
+ kv_parallel_size=2)
79
80
# Set GPU memory utilization to 0.8 for an A6000 GPU with 40GB
81
# memory. You may need to adjust the value to fit your GPU.