추측 디코딩은 연산량이 늘어도 GPU 메모리 대기 시간을 줄여 빨라진다

Hacker News13일 전

추측 디코딩(speculative decoding)은 작은 드래프트 모델이 다음에 나올 법한 토큰 여러 개를 먼저 제안하고, 큰 모델이 그 후보를 한 번의 순전파에서 병렬로 검증하는 기법이다. 기존 LLM 디코딩이 한 번에 토큰 하나씩 처리하는 것과 달리, 여러 토큰을 동시에 살펴보고 큰 모델이 생성했을 토큰만 받아들인다. 이 방식은 특히 배치 크기가 작은 추론에서 체감 속도를 높인다.

흔한 오해는 추측 디코딩이 토큰당 연산을 줄여서 빠르다는 것이다. 실제로는 같은 시퀀스를 만들 때 총 부동소수점 연산(FLOPs)이 오히려 늘어난다. 드래프트 검증 자체가 공짜가 아니고, 거부된 토큰도 연산으로 남기 때문이다. 속도의 비밀은 일의 양이 아니라 GPU가 처리하는 일의 종류가 바뀐다는 데 있다.

GPU에서 LLM이 토큰을 하나씩 디코딩할 때는 메모리 바운드에 가깝다. 각 토큰마다 활성 파라미터를 전역 메모리에서 로컬 메모리로 계속 불러와야 하고, 연산 유닛은 데이터 로딩을 기다리는 시간이 길다. 반대로 수천 토큰을 한 번에 배치로 디코딩하면 연산 바운드가 된다. 메모리 바운드 상태에서는 GPU의 연산 능력이 남는데, 추측 디코딩은 이 남는 연산을 활용해 작은 배치에서도 한 시퀀스를 더 깊게 전진시킨다.

예를 들어 한 번의 순전파에서 최적으로 처리할 토큰 수를 5라고 두면, 배치 크기 5로 다섯 시퀀스를 각각 한 토큰씩 처리하는 '넓이' 방식과 한 시퀀스에서 5토큰을 추측해 처리하는 '깊이' 방식은 GPU가 하는 일의 양이 대략 비슷하다. 전자는 여러 요청이 조금씩 나아가고, 후자는 한 요청이 크게 나아간다. 배치 크기가 최적 토큰 수보다 훨씬 작을 때 추측 디코딩이 '공짜 연산'을 써서 전체 완료 시간을 줄이는 이유다.

문제는 배치가 이미 충분히 클 때다. 최적 배치 크기에 도달한 상태에서 추측 디코딩을 계속 켜면 연산 바운드로 넘어가 오히려 느려질 수 있다. 수용률이 100%가 아니면 거부된 토큰은 그대로 낭비 연산이 된다. 그래서 vLLM은 `--speculative-disable-by-batch-size` 파라미터로 특정 배치 크기 이상에서는 추측 디코딩을 끌 수 있게 한다.

이 관점은 추측 깊이를 고정할 필요가 없다는 연구로 이어진다. 배치가 작아 연산 여유가 있을 때는 깊게, 배치가 커질수록 얕게, 필요하면 시퀀스별로 동적으로 추측하는 편이 낫다는 발상이다. dSpark 같은 아이디어가 이런 직관을 바탕으로 한다.

개발자에게 중요한 것은 추측 디코딩을 무조건 빠른 스위치로 보지 않는 것이다. 소규모 배치나 지연 시간이 중요한 서빙에서는 효과가 크지만, 처리량이 이미 높은 구간에서는 수용률과 배치 크기를 함께 보고 켜고 끄는 판단이 필요하다. 드래프트 모델의 정확도가 낮으면 이득이 줄고, 거부 토큰이 많을수록 손해가 커진다.

결국 추측 디코딩의 성능은 FLOPs 절감이 아니라 메모리 병목과 연산 병목 사이의 균형에서 나온다. 배치 크기, 수용률, 목표 지연 시간을 기준으로 추측 깊이와 활성화 여부를 조정하는 것이 실무적인 접근이다. 원문은 구체적인 벤치마크 수치를 제시하지 않으므로, 실제 도입 전에는 자신의 워크로드에서 측정해야 한다.