[vllm] vLLM Mooncake KV 오프로딩 최적화: 불필요한 KV 조회 건너뛰기

2026년 6월 18일수정: 2026년 6월 18일

PR 링크: vllm-project/vllm#45444 상태: Merged | 변경: +120 / -18

들어가며

대규모 언어 모델(LLM)의 추론 성능은 KV 캐시 관리 방식에 크게 좌우됩니다. 특히, vLLM은 KV 캐시를 효율적으로 관리하기 위해 다양한 기법을 도입하고 있으며, 그중 하나가 'Mooncake'라는 KV 오프로딩 메커니즘입니다. Mooncake는 KV 캐시를 GPU 메모리뿐만 아니라 더 저렴한 스토리지(예: CPU RAM 또는 NVMe)로 오프로딩하여 전체 모델의 메모리 용량을 확장하는 것을 목표로 합니다.

하지만 KV 캐시 오프로딩은 필연적으로 추가적인 오버헤드를 발생시킵니다. 특히, 요청된 토큰 시퀀스가 기존 KV 캐시에 존재하는지 확인하는 'lookup' 과정에서 불필요한 스토리지 접근이 발생할 수 있습니다. 이번 PR은 이러한 불필요한 KV 조회 및 스토리지 관련 오버헤드를 줄여 Mooncake의 성능을 개선하는 데 초점을 맞추고 있습니다.

주요 개선 사항은 다음과 같습니다:

불필요한 SWA(Sliding Window Attention) 블록 조회 건너뛰기: KV 캐시 그룹의 reachable_block_mask를 활용하여, 캐시 히트 가능성이 없는 SWA 블록에 대한 조회를 건너뜁니다.
전체 캐시 저장 시 리스트 생성 오버헤드 감소: store_mask에서 모든 청크가 저장 대상일 경우, True 리스트를 명시적으로 생성하는 대신 None을 사용하여 효율성을 높입니다.

이 글에서는 해당 PR의 코드 변경 사항을 상세히 분석하고, 이러한 최적화가 왜 성능 향상으로 이어지는지, 그리고 어떤 일반적인 교훈을 얻을 수 있는지 살펴보겠습니다.

코드 분석

1. `tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py` 및 `tests/v1/kv_connector/unit/test_mooncake_store_worker.py`

이 파일들은 주로 테스트 코드의 변경을 포함합니다. 핵심은 store_mask 및 lookup 함수의 동작 방식이 변경됨에 따라, 테스트 케이스들이 새로운 로직을 반영하도록 수정된 것입니다.

변경 전 (store_mask 관련 테스트):

기존에는 store_mask 함수가 반환하는 마스크가 항상 list[bool] 형태였습니다. 예를 들어, 전체 어텐션(Full Attention)의 경우 모든 청크가 저장 대상이므로 [True, True, True, True]와 같이 명시적인 True 리스트가 반환되었습니다.

-    assert masks == ([True, True, True, True],)
+    assert masks == (None,)

변경 후 (store_mask 관련 테스트):

이제 store_mask 함수는 특정 조건(예: 전체 어텐션 그룹에서 모든 청크가 저장 대상인 경우)에서 None을 반환하도록 변경되었습니다. 이는 True 리스트를 명시적으로 생성하는 오버헤드를 줄이기 위한 최적화입니다. 테스트 코드에서는 이 변경 사항을 반영하여 None을 기대하도록 수정되었습니다.

-    assert masks == ([True, True, True, True],)
+    assert masks == (None,)

또한, lookup 함수와 관련된 테스트에서도 SWA 블록의 조회 로직 변경을 반영하는 수정이 이루어졌습니다. 특히, test_lookup_checks_all_potential_swa_hit_boundaries 테스트는 SWA 청크 중에서도 실제로 캐시 히트가 발생할 가능성이 있는 경계(boundary)만 조회하도록 개선된 로직을 검증합니다.

+def test_lookup_checks_all_potential_swa_hit_boundaries():
+    """Lookup should skip SWA chunks that can never validate a hit, but still
+    check earlier aligned boundaries when sparse retention stores only the
+    current request's replay boundary.
+    """
+    # ... (테스트 코드 내용) ...
+
+    assert result == 32
+    keys = worker.store.batch_is_exist.call_args.args[0]
+    assert len(keys) == 6
+    swa_keys = [key for key in keys if "@group:1@" in key]
+    assert swa_keys == [
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1@6833",
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1@6837",
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1@683131",
+    ]

2. `vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py`

이 파일은 Mooncake KV 오프로딩의 코디네이터 로직을 담당합니다. 주요 변경 사항은 store_mask 함수가 내부적으로 _reachable_masks 함수를 호출하도록 리팩토링되고, lookup_mask라는 새로운 함수가 추가된 것입니다.

변경 전 (store_mask 함수):

-    def store_mask(
-        self,
-        aligned_token_len: int,
-        num_prompt_tokens: int | None = None,
-    ) -> tuple[list[bool], ...]:
-        """Per-group store masks: ``mask[g][i]`` is True iff chunk ``i`` of
-        group ``g`` should be written to the store so a future cache hit can
-        consume it.
-        """
-        assert aligned_token_len % self.lcm_block_size == 0, (
-            f"aligned_token_len ({aligned_token_len}) must be a multiple of "
-            f"lcm_block_size ({self.lcm_block_size})"
-        )
-        masks: list[list[bool]] = []
-        for g_idx, g in enumerate(self.kv_cache_groups):
-            spec = _unwrap_spec(g.kv_cache_spec)
-            num_chunks = aligned_token_len // spec.block_size
-            mask = _get_reachable_block_mask(
-                aligned_token_len,
-                num_chunks,
-                alignment_tokens=self.lcm_block_size,
-                kv_cache_spec=spec,
-                use_eagle=g_idx in self.eagle_group_ids,
-                retention_interval=self.retention_interval,
-                num_prompt_tokens=num_prompt_tokens,
-            )
-            masks.append([True] * num_chunks if mask is None else mask)
-        return tuple(masks)

변경 후 (_reachable_masks 및 store_mask, lookup_mask 함수):

+    def store_mask(
+        self,
+        aligned_token_len: int,
+        num_prompt_tokens: int | None = None,
+    ) -> tuple[list[bool] | None, ...]:
+        """Per-group store masks. 
+
+        ``mask[g][i]`` is True iff chunk ``i`` of group ``g`` should be
+        written to the store so a future cache hit can consume it. ``None`` is
+        the all-True sentinel.
+
+        Reuses the engine's ``SingleTypeKVCacheManager.reachable_block_mask``
+        so the store retains exactly the blocks the local prefix cache would.
+        """
+        return self._reachable_masks(
+            aligned_token_len,
+            retention_interval=self.retention_interval,
+            num_prompt_tokens=num_prompt_tokens,
+        )
+
+    def lookup_mask(
+        self,
+        aligned_token_len: int,
+    ) -> tuple[list[bool] | None, ...]:
+        """Per-group lookup masks.
+
+        ``mask[g][i]`` is True iff chunk ``i`` of group ``g`` should be
+        looked up as an aligned hit boundary. ``None`` is the all-True
+        sentinel.
+        """
+        return self._reachable_masks(
+            aligned_token_len,
+            retention_interval=None,
+            num_prompt_tokens=None,
+        )
+
+    def _reachable_masks(
+        self,
+        aligned_token_len: int,
+        *,
+        retention_interval: int | None,
+        num_prompt_tokens: int | None,
+    ) -> tuple[list[bool] | None, ...]:
         assert aligned_token_len % self.lcm_block_size == 0, (
             f"aligned_token_len ({aligned_token_len}) must be a multiple of "
             f"lcm_block_size ({self.lcm_block_size})"
         )
-        masks: list[list[bool]] = []
+        masks: list[list[bool] | None] = []
         for g_idx, g in enumerate(self.kv_cache_groups):
             spec = _unwrap_spec(g.kv_cache_spec)
             num_chunks = aligned_token_len // spec.block_size
-            mask = _get_reachable_block_mask(
+            mask = _get_reachable_block_mask(
                 aligned_token_len,
                 num_chunks,
                 alignment_tokens=self.lcm_block_size,
                 kv_cache_spec=spec,
                 use_eagle=g_idx in self.eagle_group_ids,
-                retention_interval=self.retention_interval,
-                num_prompt_tokens=num_prompt_tokens,
+                retention_interval=retention_interval,
+                num_prompt_tokens=num_prompt_tokens,
             )
-            masks.append([True] * num_chunks if mask is None else mask)
+            if mask is not None:
+                assert len(mask) == num_chunks
+            masks.append(mask)
         return tuple(masks)
 
     def block_hashes_for_spec(

_reachable_masks 함수: store_mask와 새로 추가된 lookup_mask 함수가 공통으로 사용하는 로직을 추출했습니다. 이 함수는 retention_interval과 num_prompt_tokens 인자를 받아 KV 캐시 블록의 도달 가능성(reachable) 마스크를 계산합니다. store_mask는 이 마스크를 사용하여 저장할 블록을 결정하고, lookup_mask는 이 마스크를 사용하여 조회할 블록을 결정합니다.
store_mask 함수의 None 사용: _reachable_masks 함수에서 반환된 mask가 None일 경우, 기존처럼 [True] * num_chunks를 생성하는 대신 None 자체를 masks 리스트에 추가합니다. 이는 모든 청크가 저장 대상일 때 불필요한 리스트 생성을 피하는 최적화입니다.
lookup_mask 함수의 도입: 이 함수는 store_mask와 유사하지만, retention_interval과 num_prompt_tokens를 None으로 설정하여 호출합니다. 이는 KV 캐시를 조회할 때는 저장 시점과는 다르게, 더 넓은 범위의 블록을 고려해야 할 수 있기 때문입니다. 특히, SWA의 경우 캐시 히트 가능성이 낮은 블록은 조회 대상에서 제외하여 성능을 높입니다.

3. `vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py`

이 파일은 Mooncake KV 오프로딩 워커의 핵심 로직을 구현합니다. _handle_request 및 lookup 함수에서 변경 사항이 있습니다.

변경 전 (_handle_request 함수):

기존 _handle_request 함수는 mask가 None이 아닌 경우에만 마스크를 확인했습니다. 이는 mask가 None일 때 (즉, 모든 청크가 저장 대상일 때) 마스크 검사를 건너뛰는 것을 의미했습니다.

-                    if chunk_idx >= len(mask) or not mask[chunk_idx]:
+                    if mask is not None and (
+                        chunk_idx >= len(mask) or not mask[chunk_idx]
+                    ):
                         continue

변경 후 (_handle_request 함수):

변경 후에는 mask is not None 조건을 명시적으로 추가하여, mask가 None일 경우에는 해당 if 조건문 내부의 코드가 실행되지 않도록 합니다. 즉, mask가 None이면 모든 청크에 대해 continue를 건너뛰고 처리를 진행합니다. 이는 store_mask에서 None을 반환하는 최적화와 연계됩니다.

변경 전 (lookup 함수):

기존 lookup 함수는 store_mask에서 반환된 마스크를 직접 사용하거나, 혹은 store_mask의 로직을 내부적으로 복제하여 사용했을 가능성이 있습니다. 하지만 명시적으로 lookup_mask와 같은 별도의 마스크 생성 로직은 없었습니다.

변경 후 (lookup 함수):

+        lookup_masks = self.coord.lookup_mask(token_len)
+        tp_count = min(self.tp_size, self.num_kv_head)
+        for g_idx, db in enumerate(self.token_dbs):
+            spec_block_size = db.block_size
+            lookup_mask = lookup_masks[g_idx]
+            group_hashes = self.coord.block_hashes_for_spec(
+                block_hashes, self._kv_cache_groups[g_idx].kv_cache_spec
+            )
+            # ... (이하 생략) ...
+            # if lookup_mask is not None and not lookup_mask[chunk_id]:
+            #     continue
+            # ... (이하 생략) ...

lookup 함수는 이제 coordinator.lookup_mask(token_len)를 호출하여 조회 마스크를 얻습니다.
각 KV 캐시 그룹(db)에 대해 해당 그룹의 lookup_mask를 가져옵니다.
이 lookup_mask를 사용하여 실제로 조회해야 할 청크(chunk_id)인지 판단합니다. 만약 lookup_mask가 None이 아니고 해당 청크가 False이면, 해당 청크에 대한 조회를 건너뜁니다. 이는 SWA와 같이 캐시 히트 가능성이 낮은 블록에 대한 불필요한 디스크/메모리 접근을 방지하여 lookup 과정의 오버헤드를 줄입니다.

왜 이게 좋은가?

성능 향상

PR 설명에 포함된 벤치마크 결과는 이 최적화의 효과를 명확히 보여줍니다. DeepSeek v4 모델을 4개의 GB300 GPU에서 실행했을 때, 응답 속도(latency)가 크게 개선되었습니다. 특히, 99% percentile latency가 약 15% 감소한 것을 확인할 수 있습니다. 이는 초당 처리할 수 있는 토큰 수(tokens/sec)가 증가했음을 의미하며, 더 빠르고 효율적인 추론이 가능해졌다는 것을 나타냅니다.

이러한 성능 향상의 주요 원인은 다음과 같습니다:

불필요한 KV 조회 감소: lookup_mask를 사용하여 캐시 히트 가능성이 희박한 SWA 블록에 대한 조회를 건너뜀으로써, 스토리지(CPU RAM 등) 접근 횟수가 줄어듭니다. 스토리지 접근은 GPU 메모리 접근보다 훨씬 느리기 때문에, 이러한 불필요한 접근을 제거하는 것은 전체 응답 시간 단축에 크게 기여합니다.
스토리지 오버헤드 감소: store_mask에서 모든 청크가 저장 대상일 때 None을 사용하는 최적화는, 매번 True 리스트를 생성하는 데 드는 작은 메모리 및 연산 오버헤드를 줄입니다. 비록 이 부분의 성능 기여도는 앞선 조회 감소보다는 작을 수 있지만, 누적되면 의미 있는 성능 향상으로 이어질 수 있습니다.

일반적인 교훈

이 PR은 LLM 추론 최적화에 있어 다음과 같은 중요한 교훈을 제공합니다:

데이터 구조 및 알고리즘의 중요성: KV 캐시와 같이 복잡한 데이터 구조를 다룰 때, 어떤 정보를 저장하고 어떻게 조회하는지에 대한 알고리즘적 접근은 성능에 지대한 영향을 미칩니다. reachable_block_mask와 같은 메타데이터를 활용하여 불필요한 연산을 제거하는 것은 매우 효과적인 전략입니다.
메모리 계층 구조 고려: LLM은 방대한 양의 KV 캐시를 생성하며, 이를 GPU 메모리뿐만 아니라 다른 메모리 계층(CPU RAM, NVMe 등)으로 확장합니다. 각 계층의 접근 속도 차이를 인지하고, 가장 느린 계층으로의 접근을 최소화하는 것이 성능 최적화의 핵심입니다.
'Nothing is Free' 원칙: KV 오프로딩은 메모리 용량을 늘려주지만, 추가적인 계산 및 스토리지 접근 오버헤드를 동반합니다. 따라서 이러한 오버헤드를 최소화하기 위한 지속적인 최적화 노력이 필수적입니다.
명시적인 'True' 리스트 대신 Sentinel 값 활용: 특정 조건에서 모든 요소가 동일한 값을 가질 때, 이를 명시적으로 리스트로 생성하는 대신 None과 같은 Sentinel 값을 사용하여 해당 상태를 나타내는 것은 메모리 및 연산 효율성을 높이는 좋은 패턴입니다.

결론

이번 vLLM의 Mooncake KV 오프로딩 관련 PR은 SWA 블록에 대한 불필요한 KV 조회 건너뛰기와 스토리지 관련 오버헤드 감소를 통해 실제적인 성능 향상을 이루어냈습니다. 벤치마크 결과에서 나타난 99% percentile latency의 15% 감소는 이러한 최적화의 중요성을 잘 보여줍니다. LLM 추론 성능을 극한까지 끌어올리기 위한 vLLM 팀의 지속적인 노력과 깊이 있는 엔지니어링 역량을 엿볼 수 있는 좋은 사례라고 생각합니다.

참고 자료

https://pytorch.org/docs/stable/generated/torch.compile.html

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

PR Analysis 의 다른글

이전글 [sglang] Mamba GDN의 컨볼루션 캐시 최적화: 메모리 사용량 절반으로 줄이기
현재글 : [vllm] vLLM Mooncake KV 오프로딩 최적화: 불필요한 KV 조회 건너뛰기
다음글 [ray] Ray RLlib의 비동기 학습 성능 최적화: PULL 기반 EnvRunnerStateServer 도입