본문으로 건너뛰기

[vLLM] N-gram & Suffix Decoding: 모델 프리 드래프트

들어가며

투기적 디코딩에서 드래프트 모델은 추가 GPU 메모리와 연산을 요구한다. N-gram Proposer와 Suffix Decoding은 모델 없이(model-free) 기존 토큰 시퀀스의 패턴을 활용하여 드래프트 토큰을 제안한다. 특히 코드 생성, 반복적 텍스트, 문서 요약 등 입력과 출력의 패턴이 겹치는 태스크에서 효과적이다.

소스 경로: vllm/v1/spec_decode/ngram_proposer.py, vllm/v1/spec_decode/suffix_decoding.py

논문: Suffix Decoding

공식 문서

vLLM 공식 문서: N-gram Speculative Decoding vLLM 공식 문서: Suffix Speculative Decoding

핵심 구조/코드 분석

NgramProposer

class NgramProposer:
    def __init__(self, vllm_config: VllmConfig):
        self.min_n = vllm_config.speculative_config.prompt_lookup_min
        self.max_n = vllm_config.speculative_config.prompt_lookup_max
        self.k = vllm_config.speculative_config.num_speculative_tokens
        self.max_model_len = vllm_config.model_config.max_model_len

        # Numba JIT 미리 워밍업
        self.propose([[]] * 1024, np.zeros(1024, dtype=np.int32),
                     np.zeros((1024, self.max_model_len), dtype=np.int32))

min_n부터 max_n까지의 n-gram 길이 범위에서 현재 시퀀스의 접미사와 매칭되는 가장 긴 n-gram을 찾고, 그 뒤의 k개 토큰을 드래프트로 제안한다. 초기화 시 Numba JIT 컴파일을 미리 트리거한다.

KMP 기반 매칭 알고리즘

@jit(nopython=True)
def _find_longest_matched_ngram_and_propose_tokens(
    origin_tokens, min_ngram, max_ngram, max_model_len, k
) -> np.ndarray:
    # 토큰을 뒤집어서 접미사 매칭을 접두사 매칭으로 변환
    tokens = origin_tokens[::-1]

    # KMP의 LPS(Longest Prefix Suffix) 배열 계산
    lps = np.zeros(max_ngram, dtype=np.int32)
    longest_ngram = 0
    position = 0

    prev_lps = 0
    i = 1
    while i < total_token:
        if tokens[prev_lps] == tokens[i]:
            prev_lps += 1
            if prev_lps >= longest_ngram:
                longest_ngram = prev_lps
                position = i
            if prev_lps == max_ngram:
                prev_lps = lps[max_ngram - 1]
            i += 1
        elif prev_lps != 0:
            prev_lps = lps[prev_lps - 1]
        else:
            i += 1

    start_position = total_token - 1 - position + longest_ngram
    return origin_tokens[start_position : start_position + k]

핵심은 KMP(Knuth-Morris-Pratt) 알고리즘의 LPS 배열이다. 토큰을 뒤집어서 접미사 매칭을 접두사 매칭 문제로 변환한 뒤, O(n) 시간에 가장 긴 매칭을 찾는다. max_ngram으로 LPS 배열 크기를 제한하여 메모리를 절약한다.

Numba 병렬 배치 처리

@njit(parallel=True)
def batch_propose_numba(valid_ngram_requests, num_tokens_no_spec,
                        token_ids_cpu, min_n, max_n, max_model_len, k,
                        valid_ngram_draft, valid_ngram_num_drafts):
    for i in prange(len(valid_ngram_requests)):
        idx = valid_ngram_requests[i]
        context_token_ids = token_ids_cpu[idx, :num_tokens_no_spec[idx]]
        drafter_output = _find_longest_matched_ngram_and_propose_tokens(
            origin_tokens=context_token_ids, ...
        )
        valid_ngram_num_drafts[idx] = drafter_output.shape[0]

Numba의 prange로 배치 내 요청들을 CPU에서 병렬 처리한다. 총 토큰 수가 8192 이상일 때만 멀티스레딩을 활성화하여 오버헤드를 방지한다.

SuffixDecodingProposer

class SuffixDecodingProposer:
    def __init__(self, vllm_config):
        from arctic_inference.suffix_decoding import SuffixDecodingCache
        self.suffix_cache = SuffixDecodingCache(
            max_tree_depth=config.suffix_decoding_max_tree_depth,
            max_cached_requests=config.suffix_decoding_max_cached_requests,
        )

    def propose(self, input_batch, sampled_token_ids, ...) -> list[list[int]]:
        for i, sampled_ids in enumerate(sampled_token_ids):
            req_id = input_batch.req_ids[i]
            if req_id not in self.suffix_cache.active_requests:
                prompt_token_ids = input_batch.token_ids_cpu[index, :num_prompt_tokens]
                self.suffix_cache.start_request(req_id, prompt_token_ids)

            self.suffix_cache.add_active_response(req_id, sampled_ids)

            start = max(0, num_tokens - self.max_tree_depth)
            pattern = input_batch.token_ids_cpu[i, start:num_tokens]
            draft = self.suffix_cache.speculate(
                req_id, pattern,
                max_spec_tokens=self.num_speculative_tokens,
                max_spec_factor=self.max_spec_factor,
                min_token_prob=self.min_token_prob,
            )

Suffix Decoding은 Arctic Inference 라이브러리의 SuffixDecodingCache를 활용한다. 프롬프트별 suffix tree를 구축하고, 현재 컨텍스트의 마지막 max_tree_depth 토큰을 패턴으로 사용하여 추론한다. 요청별로 동적 수의 드래프트 토큰을 생성한다.

왜 이 설계인가

  1. 제로 GPU 오버헤드: N-gram과 Suffix Decoding 모두 CPU에서 실행된다. GPU 메모리를 전혀 사용하지 않으므로 메모리가 부족한 환경에서도 투기적 디코딩을 활용할 수 있다.

  2. KMP 알고리즘: 단순 이중 루프 대신 KMP를 사용하여 O(n) 시간 복잡도로 매칭한다. 긴 컨텍스트에서도 성능이 보장된다.

  3. Numba JIT: Python의 성능 한계를 Numba로 극복한다. @njit(parallel=True)와 prange로 C 수준의 성능과 자동 병렬화를 달성한다.

  4. 적응형 드래프트: Suffix Decoding은 요청별로 다른 수의 드래프트 토큰을 생성한다. 패턴이 강한 요청은 많은 토큰을, 패턴이 약한 요청은 적은 토큰을 제안하여 효율을 극대화한다.

논문 핵심 내용

Suffix Decoding: A Model-Free Approach to Speeding Up Large Language Model Inference (2411.04975) 논문은 suffix tree를 활용한 모델 프리 투기적 디코딩 기법을 제안했다.

핵심 아이디어: 프롬프트와 이전 출력에서 긴 토큰 시퀀스를 suffix tree로 캐싱하고, 현재 컨텍스트와 매칭되는 패턴을 찾아 드래프트 토큰을 제안한다. 수용 확률이 높을 때는 더 많은 토큰을 투기하고, 낮을 때는 줄이는 적응형 깊이 조절이 핵심이다.

속도 향상 벤치마크

비교 대상 속도 향상
기본 자동회귀 디코딩 대비 최대 5.3x
모델 기반 방법 (EAGLE-2/3) 대비 2.8x
모델 프리 방법 (Token Recycling) 대비 1.9x

SWE-Bench(소프트웨어 엔지니어링)와 Text-to-SQL 같은 에이전틱 벤치마크에서 테스트됐다. 에이전틱 AI 워크로드가 특히 효과적인 이유는, 반복적인 코드 패턴, SQL 구조, 이전 대화 컨텍스트 등 재사용 가능한 패턴이 풍부하기 때문이다. NeurIPS 2025 스포트라이트 논문으로 채택됐다.

모델 기반 방법(EAGLE-2/3) 대비 2.8배 빠르다는 결과가 특히 의미 있다. 추가 GPU 메모리를 전혀 사용하지 않으면서 모델 기반 접근보다 빠른 건, 에이전틱 워크로드에서 패턴 재사용률이 매우 높기 때문이다.

참고

댓글

관련 포스트

vLLM 의 다른글