🤖 machine learning

PersistentKV: Page-Aware Decode Scheduling for Long-Context LLM Serving on Commodity GPUs

PersistentKV는 네이티브 블록 테이블 디코드 어텐션 엔진과 적응형 페이지 인식 스케줄링 정책을 도입하여, 배치 크기와 워크로드 특성에 따라 FlashInfer와 특화된 워크큐 전략 사이를 동적으로 선택함으로써 범용 GPU에서의 롱 컨텍스트 LLM 서빙을 최적화하고, 기존의 단일 커널 방식 대비 상당한 처리량 향상을 달성합니다.

원저자: Muhammad Ahmed

게시일 2026-06-26

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Muhammad Ahmed

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 거대한 도서관을 운영하는 사서라고 상상해 보세요. 단 한 명의 사서(AI)가 동시에 수많은 사람(사용자)의 질문에 답하려고 노력하고 있습니다. 이를 위해 사서는 모든 대화에 대한 방대한 사실이 담긴, 끊임없이 성장하는 거대한 노트( "KV 캐시")를 계속 작성해야 합니다.

문제는 현대의 도서관에서 이 노트들이 매우 거대하다는 점입니다. 사서는 실제로 답변을 쓰는 시간보다, 적절한 노트를 찾기 위해 페이지를 넘기고 서가로 걸어가는 데 더 많은 시간을 보냅니다. 이것이 AI를 느리게 만드는 "메모리 트래픽" 문제입니다.

PersistentKV는 사서의 작업 흐름을 조직하는 새로운 방법으로, 특히 슈퍼컴퓨터급 데이터 센터 장비가 아닌 일반적인 상용 컴퓨터(예: 게이밍 노트북)에서 더 빠르게 작동하도록 설계되었습니다.

다음은 쉬운 비유를 사용한 상세 설명입니다:

1. 문제점: "일률적인 방식"의 실수

현재 대부분의 AI 시스템은 FlashInfer라고 불리는 매우 효율적인 방법을 사용합니다. FlashInfer를 아주 숙련된 사서라고 생각해보세요. 이 사서는 짧고 간단한 질문을 하는 군중을 처리하는 데 매우 능숙합니다. 이들은 그룹 전체를 한꺼번에 매우 빠르게 처리할 수 있습니다.

하지만 이 방식은 다음과 같은 상황에서 어려움을 겪습니다:

군중은 적지만 질문이 거대한 경우: 만약 단 한 명의 사람이 매우 길고 복잡한 질문( "long-context" 쿼리)을 하고 있다면, 사서는 활용되지 못합니다. 사서는 다음 사람이 오기를 기다리며 시간을 낭비하게 됩니다. 즉, 놀고 있는 셈입니다.
군중이 섞여 있는 경우: 짧은 질문을 하는 사람과 거대한 긴 질문을 하는 사람이 섞여 있다면, 시스템은 모든 사람을 동일한 "배치(batch)"에 억지로 집어넣으려 합니다. 이는 마치 1페이지짜리 에세이를 쓰는 사람을 100페이지짜리 소설을 쓰는 사람과 같은 줄에 세우거나, 더 나아가 짧은 에세이를 소설처럼 보이게 하려고 빈 페이지를 채워 넣는 것과 같습니다. 이는 낭비되는 노력을 초래합니다.

2. 해결책: "스마트 분할" 전략 (PersistentKV)

저자들은 PersistentKV라는 새로운 시스템을 구축했습니다. 모든 사람을 하나의 큰 그룹으로 묶는 대신, 이 시스템은 각 사람의 구체적인 요구 사항을 살펴보고 작업을 다르게 나누는 스마트한 매니저처럼 행동합니다.

"분할" 비유: 읽어야 할 긴 소설이 있다고 상상해 보세요. 한 사람이 한 번에 전체를 읽는 대신, 매니저는 책을 32개의 작은 장으로 나눕니다. 그리고 여러 명의 조수에게 동시에 서로 다른 부분을 읽도록 배정합니다.
- 이것이 도움이 되는 이유: 만약 단 한 명이 긴 질문을 하고 있다면, 이 "분할" 방식은 사서의 팀이 그 하나의 긴 이야기를 동시에 읽게 함으로써 팀원들을 바쁘게 유지합니다. 이는 컴퓨터의 뇌 속에 있는 "빈 좌석"을 채워줍니다.
"작업 큐(Work Queue)" 비유: 기존 시스템에서는 8명의 사람이 서로 다른 길이의 이야기를 하고 있다면, 시스템은 (각기 다른 길이에 맞춰) 16개의 서로 다른 작은 작업들을 시작하려고 시도하여 혼란스럽고 느려질 수 있습니다.
- PersistentKV의 해결책: 이 시스템은 "압축된 작업 큐(compact work queue)"를 사용합니다. 8명의 사람을 살펴보고, 누가 정확히 무엇을 필요로 하는지 파악한 뒤, 단 하나의 효율적인 작업 목록을 만듭니다. 빈 페이지를 건너뛰고, 실제로 작업이 필요한 곳에만 일을 보냅니다.

3. "적응형 정책": 스마트 매니저

이 논문에서 가장 중요한 부분은 단순히 새로운 도구를 만든 것이 아니라, 바로 의사 결정 규칙입니다. 저자들은 "분할" 전략이 항상 더 나은 것은 아니라는 점을 깨달았습니다.

시나리오 A (적은 인원, 긴 이야기): 만로 단 1명의 사람이 긴 이야기를 하고 있다면, 새로운 "분합" 방식이 승리합니다. 이 방식은 속도를 1.4배 높여줍니다.
시나리오 B (중간 규모의 그룹, 혼합된 이야기): 8명의 사람이 다양한 길이의 이야기를 하고 있다면, "압축된 작업 큐"가 승리합니다. 이 방식은 속도를 약 1.2배 높여줍니다.
시나리오 C ("골디락스" 존 - 4명): 만약 4명의 사람이 있다면, 새로운 방식은 오히려 작업의 분할과 병합 과정에서 발생하는 오버헤드 때문에 더 느려집니다.
- 해결책: 시스템은 똑똑하게도 "잠깐, 4명일 때는 기존의 믿음직한 FlashInfer 방식을 쓰자"라고 판단합니다. 상황에 따라 자동으로 도구를 전환합니다.

4. 결과: 실제로 어떤 일이 일어났는가?

연구진은 일반적인 RTX 3060 그래픽 카드(슈퍼컴퓨터가 아닌 흔한 소비자용 GPU)에서 테스트를 진행했습니다.

정확도: 답변은 기존 방식과 비교했을 때 (아주 미세한 오차 범위 내에서) 똑같이 정확했습니다.
속도:
- 단일, 매우 긴 대화의 경우, 40% 더 빨랐습니다.
- 다양한 대화 길이를 가진 8명의 그룹의 경우, 6%에서 26% 더 빨랐습니다.
- 4명의 그룹의 경우, 새로운 방식을 시도하지 않고 기존 방식을 고수하여 속도가 느려지는 것을 방지했습니다.

핵심 요약

이 논문은 자신들의 새로운 방법이 모든 상황에서 "최고"라고 주장하는 것이 아닙니다. 대신, 작업을 어떻게 스케줄링하느냐가 수학적 계산만큼이나 중요하다는 것을 증명합니다.

AI를 거대한 작업을 조각으로 나누어 처리할지, 아니면 기존의 루틴을 따를지 결정할 줄 아는 유연한 매니저처럼 다룸으로써, 일반 컴퓨터가 길고 복잡한 AI 대화를 훨씬 더 빠르게 실행할 수 있게 만듭니다. 이는 모든 못에 똑같은 망치를 사용하는 대신, 군중의 크기에 맞는 적절한 도구를 찾는 것에 관한 이야기입니다.

기술 요약: PersistentKV

문제 정의

자기회귀적 대규모 언어 모델(LLM) 서빙은 밀집 행렬 곱셈(dense matrix multiplication)보다는 키-값(KV) 캐시 이동에 의해 점점 더 병목 현상이 발생하고 있습니다. 현대적인 페이지드 어텐션(paged-attention) 시스템(예: vLLM, FlashInfer)은 KV 캐시 파편화를 성공적으로 줄이고 네이티브 페이지드 디코드 커널을 최적화했지만, 모든 서빙 시나리오에 적합한 단 하나의 "최적의" 커널은 존재하지 않습니다.

구체적으로, 다음과 같은 긴장 관계가 존재합니다:

저활성 장기 컨텍스트 디코드(Low-active long-context decode): 단일 최적화 커널이 불충분한 병렬 작업으로 인해 범용 GPU를 충분히 활용하지 못하는 경우.
혼합된 시퀀스 길이(Mixed sequence lengths): 많은 정확한 길이(exact-length) 런치를 생성하면 오버헤드가 높고, 거친 패딩 배치(coarse padded batches)는 비효율성을 초래하는 경우.

본 논문은 네이티브 블록 테이블 디코드 엔진과 페이지 인식 스케줄링을 결면했을 때, 이미 프로덕션 품질의 페이지드 어텐션 커널이 강력한 성능을 보이는 영역에서도 범용 GPU 상의 서빙 지연 시간을 개선할 수 있는지 조사합니다.

방법론

저자들은 네이티브 블록 테이블 디코드 엔진과 페이지 인식 스케줄링 정책으로 구성된 시스템인 PersistentKV를 제안합니다.

1. 네이티브 블록 테이블 디코드 엔진

PersistentKV는 $G = H_q/H_{kv}$ 개의 쿼리 헤드가 하나의 KV 헤드를 공유하는 그룹화된 쿼리 어텐션(GQA)을 위해 설계되었습니다.

레이아웃: K, V 페이지를 연속적인 텐서로 재구성(repacking)하지 않고, 네이티브 페이지드 KV-캐시 레이아웃(블록 테이블) 위에서 직접 실행됩니다.
시퀀스 분할(Sequence Splitting): 저활성 워크로드에 충분한 병렬성을 노출하기 위해, 시퀀스 차원을 $S$ 개의 서로 소인 범위로 분할합니다. 각 분할은 부분적인 온라인 소프트맥스 상태(partial online-softmax state)를 계산하며, 별도의 머지 커널(merge kernel)이 이 상태들을 정확하게 결합합니다.
행 로컬 경계(Row-Local Bounds): 패딩된 길이를 사용하는 버킷 방식과 달리, PersistentKV는 실제 행별 시퀀스 길이( $N^*_i$ )를 사용하여 타일 개수와 분할 범위를 결정합니다. 빈 분할은 중립적인 상태를 쓰고 조기에 종료되어, 패딩된 KV 읽기를 방지합니다.
컴팩트 워크큐(Compact Workqueue): 다양한 길이의 요청(ragged batches)을 처리할 때 과도한 커널 런치를 피하기 위해, 스케줄러는 비어 있지 않은 행-KV-헤드-분할 구간에 대해서만 1D 그리드 작업을 방출합니다. 이는 시퀀스 병렬성을 유지하면서도 정확한 길이 버케팅(exact-length bucketing)의 "런치 팬아웃(launch fan-out)" 문제를 방지합니다.

2. 보정된 적응형 정책 (Calibrated Adaptive Policy)

시스템은 배치 크기( $B$ )와 최대 컨텍스트 길이( $N_{max}$ )에 따라 실행 경로를 선택하는 수동 보정 정책을 채택합니다.

$B=1$ (장기 컨텍스트): 길이 버킷과 시퀀스 분할을 사용하여 충분한 연산 스레드 배열(CTA)을 노출하는 PersistentKV를 사용합니다.
$B \ge 8$ (장기 컨텍스트): 분할 및 머지의 오버헤드를 상쇄하기 위해 컴팩트 워크큐를 사용하는 PersistentKV를 사용합니다.
$B=4$ (경계값): FlashInfer로 라우팅합니다. 이 배치 크기에서는 PersistentKV의 분할 및 머지 오버헤드가 이점보다 커져서 실제 실행 시간(wall-clock time)에서 성능 저하를 일으킵니다.
그 외의 경우: FlashInfer를 사용합니다.

3. 평가 프로토콜

하드웨어: NVIDIA RTX 3060 (28 SMs, 12GB VRAM).
베이스라인: FlashInfer 0.2.5 (주요 비교 대상), vLLM, TensorRT-LLM.
트레이스(Traces): 컨텍스트 길이, 활성 배치 크기, 길이의 불규칙성(raggedness)을 격리하기 위해 설계된 합성 트레이스(Bucketed, Bimodal, Uniform, Zipf).
메트릭: 정확성(FlashInfer 대비 FP16 허용 오차), CUDA 이벤트 타이밍(큐잉된 GPU 작업), 동기화된 월 타임(end-to-end latency, 호스트 플래닝 포함).

주요 기여

네이티브 블록 테이블 GQA 엔진: KV 캐시를 재구성하지 않고도 저활성 장기 컨텍스트 디코드에서 시퀀스 병렬 작업을 노출하는 디코드 엔진.
페이지 인식 스케줄링: 실제 요청 길이에 맞춰 분할 경계와 프리페치(prefetch)를 정렬하고 컴팩트 워크큐를 구현하여, 직사각형 형태의 빈 그리드를 피하는 행 로컬 시퀀스 길이 스케줄링.
적응형 방법론: 보정/홀드아웃 트레이스 시드, 정확성 검사, 스케줄러 유도 워크로드 카운터를 사용하는 엄격한 평가 프레임워크를 통해, 작업 할당이 어텐션 연산만큼이나 결정적인 변수임을 입증함.

결과

고정된 임계값을 가진 보정 시드(20260622)로부터 유도된 홀드아웃 시드(20260623)에 대한 결과:

단일 요청 ( $B=1$ ): 버킷형 장기 컨텍스트 트레이스에서 PersistentKV(길이 버킷 사용)는 FlashInfer 대비 동기화된 월 스루풋(synchronized wall throughput)에서 1.399배 향상된 성능을 달달성했습니다.
높은 배치 ( $B=8$ ):
- Bimodal: 1.265배 향상.
- Uniform: 1.082배 향상.
- Zipf-like: 1.063배 향상.
경계 사례 ( $B=4$ ): 적응형 정책이 FlashInfer를 올바르게 선택하여, PersistentKV가 성능 저하를 일으킬 수 있는 상황을 방지했습니다.
스케줄링 오버헤드: 컴팩트 워크큐는 bimodal $B=8$ 트레이스에서 런치 팬아웃을 step당 16.00회(정확한 버킷 기준)에서 2.00회로 줄였으며, 머지 상태 트래픽을 약 37% 감소시켰습니다.

의의 및 주장

본 논문은 다음과 같은 겸손하면서도 시스템 중심적인 주장을 합니다: PersistentKV는 적응형 페이지 인식 디코드 스케줄링을 위한 구체적인 시스템적 틈새(niche)를 식별했습니다.

보편적인 커널이 아님: 저자들은 PersistentKV가 항상 더 빠른 것은 아니라고 명시합니다. 분할 및 머지 오버헤드가 지배적인 작은 활성 배치에서는 FlashInfer가 여전히 우수합니다.
작업 할당의 중요성: 결과는 "어텐션 연산뿐만 아니라 작업 할당(work assignment) 또한 결정적인 서빙 시스템 변수"임을 보여줍니다.
워크숍 수준의 기여: 본 논문은 재현 가능한 방법론과 절제 실험(ablation)을 제공하는 워크숍 수준의 기여를 목표로 합니다. 저자들은 저활성 장기 컨텍스트 환경의 범용 GPU에서 시퀀스를 분할하고 컴팩트 워크큐를 사용하는 것이 단일 최적화 커널보다 훨씬 나은 성능을 낼 수 있음을 강조합니다.
향후 방향: 저자들은 프로덕션급 시스템을 위해서는 학습된 비용 모델(learned cost model), 전체 서빙 런타임(MLP 및 샘플링 포함)으로의 통합, 그리고 다른 성능 특성을 가진 Hopper 세대 하드웨어에서의 테스트가 필요하다고 언급합니다.