[vllm] vLLM 스케줄러 최적화: Cache-aware Admission으로 KV 캐시 효율 극대화하기

PR 링크: vllm-project/vllm#54625 상태: Merged | 변경: +322 / -0

들어가며: FCFS와 Prefix Caching의 충돌

대규모 언어 모델(LLM) 서빙에서 Prefix Caching은 시스템 성능을 결정짓는 핵심 요소입니다. 하지만 기존의 엄격한 FCFS(First-Come-First-Served) 입장(Admission) 정책은 캐시 효율성을 저해하는 고질적인 문제를 안고 있었습니다.

문제의 핵심은 이렇습니다. 대기 큐(Waiting Queue)의 맨 앞에 있는 요청이 캐시에 데이터가 없는 'Cold' 요청이고, 바로 뒤에 있는 요청이 이미 캐시에 데이터가 있는 'Warm' 요청일 때, FCFS는 무조건 Cold 요청을 먼저 처리합니다. 이때 Cold 요청이 새로운 KV Cache 블록을 할당받으면서, 정작 바로 다음에 들어올 Warm 요청이 필요로 하는 캐시 블록을 밀어내는(Eviction) 현상이 발생합니다. 결과적으로 캐시 히트가 될 수 있었던 요청이 재연산(Recomputation)으로 변하며 시스템 전체의 처리량이 저하됩니다.

이번에 분석할 PR([Core] Add cache-aware admission ordering)은 대기 큐의 특정 범위(Window) 내에서 캐시 상태를 미리 확인하여, 캐시 히트 가능성이 높은 요청을 먼저 입장시키는 Cache-aware Admission Ordering을 도입합니다.

코드 분석: 핵심 변경 사항

1. SchedulerConfig: 새로운 제어 파라미터 도입

먼저 스케줄러 설정에 두 가지 핵심 옵션이 추가되었습니다. cache_aware_admission_window는 대기 큐를 얼마나 멀리 내다볼지 결정하고, cache_aware_admission_threshold는 캐시 사용량이 어느 정도일 때 이 로직을 활성화할지 결정합니다.

Before & After: SchedulerConfig

# Before (기존에는 FCFS 순서만 존재)
class SchedulerConfig:
    def __init__(self, ...):
        self.scheduling_policy = "fcfs"

# After (캐시 인지형 파라미터 추가)
class SchedulerConfig:
    cache_aware_admission_window: int = Field(default=0, ge=0)
    """Look-ahead window for cache-aware admission. ..."""

    cache_aware_admission_threshold: float = Field(default=0.5, ge=0.0, le=1.0)
    """KV cache usage below which the window is not applied. ..."""

이 변경으로 인해 사용자는 window 값을 설정하여 FCFS의 엄격함을 완화하고 성능을 최적화할 수 있게 되었습니다. 특히 threshold 기본값이 0.5로 설정되어, 캐시가 충분히 여유로울 때는 불필요한 재정렬을 피하도록 설계되었습니다.

2. KVCacheManager: 부작용 없는 캐시 상태 조사

재정렬을 위해서는 대기 중인 요청이 캐시에 얼마나 있는지 확인해야 합니다. 하지만 기존의 get_computed_blocks 함수는 호출 시 BlockStored 이벤트를 발생시키는 등 부작용(Side-effect)이 있었습니다. 이를 해결하기 위해 새로운 프로브(Probe) 함수가 추가되었습니다.

Implementation Note: PR 설명에 따르면 KVCacheManager.get_num_cached_tokens는 레퍼런스를 가져오거나 블록을 할당하지 않고 순수하게 캐시 상태만 조회합니다. 이는 스케줄링되지 않을 수도 있는 요청에 대해 안전하게 호출할 수 있음을 의미합니다.

3. 테스트 코드로 보는 동작 원리

이 PR의 진가는 테스트 코드에서 명확히 드러납니다. warm 요청의 프롬프트가 이미 캐싱된 상태에서 cold 요청이 먼저 도착했을 때의 시나리오를 검증합니다.

Test Case: test_cache_aware_admission_protects_resident_prefix_from_eviction

# 테스트 코드의 핵심 로직 요약
def test_cache_aware_admission_protects_resident_prefix_from_eviction(window: int):
    # 1. 캐시가 부족한 상황 설정 (5개 블록만 존재, 각 요청은 3개 블록 필요)
    scheduler, warm, cold = _setup_cache_aware(num_blocks=5, ...)
    
    scheduler.add_request(cold) # 먼저 도착
    scheduler.add_request(warm) # 나중에 도착
    output = scheduler.schedule()

    if window > 0:
        # After: Window가 설정되면 warm 요청이 먼저 입장하여 캐시를 활용함
        assert [r.req_id for r in output.scheduled_new_reqs][0] == "warm"
        assert output.num_scheduled_tokens["warm"] < warm.num_prompt_tokens
    else:
        # Before: Window가 0이면 cold가 먼저 입장하여 warm의 캐시를 밀어냄
        assert [r.req_id for r in output.scheduled_new_reqs][0] == "cold"

왜 이게 좋은가? (성능 및 교훈)

1. 실질적인 성능 향상

PR 작성자가 MI355X GPU 환경에서 테스트한 결과에 따르면, 이 최적화는 다음과 같은 이득을 가져왔습니다.

• Interactivity p90: 약 +1.96% ~ +4.95% 향상 (대기 큐가 깊을수록 효과적)
• Throughput (tok/s/GPU): 약 +4.47% 향상
• Prefix-cache hit rate: 90.7%에서 91.5%로 상승

특히 #54627(Prefill cadence gate)과 결합했을 때, p90 Interactivity가 최대 44%까지 개선되는 놀라운 결과를 보여주었습니다.

2. 안전한 최적화 설계

이 최적화는 단순히 순서를 바꾸는 것이 아니라 여러 안전 장치를 두었습니다.

• Peer-only Reordering: WAITING 상태인 요청들 사이에서만 순서를 바꿉니다. 이미 일부 연산이 진행되었거나(partially computed), 외부 상태를 기다리는(blocked) 요청의 위치는 고정하여 시스템의 일관성을 유지합니다.
• Policy Compatibility: fcfs 정책에서만 동작하며, 사용자가 명시적으로 우선순위를 부여한 priority 정책에서는 비활성화되어 충돌을 방지합니다.
• Bounded Cost: 윈도우 크기만큼만 루프를 돌기 때문에 스케줄링 단계에서의 오버헤드가 극히 적습니다.

결론

이번 PR은 "무조건 먼저 온 순서대로"라는 단순한 규칙이 복잡한 캐시 시스템에서 어떻게 독이 될 수 있는지, 그리고 이를 해결하기 위해 제한된 범위 내의 유연성(Bounded Look-ahead)을 도입하는 것이 얼마나 효과적인지 잘 보여줍니다.

시니어 엔지니어로서 배울 수 있는 교훈은, 최적화를 위해 기존의 대원칙(FCFS)을 완전히 뒤엎는 대신, 부작용 없는 프로브 함수와 명확한 활성화 임계값(Threshold)을 통해 점진적이고 안전한 개선을 이끌어내는 설계 능력입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://github.com/vllm-project/vllm/blob/main/vllm/config/scheduler.py

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