Medusa 심층 아키텍처: 별도 초안 모델을 제거한 다중 헤드(Multi-Head) 추측 디코딩과 Tree Attention 3배 가속 원리

대규모 언어 모델(LLM)을 활용한 실시간 대화형 서비스, 인라인 코드 완성(Code Autocomplete), 복합 자율 에이전트 환경에서 최종 사용자의 만족도를 결정하는 가장 결정적인 성능 지표는 바로 토큰 생성 지연시간(Inter-Token Latency, ITL)입니다. 아무리 거대하고 똑똑한 모델이라 할지라도 한 글자씩 출력되는 속도가 인간의 읽기 속도보다 느리다면 실무 적용에 막대한 제약이 따릅니다.

트랜스포머 디코더 모델의 생성 과정은 이전의 모든 토큰을 기반으로 다음 단 1개의 토큰을 순차적으로 예측하는 자가회귀(Autoregressive, AR) 방식으로 작동합니다. 이로 인해 GPU 가속기 내부에서는 심각한 하드웨어 불균형이 발생합니다. 배치 크기가 1인 실시간 추론 환경에서 단 1개의 토큰을 생성할 때마다 수십 기가바이트(GB)에 달하는 모델 가중치 전체를 고대역폭 메모리(HBM)에서 온칩 레지스터로 끌어와야 하기 때문입니다.

이 과정에서 텐서 코어가 수행하는 부동소수점 곱셈-누적(MAC) 연산량에 비해 메모리 버스에서 전송해야 하는 바이트 수가 압도적으로 크기 때문에, 현대 가속기는 연산 장치가 아니라 메모리 대역폭에 발이 묶이는 메모리 병목(Memory-Bound) 상태에 빠집니다. 그 결과 최신 H100이나 A100 GPU의 막강한 연산 코어 중 95% 이상이 데이터가 메모리 버스를 통과하기만을 기다리며 공회전(Idle)하는 기이한 현상이 반복됩니다.

이 고질적인 자가회귀 병목을 타파하기 위해 컴퓨터 과학계가 고안한 혁신이 바로 추측 디코딩(Speculative Decoding)입니다. 그리고 이 추측 디코딩의 복잡성을 한 차원 높은 수준으로 단순화하면서도 가속 효율을 극대화한 오픈소스 프로젝트가 바로 프린스턴 대학과 Together AI 연구진이 주도한 Medusa(메두사)입니다.

Medusa는 추가적인 별도 초안 모델(Draft Model)을 서빙해야 하는 전통적 추측 디코딩의 거대한 운영 비용을 완전히 제거했습니다. 대신 원본 모델의 마지막 은닉 상태(Hidden State) 상단에 가벼운 단일 계층 MLP 헤드들을 여러 개 장착하고, 2차원 트리 어텐션(Tree Attention) 마스크를 통해 단 한 번의 검증 순방향 패스(Forward Pass)에서 복수 개의 미래 토큰을 병렬로 획득합니다. 이번 글에서는 Medusa의 수학적 검증 원리, 기각 샘플링과 엔트로피 기반 탐색, 그리고 트리 어텐션 커널의 하드웨어 최적화 메커니즘을 심층 분석합니다.

Medusa 다중 헤드 추측 디코딩 아키텍처

자가회귀 디코딩의 물리적 한계: 산술 강도(Arithmetic Intensity)의 붕괴

자가회귀 디코딩의 성능 병목을 정확히 이해하기 위해서는 고성능 컴퓨팅(HPC)의 척도인 산술 강도(Arithmetic Intensity, FLOPs per Byte) 개념을 살펴보아야 합니다.

산술 강도는 연산 장치가 메모리에서 1바이트의 데이터를 읽어올 때 몇 번의 부동소수점 연산(FLOPs)을 수행할 수 있는지를 나타내는 비율입니다:

$$ ext{Arithmetic Intensity} = rac{ ext{Total Floating Point Operations (FLOPs)}}{ ext{Total Memory Access (Bytes)}}$$

배치 크기 $B=1$인 환경에서 매개변수 수 $P$를 가진 모델이 16비트 정밀도(FP16/BF16, 파라미터당 2바이트)로 단 1개의 토큰을 디코드할 때:

메모리에서 로드하는 가중치 크기: $2P$ 바이트
수행하는 행렬-벡터 곱셈(GEMV) 연산량: $2P$ FLOPs
산술 강도: $2P / 2P = 1 ext{ FLOP/Byte}$

NVIDIA H100 SXM 가속기의 하드웨어 스펙을 대입해 보겠습니다:

FP16/BF16 텐서 코어 연산 성능: 약 1,979 TFLOPS
HBM3 메모리 대역폭: 약 3.35 TB/s
하드웨어 임계 산술 강도: $1,979 imes 10^{12} / (3.35 imes 10^{12}) pprox 590 ext{ FLOPs/Byte}$

즉, H100 GPU의 연산 코어가 100% 포화되어 최대 효율을 내기 위해서는 메모리에서 1바이트를 읽을 때마다 최소 590번 이상의 부동소수점 연산을 쏟아부어야 합니다. 하지만 1토큰 자가회귀 디코딩의 산술 강도는 고작 1에 불과합니다.

결과적으로 H100 텐서 코어의 실제 연산 점유율(MFU)은 채 1%에도 미치지 못하며, 가중치를 읽어오는 메모리 버스 대역폭이 디코딩 속도의 절대적인 병목 상한선을 결정짓게 됩니다.

전통적 추측 디코딩(Speculative Decoding)의 등장과 현실적 장벽

2023년 Leviathan 등과 Chen 등이 독립적으로 제안한 표준 추측 디코딩(Speculative Decoding)은 이 문제를 해결하기 위해 초안 모델(Draft Model, $M_q$)과 목표 모델(Target Model, $M_p$)이라는 두 가지 이종 신경망의 협업 구조를 제시했습니다.

초안 생성 단계: 훨씬 작고 빠른 소형 초안 모델(예: 70B 목표 모델 대비 1B 소형 모델)이 메모리를 적게 소모하며 미래의 후보 토큰 시퀀스 $[x_1, x_2, \dots, x_K]$를 빠르게 순차 생성합니다.
병렬 검증 단계: 거대한 목표 모델이 한 번의 순방향 패스(GEMM 연산)로 프롬프트와 초안 토큰들을 한꺼번에 병렬 입력으로 받아들여, 각 위치에서의 정확한 확률 분포를 동시에 산출합니다.
기각 샘플링(Rejection Sampling): 목표 모델은 확률 비율을 기반으로 초안 토큰들을 앞에서부터 차례대로 검증하여, 일치하는 토큰들을 한 번에 대량 수용(Accept)하고 불일치하는 첫 번째 토큰 위치에서 수정을 가합니다.

한 번의 거대 모델 메모리 로드(Forward Pass)로 복수 개의 토큰(예: 3개 ~ 5개)을 동시에 통과시킬 수 있으므로 산술 강도가 $K$배로 상승하고, 전체 생성 지연시간이 2배에서 3배까지 비약적으로 단축됩니다.

전통적 추측 디코딩이 안고 있는 3가지 운영 장벽

그러나 표준 추측 디코딩을 실제 대규모 프로덕션 환경에 배포하려던 인프라 팀들은 심각한 현실적 장벽에 부딪혔습니다:

초안 모델 서빙을 위한 추가 VRAM 점유: 소형 모델이라 할지라도 1B~7B 모델을 GPU 메모리에 함께 상주시켜야 하므로 귀중한 VRAM이 수 GB 이상 추가 낭비됩니다.
토크나이저 및 어휘 사전 불일치: 목표 모델과 완벽히 동일한 어휘 사전, 동일한 바이트 페어 인코딩(BPE) 규칙, 동일한 특수 토큰 처리를 공유하는 초안 모델을 찾는 것은 극히 어렵습니다. 토크나이저가 미세하게 다르면 추측 자체가 불가능합니다.
파이프라인 스케줄링 오버헤드: 서로 다른 두 개의 모델 간에 텐서를 주고받는 과정에서 추가적인 IPC(프로세스 간 통신)나 CUDA 스트림 동기화 지연시간이 발생하여 실제 체감 가속율이 깎여나갑니다.

Medusa의 핵심 철학: 초안 모델을 완전히 소거한 다중 헤드 구조

Together AI와 프린스턴 대학 연구진이 발표한 Medusa는 이 모든 복잡성을 단 하나의 우아한 설계로 해체했습니다: “초안 모델을 별도로 두지 말고, 원래 모델의 마지막 레이어 위에 가벼운 예측 헤드들을 여러 개 얹어라.”

트랜스포머 디코더의 마지막 트랜스포머 블록을 통과한 출력 은닉 상태(Hidden State) 벡터 $h_t$는 현재 토큰 $t$뿐만 아니라 바로 뒤이어 등장할 문맥에 대한 풍부한 언어적 표상(Representation)을 이미 내포하고 있습니다.

Medusa는 $h_t$를 입력으로 받아 각각 서로 다른 미래 시점을 예측하는 $K$개의 독립적인 Medusa 헤드(Medusa Heads)를 부착합니다:

Medusa 헤드 0: 원본 언어 모델 헤드(LM Head), 시점 $t+1$ 토큰 예측
Medusa 헤드 1: 단일 계층 잔차 MLP, 시점 $t+2$ 토큰 예측
Medusa 헤드 2: 단일 계층 잔차 MLP, 시점 $t+3$ 토큰 예측
Medusa 헤드 $k$: 단일 계층 잔차 MLP, 시점 $t+k+1$ 토큰 예측

잔차 블록(Residual Block)을 포함한 Medusa 헤드 구조

각 Medusa 헤드는 파라미터 수가 매우 적은 경량 모듈입니다. 은닉 차원 $d$에 대해 단 하나의 잔차 피드포워드 계층과 언임베딩 선형 프로젝션으로 구성됩니다:

$$h_t^{(k)} = ext{SiLU}(h_t W_{ ext{res}}^{(k)}) + h_t$$

$$p_k = ext{softmax}(h_t^{(k)} W_{ ext{head}}^{(k)})$$

여기서 가중치 $W_{ ext{head}}^{(k)}$는 원본 모델의 언임베딩 가중치($W_{ ext{lm_head}}$)를 동결(Freeze)한 채 그대로 공유하거나 미세조정할 수 있으므로, 헤드 하나당 추가되는 메모리는 고작 수십 메가바이트(MB)에 불과합니다. 원본 70B 모델 대비 0.1%도 안 되는 극소량의 오버헤드로 다중 토큰 예측 능력을 획득하는 것입니다.

트리 어텐션(Tree Attention): 2차원 마스크 기반의 단일 패스 병렬 검증

단순히 선형적으로 $t+1, t+2, t+3$의 최상위(Top-1) 토큰 하나씩만 뽑아서 검증하면, 중간에 단 한 토큰만 틀려도 그 뒤의 모든 예측이 즉시 폐기됩니다. 따라서 수용률(Acceptance Rate)이 급격히 떨어집니다.

Medusa는 이를 극복하기 위해 각 헤드에서 확률이 높은 상위 복수 개의 후보(Top-s Candidates)를 추출하고, 이를 나뭇가지처럼 갈라지는 트리 구조(Tree of Tokens)로 확장하여 한꺼번에 검증하는 트리 어텐션(Tree Attention) 아키텍처를 도입했습니다.

고정 트리 토폴로지(Tree Topology)와 2차원 어텐션 마스크

예를 들어 3개의 Medusa 헤드를 사용하여 다음과 같은 트리 분기를 구성할 수 있습니다:

헤드 0 (t+1): Top-2 후보 선택 ($A_1, A_2$)
헤드 1 (t+2): 각 $A_i$ 뒤에 붙을 Top-3 후보 선택 ($B_1, B_2, B_3$)
헤드 2 (t+3): 각 $B_j$ 뒤에 붙을 Top-2 후보 선택 ($C_1, C_2$)

이 조합은 총 수십 개의 서로 다른 문장 후보 경로(Path)를 형성합니다.

일반적인 어텐션은 인과적 마스크(Causal Mask, 하삼각 행렬)를 사용하므로 토큰들이 일렬로 나열되어 있어야 합니다. 하지만 트리의 서로 다른 나뭇가지에 있는 토큰들은 서로를 참조(Attend)해서는 안 됩니다. 오직 자신의 조상 노드(Ancestors)만을 참조해야 합니다.

Medusa는 이를 위해 특별한 2차원 트리 어텐션 마스크(Tree Mask)를 동적으로 구성합니다:

$$ ext{TreeMask}[i][j] = egin{cases} 0 & ext{if token } j ext{ is an ancestor of token } i ext{ or } i = j \ -\infty & ext{otherwise} \end{cases}$$

이 커스텀 마스크를 FlashAttention이나 Triton 어텐션 커널에 주입하면, 트리에 포함된 수십 개의 토큰 후보 전체를 단 1회의 GEMM 순방향 패스로 동시에 어텐션 연산할 수 있습니다. 각 나뭇가지는 서로 간섭하지 않으면서 독립된 문맥으로 완벽하게 평가됩니다.

수학적 무손실 보장: 기각 샘플링과 가중치 분포 보존

추측 디코딩이 산업계에서 절대적인 신뢰를 얻는 이유는 연산 속도를 3배 가속하면서도 모델의 최종 수학적 출력 분포를 0.0001%의 왜곡도 없이 100% 동일하게 보존하기 때문입니다.

이를 가능하게 하는 핵심 알고리즘이 바로 무손실 기각 샘플링(Lossless Rejection Sampling)입니다.

목표 모델이 토큰 $x$를 생성할 확률을 $p(x)$, Medusa 헤드가 예측한 확률을 $q(x)$라고 할 때, 후보 토큰 $x$에 대한 수용 확률(Acceptance Probability)은 다음과 같이 정의됩니다:

$$lpha(x) = \min \left( 1, rac{p(x)}{q(x)} ight)$$

균등 난수 $r \sim U(0, 1)$을 추출하여 $r < lpha(x)$이면 해당 토큰을 최종 출력으로 즉각 수용(Accept)합니다.
만약 $r \ge lpha(x)$이면 해당 토큰을 기각(Reject)하고, 목표 모델의 수정된 잔여 확률 분포(Residual Distribution) $p’(x)$에서 새로운 토큰을 다시 샘플링합니다:

$$p’(x) = rac{\max(0, p(x) - q(x))}{\sum_y \max(0, p(y) - q(y))}$$

수학적 증명에 따르면, 이 샘플링 절차를 거친 최종 토큰의 주변 확률 분포(Marginal Distribution)는 원래 거대 모델 하나만을 켜두고 순수 자가회귀로 샘플링했을 때의 확률 분포 $p(x)$와 비트 단위로 완벽히 일치합니다. 즉, 어떠한 근사 오차나 품질 저하도 발생하지 않는 완전 무손실 가속입니다.

소스코드 레벨 구조 분석: Medusa 구조체 및 트리 마스킹 구현체

Medusa의 핵심 헤드 및 트리 어텐션 검증 로직이 어떻게 구현되는지 단순화된 파이썬/PyTorch 코드로 살펴보겠습니다:

import torch
import torch.nn as nn
from typing import List, Tuple, Optional

class MedusaHead(nn.Module):
    # Single-layer residual MLP predicting future token offset
    def __init__(self, hidden_size: int, vocab_size: int):
        super().__init__()
        self.linear = nn.Linear(hidden_size, hidden_size, bias=False)
        self.act = nn.SiLU()
        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        # Residual connection preserves deep representation
        residual = hidden_states
        out = self.act(self.linear(hidden_states))
        out = out + residual
        logits = self.lm_head(out)
        return logits


class MedusaModel(nn.Module):
    def __init__(self, base_model: nn.Module, num_heads: int = 4):
        super().__init__()
        self.base_model = base_model
        self.hidden_size = base_model.config.hidden_size
        self.vocab_size = base_model.config.vocab_size
        self.num_heads = num_heads

        # Attach lightweight Medusa heads
        self.medusa_heads = nn.ModuleList([
            MedusaHead(self.hidden_size, self.vocab_size)
            for _ in range(num_heads)
        ])

    def generate_candidates(self, hidden_states: torch.Tensor, top_k: int = 4) -> List[torch.Tensor]:
        # Emits parallel token guesses for offsets t+1, t+2, t+3 ...
        candidates = []
        for head in self.medusa_heads:
            logits = head(hidden_states)
            topk_tokens = torch.topk(logits, top_k, dim=-1).indices
            candidates.append(topk_tokens)
        return candidates

    @staticmethod
    def build_tree_attention_mask(tree_structure: List[Tuple[int, int]], total_tokens: int) -> torch.Tensor:
        # Constructs 2D attention mask enforcing ancestor-only visibility
        mask = torch.full((total_tokens, total_tokens), float("-inf"))
        
        # Diagonal elements self-attend
        for i in range(total_tokens):
            mask[i, i] = 0.0

        # Ancestor visibility paths wired deterministically
        for child_idx, parent_idx in tree_structure:
            mask[child_idx, parent_idx] = 0.0
            # Propagate transitive ancestral relations
            for k in range(total_tokens):
                if mask[parent_idx, k] == 0.0:
                    mask[child_idx, k] = 0.0

        return mask

위 구현체에서 볼 수 있듯이, 모델의 기본 백본(Llama, Mistral 등) 가중치는 전혀 손대지 않으며 오직 상단에 부착된 경량 헤드들과 2차원 트리 마스크 제어만으로 전체 시스템이 완성됩니다.

벤치마크 실측 지표: 도메인별 속도 향상 및 수용률

Medusa 공식 벤치마크 및 다양한 다운스트림 작업에서 실측된 가속 지표는 실시간 LLM 추론의 판도를 바꾸어 놓았습니다.

다음은 단일 NVIDIA A100 (80GB VRAM) 환경에서 Vicuna-7B 및 Llama-3 8B 모델을 Medusa(4개 헤드, 64개 트리 분기)로 구동했을 때의 실측 결과입니다:

인간 평가 및 일반 대화 (Alpaca / MT-Bench):
- 순수 자가회귀 생성 속도: 초당 38.4 토큰
- Medusa 가속 생성 속도: 초당 98.6 토큰 (2.57배 가속 달성)
- 스텝당 평균 수용 토큰 수: 2.85개
정형화된 코드 생성 (HumanEval):
- 프로그래밍 언어는 문법적 구조(키워드, 괄호, 들여쓰기)가 명확하므로 Medusa 헤드의 예측 정확도가 극대화됩니다.
- 가속 배율: 3.15배 가속 (스텝당 평균 3.32개 토큰 동시 수용)
복합 추론 및 수학 문제 (GSM8k):
- 가속 배율: 2.38배 가속 (스텝당 평균 2.54개 토큰 동시 수용)
추가 훈련 비용 (Training Overhead):
- 원본 모델의 사전 훈련 토큰 수의 0.1% 미만인 소규모 데이터셋으로 단 몇 시간 만에 Medusa 헤드 학습 완료 (LoRA 파인튜닝과 유사한 초저비용)

엔터프라이즈 도입 시 고려해야 할 설계 트레이드오프

Medusa가 제공하는 강력한 지연시간 단축에도 불구하고, 프로덕션 서빙 시스템에 통합할 때 아키텍트가 반드시 고려해야 할 트레이드오프들이 존재합니다.

높은 동시 요청(High Concurrency) 환경에서의 이득 감소

배치 크기가 1인 저지연(Low Latency) 환경에서는 메모리 대역폭이 남아돌기 때문에 Medusa의 트리 검증이 공짜에 가까운 효율을 냅니다. 하지만 배치 크기가 64나 128 이상으로 커지면 시스템은 이미 연산 집약적(Compute-bound) 상태로 전환됩니다. 이때 트리 어텐션으로 검증 토큰 수가 수십 배로 불어나면 텐서 코어 연산량 자체가 포화되어 순수 자가회귀 대비 가속 배율이 1.2배 이하로 둔화될 수 있습니다. 따라서 Medusa는 배치 크기가 작고 반응 속도가 극도로 중요한 실시간 대화형 서비스나 에이전트 인라인 추론에 가장 적합합니다.

트리 크기(Tree Size)와 검증 오버헤드의 균형

트리 분기 수를 16개, 64개, 128개로 무작정 늘린다고 해서 수용률이 비례하여 증가하지는 않습니다. 너무 큰 트리는 검증 순방향 패스 자체의 GEMM 연산 시간을 증가시켜 오히려 배율을 떨어뜨립니다. 따라서 서빙 대상 모델의 크기와 GPU 아키텍처에 맞추어 최적의 트리 토폴로지(보통 32개 ~ 64개 노드)를 프로파일링을 통해 결정해야 합니다.

결론: 자가회귀의 굴레를 벗어나는 추론 엔진의 미래

Medusa의 혁신은 거대 언어 모델의 지연시간을 단축하기 위해 하드웨어 가속기만 무작정 증설하던 낡은 접근법을 지성적으로 극복한 쾌거입니다.

별도의 복잡한 초안 모델을 완전히 걷어내고, 원본 모델 내부의 풍부한 은닉 잠재력을 다중 헤드로 이끌어내어 트리 어텐션이라는 단일 패스 병렬 검증으로 완성해 냈습니다.

실시간 고성능 인공지능 인프라를 설계하는 시스템 엔지니어라면 다음 세 가지 핵심 원칙을 가슴에 새겨야 합니다:

메모리 바운드 워크로드의 산술 강도 복원 원칙: 단 1개의 토큰을 뽑기 위해 전체 가중치를 읽지 마십시오. 한 번 메모리를 읽어올 때 최대한 많은 유효 토큰을 병렬로 회수하십시오.
단일 모델 자율 추측의 원칙: 복잡한 이종 모델 협업 파이프라인 대신, 원본 모델의 잠재 벡터를 직접 활용하는 경량 다중 헤드 아키텍처로 엔지니어링 복잡성을 통제하십시오.
무손실 수학적 검증의 원칙: 속도를 위해 모델의 지능을 훼손하는 근사법 대신, 기각 샘플링을 결합하여 출력 분포를 온전히 보존하는 신뢰할 수 있는 가속을 구축하십시오.

알고리즘의 영리한 조율과 GPU 아키텍처의 정밀한 이해가 결합될 때, 우리는 지연 없는 완벽한 실시간 생성형 인공지능의 미래를 앞당길 수 있습니다.