[sglang] Whisper 모델 추론 성능 극대화: 동시 Prefill 요청을 위한 배치 인코더 최적화
PR 링크: sgl-project/sglang#22361 상태: Merged | 변경: +None / -None
들어가며
대규모 언어 모델(LLM) 기반의 음성 인식 시스템에서 추론 성능은 사용자 경험과 직결되는 핵심 요소입니다. 특히, OpenAI의 Whisper 모델은 뛰어난 음성 인식 성능으로 널리 사용되고 있지만, 동시 요청 처리 시 발생하는 성능 병목 현상은 여전히 해결해야 할 과제입니다.
이번 분석 글에서는 sglang 레포지토리의 PR ([Whisper] Batch encoder forward for concurrent prefill requests)을 통해 Whisper 모델의 Prefill 단계에서 발생하는 성능 저하 문제를 어떻게 해결했는지 심층적으로 살펴보겠습니다. 이 PR은 여러 개의 새로운 Whisper 요청이 동시에 들어왔을 때, 기존에 각 요청마다 개별적으로 수행되던 인코더 연산을 하나의 배치로 묶어 처리함으로써 전체적인 추론 처리량(throughput)을 크게 향상시키는 것을 목표로 합니다.
기존 코드에서는 N개의 새로운 Whisper 요청이 Prefill 배치에 포함될 경우, 인코더가 N번 순차적으로 호출되었습니다. GPU 종류에 따라 다르지만, 각 인코더 호출에 약 7ms ~ 18ms가 소요됩니다. 만약 5개의 요청이 동시에 들어온다면, 디코더(decode) 단계에서 처리할 수 있는 다른 요청들을 기다리게 만들면서 35ms ~ 90ms의 시간을 허비하게 됩니다. 프로파일링 결과에 따르면, 이러한 순차적인 인코더 호출은 Prefill 단계에서 전체 스케줄러 시간의 35%를 차지하며, 동시성(concurrency)이 높아질수록 성능 확장성을 심각하게 저해하는 주요 원인이었습니다.
이 PR은 이러한 문제를 해결하기 위해, 캐시되지 않은 모든 오디오 특징(audio features)들을 하나의 배치 텐서로 수집하고, 인코더를 요청마다 호출하는 대신 Prefill 배치당 단 한 번만 실행하도록 변경했습니다. 이 변경이 어떻게 이루어졌고, 왜 성능 향상으로 이어졌는지 자세히 알아보겠습니다.
코드 분석
python/sglang/srt/models/whisper.py
이 PR의 핵심 변경 사항은 whisper.py 파일 내 forward 함수의 prefill 로직 부분에 집중되어 있습니다.
1. 장치(Device) 이동 로직 추가
PR에서는 forward 함수 시작 부분에 다음과 같이 장치 이동 로직을 추가했습니다. 이는 모델의 가중치가 있는 장치로 입력 텐서들을 명시적으로 이동시켜 연산이 올바른 장치에서 수행되도록 보장합니다.
Before:
# No explicit device placement for input_features and position_ids
After:
device = self.conv1.weight.device
input_features = input_features.to(device=device)
position_ids = position_ids.to(device=device)
이 변경은 직접적인 성능 최적화라기보다는, 모델이 여러 장치에서 올바르게 동작하도록 보장하는 안정성 개선에 가깝습니다. 특히 비동기 연산이나 데이터 로딩 시 발생할 수 있는 장치 불일치 문제를 방지합니다.
2. 인코더 순차 호출 제거 및 배치 처리 도입
가장 중요한 변경은 Prefill 단계에서 여러 요청의 인코더 연산을 묶어서 한 번에 처리하는 부분입니다. 기존에는 mm_inputs_list와 encoder_cached_list를 순회하며 각 요청에 대해 개별적으로 인코더를 호출하고 그 결과를 encoder_list에 저장했습니다.
Before:
encoder_list = []
for i, (mm_input, cached) in enumerate(
zip(mm_inputs_list, encoder_cached_list)
):
if cached or mm_input is None or not mm_input.mm_items:
continue
features = mm_input.mm_items[0].feature
if features.ndim == 2:
features = features.unsqueeze(0)
encoder_len = features.shape[-1] // 2
encoder_position_ids = torch.arange(encoder_len).to(
features.device, non_blocking=True
)
req_encoder_output = self.encoder(
features.to(dtype), encoder_position_ids, forward_batch
)
req_encoder_output = req_encoder_output.squeeze(0)
encoder_list.append(req_encoder_output)
if encoder_list:
encoder_hidden_states = torch.cat(encoder_list, dim=0)
위 코드에서는 각 요청(mm_input)의 오디오 특징(features)을 가져와 encoder_position_ids를 생성하고, self.encoder를 호출한 뒤 결과를 encoder_list에 추가합니다. 마지막으로 encoder_list에 쌓인 결과들을 torch.cat으로 합칩니다. 이 과정이 각 요청마다 반복되어 상당한 오버헤드를 발생시킵니다.
After:
# Collect features from all uncached requests for batched encoding
features_to_encode = []
for mm_input, cached in zip(mm_inputs_list, encoder_cached_list):
if cached or mm_input is None or not mm_input.mm_items:
continue
features = mm_input.mm_items[0].feature
if features.ndim == 2:
features = features.unsqueeze(0)
features_to_encode.append(features.to(dtype))
if features_to_encode:
# Batch all features and run encoder once instead of sequentially
features_batch = torch.cat(features_to_encode, dim=0)
encoder_len = features_batch.shape[-1] // 2
encoder_position_ids = torch.arange(
encoder_len, device=features_batch.device
)
batched_output = self.encoder(
features_batch, encoder_position_ids, forward_batch
)
# Flatten [N, seq_len, dim] → [N*seq_len, dim] for cross-attention
encoder_hidden_states = batched_output.reshape(
-1, batched_output.shape[-1]
)
변경된 코드에서는 먼저 features_to_encode 리스트에 캐시되지 않은 모든 요청의 오디오 특징들을 수집합니다. 그 후, if features_to_encode: 블록 안에서 이 특징들을 torch.cat을 사용하여 하나의 features_batch로 합칩니다. 이 features_batch에 대해 self.encoder를 단 한 번만 호출합니다. 인코더의 각 레이어(Conv1d, 임베딩, 트랜스포머 레이어, LayerNorm 등)는 배치 차원을 지원하므로, 이러한 배칭 연산은 문제없이 수행됩니다. 마지막으로, 인코더의 출력(batched_output)은 [N, seq_len, dim] 형태에서 [N*seq_len, dim] 형태로 변환되어 다운스트림 크로스 어텐션(cross-attention) 레이어의 KV 캐시에 사용됩니다.
이 변경은 GPU 커널 실행 오버헤드와 연산 준비 시간을 크게 줄여주며, 특히 배치 크기가 작을 때 효과적입니다. 또한, GPU가 더 큰 배치 작업을 효율적으로 처리할 수 있게 하여 연산 밀도를 높입니다.
왜 이게 좋은가?
성능 향상
PR에서 제공된 벤치마크 결과는 이 최적화의 효과를 명확하게 보여줍니다.
B200 (183 GB, 1 GPU):
| Concurrency | Baseline (req/s) | Batched (req/s) | Change |
|---|---|---|---|
| 1 | 10.02 | 10.01 | 0% |
| 4 | 23.33 | 23.75 | +2% |
| 16 | 37.43 | 38.12 | +2% |
| 64 | 43.28 | 44.63 | +3% |
B200 GPU에서는 배치 크기가 작을 때 인코더 연산 자체가 차지하는 시간이 상대적으로 짧기 때문에(약 7ms/req), 성능 향상이 미미합니다. 하지만 동시성이 증가함에 따라 약간의 개선이 관찰됩니다.
GB300 (GKE, 1 GPU):
| Concurrency | Baseline (req/s) | Batched (req/s) | Change |
|---|---|---|---|
| 1 | 7.43 | 7.80 | +5.0% |
| 4 | 13.87 | 14.99 | +8.0% |
| 16 | 21.61 | 26.50 | +22.6% |
| 64 | 24.96 | 32.17 | +28.9% |
GB300 GPU에서는 상황이 훨씬 극적입니다. 특히 동시 요청 수가 16개와 64개일 때, 각각 **22.6%**와 **28.9%**의 처리량 증가를 보였습니다. 이는 GB300 GPU에서 인코더 호출 오버헤드가 B200보다 크기 때문에, 배치 처리를 통한 오버헤드 감소 효과가 더 두드러지기 때문입니다. 인코더 타이밍 로그에서도 이를 확인할 수 있습니다:
Encoder timing (from server logs, GB300):
Before (sequential):
n_requests=4 total=72ms avg=18ms/req (4 x 18ms serial)
After (batched):
batch_size=4 encoder=19.1ms
batch_size=3 encoder=15.8ms
batch_size=2 encoder=16.0ms
4개의 요청을 순차적으로 처리할 때 총 72ms가 걸렸던 것이, 배치 처리 시에는 배치 크기에 따라 15.8ms ~ 19.1ms로 크게 단축되었습니다. 이는 배치 처리로 인한 연산 효율성 증대와 오버헤드 감소 효과를 명확히 보여줍니다.
또한, 모든 테스트에서 WER(Word Error Rate)은 12.77-12.78%로 거의 변하지 않아, 성능 향상이 음성 인식 정확도를 저하시키지 않았음을 알 수 있습니다.
일반적 교훈
- 순차 연산의 위험성 인지: LLM 추론 파이프라인에서 여러 요청을 개별적으로 처리하는 순차적인 루프는 심각한 성능 병목을 야기할 수 있습니다. 특히, 각 단계의 연산 비용이 크지 않더라도 반복 호출 자체의 오버헤드가 누적되면 전체 처리량에 큰 영향을 미칩니다.
- 배치 연산의 힘 활용: 딥러닝 모델의 많은 연산(컨볼루션, 행렬 곱셈, 어텐션 등)은 배치 연산에 최적화되어 있습니다. 가능한 경우, 여러 입력을 묶어 배치로 처리하면 GPU 활용률을 높이고 커널 실행 오버헤드를 줄여 성능을 크게 향상시킬 수 있습니다.
- 모델 아키텍처 이해: Whisper의 인코더와 같이 배치 차원을 자연스럽게 지원하는 모듈은 배치 처리에 매우 적합합니다. 모델의 각 구성 요소가 배치 연산을 어떻게 지원하는지 이해하는 것이 최적화 전략 수립에 중요합니다.
- 벤치마킹의 중요성: 다양한 하드웨어와 부하 조건에서 성능을 측정하는 것은 최적화의 효과를 검증하고, 특정 환경에서의 성능 개선 정도를 파악하는 데 필수적입니다. B200과 GB300에서의 결과 차이는 하드웨어 특성과 오버헤드 민감도를 보여줍니다.
결론
sglang 레포지토리의 이 PR은 Whisper 모델의 Prefill 단계에서 발생하는 인코더 순차 호출 문제를 효과적으로 해결했습니다. 여러 요청의 오디오 특징을 하나의 배치로 묶어 인코더를 단 한 번만 실행함으로써, 특히 동시 요청 수가 많을 때 추론 처리량을 최대 28.9%까지 향상시켰습니다. 이는 딥러닝 모델 추론 최적화에서 배치 연산의 중요성을 다시 한번 강조하며, LLM 기반 서비스의 성능을 개선하는 데 중요한 기여를 했습니다.
참고 자료
- https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html
- https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html
- https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.
관련 포스트
PR Analysis 의 다른글
- 이전글 [vllm] vLLM 성능의 한계를 넘다: MXFP8 양자화 지원 및 MoE 최적화 분석
- 현재글 : [sglang] Whisper 모델 추론 성능 극대화: 동시 Prefill 요청을 위한 배치 인코더 최적화
- 다음글 [sglang] SGLang 토크나이저 매니저: O(n²) 복사 비용 제거를 통한 스트리밍 성능 최적화
댓글