[vllm] NGram GPU 구현 - 비동기 스케줄러 호환 GPU 기반 N-gram Drafting

PR 링크: vllm-project/vllm#29184 상태: Merged | 변경: +940/-12

들어가며

N-gram speculative decoding은 별도의 draft 모델 없이 이전 출력에서 반복 패턴을 찾아 다음 토큰을 예측하는 방식이다. 기존 구현은 CPU에서 동작하여 비동기 스케줄링과 호환되지 않았다. 이 PR은 N-gram 매칭 로직을 GPU로 이전하여 비동기 스케줄러와의 호환성을 확보하고 성능을 개선한다.

핵심 코드 분석

GPU 기반 N-gram 매칭

기존 CPU 구현에서는 토큰 시퀀스를 CPU로 복사한 후 Python에서 N-gram 매칭을 수행했다. GPU 구현에서는 이 과정이 CUDA 커널로 대체되어 GPU 메모리에서 직접 동작한다.

비동기 스케줄링 호환 테스트

def test_with_ngram_gpu_spec_decoding(monkeypatch):
    ngram_gpu_config = {
        "method": "ngram_gpu",
        "num_speculative_tokens": 3,
        "prompt_lookup_max": 3,
        "prompt_lookup_min": 2,
    }

    test_configs = [
        (False, "mp", False, None, False),
        (False, "mp", False, ngram_gpu_config, False),
        (True, "mp", False, ngram_gpu_config, True),     # preemption + chunking
        (False, "mp", True, ngram_gpu_config, False),     # async scheduling
        (True, "mp", True, ngram_gpu_config, False),      # async + preemption
        (True, "uni", True, ngram_gpu_config, False),     # uni executor
        (True, "mp", True, ngram_gpu_config, True),       # all combined
    ]

7개의 테스트 구성으로 preemption, async scheduling, prefill chunking과의 모든 조합을 검증한다.

왜 이게 좋은가

1. CPU-GPU 동기화 제거: N-gram 매칭이 GPU에서 수행되어 데이터 전송 오버헤드 제거
2. 비동기 스케줄링 호환: CPU 기반 구현의 가장 큰 제약이었던 동기화 문제 해결
3. 드래프트 모델 불필요: 추가 GPU 메모리 없이 speculative decoding 적용 가능
4. 다양한 환경 검증: preemption, chunking, 멀티프로세스 등 실제 운영 조건 모두 테스트

정리

N-gram GPU 구현은 "무료 점심"에 가까운 최적화다. 추가 모델 없이, 기존 출력 패턴만으로 토큰 생성을 가속하며, GPU에서의 실행으로 비동기 파이프라인에 자연스럽게 통합된다.

참고 자료

• PR #29184
• Async Scheduling #27614

---

이 글은 AI(Claude)의 도움을 받아 작성되었습니다. 코드 분석 내용은 실제 PR diff를 기반으로 합니다.