추측을 통한 Sonar 가속화

추측형 디코딩(Speculative decoding)은 작고 빠른 초안 모델을 사용하여 검증을 거치는 후보 완성본을 생성함으로써 거대 언어 모델(LLM)의 생성 속도를 높입니다. 이 방식에서는 값비싼 대상 모델이 한 번에 하나의 토큰만 생성하는 대신, 단일 단계에서 여러 토큰을 방출합니다. 본 문서에서는 Perplexity에서 Sonar 모델의 토큰 간 지연 시간을 줄이기 위해 적용한 다양한 추측형 디코딩의 구현 세부 사항을 설명합니다.

추측형 디코딩

추측형 디코딩은 자연어의 구조와 트랜스포머의 자기회귀적 성질을 활용하여 토큰 생성 속도를 높입니다. Llama-70B와 같은 대형 모델이 Llama-1B와 같은 소형 모델보다 더 많은 지식을 보유하고 있더라도, 일부 단순한 작업에서는 비슷한 성능을 보입니다. 이러한 중복성은 특정 시퀀스는 비용이 적게 드는 모델이 생성하는 것이 더 효율적이며, 복잡한 문제는 대형 모델이 처리하도록 남겨둘 수 있음을 시사합니다. 관건은 어떤 완성본이 더 나은지, 그리고 소형 모델이 생성한 결과물이 대형 모델의 품질과 동일한지 결정하는 것입니다.

다행히 LLM은 자기회귀 트랜스포머이므로, 토큰 시퀀스가 주어지면 다음 토큰에 대한 확률 분포를 출력합니다. 또한 입력 시퀀스의 토큰과 관련된 중간 특징에서 도출된 로짓(logit)은 해당 모델이 정확히 그 토큰을 방출할 가능성을 나타냅니다. 이 속성을 통해 추측이 가능해집니다. 입력 접두사로부터 소형 모델이 토큰 시퀀스를 생성하면, 대형 모델을 통해 해당 시퀀스가 대상 모델과 얼마나 일치하는지 확인할 수 있습니다. 후보 시퀀스의 각 접두사는 확률 점수가 매겨지며, 허용 임계값을 넘는 가장 긴 시퀀스가 선택됩니다. 추가적인 이점으로, 대상 모델은 후속 토큰을 무료로 제공합니다. 즉, 초안 모델이 n개의 토큰을 생성하면 한 단계에서 최대 n + 1개의 토큰이 방출될 수 있습니다.

확대

추론 시점에 추측형 샘플링 과정은 대략 4단계로 나눌 수 있습니다:

프리필(Prefill): 대상 모델과 초안 모델 모두 입력 시퀀스에서 실행되어 KV 캐시 항목을 채워야 합니다. Medusa와 같은 일부 방식은 예측을 위해 단순한 밀집 층(dense layer)을 사용하지만, 본 게시물에서는 자체 KV 캐시가 필요한 트랜스포머 기반 초안 모델에 중점을 둡니다.

초안 생성: 초안 모델이 반복적으로 고정된 수의 토큰을 생성합니다. 초안 시퀀스는 선형일 수도 있고, 모델이 주어진 깊이까지 트리 형태의 구조를 탐색할 수도 있습니다(EAGLE, Medusa). 여기서는 선형 시퀀스에 중점을 둡니다.

허용(Acceptance): 대상 모델이 초안 시퀀스에서 실행되며 각 초안 토큰에 해당하는 로짓을 생성합니다. 허용 가능한 가장 긴 시퀀스의 길이가 결정됩니다.

대상 생성: 대상 모델이 로짓을 생성했으므로, 일치하지 않는 위치나 시퀀스의 끝부분에서 로짓은 후속 토큰에 대응하게 됩니다. 이 로짓들을 샘플링하여 대상 모델로부터 견고한 토큰을 제공하고 시퀀스를 마무리할 수 있습니다.

추측형 디코딩을 구현하는 다양한 방법이 존재합니다. 본 게시물에서는 자체 개발한 1B 모델을 사용하여 Sonar 모델을 가속화하는 데 사용한 방식과, DeepSeek 규모의 모델 속도를 높이기 위해 구축 중인 예측 메커니즘에 중점을 둡니다.

대상-초안(Target-Draft)

추측형 디코딩은 기존의 작은 LLM을 초안 모델로 대상 모델과 결합하여 후보 시퀀스를 생성함으로써 달성할 수 있습니다. 운영 환경에서 당사는 대상 모델과 동일한 데이터셋으로 미세 조정된 Llama-1B 모델을 사용하여 Sonar를 가속화했습니다. 이 방식은 초안 모델을 처음부터 학습시킬 필요는 없으나, 소형 모델이 상당한 KV 캐시 용량을 사용하고 약간의 프리필 오버헤드를 발생시켜 TTFT(첫 토큰 생성 시간)를 증가시킨다는 단점이 있습니다.

이 체계에서 디코더는 프리필 또는 혼합된 프리필-디코드 배치 동안 표준 샘플링을 통해 토큰을 생성하며, 디코드 전용 배치에서만 추측을 수행합니다. 프리필 단계에서 대상 로짓은 즉시 샘플링되어 초안의 KV 캐시 내 새로 생성된 토큰을 프리필하는 데도 사용됩니다. 초안은 아직 샘플링되지 않지만, 생성된 로짓은 디코드 단계로 이월됩니다.

확대

디코드 단계에서 초안 모델이 진행되며 각 단계에서 상위 토큰을 샘플링합니다. 원하는 초안 길이에 도달하면, 토큰들은 대상 모델을 통과하여 샘플러가 허용된 시퀀스 길이를 식별하는 기준이 되는 로짓을 생성합니다. 허용 여부는 초안과 대상으로부터 나온 전체 확률 분포를 비교하여 결정됩니다. 대상 모델은 항상 허용된 초안 시퀀스 다음에 하나의 로짓 세트를 출력하므로, 이를 샘플링하여 추가적인 출력을 생성합니다. 초안 모델은 아직 해당 허용 토큰을 확인하지 못했으므로, 다음 디코드 단계를 준비하기 위해 KV 캐시 항목을 채우도록 재실행되며 로짓을 다시 이월합니다.

EAGLE

EAGLE은 확률적 초안 토큰의 트리 탐색을 통해 생성된 다수의 초안 시퀀스를 탐색하는 추측형 디코딩 체계입니다. 고정(EAGLE) 또는 동적 형태(EAGLE-2)의 트리를 초안 토큰의 연속적인 실행을 통해 탐색하며, 선형 시퀀스에서 가장 높은 점수의 토큰을 따르는 대신 각 노드에서 상위 K(Top-K) 후보를 고려합니다. 시퀀스들에 점수가 매겨지면 가장 적절하고 긴 시퀀스가 선택되어 이어지며, 대상 모델의 추가 토큰도 덧붙여집니다.

확대

더 정확한 예측을 위해 EAGLE 초안 모델은 토큰뿐만 아니라 대상 모델의 대상 특징(마지막 층 은닉 상태)을 사용하여 예측합니다. EAGLE의 단점은 낮은 지연 시간 예산 내에서 적절한 후보를 생성할 만큼 정확한 맞춤형 소형 초안 모델을 학습시켜야 한다는 점입니다. 일반적으로 초안 모델은 원래 모델의 디코더 층과 동일한 단일 트랜스포머 층으로 구성되며, 임베딩 및 lm_head 투영에 연결되어 대상 모델과 밀접하게 결합됩니다. 이는 더 적은 KV 캐시 용량을 필요로 하므로 EAGLE은 메모리 점유율이 낮습니다.

대상 모델에서 트리 형태의 시퀀스를 검증하려면 맞춤형 어텐션 마스크를 사용해야 합니다. 안타깝게도 전체 시퀀스에 맞춤형 어텐션 마스크를 사용하면 실제 입력 길이에 대해 어텐션 속도가 크게 느려져(최대 50%), 추측을 통해 얻을 수 있는 속도 향상이 상쇄됩니다. 이러한 이유로 당사는 아직 전체 트리 탐색을 운영 환경에 배포하지 않았으며, 대신 DeepSeek-V3 기술 보고서에 제시된 MTP 유사 체계를 통한 단일 토큰 예측의 특수한 경우에 집중하고 있습니다.

MTP

이 체계는 예측을 위해 토큰과 함께 은닉 상태를 사용한다는 점을 제외하면 초안-대상 디코딩과 유사합니다. 일반적인 초안-대상 추측보다 프리필 및 디코드 단계에서 조금 더 많은 작업이 수행되어야 합니다. 초안 모델은 토큰과 은닉 상태를 모두 사용합니다. 토큰 t_{i+1}은 토큰 t_i에 해당하는 로짓 L_i로부터 샘플링되며, 이는 다시 은닉 상태 H_i에서 파생됩니다. 결과적으로, 입력 토큰 버퍼는 대상이 출력한 은닉 상태 벡터에 대해 왼쪽으로 한 단계 이동해야 합니다. 아래 그림은 학습에 사용된 대응 관계와 추론 중의 이동을 나타냅니다.

확대

디코딩 흐름은 초안-대상 디코딩과 매우 유사하며, 은닉 상태와 로짓이 모두 이월된다는 점만 다릅니다. 당사의 구현은 관련된 모든 샘플링 및 로짓 처리 로직을 공유하며 모델 순방향 호출만 특화합니다. 여러 토큰이 예측될 때, 초안 모델은 예측을 위해 초안 은닉 상태를 사용하며 자체 특징에 기반하여 KV 캐시 항목을 채웁니다. 장기적으로 이는 정확도를 저하시킬 수 있습니다. 이후 대상 예측을 위한 KV 캐시 항목을 채우기 위해 초안 모델을 실행할 때, 더 정확한 대상 은닉 상태를 입력으로 취하여 전체 시퀀스에 대해 실행합니다. 이러한 초안 모델은 작기 때문에 추가 토큰 처리에 따른 비용은 무시할 수 있는 수준입니다.

MTP 헤드 학습

MTP의 이점을 활용하기 위해 당사는 8xH100 장치가 탑재된 단일 노드에서 Perplexity의 데이터셋으로 미세 조정된 모델에 부착할 MTP 헤드를 학습시키는 인프라를 구축했습니다. 약 하루 만에 Llama-1B에서 Llama-70B 및 DeepSeek V2-Lite에 이르는 모델의 헤드를 구축할 수 있습니다. 더 큰 모델의 경우 미세 조정 과정에서 구축된 MTP 헤드에 의존합니다.

MTP 학습의 목표는 대상 은닉 상태에서 외삽된 초안 은닉 상태 및 로짓을 대상의 다음 토큰 로짓 및 은닉 상태와 일치시키는 것입니다. 은닉 상태에 대한 추론 비용이 높기 때문에, 학습 중에 사용할 대상 모델의 추론 최적화 구현을 사용하여 이를 사전 계산합니다. 그러나 추론 MTP 구현을 검증하고 양자화나 최적화로 인한 수치적 차이가 결과를 방해하지 않도록, 검증 손실 및 정확도 추정을 위해서는 대상 및 초안 모델의 추론 구현을 완전히 재사용합니다.

원래 논문에서 사용된 ShareGPT 데이터셋에서 더 큰 샘플로 확장할 때, EAGLE 논문에 명시되고 구현된 MTP 헤드 아키텍처가 70B 규모의 모델에서는 학습되지 않는다는 것을 발견했습니다. 더 많은 수의 짧은 시퀀스를 포함하는 ShareGPT와 달리, 당사는 더 적은 수의 훨씬 긴 프롬프트로 학습합니다. 원래의 EAGLE 헤드는 일반적인 트랜스포머와 구조가 약간 달랐기 때문에, 제거되었던 일부 RMS 정규화(Normalization) 층을 재도입했습니다. 이는 학습 수렴을 가능하게 했을 뿐만 아니라 헤드의 정확도를 몇 퍼센트 포인트 높였습니다.

확대

층 정규화는 학습을 원활하게 할 뿐만 아니라 수학적으로도 직관적입니다. MTP 헤드는 대상 모델의 임베딩과 로짓 투영을 재사용하는데, 이는 크기가 상당할 수 있습니다(Llama 70B의 경우 약 2GB). 학습 중에 이들은 고정되며, MTP 층이 원래 모델의 투영 층이 학습 중에 배운 것과 동일한 벡터 공간에 예측을 임베딩하는 법을 배우는 것이 기대됩니다. 정규화를 제거하면 단일 MLP가 정규화가 뒤따르는 MLP와 동일한 함수를 학습해야 하므로, 초안 모델과 대상 모델의 은닉 상태 간 일치를 방해합니다.

추측형 디코딩을 통한 추론

추론 엔진에서 입력 시퀀스에 대한 토큰을 생성하려면, 먼저 합리적인 크기의 배치로 그룹화한 다음 KV 캐시에서 다음 토큰을 위한 페이지를 할당해야 합니다. 입력 토큰과 KV 페이지 정보는 모델을 실행하는 모든 병렬 순위(parallel ranks)에 브로드캐스트되는 버퍼로 패킹됩니다. 마지막으로 메타데이터가 GPU 메모리로 복사되고 모델이 실행되어 다음 토큰이 샘플링되는 로짓이 생성됩니다.

요청을 조율하는 래퍼(wrapper)를 통해 초안 및 대상 추론 서버를 느슨하게 결합하는 특정 구현과 달리, 당사의 초안-대상 쌍은 밀접하게 결합되어 생성 단계에서 일제히 움직입니다. 배치 스케줄링과 KV 페이지 할당은 모든 형태의 추측형 디코딩에서 모델 간 공유됩니다. 이는 모든 모델이 동일한 인터페이스를 노출하므로 모델과 상위 추론 서버를 연결하는 로직을 통합합니다.

Perplexity의 추론 런타임은 FlashInfer를 중심으로 구성되어 있으며, 이는 어텐션 커널을 구성하고 스케줄링하는 데 필요한 메타데이터를 결정합니다. 배치를 형성하는 일부 입력 시퀀스가 주어지면, 프리필, 디코드 또는 검증을 위해 CPU 측에서 중간 버퍼를 할당하고 어텐션에 사용되는 특정 상수 버퍼를 채우는 작업이 수행되어야 합니다. 이 작업은 GPU 활용도를 극대화하기 위해 숨겨야 할 지연 시간을 초래하는 배치 스케줄링 및 KV 페이지 할당 비용 외에 추가로 발생합니다.

추측 없이 추론을 위해 CPU 및 GPU 측 작업을 완전히 병렬화했지만, 추측형 디코딩을 위한 CPU-GPU 균형은 더 복잡하다는 것을 알게 되었습니다. 주요 문제는 허용된 토큰 수가 후속 실행의 시퀀스 길이를 결정하여, 피하기 어려운 GPU-CPU 동기화 지점이 발생한다는 점입니다. 당사는 CPU 작업의 지연 시간을 최적으로 숨기기 위해 다양한 스케줄링 체계를 실험했습니다.

대상-초안 스케줄

전체 LLM을 초안으로 사용하면 대상 모델보다 작음에도 불구하고 GPU에서 상당한 지연 시간이 발생하며, 비용이 많이 드는 CPU 작업을 숨길 여유가 생깁니다. 소형 모델은 텐서 병렬화의 이점을 얻지 못하므로 대상 모델과 초안 모델이 분할되는 순위(rank) 수 사이에 불일치가 발생합니다. 당사의 구현에서 초안 모델은 TP 그룹의 리더 순위에서만 실행됩니다.

확대

앞서 언급했듯이 디코드 단계는 로짓을 다음 실행으로 이월합니다. 이를 통해 초안 모델의 한 번의 실행과 CPU 측 배치 스케줄링 작업을 겹칠 수 있습니다. 배치가 구성된 후 샘플러와 초안을 반복적으로 호출하여 초안 토큰을 생성합니다. 이와 병렬로 대상 모델을 위한 검증용 배치가 구성되어 병렬 작업자와 동기화됩니다. 대상 로짓은 검증 및 샘플링되어 허용된 시퀀스 길이가 결정됩니다. 이 시점에서 후속 시퀀스 길이를 결정하기 위해 GPU-CPU 동기화가 필요합니다. 초안 모델은 리더 노드에서만 실행되므로, 배치가 순차적으로 설정되고 대상 모델이 생성한 추가 토큰으로 KV 캐시 항목을 채우기 위해 실행이 시작됩니다. 현재 실행 중에 이 초안 실행에서 생성된 로짓은 후속 실행에서 첫 번째 초안 토큰을 샘플링하는 데 사용됩니다. 가장 중요한 점은 초안이 실행되는 동안 다음 배치를 스케줄링할 수 있다는 것입니다.

단일 토큰을 위한 MTP 스케줄

런타임이 아직 Eagle 스타일의 초안 트리 탐색을 제공하지는 않지만, 당사는 단일 트랜스포머 디코더 층 크기의 모델이 생성한 선형 초안 토큰 시퀀스를 고려하는 이 체계의 특수한 경우를 구현했습니다. 이 체계는 DeepSeek R1의 오픈 소스 가중치를 사용하여 초안 예측에 사용할 수 있습니다. 단일 토큰을 예측하는 하위 사례는 대형 MTP 층이 오버헤드를 정당화할 만큼 충분히 높은 허용률을 달성하므로 흥미롭습니다.

초안 모델이 훨씬 빠르고 CPU 측 지연 시간을 덜 숨기기 때문에 MTP 스케줄링은 다소 복잡합니다. 또한 초안은 대상 모델과 함께 분할(sharded)되어 배치 정보를 위해 공유 메모리 전송이 필요합니다. 이전 체계와 유사하게, 배치 정보를 전송하고 이월된 로짓에서 첫 번째 토큰을 샘플링하며 실행이 시작됩니다. 다음으로, 디코드 배치 크기를 D라고 할 때 2 * D 토큰을 처리하며 토큰을 검증하기 위해 대상 모델이 실행됩니다. 배치가 두 개의 절반으로 균등하게 분할되므로 InfiniBand와 같은 느린 인터커넥트를 통해 Mixture-of-Experts(MoE) 모델에서 마이크로 배치 처리를 수행하는 데 이상적입니다. 대상 모델의 은닉 상태는 다음 초안 실행으로 이월되고, 로짓은 검증을 위해 샘플러로 전달됩니다.

확대

GPU에서 제한된 추가 작업을 수행함으로써 초안 시퀀스 허용 후의 CPU-GPU 동기화를 피합니다. 대상 모델의 입력 토큰이 이동된 후, 커널이 다음 대상 토큰을 해당 위치에 삽입합니다. 그런 다음 초안 모델은 대상과 동일한 배치 정보로 재실행되어 KV 캐시 항목을 채우고 다음 실행을 위한 로짓과 은닉 상태를 구축하며, 허용되지 않은 토큰에 대해 일부 중복 작업을 수행합니다. 이러한 상황에서 사용되지 않은 작업의 지연 시간은 초안 모델의 작은 크기 덕분에 거의 측정되지 않습니다. 초안 실행과 병렬로 시퀀스 길이는 CPU에서 결정되며, GPU 작업이 종료될 때까지 기다릴 필요 없이 다음 배치의 스케줄링이 시작됩니다.

초안 층에서의 추가 작업으로 인한 오버헤드는 어텐션에서는 눈에 띄지 않으나 MLP 층은 더 문제가 됩니다. 행렬 곱셈 명령어가 토큰 수 차원을 따라 64의 경계로 패딩되기 때문에, 두 배로 늘리는 작업에 상당한 블록이 필요하지 않다면 오버헤드는 숨겨집니다. 더 긴 초안 시퀀스의 경우 오버헤드 비용이 더 커지며 일반적인 초안-대상 모델에 사용되는 체계가 더 효과적입니다.

참고 문헌

추측형 디코딩을 통한 트랜스포머의 고속 추론

EAGLE: 추측형 샘플링은 특징 불확실성에 대한 재고를 요구함

EAGLE-2: 동적 초안 트리를 사용한 언어 모델의 더 빠른 추론

EAGLE-3: 학습 시간 테스트를 통한 거대 언어 모델의 추론 가속화 확장

Medusa: 다중 디코딩 헤드를 갖춘 단순 LLM 추론 가속화 프레임워크

층 정규화

FlashInfer: LLM 추론 서비스를 위한 효율적이고 사용자 정의 가능한 어텐션 엔진

DeepSeek-V3 기술 보고서

인용

@online{perplexity2025AcceleratingSonarThroughSpeculation,
  author = {{Perplexity} and Nandor Licker and Erin McKnight},
  title = {{Accelerating Sonar through speculation}},
  date = {2025-06-10},
  url = {https://perplexity.coolkeypoint.com/hub/blog/accelerating-sonar-through-speculation}
}