내 프리픽스 캐시가 동일한 요청을 놓친 이유

저는 macOS용 Rust LLM 추론 엔진인 mini-vllm-rs에 프리픽스 캐싱(prefix caching) 기능을 활성화했고, 콜드 캐시(cold cache)에 동일한 요청 두 개를 동시에 제출했습니다.

두 요청 모두 단 하나의 토큰도 재사용하지 못했습니다. 둘 다 같은 프롬프트를 계산했기 때문입니다.

왜일까요? 제 캐시는 요청이 완료되었을 때만 재사용 가능한 페이지를 게시(publish)했습니다. 두 요청 모두 어떤 것도 게시하기 전에 도착했기 때문에, 어느 쪽도 재사용할 프리픽스를 찾지 못했습니다. 나중에 도착한 요청들은 완성된 프리픽스를 올바르게 재사용했습니다.

동시 재사용을 지원하려면 더 많은 결정이 필요했습니다: 캐시된 페이지가 언제 가시화되어야 하는지, 다른 요청이 그것들을 기다려야 하는지, 그리고 그 대기가 언제 끝나야 하는지 말입니다?

벤치마크를 통해 저는 증분 게시(incremental publication)와 싱글-플라이트 조정(single-flight coordination)을 추가하게 되었습니다. 이때 기다리는 요청들을 언제 해제할지 결정하는 과정에서 또 다른 문제가 드러났습니다: 점점 복잡해지는 AI 생성 상태 기계였습니다.

따뜻한 캐시에서의 성공이 콜드 캐시 문제를 남겼다

저는 Mac mini의 Metal GPU에서 4비트 양자화된 Qwen2.5 0.5B 모델을 사용하고, 16토큰 캐시 페이지와 3번의 반복 실행으로 테스트했습니다. 저의 첫 재사용 벤치마크는 따뜻한(warm) 재사용과 동시 콜드 도착을 비교했습니다:

그 콜드 캐시 결과가 바로 이 글에서 제가 마주친 문제였습니다.

소유자가 아직 생성 중일 때 페이지 게시하기

저는 **증분 게시(incremental publication)**를 추가했습니다. 각 성공적인 배치 모델 포워드(batched model forward) 이후, 엔진은 커밋된 토큰을 포함하는 새로 완성된 모든 페이지를 즉시 게시하여, 이 요청이 완료되기 전에 재사용 가능하도록 했습니다.

이 변경 사항을 테스트하기 위해 요청 하나를 시작하고 첫 번째 출력 토큰이 반환될 때까지 기다렸습니다. 나머지 128개 토큰 출력이 계속 생성되는 동안, 캐시된 프리픽스를 재사용할 수 있는지 확인하기 위해 동일한 프롬프트로 더 많은 요청을 제출했습니다.

증분식 발행(incremental publication)은 각 팔로워가 493개의 프롬프트 토큰 중 480개를 복구하도록 했습니다. 2개, 4개, 8개 그룹의 동시 요청에 걸쳐, 평균 TTFT는 완료 전용 발행(completion-only publication)과 비교하여 48–93% 감소했습니다. 이 평균값에는 제출부터 첫 번째 스트리밍 출력까지 모든 요청의 대기 시간이 포함됩니다. 자세한 측정 내용은 follow-up benchmark report를 참조하십시오.

하지만, 증분식 발행은 동시에 발생하는 콜드 도착(cold arrivals)에 대한 중복 작업을 제거하지는 못했습니다. 요청이 실행을 시작할 때 다른 요청이 여전히 프리필링(prefilling) 중이더라도 이미 발행된 완전한 프리픽스 페이지를 재사용할 수 있습니다. 그러나 프롬프트의 일부만 찾거나, 아예 아무것도 찾지 못하여 나머지 부분을 스스로 계산해야 할 수도 있습니다. 따라서 재사용은 실행 순서와 발행 시기에 따라 달라집니다.

이러한 경쟁 상황(race)을 피하려면 누가 프리픽스를 계산할지 결정해야 했습니다.

요청 하나가 프리픽스를 계산하게 하고, 다른 요청들은 얼마나 기다려야 할까요?

저는 정확 입력 단일 비행(exact-input single flight) 기능을 추가했습니다. 한 요청이 리더가 되고, 동일한 입력 토큰을 가진 다른 요청들이 팔로워로서 기다립니다. 이것은 답변이 아니라 프롬프트 계산을 조정합니다. 모든 팔로워는 여전히 예상대로 자체 출력을 생성합니다. 즉, 동일한 프롬프트라도 온도(temperature)와 같은 서로 다른 생성 설정을 사용할 수 있습니다.

제가 처음 구현했을 때는 팔로워들이 리더의 전체 요청이 끝날 때까지 기다리도록 했습니다. 캐시 재사용은 예측 가능해졌지만, 이 정책 때문에 디코드(decode) 과정 동안에도 기다려야 했습니다.

각각 128 토큰의 출력을 생성하는 8개의 동시 요청을 처리했을 때, 이 정책은 독립 실행과 비교하여 평균 TTFT와 가장 느린 요청의 완료 지연 시간을 모두 줄였습니다. 하지만 단지 2개의 요청만 있을 경우, 완료 지연 시간은 개선되었음에도 불구하고 평균 TTFT가 독립 실행보다 나빠졌습니다. 중복 계산을 피하려다 새로운 대기 비용이 발생한 것입니다.

프롬프트 캐시가 이미 준비되어 있는데 왜 디코딩(decode) 과정 동안 기다려야 할까요? 저는 두 번째 릴리스 정책을 추가했습니다. 리더(leader)가 프리필(prefill)을 완료하면, 엔진은 완성된 접두사 페이지를 게시하고 팔로워(follower)들이 진입할 수 있도록 만듭니다:

sequenceDiagram
    participant S as Scheduler
    participant C as Prefix cache
...

각각 128개의 출력 토큰을 생성하는 동시 요청의 경우, 두 정책 모두 동일한 수의 캐시된 토큰을 복원했지만, 프리필 완료 후 팔로워를 해제하는 방식이 평균 TTFT를 줄였습니다:

두 정책을 독립 실행과 비교하기 위해, 아래 표는 8개 요청 버스트에 초점을 맞춥니다. 각 요청은 532 토큰의 프롬프트와 128 토큰의 출력을 가졌습니다:

정기적인 증분 게시(Regular incremental publication)는 모든 행에서 비활성화됩니다. 조기 릴리스 정책은 프리필이 완료될 때 프롬프트 접두사를 명시적으로 게시합니다.

하지만, 더 이른 릴리스가 모든 지표를 개선하지는 못했습니다. 별도의 테스트에서는 원래 요청이 디코딩 단계에 진입한 직후 두 번째 요청이 도착했습니다. 완료 시 릴리스 단일 플라이트(single flight)와 비교했을 때, 조기 릴리스는 평균 TTFT를 72% 줄였지만, 가장 느린 요청의 완료 지연 시간을 14% 증가시켰습니다.

왜일까요? 원래의 플라이트(flight)가 이미 종료되었기 때문에 두 번째 요청은 새로운 플라이트를 시작하고 따뜻한 프리픽스(warm prefix)를 재사용했습니다. 그 결과 두 요청이 순차적으로 디코드(decode)하는 대신 동시에 디코드하면서 활성 배치(active batch)가 변경되었습니다. 이것이 더 긴 완료 시간의 원인이 되었을 수 있습니다.

초기 릴리스 정책은 필요한 체크포인트 자체를 게시합니다. 추가적인 증분 발행(incremental publications)은 이러한 조정된 정확한 입력 워크로드(coordinated exact-input workloads)에는 큰 도움이 되지 않았습니다. 이들은 프리픽스의 일부만 공유하는 요청을 포함하여, 조정 없이도 유용하게 남아 있습니다.

발행(Publication)은 데이터가 언제 재사용 가능한지를 결정합니다. 조정(Coordination)은 누가, 그리고 얼마나 오랫동안 기다릴지를 결정합니다.

이 작은 구현이 남겨두는 열린 문제들

이것들은 하나의 백엔드에서 소규모 모델을 사용한 합성 워크로드입니다. 정확한 입력 조정은 이 제어(control)에서 서로 다른 입력에 대해 의미 있는 완료 시간 회귀를 보여주지 않았지만, 이 벤치마크는 프로덕션 트래픽 하에서의 동작을 확립하지 못합니다. 몇 가지 미래 방향은 다음과 같습니다:

1. 부분 프리픽스 조정(Partial-prefix coordination). 요청들이 긴 문서를 공유하지만 다른 질문으로 끝날 수 있습니다. 얼마나 많은 중복이 기다릴 가치가 있을까요?
2. 기다릴 것인가, 재계산할 것인가(Wait or recompute). 팔로워의 지연 시간 목표와 리더의 남은 프리필(prefill)이 이 선택을 바꿀 수 있습니다.
3. 다른 어텐션 상태(Different attention states). 슬라이딩 윈도우 및 하이브리드 모델은 유지(retention)와 재사용을 변경합니다. vLLM's hybrid KV cache manager가 이러한 제약 조건을 탐구합니다.
4. 백엔드 및 메모리 계층 전반의 재사용(Reuse across backends and memory tiers). 라우팅은 프리픽스 지역성(prefix locality)과 속도, 부하를 저울질해야 합니다. SGLang HiCache는 GPU, 호스트 및 외부 스토리지 전반에 걸쳐 캐싱을 확장합니다.

릴리스 변경은 또한 조정자(coordinator) 자체의 설계 문제도 노출했습니다.

오류를 수정하고, 상태 기계(state machine)에 의문을 제기하기

AI의 도움을 받아 이러한 기능들을 구현했습니다. 제 첫 번째 글에서처럼, 그 디자인을 검토하는 것은 제가 이해하고 유지할 수 있는 시스템을 구축하는 과정의 일부였습니다.

처음에 AI는 팔로워들을 방출했을 때 플라이트를 삭제했습니다. 이는 대기 시간이 끝났기 때문에 합리적으로 보였습니다.

하지만 첫 번째 싱글-플라이트 벤치마크에서 문제가 드러났습니다. 나중에 완료된 팔로워들이 코디네이터에게 더 이상 인식하지 못하는 ID를 통보했기 때문입니다. AI는 알 수 없는 요청(unknown-request) 검사를 제거할 것을 제안했습니다.

저는 그 수정 방식을 거부했습니다. 왜냐하면 그것은 실제로 유효하지 않은 요청 ID를 담고 있는 종료 알림까지 숨길 것이었기 때문입니다. 원래 디자인에는 FlightState가 없었습니다. 검증을 유지하기 위해, 우리는 Active와 Draining 상태를 추가하고 방출된 팔로워들이 종료될 때까지 추적했습니다. 메서드는 생략하고 첫 번째 수정본은 다음과 같습니다:

enum FlightState {
    Active { leader_request_id: u64 },
    Draining,
...

리더는 이미 끝났기 때문에, Draining은 방출된 팔로워들만 추적했습니다.

이전 릴리스는 그 표현을 더 복잡하게 만들었습니다. 리더가 팔로워들이 방출된 후에도 여전히 디코딩할 수 있었습니다. AI는 동일한 Flight 필드를 유지했지만, 열거형(enum)을 변경했습니다:

enum FlightState {
    Active { leader_request_id: u64 },
    Draining { leader_request_id: Option<u64> },
...

선택적 리더(Option)는 저를 멈칫하게 만들었습니다. 다른 방출 정책들은 서로 다른 조합을 요구했습니다. 만약 어떤 정책이 리더를 유지해서는 안 되지만 상태에 리더가 포함되어 있다면, 우리는 또 다른 검사가 필요할까요?

잠시 멈춰 생각했습니다: 정말로 Draining 상태가 필요했을까요? 저희는 요청이 결국 종료될 때 이를 인식하기 위해 이 상태를 추가했었습니다. 이를 위해서는 해당 요청의 ID를 기억해야 했지만, 이전 플라이트 멤버십을 보존할 필요는 없었습니다. 왜 이 ID들을 플라이트별로 따로 모으지 않았을까요?

저는 AI에게 이러한 책임 분리를 요청했습니다. 저희는 FlightState를 제거하고 해제된 요청 ID를 별도로 추적하게 했습니다:

pub(super) struct PrefixCacheFlightCoordinator {
    flights: HashMap<FlightId, Flight>,
    released_request_ids: HashSet<u64>,
...

요청이 해제될 때, 저희는 플라이트와 그 인덱스들을 제거하고, 아직 종료되지 않은 요청들의 ID만 남겼습니다. 이렇게 하면 실행에 더 이상 영향을 미치지 않는 리더/팔로워 관계를 유지하지 않으면서도 종료 검증을 보존할 수 있습니다. 알려지지 않은 요청 ID는 여전히 오류를 발생시키며, 동일한 프롬프트로 새로운 요청이 들어오면 새로운 플라이트를 시작합니다.

제가 삭제에 이의를 제기하자 AI는 검증을 유지했습니다. 나중에는 제가 추상화 자체에 의문을 제기할 때까지 그 해결책을 계속 확장했습니다.

인간의 검토는 불변성(invariant)을 보존하는 것과, 그것을 강제하기 위해 사용된 구조를 재고하는 것을 모두 의미했습니다. 벤치마크는 요청이 언제 기다림을 멈출 수 있는지 보여주었습니다. 코드를 검토한 것은 플라이트가 언제 존재할 수 없는지 보여주었습니다.