Disaggregated Prefill V1

Disaggregated Prefill V1 (Example Connector)

vLLM의 오프라인 설정에서 disaggregated prefill(prefill과 decode를 분리)을 시연하는 예제 스크립트 모음입니다. ExampleConnector라는 로컬 공유 스토리지 기반 KV 커넥터를 사용해, 전용 prefill 인스턴스가 KV 상태를 저장하고 전용 decode 인스턴스가 이를 불러와 이어서 생성합니다.

출처: 문서

본문

소스: https://github.com/vllm-project/vllm/tree/main/examples/disaggregated/example_connector

이 예제는 vLLM의 오프라인 설정에서 disaggregated prefill을 시연하는 스크립트들을 담고 있습니다.

파일 (Files)

• run.sh — prefill_example.py와 decode_example.py를 순차 실행하는 헬퍼 스크립트. run.sh 실행 전에 examples/disaggregated/example_connector 디렉토리 안에 있어야 합니다.
• prefill_example.py — prefill만 수행하고, KV 상태를 local_storage 디렉토리에, 프롬프트를 output.txt에 저장하는 스크립트.
• decode_example.py — decode만 수행하고, local_storage 디렉토리에서 KV 상태를, output.txt에서 프롬프트를 불러오는 스크립트.

예제 자료 (Example materials)

decode_example.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from vllm import LLM, SamplingParams
from vllm.config import KVTransferConfig

def read_prompts():
    """Read prompts from output.txt."""
    prompts = []
    try:
        with open("output.txt") as f:
            for line in f:
                prompts.append(line.strip())
        print(f"Loaded {len(prompts)} prompts from output.txt")
        return prompts
    except FileNotFoundError:
        print("Error: output.txt file not found")
        exit(-1)

def main():
    prompts = read_prompts()
    sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)

    llm = LLM(
        model="meta-llama/Llama-3.2-1B-Instruct",
        enforce_eager=True,
        gpu_memory_utilization=0.8,
        max_num_batched_tokens=64,
        max_num_seqs=16,
        kv_transfer_config=KVTransferConfig(
            kv_connector="ExampleConnector",
            kv_role="kv_both",
            kv_connector_extra_config={"shared_storage_path": "local_storage"},
        ),
    )  # , max_model_len=2048, max_num_batched_tokens=2048)

    # 1ST generation (prefill instance)
    outputs = llm.generate(prompts, sampling_params)

    print("-" * 30)
    for output in outputs:
        prompt = output.prompt
        generated_text = output.outputs[0].text
        print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
        print("-" * 30)

if __name__ == "__main__":
    main()

prefill_example.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from vllm import LLM, SamplingParams
from vllm.config import KVTransferConfig

def read_prompts():
    context = "Hi " * 1000
    context2 = "Hey " * 500
    return [
        context + "Hello, my name is",
        context + "The capital of France is",
        context2 + "Your name is",
        context2 + "The capital of China is",
    ]

def main():
    prompts = read_prompts()

    sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1)

    llm = LLM(
        model="meta-llama/Llama-3.2-1B-Instruct",
        enforce_eager=True,
        gpu_memory_utilization=0.8,
        kv_transfer_config=KVTransferConfig(
            kv_connector="ExampleConnector",
            kv_role="kv_both",
            kv_connector_extra_config={"shared_storage_path": "local_storage"},
        ),
    )  # , max_model_len=2048, max_num_batched_tokens=2048)

    # 1ST generation (prefill instance)
    outputs = llm.generate(
        prompts,
        sampling_params,
    )

    new_prompts = []
    print("-" * 30)
    for output in outputs:
        prompt = output.prompt
        generated_text = output.outputs[0].text
        new_prompts.append(prompt + generated_text)
        print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
        print("-" * 30)

    # Write new_prompts to output.txt
    with open("output.txt", "w") as f:
        for prompt in new_prompts:
            f.write(prompt + "\n")
    print(f"Saved {len(new_prompts)} prompts to output.txt")

if __name__ == "__main__":
    main()

run.sh

rm -rf local_storage/

if [ -f "output.txt" ]; then
    rm output.txt
fi

# The directory of current script
SCRIPT_DIR=$(dirname "$(readlink -f "$0")")

VLLM_ENABLE_V1_MULTIPROCESSING=0 CUDA_VISIBLE_DEVICES=0 python3 "$SCRIPT_DIR/prefill_example.py"
VLLM_ENABLE_V1_MULTIPROCESSING=0 CUDA_VISIBLE_DEVICES=0 python3 "$SCRIPT_DIR/decode_example.py"

핵심 포인트:

• prefill 인스턴스는 max_tokens=1로 짧은 생성만 하고, KV 상태를 local_storage에 저장하며 새 프롬프트를 output.txt에 기록합니다.
• decode 인스턴스는 output.txt에서 프롬프트를, local_storage에서 KV 상태를 불러와 max_tokens=10으로 이어서 생성합니다.
• 두 인스턴스 모두 KVTransferConfig에 kv_connector="ExampleConnector", kv_role="kv_both", shared_storage_path="local_storage"를 지정합니다.

더 알아보기 (Learn more)

• Disaggregated Serving — 온라인 disaggregated 서빙
• KVTransferConfig — KV 커넥터·역할 설정