본문으로 건너뛰기

[vllm] vLLM 하이브리드 모델을 위한 혁신: Partial Prefix Cache Hit 구현 분석

PR 링크: vllm-project/vllm#46384 상태: Merged | 변경: +1651 / -175

들어가며

대규모 언어 모델(LLM) 추론 엔진인 vLLM에서 Prefix Caching은 다회차 대화(Multi-turn)나 시스템 프롬프트 재사용 시 성능을 결정짓는 핵심 요소입니다. 기존의 Prefix Caching은 주로 '물리적 KV-cache 블록' 단위로 작동했습니다. 하지만 Full-Attention과 Mamba가 혼합된 하이브리드 모델(예: Jamba, Qwen2-Mamba)에서는 문제가 발생합니다. Mamba의 Recurrent State는 토큰 단위로 상태가 변하며, 물리 블록 크기(예: 16)와 해시 매칭 단위가 일치하지 않을 경우 캐시 효율이 급격히 떨어지기 때문입니다.

이번 PR([2/N][Core] support partial prefix cache hit for hybrid model)은 이러한 문제를 해결하기 위해 물리 블록 크기보다 더 세밀한(fine-grained) 단위의 캐시 히트를 지원하는 메커니즘을 도입했습니다. 이를 통해 하이브리드 모델에서 불필요한 재계산을 줄이고 지연시간을 획기적으로 개선했습니다.

코드 분석: 핵심 변경 사항

1. prefix_match_unit 도입과 세밀한 해시 뷰

기존에는 물리 블록 크기가 곧 캐시 매칭의 최소 단위였습니다. 이번 변경에서는 prefix_match_unit이라는 독립적인 단위를 도입하여, 물리 블록 내부에서도 해시 매칭 여부를 판단할 수 있게 했습니다.

Before (Conceptual):

# 물리 블록 크기(block_size)가 16이면, 16토큰 단위로만 캐시 히트 판단
# 20토큰이 겹쳐도 16토큰만 히트 처리

After (Implementation in KVCacheConfig):

# vllm/config/cache.py (추정)
# prefix_match_unit을 통해 물리 블록(예: 16)보다 작은 단위(예: 2)로 해시 계산
# Request.block_hashes는 이제 prefix_match_unit 단위로 생성됨

이 변경 덕분에 FullAttentionManager는 물리 블록을 다 채우지 못한 'Partial Tail' 영역에서도 해시를 대조하여 더 긴 접두사를 찾아낼 수 있습니다.

2. 스케줄러의 전략적 중단 (Mamba Block-Aligned Split)

Mamba 모델은 특정 지점의 상태(State)를 저장해야 나중에 그 지점부터 재개할 수 있습니다. 이를 위해 스케줄러가 해시 경계면에서 정확히 멈추도록 로직이 수정되었습니다.

Before/After (Scheduler Logic):

# tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py 내 테스트 코드 인용
def test_mamba_align_split_partial_tail_schedule():
    # block_size=512, hash_block_size=32
    # 10000 토큰 프롬프트 처리 시
    
    # Before: 물리 블록 단위로만 끊김
    # After: 마지막 해시 경계(9984)에서 한 번 더 멈춤 (Partial Tail 등록을 위해)
    assert split(self=mock, request=req, num_new_tokens=272) == 256 # 9728 -> 9984
    req.num_computed_tokens = 9984
    assert split(self=mock, request=req, num_new_tokens=16) == 16   # 9984 -> 10000

이처럼 9984 토큰 지점에서 의도적으로 스케줄링을 끊음으로써, Mamba의 Recurrent State를 해당 해시 경계에 정확히 저장(Register)할 수 있게 됩니다.

3. 하이브리드 코디네이터의 수렴 (Convergence)

Attention과 Mamba는 서로 다른 캐시 구조를 가집니다. HybridCoordinator는 각 그룹의 히트 길이를 비교하여 공통된 최소 길이를 도출합니다.

Code Snippet (v1/core/kv_cache_coordinator.py):

# 여러 캐시 그룹 간의 히트 길이를 동기화
for group_id, hit_length in hit_length_by_group.items():
    # 모든 그룹이 커버할 수 있는 최소 공통 접두사 길이로 수렴
    curr_hit_length = min(curr_hit_length, hit_length)

# 물리 블록 리스트도 이 길이에 맞춰 트리밍(Ceiling division)

예를 들어 Attention은 6토큰 히트인데 Mamba가 4토큰 히트라면, 전체 히트 길이를 4로 맞추어 데이터 일관성을 보장합니다.

4. Copy-on-Write (CoW)를 통한 공유 블록 보호

공유된 Partial Block에 새로운 토큰을 추가해야 할 때, 기존 캐시를 오염시키지 않기 위해 CoW를 수행합니다. 특히 Mamba의 경우 워커의 블록 테이블이 Append-only이므로, 기존 블록을 유지하되 캐시 엔트리를 새로운 블록으로 복사하는 'Reversed CoW' 전략을 사용합니다.

Code Snippet (Test Case):

# req0가 사용하던 partial_mamba_block_id를 유지하면서
# 캐시 엔트리(hash)만 새로운 dst_block_id로 이동
cow_copy = next(c for c in copies if c.src_block_id == partial_mamba_block_id)
assert cow_copy.dst_block_id != partial_mamba_block_id

왜 이게 좋은가?

1. 압도적인 성능 향상

PR 설명에 포함된 벤치마크 결과에 따르면, Qwen3.5-35B-A3B 모델에서 10,000토큰 입력 후 추가 입력을 주었을 때의 지연시간이 다음과 같이 개선되었습니다.

  • Partial Cache Hit 미적용: 0.1337s
  • Partial Cache Hit 적용: 0.0955s (약 28.5% 감소)

2. 메모리 효율과 유연성

물리 블록 크기를 작게 줄이면 관리 오버헤드가 커집니다. 이 방식은 물리 블록은 크게 유지(예: 16 or 32)하여 관리 효율을 높이면서도, 논리적인 매칭 단위(prefix_match_unit)만 작게 설정하여 두 마리 토끼를 모두 잡았습니다.

3. 기술적 교훈: 리뷰어 피드백 반영

리뷰 과정에서 njhill은 데이터 복사 시 효율성을 위해 async_tensor_h2d와 NumPy를 사용할 것을 제안했습니다.

# njhill의 제안: 비동기 복사를 통한 성능 최적화
indices_np = np.array(kv_cache_block_copies, dtype=np.int64)
indices = async_tensor_h2d(indices_np, device=storage_tensors[0].device)
src_indices, dst_indices = indices.unbind(dim=1)

또한, Mamba 블록은 초기화(Zeroing)가 필수적이지만 Full-Attention 블록은 선택적일 수 있다는 논의를 통해 하이브리드 모델 특유의 메모리 관리 전략을 정교화했습니다.

결론

이번 PR은 vLLM이 단순한 Transformer 모델을 넘어 Mamba와 같은 차세대 아키텍처에서도 업계 최고 수준의 추론 효율을 낼 수 있음을 증명했습니다. prefix_match_unit이라는 추상화 계층을 통해 물리적 제약을 극복한 설계는 복잡한 하이브리드 모델 최적화의 모범 사례로 남을 것입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글