Prefill and Decode for Concurrent Requests - Optimizing LLM Performance

TNG에서는 24 개의 H100 GPU 클러스터에 여러 개의 대형 언어 모델 (LLM) 을 자체 호스팅하고 있습니다. 이는 50 개의 다른 애플리케이션을 지원하며, 시간당 5,000 회 이상의 추론을 처리하고 하루에 1,000 만 개 이상의 토큰을 생성합니다.

대부분의 LLM 은 토큰마다 텍스트를 생성하며, 이는 모든 새로운 토큰이 이전 토큰들에 기반하여 계산됨을 보장합니다 (이 모델 속성은 auto-regressive라고 불립니다). 첫 번째 출력 토큰은 모든 프롬프트 토큰에 의존하지만, 두 번째 출력 토큰은 모든 프롬프트 토큰뿐만 아니라 첫 번째 출력 토큰에도 의존하며, 그 이후로 계속됩니다. 따라서 토큰 생성은 개별 요청 수준에서 병렬화할 수 없습니다.

주의 메커니즘을 가진 LLM 에서 새로운 토큰을 계산하려면 이전 각 토큰에 대해 key, value, 및 query 벡터를 계산해야 합니다. 다행히도 일부 특정 계산의 결과는 후속 토큰을 위해 재사용할 수 있습니다. 이 개념은 key-value (KV) 캐시라고 알려져 있습니다. 추가적인 출력 토큰에 대해, 하나의 더 많은 세트의 key 및 value 벡터만 계산되고 KV 캐시에 추가되어야 합니다. 그러나 첫 번째 출력 토큰에 대해서는, 우리는 초기로 빈 KV 캐시를 시작하며 입력 프롬프트의 토큰 수만큼의 key 및 value 벡터 세트를 계산해야 합니다. 다행히도, 그리고 후속 토큰 생성과 달리, 모든 입력 토큰은 처음부터 알려져 있으며, 각각의 key 및 value 벡터 계산을 병렬화할 수 있습니다. 이 차이는 prefill (첫 번째 출력 토큰 계산) 과 decode 단계 (후속 출력 토큰 계산) 사이의 구분을 동기를 부여합니다.

prefill 단계에서는 모든 입력 토큰에 대한 계산을 병렬로 실행할 수 있지만, decode 단계에서는 개별 요청 수준에서 병렬화가 불가능합니다.

prefill 및 decode 단계의 차이점은 텍스트 생성을 위한 두 가지 주요 지표인 Time to first token 과 time per output token 에서도 반영됩니다. * time to first token* 은 prefill 단계의 지연 시간으로, *is the latency of a single decode step.*는 단일 decode 단계의 지연 시간입니다. prefill 단계가 또한 단 하나의 토큰만 생성하지만, 모든 입력 토큰이 처리되어야 하므로 단일 decode 단계보다 훨씬 더 오래 걸립니다. 반면에 동일한 출력 토큰 수를 가진 decode 단계에 비해 prefill 단계는 입력 토큰의 수에 대해 훨씬 더 빠릅니다 (이 차이는 상업적 LLM API 가 입력 토큰을 출력 토큰보다 훨씬 낮은 비율로 부과하는 이유입니다).

time per output token두 지연 시간은 채팅 봇과 같은 인터랙티브 애플리케이션에 관련된 지표입니다. 사용자가 응답을 보기 전에 5 초 이상 기다려야 한다면, 그들은 애플리케이션이 고장 났다고 생각하여 떠날 수 있습니다. 또한 텍스트 생성이 초당 1 토큰만큼 느리다면, 그들은 완료될 때까지 기다릴 만큼 충분히 인내할 수 없습니다. 인터랙티브 애플리케이션의 일반적인 지연 시간 목표는 출력 토큰 당 100-300ms (즉, 초당 3-10 토큰의 생성 속도, 최소로 읽기 속도와 동일하며, 이상적으로는 생성되는 출력 텍스트를 스크리밍할 수 있도록 허용) 및 time to first token 3 초 또는 그 이하입니다. 이 두 지연 시간 목표는 모델 크기, 하드웨어, 프롬프트 길이, 및 동시 부하에 따라 달성하기가 매우 어려울 수 있습니다.

다른 상호작용이 없는 사용 사례는 개별 요청의 지연 시간보다 모든 요청에 걸친 총 토큰 처리량 (초당 토큰 수, 동시 실행 중인 모든 요청의 합계) 에만 관심이 있을 수 있습니다. 이는 책을 번역하거나 대규모 저장소에서 코드 파일을 요약하고 싶을 때 관련이 있을 수 있습니다.

후반 섹션에서 보겠지만, 총 처리량을 최대화하고 개별 요청의 지연 시간을 최소화하는 것 사이에는 일반적으로 트레이드오프가 존재합니다.

모든 입력 토큰에 대한 병렬 계산으로 인해 prefill 단계는 매우 GPU 연산 집약적입니다. 반면, 개별 출력 토큰을 디코딩하는 단계는 매우 적은 컴퓨팅 파워를 활용하며, 여기서는 속도는 일반적으로 GPU 메모리 대역폭에 의해 제한됩니다. 즉, 모델 가중치 및 활성화 (포함 키 와 값 벡터) 를 GPU 메모리에서 로드하고 액세스할 수 있는 속도입니다.

일반적으로 GPU 활용도 (컴퓨팅 파워 기준) 에 포화될 때까지 토큰 처리량을 증가시킬 수 있습니다. prefill 단계에서는 긴 프롬프트를 가진 단일 요청으로도 최대 GPU 활용도를 달성할 수 있습니다. decode 단계에서는 여러 요청의 배치 처리를 통해 GPU 활용도를 높일 수 있습니다. 따라서, 동시 실행 중인 요청 수에 대한 토큰 처리량을 그래프로 그리면, 낮은 병렬성에서 거의 선형적인 처리량 증가를 볼 수 있으며, 이는 메모리 제한 영역이 더 큰 배치 크기를 이득으로 삼기 때문입니다. GPU 활용도가 포화되고 계산 제한 영역에 진입한 후 처리량은 병렬성 증가에 무관하게 일정하게 유지됩니다.

이제 짧은 시간 간격 내에 도착하는 여러 요청을 추론 엔진이 어떻게 처리하는지 고려하겠습니다.

prefill 및 decode 단계 모두 동일한 연산을 다른 요청에 적용하기 위해 배치 전략을 사용할 수 있습니다. 그러나 서로 다른 요청의 prefill 및 decode를 동시에 실행했을 때 어떤 결과가 발생하는지 궁금합니다.

가장 단순한 배치는 **정적 배치 (static batching)**입니다. (1) 빈 배치를 시작하고, (2) 기다리는 항목이 많고 배치에 들어갈 수 있는 만큼 채우며, (3) 모든 배치된 항목을 처리하기까지 배치를 처리하고, (4) 새로운 빈 배치를 반복합니다.

모든 요청은 prefill 단계를 동시에 시작합니다. prefill 은 단일이지만 매우 병렬화된 GPU 연산이기 때문에 (거대한 행렬 곱셈이라고 생각해보세요), 모든 동시 요청의 prefill 단계는 동시에 완료됩니다. 그런 다음, 모든 decode 단계가 동시에 시작됩니다. 출력 토큰이 적은 요청은 먼저 종료되지만, 정적 배치로 인해 다음 기다리는 요청은 가장 긴 배치된 요청이 완료될 때까지만 시작할 수 있습니다.

정적 배치는 출력 토큰당 시간을 최적화합니다. 왜냐하면 decode 단계가 중단되지 않기 때문입니다. 단점은 매우 비효율적인 자원 활용입니다. 단일 긴 프롬프트는 prefill 동안 컴퓨팅 파워를 포화시킬 수 있으므로, 여러 prefill 을 병렬로 처리하는 것은 속도 향상을 가져오지 않으며 GPU 활용도를 최대치로 만들 것입니다. 반면, decode 단계에서는 GPU 가 충분히 활용되지 않을 가능성이 높습니다. 왜냐하면 긴 프롬프트의 prefill 보다 계산 집약적이지 않은 많은 동시 디코드가 있기 때문입니다.

그러나 가장 큰 단점은 잠재적으로 긴 *첫 토큰 생성 시간 (time to first token)*입니다. 일부 짧은 요청이 조기 종료되더라도 다음 대기열 요청은 배치에서 가장 긴 디코딩이 완료될 때까지 기다려야 합니다. 이러한 정적 배치의 결함으로 인해 추론 엔진은 일반적으로 연속 배치 (continuous batching) 전략을 구현합니다. 여기서는 완료된 요청은 즉시 배치에서 제거되고, 배치 공간은 다음 요청으로 채워집니다. 결과적으로 모든 연속 배치 전략은 prefill 와 decode 단계 간의 병행 처리를 처리해야 합니다.

요청의 대기 시간을 줄이기 위해 추론 엔진 (예: vLLM, TGI) 은 새로운 요청이 도착할 때 즉시 prefill 단계를 스케줄링하고 현재 배치에 적합하게 배치합니다. 모든 이전 요청 각각의 단일 디코딩 단계와 병행하여 새로운 요청의 prefill 을 실행하는 것이 가능합니다. 그러나 모든 작업이 동일한 GPU 연산에서 실행되므로, 그 지속 시간은 prefill 에 의해 지배되며, 디코딩 단계의 모든 요청마다 해당 시간 동안만 하나의 출력 토큰을 생성할 수 있습니다. 따라서 이러한 prefill 우선 (prioritization of prefills) 은 첫 토큰 생성 시간을 최소화하지만 이미 실행 중인 요청의 decode 단계를 중단시킵니다. 채팅 애플리케이션에서는 다른 사용자가 긴 프롬프트를 제출할 때 스트림화된 토큰 생성이 일시적으로 중단되는 것을 경험할 수 있습니다.

다음 측정에서 prefill-first 전략을 가진 연속 배치의 영향을 볼 수 있습니다.

실행 중인 디코딩에 대한 중단형 prefill 의 영향 완화를 위한 접근법은 분할 prefill (chunked prefill) 입니다. 전체 프롬프트를 단일 prefill 단계로 처리하는 대신, 여러 청크로 분배할 수 있습니다. 그러면 prefill 동안 하나의 디코딩 단계보다 더 많은 병행 디코딩 단계를 가질 수 있습니다 (전체 prefill 동안의 각 병행 요청당 단일 디코딩 단계와 대조). 분할 prefill 단계는 여전히 고립된 디코딩 단계보다 오래 걸리지만, 작은 청크 크기의 경우 사용자는 토큰 생성의 완전한 중단이 아닌 지연만 경험합니다. 이는 평균 *출력 토큰당 시간 (time per output token)*을 줄입니다. 중단 요청의 관점에서 분할 prefill 은 일부 오버헤드와 함께 고립된, 연속적인 prefill 보다 약간 더 오래 걸리므로 *첫 토큰 생성 시간 (time to first token)*에 작은 증가가 있습니다. 현재 청크 크기는 **첫 토큰 생성 시간 또는 출력 토큰당 시간을 우선시하는 튜닝 핸들 (tuning knob)**이 됩니다. 일반적인 청크 크기는 512 에서 8192 토큰 사이입니다 (vLLM 의 기본값은 chunked prefill 이 처음 구현되었을 때 512 였으며, 이후 더 높은 값으로 업데이트되었습니다).

그러나 이 전략의 가장 큰 장점은 분할 prefill 은 자원 효율성을 최대화한다는 것입니다. Prefill 은 계산 집약적이지만 decode 는 메모리 제한적입니다. 두 작업을 병행하여 실행하면 GPU 자원에 의해 제한되지 않고 전체 트루풋을 증가시킬 수 있습니다. 물론, 최대 효율성은 특정 청크 크에서만 달성되며, 이는 정확한 로드 패턴에 의존합니다.

표준 vLLM 배포 환경과 균일한 크기의 요청에서, 우리는 Chunked Prefill이 전체 토큰 처리량을 +50% 증가시켰음을 관찰했습니다. 이는 이제 TNG 에서 자체 호스팅 LLM 의 모든 vLLM 배포에 활성화되었습니다. 전반적으로 Chunked Prefill 은 대부분의 사용 사례에 대한 좋은 기본 전략입니다. 그러나 Chunk Size 를 최적화하는 것은 예측 불가능한 로드 패턴이 있는 환경 (예: 다양한 애플리케이션이 많은 TNG) 에서 매우 어렵습니다; 일반적으로 당신은 기본값을 유지합니다.

구체적인 Chunk Size 설정과 무관하게, Chunked Prefill 과 함께 동시 처리는 다음 기사에서 해결할 두 가지 도전 과제를 동반합니다.