[vllm] vLLM KV Offload 최적화: Tensor Parallelism 환경에서 MLA KV 캐시 복제본 제거하기

2026년 7월 26일수정: 2026년 7월 26일

PR 링크: vllm-project/vllm#48906 상태: Merged | 변경: +1135 / -498

들어가며

대규모 언어 모델(LLM)을 서빙할 때, KV 캐시(Key-Value Cache)는 추론 성능에 결정적인 영향을 미칩니다. 특히 긴 컨텍스트 길이의 시퀀스를 처리할 때 KV 캐시가 GPU 메모리를 빠르게 소모하여 배치 크기를 제한하는 병목 현상이 발생하곤 합니다. 이를 해결하기 위해 vLLM은 KV 캐시 오프로딩(Offloading) 기능을 제공하여, 사용 빈도가 낮은 KV 캐시 블록을 CPU 메모리로 옮겨 GPU 메모리 부담을 줄입니다.

하지만 Tensor Parallelism(TP) 환경에서 Multi-Layer Attention (MLA) 모델을 사용할 경우, 각 TP 랭크(rank)가 동일한 KV 캐시 페이로드의 복제본을 CPU 호스트 메모리에 저장하는 비효율적인 문제가 있었습니다. 이는 TP 스케일(N)에 비례하여 CPU 메모리 사용량과 GPU에서 CPU로의 데이터 전송(D2H) 트래픽을 불필요하게 증가시키는 원인이 됩니다. 이번 PR([KV Offload] Deduplicate replicated MLA KV in the shared CPU region)은 이러한 문제를 해결하기 위해 공유 CPU 영역에서 MLA KV 캐시 복제본을 제거하는 중요한 최적화를 구현했습니다.

이 글에서는 해당 PR의 코드 변경사항을 자세히 분석하고, 이 최적화가 왜 좋은 개선인지 설명하며, 실제 프로덕션 환경에서의 성능 향상 사례와 기술적 교훈을 공유하고자 합니다.

코드 분석: 복제본 제거의 핵심 로직

이 PR의 핵심은 Tensor Parallelism 환경에서 MLA KV 캐시 블록이 여러 TP 랭크에 걸쳐 복제되는 것을 방지하고, 단일 물리적 복사본만 CPU에 저장하도록 하는 것입니다. 이를 위해 OffloadingConfig에 replicated_layout이라는 새로운 플래그를 도입하고, 이 플래그에 따라 KV 캐시 저장 및 로드 방식, 그리고 CPU 메모리 할당 로직을 변경합니다.

1. `replicated_layout` 결정 로직 (`vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py`)

가장 먼저, 시스템은 현재 KV 캐시 구성과 병렬 처리 설정에 따라 replicated_layout을 사용할지 여부를 결정해야 합니다. 이 로직은 build_offloading_config 함수에 추가되었습니다.

Before (개념):

# OffloadingConfig에 replicated_layout 필드가 없었으며, 항상 False로 간주되었습니다.
# build_offloading_config는 단순히 OffloadingConfig를 생성했습니다.

After:

# vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
# ...
def build_offloading_config(
    vllm_config: VllmConfig,
    kv_cache_config: KVCacheConfig,
) -> OffloadingConfig:
    # ...
    replicated_layout = (
        spec_name == "TieringOffloadingSpec"
        and vllm_config.parallel_config.tensor_parallel_size > 1
        and vllm_config.parallel_config.pipeline_parallel_size == 1
        and vllm_config.parallel_config.prefill_context_parallel_size == 1
        and vllm_config.parallel_config.decode_context_parallel_size == 1
        and vllm_config.parallel_config.distributed_executor_backend == "mp"
        and vllm_config.parallel_config.nnodes == 1
        and len(kv_cache_config.kv_cache_groups) == 1
        and isinstance(kv_cache_config.kv_cache_groups[0].spec,
                       MLAAttentionSpec)
    )

    return OffloadingConfig(
        worker_kv_bytes_per_block=worker_kv_bytes_per_block,
        parallel=parallel_config,
        spec_name=spec_name,
        cpu_bytes_to_use=cpu_bytes_to_use,
        eviction_policy=eviction_policy,
        replicated_layout=replicated_layout, # 새로운 필드
    )

replicated_layout은 TieringOffloadingSpec을 사용하고, TP 사이즈가 1보다 크며, 다른 병렬 처리 축(Pipeline, Prefill Context, Decode Context)이 1이고, 단일 노드 mp executor 환경에서, 그리고 단일 MLAAttentionSpec 그룹일 때만 True로 설정됩니다. 이는 최적화가 적용될 수 있는 보수적인 조건을 정의하여 안정성을 확보합니다.

2. KV 캐시 저장 로직 최적화 (`vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py`)

replicated_layout이 True일 때, 오직 TP 랭크 0만이 GPU에서 CPU로 KV 캐시 블록을 실제로 저장합니다. 다른 랭크들은 저장 작업을 건너뛰고 완료 신호만 보냅니다.

Before:

# vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
# ...
def _store_kv_cache_block_impl(self, block_id: int,
                                    kv_cache_block: KVCacheBlock):
    # 모든 랭크가 물리적으로 블록을 저장합니다.
    self.offloading_backend.store_block(
        block_id,
        kv_cache_block,
        self.offloading_config.parallel.rank,
        self.offloading_config.parallel.world_size,
    )

After:

# vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
# ...
def _store_kv_cache_block_impl(self, block_id: int,
                                    kv_cache_block: KVCacheBlock):
    # ...
    if (self.offloading_config.replicated_layout
            and self.offloading_config.parallel.rank != 0):
        # 비-작성자 랭크는 D2H 전송 없이 저장 완료를 알립니다.
        self.offloading_backend.acknowledge_store(block_id)
        return

    self.offloading_backend.store_block(
        block_id,
        kv_cache_block,
        self.offloading_config.parallel.rank,
        self.offloading_config.parallel.world_size,
    )

이 변경으로 인해 D2H 트래픽이 TP 스케일(N)만큼 감소합니다.

3. KV 캐시 로드 로직 최적화 (`vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py`)

저장과 달리, 로드 시에는 모든 TP 랭크가 공유 CPU 메모리에서 동일한 슬롯의 데이터를 로드합니다.

Before:

# vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
# ...
def _load_kv_cache_block_impl(self, block_id: int,
                                    kv_cache_block: KVCacheBlock):
    # 각 랭크는 자신의 랭크에 해당하는 슬롯에서 로드합니다.
    self.offloading_backend.load_block(
        block_id,
        kv_cache_block,
        self.offloading_config.parallel.rank,
        self.offloading_config.parallel.world_size,
    )

After:

# vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
# ...
def _load_kv_cache_block_impl(self, block_id: int,
                                    kv_cache_block: KVCacheBlock):
    # ...
    self.offloading_backend.load_block(
        block_id,
        kv_cache_block,
        # replicated_layout의 경우, 모든 랭크는 동일한 공유 슬롯에서 로드합니다.
        # 백엔드는 이를 슬롯 0으로 매핑합니다.
        self.offloading_config.parallel.rank,
        self.offloading_config.parallel.world_size,
        replicated_layout=self.offloading_config.replicated_layout, # 새로운 인자
    )

load_block 함수에 replicated_layout 인자가 추가되었고, 이 인자를 통해 백엔드(shared_offload_region)가 올바른 슬롯을 결정하게 됩니다.

4. 공유 CPU 메모리 할당 및 슬롯 인덱싱 (`vllm/v1/kv_offload/cpu/shared_offload_region.py`)

공유 메모리 영역(SharedOffloadRegion)은 replicated_layout에 따라 블록당 할당되는 메모리 크기와 슬롯 인덱스를 다르게 계산합니다.

Before (_get_block_row_size):

# vllm/v1/kv_offload/cpu/shared_offload_region.py
# ...
def _get_block_row_size(
    block_size: int,
    num_kv_heads: int,
    head_size: int,
    dtype_size_in_bytes: int,
    world_size: int,
) -> int:
    # 블록당 물리적 크기를 그대로 반환합니다.
    return block_size * num_kv_heads * head_size * dtype_size_in_bytes

After (_get_block_row_size):

# vllm/v1/kv_offload/cpu/shared_offload_region.py
# ...
def _get_block_row_size(
    block_size: int,
    num_kv_heads: int,
    head_size: int,
    dtype_size_in_bytes: int,
    replicated_layout: bool, # 새로운 인자
    world_size: int,
) -> int:
    # ...
    row_size = (block_size * num_kv_heads * head_size *
                dtype_size_in_bytes)
    if replicated_layout:
        # replicated_layout의 경우, 구성된 CPU 용량은 약 N배 더 많은
        # 복제된 MLA 블록을 수용할 수 있습니다.
        # 이는 row_size를 world_size로 나누어, 각 "논리적" 블록이
        # 물리적으로는 하나의 복사본만 저장되므로 N배 작아지도록 합니다.
        row_size //= world_size
    return row_size

replicated_layout이 True일 때, _get_block_row_size는 실제 블록 크기를 world_size로 나눕니다. 이는 CPU 캐시 용량을 계산할 때, 각 블록이 world_size만큼 더 많은 논리적 블록을 저장할 수 있도록 하여, 실제로는 하나의 물리적 복사본만 저장되므로 N배 더 많은 블록을 저장할 수 있음을 의미합니다.

Before (_get_slot_idx):

# vllm/v1/kv_offload/cpu/shared_offload_region.py
# ...
def _get_slot_idx(rank: int, world_size: int) -> int:
    # 각 랭크는 자신의 랭크에 해당하는 슬롯을 사용합니다.
    return rank

After (_get_slot_idx):

# vllm/v1/kv_offload/cpu/shared_offload_region.py
# ...
def _get_slot_idx(rank: int,
                    world_size: int,
                    replicated_layout: bool) -> int: # 새로운 인자
    if replicated_layout:
        # replicated_layout의 경우, 모든 워커는 슬롯 0으로 매핑됩니다.
        return 0
    return rank

_get_slot_idx 함수는 replicated_layout이 True일 때 항상 슬롯 0을 반환하여, 모든 TP 랭크가 동일한 물리적 메모리 위치를 참조하도록 합니다.

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.

PR Analysis 의 다른글

이전글 [vllm] vLLM: PyNvVideoCodec 하드웨어 디코더 동시성 최적화로 비디오 처리량 94% 향상
현재글 : [vllm] vLLM KV Offload 최적화: Tensor Parallelism 환경에서 MLA KV 캐시 복제본 제거하기
다음글 [vllm] vLLM DeepSeek-OCR-2 모델의 TTFT 성능 최적화 분석

[vllm] vLLM KV Offload 최적화: Tensor Parallelism 환경에서 MLA KV 캐시 복제본 제거하기

들어가며

코드 분석: 복제본 제거의 핵심 로직

1. `replicated_layout` 결정 로직 (`vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py`)

2. KV 캐시 저장 로직 최적화 (`vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py`)

3. KV 캐시 로드 로직 최적화 (`vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py`)

4. 공유 CPU 메모리 할당 및 슬롯 인덱싱 (`vllm/v1/kv_offload/cpu/shared_offload_region.py`)

5. 테스트 및 CI 통합

왜 이게 좋은가?

1. 리소스 효율성 극대화 (CPU 메모리 및 D2H 트래픽 감소)

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글