DLite가 에이전트 추론에서 초안 모델의 과거 KV 캐시 유지 비용을 없앤다

무엇인가

이 논문은 LLM 에이전트가 코딩, 웹 검색, 의료 같은 작업에서 빈번히 호출되면서 추론 효율이 병목이 되는 상황을 다룬다. 투기적 디코딩은 가벼운 초안 모델이 후보 토큰을 만들고 타깃 모델이 한 번에 검증하는 방식인데, 에이전트 워크로드는 긴 입력, 짧은 출력, 잦은 멀티턴, 계속 커지는 문맥을 특징으로 한다. 저자들은 이런 환경에서 기존 초안 모델의 수용 길이가 1~2토큰까지 떨어져 80% 넘게 악화될 수 있고, 초안 모델의 과거 KV 캐시를 유지·갱신하는 비용이 전체 추론 지연의 10~20%를 차지할 수 있다고 관찰한다. 두 문제의 공통 원인은 타깃 모델이 이미 처리한 이력을 초안 모델이 독립적으로 다시 인코딩하고 유지해야 한다는 점이다.

어떻게 동작하나

DLite의 핵심은 초안 모델이 긴 이력을 스스로 재구성하지 않게 만드는 것이다. 전역 문맥 모델링은 타깃 모델에 맡기고, 초안 쪽은 타깃 모델이 만든 히든 상태를 재사용해 현재 디코딩 위치 주변만 국소적으로 예측한다. 이렇게 하면 초안 모델이 다루는 위치 범위와 구조가 누적 시퀀스 길이와 무관해지고, 과거 초안 KV 캐시를 유지할 필요가 없어진다. 저자들은 이를 성장하는 이력 모델링 문제를 국소 예측 문제로 바꾸는 설계라고 설명한다.

무엇과 다른가

구조는 히스토리 없는 병렬 백본과 가벼운 순차 헤드로 나뉜다. 타깃 모델은 입력 문맥을 처리해 원래 KV 캐시를 유지하고, 각 투기 라운드에서 DLite는 타깃의 문맥 히든 상태, 현재 앵커 토큰, K개의 드래프트 자리표시자를 입력받는다. 병렬 백본은 모든 투기 위치의 히든 표현을 한 번에 만들고, 학습 가능한 깊이 임베딩으로 타깃 히든 상태와 일반 토큰 임베딩, 여러 타깃 레이어의 출처를 구분한다. 위치도 전역 누적 위치가 아니라 국소 위치를 쓴다. 병렬 계산이 잃는 블록 내부 인과 의존성은 저랭크 RNN 형태의 순차 헤드가 이전 드래프트 토큰을 순환 상태로 압축해 복원한다. 이 순차 모델링은 히든 공간에서 수행되고, 어휘 투영은 블록 내 의존성을 반영한 뒤에만 적용해 기존 준자기회귀 방식의 반복적인 D×V 투영 비용을 피한다.

어떻게 쓰나

훈련은 두 단계다. 타깃 모델에서 병렬 백본으로 직접 KL 증류를 하면 두 모델의 예측 조건이 너무 달라 성능이 나쁘다고 저자들은 보고한다. 그래서 워밍업 단계에서는 DLite와 같은 백본·예측 구조를 쓰되 과거 KV 캐시를 유지하는 동형의 히스토리 인지 교사를 토큰 단위 교차엔트로피로 먼저 훈련하고, 이 교사의 분포를 DLite에 전방 KL로 증류한다. 이때 가중치 w_i=exp(-(i-1)/γ)로 앞쪽 위치를 더 강조한다. 안정화 단계에서는 교사를 버리고 타깃 모델에 직접 맞추기 위해 토큰 단위 교차엔트로피와 총변동 거리 손실을 함께 쓴다. 이 교사는 자기회귀 타깃과 히스토리 없는 드래프터 사이의 최적화 다리 역할을 한다.

전제와 한계

실험은 Qwen3-4B와 Qwen3-8B에서 NVIDIA H800으로 수행됐다. 투기 블록 크기는 K=7이고, DLite는 12/14번째 타깃 히든 상태를 시퀀스 방향으로 연결하며 순차 헤드 랭크는 Qwen3-4B에서 320, Qwen3-8B에서 512를 쓴다. 비교 대상은 자기회귀 디코딩과 EAGLE-3, DFlash, DFlash2, DSpark, DFlow다. 평가는 일반 과제의 수학·코드·채팅, LongBench_v2의 장문맥 과제, AgentMemoryBench·AgentLongBench·BFCL-v3·SWE-Bench의 에이전트 과제로 나뉜다. 훈련 데이터는 130만 샘플의 Open-PerfectBlend와 60만 샘플의 NVIDIA Nemotron Post-Training Dataset V2 프롬프트를 쓰고 응답은 타깃 모델로 생성했다.

본문이 보고한 핵심 수치는 수용 길이 2~6, DSpark 대비 최대 70% 개선, SWE-Bench에서 약 2배의 종단간 속도 향상이다. 일반 과제에서는 기존 드래프터와 동등하거나 더 나았고, 장문맥에서는 문맥이 길어질수록 DLite의 장점이 더 뚜렷해졌다. 에이전트 과제에서는 상대적 이득이 가장 컸다. 동시성 1~64 구간의 처리량 실험에서도 DLite는 DSpark보다 높은 처리량을 보였는데, 수용 길이가 비슷한 경우에도 차이가 남아 순차 헤드가 중간 어휘 투영을 생략해 드래프팅 오버헤드를 줄인 효과가 있다고 해석한다. 절제 실험에서는 두 단계 훈련이 단일 단계보다 꾸준히 높은 수용 길이를 냈고, 128·512토큰 슬라이딩 윈도 KV 캐시를 둔 변형도 충분히 훈련하면 DLite와 비슷한 수용 길이를 보여 과거 KV 상태가 추론에 필수는 아니라는 주장을 뒷받침한다. 다만 발췌 본문에는 표 1~4의 셀별 수치가 모두 제시되어 있지 않아 벤치마크별 정확한 값은 원문 표를 확인해야 한다.

개발자 관점에서 DLite는 에이전트 서빙처럼 문맥이 계속 길어지고 멀티턴 호출이 잦은 환경에서 초안 모델의 KV 캐시 메모리와 갱신 오버헤드를 줄이는 설계로 참고할 만하다. 특히 타깃 모델의 히든 상태를 꺼내 쓸 수 있는 서빙 스택이라면 드래프터를 별도 이력 모델로 유지하는 대신 국소 예측기로 바꾸는 선택지를 검토할 수 있다. 도입 전에는 타깃 모델의 히든 상태 추출 지원, 블록 크기와 순차 헤드 랭크 튜닝, 두 단계 훈련 비용, 기존 SGLang·vLLM류 배포와의 호환성을 확인해야 한다.

저자들은 별도의 한계 절을 두지 않는다. 논문의 전제는 타깃 모델이 원래 KV 캐시를 유지하고 히든 상태를 제공한다는 것이며, 실험은 Qwen3-4B와 Qwen3-8B, H800 환경에 한정된다. 또한 타깃 모델에 직접 분포 수준 증류를 하기 어렵다는 구조적 간극을 인정하고 교사 모델을 거치는 두 단계 훈련으로 우회한다. 다른 모델 계열이나 하드웨어, 더 큰 블록 크기에서의 일반화는 발췌 본문만으로는 확인할 수 없다.