KV 캐시 리셋

KV 캐시 리셋 (Reset KV Cache)

이 예제는 LLMEngine을 직접 다루면서 프리픽스 캐시(prefix cache)를 동적으로 리셋하는 방법을 보여줍니다. 중간 단계에서 reset_prefix_cache()를 호출해 실행 중인 요청의 KV 캐시를 비우고, 메모리와 캐시 상태를 재설정하는 흐름을 이해할 수 있습니다.

출처: 문서

본문

LLMEngine은 로우 레벨 API로, 요청을 하나씩 추가하고 step()으로 처리합니다. 아래 코드는 10번째 스텝에서 reset_prefix_cache(reset_running_requests=True)를 호출해 캐시를 비우는 예시입니다.

reset_kv_offline.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""This file demonstrates preempt requests when using the `LLMEngine`
for processing prompts with various sampling parameters.
"""

import argparse

from vllm import EngineArgs, LLMEngine, RequestOutput, SamplingParams
from vllm.utils.argparse_utils import FlexibleArgumentParser


def create_test_prompts() -> list[tuple[str, SamplingParams]]:
    """Create a list of test prompts with their sampling parameters."""
    return [
        (
            "A robot may not injure a human being " * 50,
            SamplingParams(
                temperature=0.0, logprobs=1, prompt_logprobs=1, max_tokens=16
            ),
        ),
        (
            "A robot may not injure a human being " * 50,
            SamplingParams(
                temperature=0.0, logprobs=1, prompt_logprobs=1, max_tokens=16
            ),
        ),
        (
            "To be or not to be,",
            SamplingParams(
                temperature=0.8, top_k=5, presence_penalty=0.2, max_tokens=128
            ),
        ),
        (
            "What is the meaning of life?",
            SamplingParams(
                n=2, temperature=0.8, top_p=0.95, frequency_penalty=0.1, max_tokens=128
            ),
        ),
    ]


def process_requests(engine: LLMEngine, test_prompts: list[tuple[str, SamplingParams]]):
    """Continuously process a list of prompts and handle the outputs."""
    request_id = 0

    print("-" * 50)
    step_id = 0
    while test_prompts or engine.has_unfinished_requests():
        print("-" * 50)
        import os

        print(f"Step {step_id} (pid={os.getpid()})")

        if test_prompts:
            prompt, sampling_params = test_prompts.pop(0)
            engine.add_request(str(request_id), prompt, sampling_params)
            request_id += 1

        if step_id == 10:
            print(f"Resetting prefix cache at {step_id}")
            engine.reset_prefix_cache(reset_running_requests=True)

        request_outputs: list[RequestOutput] = engine.step()

        for request_output in request_outputs:
            if request_output.finished:
                print("-" * 50)
                print(request_output)
                print("-" * 50)
        step_id += 1


def initialize_engine(args: argparse.Namespace) -> LLMEngine:
    """Initialize the LLMEngine from the command line arguments."""
    engine_args = EngineArgs.from_cli_args(args)
    return LLMEngine.from_engine_args(engine_args)


def parse_args():
    parser = FlexibleArgumentParser(
        description="Demo on using the LLMEngine class directly"
    )
    parser = EngineArgs.add_cli_args(parser)
    return parser.parse_args()


def main(args: argparse.Namespace):
    """Main function that sets up and runs the prompt processing."""
    engine = initialize_engine(args)
    test_prompts = create_test_prompts()
    process_requests(engine, test_prompts)


if __name__ == "__main__":
    args = parse_args()
    main(args)

더 알아보기 (Learn more)