[Paper] DREX를 이용한 효율적인 Early-Exit Inference를 위한 Dynamic Rebatching

Source: arXiv - 2512.15705v1

개요

Early‑Exit (EE) 기법은 대형 언어 모델(LLM)이 “쉬운” 토큰에 대해 불필요한 레이어를 건너뛰게 하여 추론 지연 시간을 줄입니다. 논문 Dynamic Rebatching for Efficient Early‑Exit Inference with DREX는 기존 배치 파이프라인이 배치를 하나의 단일 객체로 취급하기 때문에 이 기회를 놓친다고 지적합니다—모든 요청이 동시에 종료하거나 전혀 종료되지 않습니다. 저자들은 Dynamic Rebatching을 도입했으며, 이는 실행 중에 배치를 실시간으로 재구성하여 각 요청이 처리량이나 품질을 손상시키지 않으면서 최적 레이어에서 종료될 수 있게 합니다.

주요 기여

• Dynamic Rebatching concept – 매 조기 종료 체크포인트마다 배치를 재구성하여, 종료 조건을 만족하는 토큰을 즉시 확정하고 나머지는 더 깊은 처리를 위해 버퍼링합니다.
• DREX system – 동적 재배치를 구현한 프로덕션 수준 추론 엔진으로, 두 가지 성능 향상 트릭을 포함합니다:
  1. Copy‑free rebatching buffer – 요청을 재배열할 때 비용이 많이 드는 메모리 복사를 방지합니다.
  2. EE‑and SLA‑aware scheduler – 재배치 단계가 전체 지연시간/처리량을 개선할지 분석적으로 예측하여 해로운 재구성을 방지합니다.
• Efficient KV‑cache handling – 건너뛴 레이어에 대한 누락된 키‑값 캐시 엔트리를 재구성하는 메모리 경량 방법을 제공하여 트랜스포머 캐시의 속도 이점을 유지합니다.
• Guarantee of no involuntary exits – DREX는 요청이 EE 모델 자체의 신뢰 임계값보다 일찍 종료되지 않도록 보장하여 출력 품질을 보호합니다.
• Empirical validation – 기존 EE 배치 기준에 비해 2–12 % 높은 처리량을 보여주면서도 동일한 생성 품질을 유지합니다.

방법론

1. Early‑Exit 체크포인트 – 모델은 선택된 트랜스포머 레이어 뒤에 여러 개의 “exit heads”를 삽입한다. 각 헤드는 신뢰도 점수를 생성하며, 이 점수가 사전 조정된 임계값을 초과하면 토큰을 조기에 출력할 수 있다.

2. Dynamic rebatching loop:

   • 종료 지점에서 런타임은 현재 배치를 스캔한다.
   • 신뢰도 기준을 만족하는 토큰은 finalized(최종 처리)되어 이후 처리에서 제거된다.
   • 남은 토큰은 rebatching buffer에 배치되며, 이 버퍼는 원래 위치를 추적하지만 기본 텐서 데이터를 복사하지 않는다.
   • 버퍼는 보류 중인 토큰을 새로운 배치(크기가 달라질 수 있음)로 묶어 다음 깊은 레이어로 전달한다.

3. Copy‑free buffer design – 인덱스 기반 간접 참조(예: 포인터/오프셋 벡터)를 사용하여 동일한 활성화 메모리를 재배치 단계마다 재사용한다. 이를 통해 O(N) 데이터 이동을 제거한다.

4. Scheduler analytics – 잠재적인 rebatch마다 DREX는 레이어 지연 시간, 배치 크기 스케일링, SLA 제약(예: 토큰당 최대 지연 시간) 등을 경량 모델로 활용해 비용/이익 트레이드‑오프를 추정한다. 예측된 이익이 음수이면 스케줄러는 rebatching을 연기하고 현재 배치를 그대로 유지한다.

5. KV‑cache reconstruction – 토큰이 레이어를 건너뛰면 DREX는 가장 가까운 캐시된 상태를 복사하고 저비용 선형 프로젝션을 적용해 누락된 key‑value 엔트리를 합성한다. 이를 통해 캐시 크기를 제한한다.

결과 및 발견

• 처리량 향상은 배치 크기가 커지고 종료 지점이 깊어질수록 증가합니다. DREX는 배치를 충분히 채워 GPU를 계속 바쁘게 유지하면서도 초기 종료 토큰이 신속히 나갈 수 있게 합니다.
• **비자발적 종료 0%**는 이전 EE 시스템에서 배치를 균일하게 유지하기 위해 토큰이 조기에 강제로 종료되는 미묘한 품질 버그를 제거합니다.
• 스케줄러 정확도: 분석적 이익 모델은 95 % 이상의 경우에 수익성 있는 재배치를 정확히 예측하여, 불필요한 재배치를 방지하고 오버헤드를 최소화합니다.

Practical Implications

• LLM API 배포자는 DREX를 통합하여 클라우드 GPU 비용을 절감하면서 조기 종료 모델이 약속하는 품질 보장을 유지할 수 있습니다.
• 지연 민감 애플리케이션(예: 실시간 코드 완성, 대화형 에이전트)은 쉬운 토큰이 배치 내 가장 느린 요청을 기다리지 않고 즉시 처리되므로 꼬리 지연이 낮아지는 혜택을 받습니다.
• 프레임워크 개발자(PyTorch, TensorFlow, Triton)는 메모리 복사 비용 없이 동적 배치 크기를 지원하기 위해 복사 없는 재배치 버퍼 패턴을 채택할 수 있습니다.
• SLA 인식 스케줄링은 일부 요청은 엄격한 지연 제한을, 다른 요청은 처리량을 우선시하는 하이브리드 워크로드를 가능하게 하며, DREX는 두 가지를 자동으로 균형 맞출 수 있습니다.
• KV 캐시 재구성은 레이어가 건너뛰어도 트랜스포머 캐싱을 효과적으로 유지하는 실용적인 방법을 보여주며, 이 패턴은 다른 조건부 실행 모델(예: Mixture‑of‑Experts)에도 재사용될 수 있습니다.

제한 사항 및 향후 연구

• 모델별 튜닝 – 종료 임계값과 스케줄러 비용 모델은 모델 아키텍처와 하드웨어마다 보정이 필요하며, 일괄 적용 가능한 설정은 제공되지 않는다.
• GPU 메모리 단편화 – 복사 없이 동작하지만, 간접 버퍼가 메모리 접근을 분산시켜 제한된 캐시를 가진 GPU에서 성능에 영향을 줄 수 있다.
• 대규모 배치 크기에 대한 확장성 – 보고된 성능 향상은 약 256 토큰 배치를 초과하면 정체되며, 극한 규모에서도 이점을 유지하기 위한 추가 연구가 필요하다.
• 멀티노드 추론으로의 확장 – DREX는 현재 단일 노드 GPU 환경을 목표로 하며, 분산 시나리오에서는 노드 간 재배치를 조정해야 한다.
• 보다 폭넓은 EE 전략 – 본 논문은 신뢰도 기반 종료에 초점을 맞추고 있으나, 향후 연구에서는 토큰 수준 난이도 추정기와 같은 혼합 기준을 탐색하고, 적응 깊이 또는 혼합 전문가와 같은 다른 조건부 실행 기법과 통합할 수 있다.

저자

• Xuting Liu
• Daniel Alexander
• Siva Kesava Reddy Kakarla
• Behnaz Arzani
• Vincent Liu

논문 정보

• arXiv ID: 2512.15705v1
• 분류: cs.DC, cs.LG
• 출판일: 2025년 12월 17일
• PDF: PDF 다운로드