기술 요약: PARS – 저지연 LLM 서빙을 위한 프롬프트 인지형 랭킹 스케줄러 (Prompt-Aware Ranking Scheduler)

문제 정의

대규모 언어 모델(LLM), 특히 최근 등장하는 추론 가능 모델(예: DeepSeek R1, GPT-o1)의 배포는 추론 스케줄링에서 결정적인 과제인 헤드 오브 라인(Head-of-Line, HOL) 블로킹 문제에 직면해 있습니다. vLLM 및 Orca와 같은 기존의 서빙 시스템은 주로 선입선출(First Come, First Serve, FCFS) 정책에 의존합니다. FCFS 방식에서는 긴 실행 시간이 소요되는 추론 작업(종종 매우 가변적이고 광범한 사고 사슬(Chain-of-Thought) 추론 흔적을 특징으로 함)이 뒤에 대기 중인 짧은 작업들을 차단하여, 심각한 지연 시간 급증과 처리량 감소를 초래합니다.

최단 작업 우선(Shortest-Job-First, SJF) 정책은 평균 지연 시간을 최소화하는 데 이론적으로 최적이지만, LLM의 확률적 및 자기회귀적(auto-regressive) 특성으로 인해 적용이 어렵습니다. 정확한 출력 길이는 실행 전에는 알 수 없기 때문입니다. 길이를 예측하여 SJF를 근사하려는 기존의 시도들은 세 가지 주요 한계점을 가집니다:

1. 노이즈 민감도: 점수 기반(Pointwise, 회귀) 및 리스트 기반(Listwise) 랭킹 방식은 유사한 프롬프트가 실행 시마다 변동하는 길이를 생성하는 LLM 출력의 내재적 가변성을 처리하는 데 어려움을 겪으며, 이는 노이즈가 섞인 학습 신호를 유발합니다.
2. 추론 복잡성: 표준 예측 모델들은 현대 모델에서 흔히 나타나는 다단계 추론 과정의 전체 복잡성을 반영하지 못하는 경우가 많습니다.
3. 일반화 능력 부족: 많은 접근 방식이 모델별 재학습을 요구하므로, 다양하거나 독점적인 모델 배포 환경에서는 실용성이 떨어집니다.

방법론: PARS 프레임워크

저자들은 **쌍체 학습-투-랭크(pairwise learning-to-rank)**를 통해 SJF 스케줄링을 근사하도록 설계된 경량 프레임워크인 **PARS (Prompt-Aware Ranking Scheduler)**를 소개합니다. PARS는 두 가지 주요 구성 요소로 이루어져 있습니다.

1. 쌍체 학습된 예측기 (Pairwise-Trained Predictor)

절대적인 출력 길이를 예측(Pointwise)하거나 전체 리스트를 랭킹(Listwise)하는 대신, PARS는 **마진 랭킹 손실(margin ranking loss)**을 사용하는 **쌍체 비교(pairwise comparisons)**에 집중합니다.

• 학습 목표: 모델은 어떤 프롬프트가 더 긴 응답을 생성할 것으로 예상되는지를 나타내는 이진 레이블이 있는 프롬프트 쌍 $(A, B)$를 학습합니다.
• 노이즈 필터링: 확률성을 해결하기 위해, 실제 정답 응답 길이 간의 차이가 유의미한(상대적 분산 임계값 $\delta = 0.2$로 정의됨) 프롬프트 쌍만을 포함하도록 학습 데이터를 필터링합니다. 이를 통해 모델이 샘플링 노이즈에 취 의해 순위가 뒤바뀔 수 있는 근소한 차이의 비교에 매달리는 대신, 상대적 순서가 견고한 "정보가 풍부한" 쌍으로부터 학습하도록 보장합니다.
• 아키텍처: 예측기는 경량 BERT-base-uncased 백본을 활용합니다. 프롬프트로부터 [CLS] 토큰 임베딩을 추출하여 선형 레이어에 통과시킨 후, 예상 응답 비용을 나타내는 스칼라 점수를 생성합니다.
• 손실 함수: 모델은 마진 랭킹 손실 $L(s_A, s_B, y) = \max(0, -y \cdot (s_A - s_B) + \text{margin})$을 사용하여 최적화되며, 이는 올바르게 랭킹된 쌍들 사이의 점수 격차를 강제합니다.

2. 예측기 가이드 스케줄러 (Predictor-Guided Scheduler)

학습된 예측기는 **대기 큐(Waiting Queue, W)**와 **실행 큐(Running Queue, R)**를 관리하기 위해 vLLM 서빙 시스템에 통합됩니다.

• 스케줄링 로직: 요청이 도착하면 예측기에 의해 점수가 매겨집니다. 대기 큐는 예상 길이의 오름차순으로 정렬되어 짧은 작업을 우선시함으로써 효과적으로 SJF를 근사합니다.
• 기아 현상 방지 (Starvation Prevention): 긴 실행 시간이 소요되는 작업이 무기한 연기되는 것을 방지하기 위해, 특정 임계값(기본값: 2분) 이상 대기한 모든 요청의 우선순위를 높이는 메커니즘을 갖추고 있습니다.
• 구현: 스케줄러는 반복(iteration) 단위로 작동하며, 연속 배치(continuous batching)와 호환됩니다. 또한 vLLM의 핵심 C++/CUDA 백엔드를 수정할 필요가 없습니다.

주요 기여

1. 새로운 쌍체 접근 방식: 본 논문은 마진 랭킹 손실과 데이터 필터링을 사용하는 쌍체 학습-투-랭크 방식을 제안하며, 이는 LLM 생성의 확률적 특성에 대해 점수 기반(pointwise) 및 리스트 기반(listwise) 베이스라인보다 더 높은 안정성을 입증했습니다.
2. 추론 인지 설계: 본 프레임워크는 응답 길이의 정의에 전체 추론 흔적(중간 사고 및 정당화 과정)을 명시적으로 포함하여, 현대의 추론형 LLM에 적합하도록 설계되었습니다.
3. 원활한 통합: PARS는 vLLM 내의 오픈 소스 모듈(약 1,000 라인의 코드)로 구현되었으며, 핵심 추론 엔진의 변경 없이도 사용 가능합니다.
4. 교차 모델 일반화: 특정 모델(예: GPT-4)에서 학습된 예측기가 재학습 없이도 다른 다양한 모델(예: Llama 3.1, DeepSeek-R1)을 위한 작업을 효과적으로 스케줄링할 수 있음을 보여줍니다.

실험 결과

다양한 데이터셋(Alpaca, LMSYS-Chat-1M, MATH-plus, Evolved Codealpaca)과 모델(Llama 3.1, GPT-4, DeepSeek-R1)에 대해 광범위한 평가를 수행했습니다.

• 예측 정확도: PARS는 모든 데이터셋에서 가장 높은 Kendall's Tau ($\tau_b$) 랭킹 정확도를 달성했으며, Alpaca 데이터셋에서는 최대 0.96에 도달했습니다. 이는 점수 기반(pointwise) 및 리스트 기반(listwise) 베이스라인보다 성능이 우수했으며, 특히 추론 모델에서 뛰어난 성능을 보였습니다(LMSYS-Chat-1M의 DeepSeek-R1에 대해 점수 기반 SJF 대비 5.5배 향상).
• 지연 시간 감소: 실제 워크로드에서 PARS는 Llama 3.1에 대해 기본 vLLM FCFS 스케줄러 대비 평균 토큰당 지연 시간을 최대 15.7배 감소시켰습니다. 버스트 부하 조건(2,000개의 동시 요청)에서 Llama에서는 최대 7.7배, DeepSeek-R1에서는 2배 이상의 속도 향상을 달ach했습니다.
• 꼬리 지연 시간 (Tail Latency): PARS는 p90 지연 시간을 크게 개선하여 FCFS 대비 최대 12.3배 감소시켰으며, 이는 이상적인 "Oracle SJF"(정답 길이를 사용하는 방식)의 성능에 근접하는 수준입니다.
• 처리량 및 오버헤드: PARS는 배치 동질성(batch homogeneity)을 높임으로써 시스템 처리량(초당 요청 수)을 유지하거나 약간 향상시켰습니다. 점수 계산 및 정렬을 위한 런타임 오버헤드는 약 **0.08%**로 미미했으며, 요청 실행에 추가적인 지연을 더하지 않았습니다.
• 일반화: GPT-4 데이터로만 학습된 "Cross-Model PARS" 변형은 Llama 및 DeepSeek-R1을 성공적으로 스케줄링했으며, 모델별 점수 기반 베이스라인보다 우수한 성능을 보였고 리스트 기반 방식과 대등한 성능을 보였습니다.

의의 및 주장

본 논문은 추론 능력을 갖춘 LLM이 보편화됨에 따라, 실행 시간의 가변성이 FCFS를 점점 더 비효율적으로 만들 것이라고 주장합니다. PARS의 의의는 무겁거나 모델 특정적인 스케줄링 솔루션에 대한 실용적이고 배포 가능하며 경량화된 대안을 제공한다는 점에 있습니다.

저자들의 주장은 다음과 같습니다:

• 쌍체 랭킹은 LLM 스케줄링에 더 우수합니다. 이는 불안정한 절대적 예측 대신 견고한 상대적 비교에 집중하기 때문입니다.
• 재학습 없이 일반화가 가능합니다. 따라서 단일 예측기로 이질적인 LLM 플릿(fleet)을 관리할 수 있습니다.
• 최소한의 오버헤드를 통해 응답성이 중요한 프로덕션 환경에서도 이 방식이 실용적임을 입증했습니다.

결론적으로, 이 연구는 이득의 크기가 작업 길이의 가변성에 따라 달라질 수 있지만, PARS가 상대적 실행 비용을 포착하는 강력한 메커니즘을 제공하여 HOL 블로킹을 완화하고 차세대 AI 애플리케이션을 위한 저지연 서빙을 가능하게 한다고 밝히고 있습니다. 소프트웨어와 데이터셋은 연구 및 실제 배포에서의 채택을 촉진하기 위해 오픈 소스로 공개되었습니다.