LLM 추론(inference)에서 가장 큰 비용 중 하나는 디코딩, 즉 토큰을 생성하는 과정이다.
모델이 커질수록, 그 생성 속도는 점점 더 하드웨어의 물리적 한계에 가로막히고 있다.
이 글에서는 모델의 구조나 품질을 전혀 건드리지 않으면서도 디코딩 속도를 2~2.5배 끌어올린 아이디어, Speculative Decoding에 대해 다룬다.
Motivation: LLM 디코딩은 왜 이렇게 느린가?
LLM은 텍스트를 생성할 때 토큰을 하나씩 순차적으로 만들어낸다.
"I love you"를 생성한다면, "I", "love", "you"를 각각 별개의 단계에서 생성하는 것이다.
문제는 이 각각의 단계마다 수십~수백억 개에 달하는 모델의 전체 파라미터를 GPU 메모리에서 읽어와 연산을 수행해야 한다는 점이다. GPU의 연산 능력(Compute)은 충분히 남아있는데, 데이터를 실어 나르는 **메모리 대역폭(Memory Bandwidth)**이 이를 따라가지 못한다.
FlashAttention을 다룰 때 언급했던 '메모리 벽(Memory Wall)'이 여기서도 등장한다. 다만 그때는 어텐션 행렬의 읽기/쓰기가 문제였다면, 여기서는
토큰 하나를 생성할 때마다 모델 전체를 메모리에서 끌어오는 비용
이 병목이다.
결과적으로 토큰 하나를 생성하는 데 걸리는 평균 시간, 즉 **TPOT(Time Per Output Token)**은 모델 파라미터의 크기에 거의 비례하는 물리적 한계를 갖게 된다. 모델이 크면 클수록, 한 토큰을 생성하는 데 걸리는 시간은 늘어날 수밖에 없다.
그렇다면,
"어차피 한 번 forward pass를 돌리는 김에, 토큰을 여러 개 한꺼번에 처리할 수 있다면 어떨까?"
이 질문이 Speculative Decoding의 출발점이다.
핵심 관찰: 1개 생성과 K개 검증은 시간이 거의 같다
Speculative Decoding이 가능한 이유는 Transformer의 특성에서 비롯된다.
Transformer는 학습(Training) 시에는 다음 토큰을 한꺼번에 병렬로 예측할 수 있다. 이 구조 덕분에 이미 토큰들이 주어져 있다면, 각 위치에서의 확률 분포를 한 번의 forward pass로 동시에 계산할 수 있다. 이것을 **Scoring(병렬 검증)**이라고 부른다.
핵심은 이 Scoring의 비용이 단일 토큰 생성과 크게 다르지 않다는 점이다.
그 이유를 살펴보면,
- Linear Layer: 단일 토큰을 생성할 때는 행렬-벡터 곱(GEMV) 연산이 수행된다. K개의 토큰을 병렬로 검증할 때는 이것이 행렬-행렬 곱(GEMM)으로 바뀐다. GEMM은 GPU가 훨씬 효율적으로 처리하는 연산이지만, K가 작은 범위에서는 그 차이가 크지 않아 전체 소요 시간은 비슷하게 유지된다. 오히려 하드웨어 활용률이 소폭 올라가는 효과도 있다.
- Attention의 KV Cache: 이전 토큰들의 Key/Value를 저장해두는 캐시 크기는 K가 늘어도 변하지 않는다. 따라서 어텐션 연산의 추가 비용도 미미하다. 오히려 여러번 재활용하기 때문에 compute intensity가 증가하는 효과를 얻을 수 있다.
- 분산 환경의 All-reduce 통신: 모델을 여러 칩에 분산할 때 발생하는 통신 비용 역시, K가 작은 범위에서는 거의 증가하지 않는다.
즉, 작은 K에 대해서 **"K개 토큰을 검증하는 비용 ≈ 1개 토큰을 생성하는 비용"**이라는 등가 관계가 성립한다. Speculative Decoding은 이 등가 관계를 적극적으로 활용한다.
동작 원리: 추측하고, 검증하고, 보정한다
전체 알고리즘은 세 단계로 구성된다.
1단계 - Drafting: 빠르게 초안을 만든다
먼저 draft model이라 불리는 작고 빠른 모델이 K개의 토큰을 순차적으로 자동 회귀 생성한다.
부하를 시켜서 미리 추론을 끝내도록 명령하는 것이다.
draft model은 파라미터 수가 적어 토큰 하나를 생성하는 데 걸리는 시간이 매우 짧다. 논문에서는 70B짜리 target model(Chinchilla)에 대해 4B짜리 draft model을 사용했는데, 토큰 생성 속도가 약 7.8배 빠르다.
이 단계의 목적은 정확한 답을 내놓는 것이 아니라, 검증받을 후보 초안을 빠르게 마련하는 것이다.
2단계 - Scoring: target model이 한 번에 검증한다
draft model이 만들어낸 K개의 토큰을 target model이 한 번의 forward pass로 병렬 검증한다.
각 위치에서 target model이 생각하는 확률 분포를 동시에 계산하는 것이다.
앞서 설명한 핵심 관찰이 여기서 작동한다. K개를 검증하는 비용이 1개를 생성하는 비용과 거의 같으므로, 이 단계의 오버헤드는 크지 않다.
3단계 - Modified Rejection Sampling: 수용하거나, 보정한다
이제 draft model의 초안을 그대로 쓸 수 있을지 판단해야 한다. 이 과정이 Modified Rejection Sampling이며, Speculative Decoding의 수학적 핵심이다.
판단 기준은 직관적이다.
- 수용(Accept): target model이 해당 토큰에 부여한 확률이 draft model의 확률보다 크거나 같다면, 그 토큰은 무조건 통과된다. 큰 모델이 작은 모델보다 그 토큰을 더 선호한다는 뜻이므로, 그대로 써도 문제없다.
- 거부(Reject): draft model이 과대평가한 토큰은 두 확률의 비율에 따라 확률적으로 거부된다. 거부가 발생하면 이후의 draft 토큰은 전부 폐기하고, target model이 해당 위치에서 직접 보정 토큰 하나를 생성한 뒤 다음 루프를 새로 시작한다.
마치 CPU의 Branch Misprediction 처리와 유사한 방식이다. 예측이 맞으면 미리 해둔 작업을 그대로 활용하고, 틀리면 파이프라인을 비우고 다시 시작하는 것이다.
이 수용/거부 메커니즘의 핵심 보장은, 수용된 토큰과 거부 후 보정된 토큰을 합산하면 정확히 target model의 원래 확률 분포가 복원된다는 것이다. 즉, draft model을 거쳐도 최종 출력은 target model이 직접 생성한 것과 통계적으로 동일하다.
왜 TPOT가 줄어드는가?
성능 개선의 본질을 이해하려면 올바른 지표를 비교해야 한다.
Auto-regressive 방식에서 3개의 토큰을 생성하려면 3번의 iteration이 필요하다. 각 iteration마다 1개의 토큰이 생성되므로, 총 비용은 3 × ITL이다.
Speculative Decoding에서는 한 번의 루프(draft K개 + target scoring 1회)가 하나의 iteration에 해당한다. 이 루프는 단일 토큰 생성보다 시간이 더 걸리므로 ITL은 오히려 소폭 증가한다. 그러나 이 루프 하나에서 acceptance length, 즉 평균적으로 수용되는 토큰 수만큼 여러 토큰이 한꺼번에 확정된다.
결국 비교해야 할 overhead는 다음과 같다.
Auto-regressive의 3 × ITL vs Speculative Decoding의 1 루프 시간 (≈ 2~3 TPOT 분량)
acceptance length가 충분히 크다면, 루프 하나에서 수용되는 토큰 수가 auto-regressive에서 같은 수의 토큰을 개별 생성하는 총 시간보다 훨씬 짧다. 이것이 TPOT를 줄이는 핵심 메커니즘이다.
Draft Model 설계와 현실적인 한계
draft model의 선택 기준은 세 가지다. logit을 외부에 노출할 수 있어야 하고, target model과의 분포 일치도가 충분히 높아야 하며, 레이턴시가 낮아야 한다.
가장 현실적인 선택은 target model과 동일한 아키텍처를 가진 소규모 버전을 draft model로 사용하는 것이다. 같은 토크나이저와 데이터셋으로 학습되었기 때문에 분포 일치도가 높고, 기존 인프라를 그대로 재활용할 수 있다는 엔지니어링 이점도 크다.
다만 독립 모델을 추가로 운용하는 비용은 현실적으로 존재한다. GPU 메모리를 추가로 점유하고, draft model 자체의 KV Cache도 별도로 관리해야 한다. K값 역시 트레이드오프다. K가 클수록 한 루프에서 더 많은 토큰을 얻을 가능성이 생기지만, 뒤쪽 토큰일수록 앞 토큰의 수용에 의존하기 때문에 평균 수용률이 떨어지고 루프당 시간도 선형으로 늘어난다. 논문에서는 도메인에 따라 최적 K가 다르며, 무작정 크게 설정하는 것이 오히려 역효과를 낳을 수 있음을 실험으로 보였다.
결론: 품질의 손실 없이 속도를 얻다
Speculative Decoding을 Chinchilla 70B에 적용한 결과, 생성 품질의 손실 없이 2~2.5배의 디코딩 속도 향상을 달성했다. 특히 코드 생성 태스크에서는 이론적인 메모리 대역폭 한계를 뛰어넘는 속도를 기록했는데, 코드는 반복적인 패턴이 많아 draft model의 acceptance length가 자연어보다 길기 때문이다.
이 논문이 인상적인 이유는 단순히 빠르다는 것이 아니다. 모델 구조를 전혀 바꾸지 않고, 양자화나 지식 증류 같은 품질 손실을 감수하지도 않으면서, 알고리즘 수준에서 하드웨어의 특성을 파고들어 속도를 끌어냈다는 점이다.
FlashAttention이 "메모리 접근 횟수를 줄이기 위해 연산을 재활용"했다면, Speculative Decoding은 "target model의 forward pass 한 번으로 여러 토큰을 확정"한다. 두 아이디어 모두 하드웨어의 병목을 소프트웨어 알고리즘으로 우회했다는 점에서 같은 철학 위에 서 있다.
결국 진짜 최적화는 모델 자체를 바꾸는 것이 아니라, 시스템이 작동하는 방식을 근본적으로 다시 설계하는 것에서 나온다는 선례를 다시 한번 확인할 수 있었다.
다음주에는 이 draft model을 활용하는 speculative decoding를 발전시킨 EAGLE 방법론에 대해 다뤄보겠다.
참고 문헌
vLLM 공식 문서 - https://docs.vllm.ai/en/stable/features/speculative_decoding/
Chen, Charlie, et al. "Accelerating large language model decoding with speculative sampling."(2023). -https://arxiv.org/pdf/2302.01318