[sglang] SGLang Whisper 모델의 CUDA Graph 도입 및 성능 최적화 분석

PR 링크: sgl-project/sglang#21190 상태: Merged | 변경: +388 / -123

들어가며

최근 SGLang 프로젝트에서 Whisper 모델의 추론 성능을 비약적으로 향상시키는 업데이트가 있었습니다. 기존 Whisper 구현체는 Python 레벨에서 관리되는 _encoder_cache와 커스텀 BMM(Batch Matrix Multiplication) 기반의 Cross-attention을 사용하고 있었는데, 이는 CUDA Graph 캡처와 호환되지 않는다는 치명적인 단점이 있었습니다. 이번 PR은 이를 RadixAttention으로 교체하고 CUDA Graph를 활성화함으로써, 정확도 손실 없이 처리량(Throughput)을 36% 향상시켰습니다.

코드 분석

1. Cross-attention의 현대화 (FlashInfer Backend)

기존의 수동 BMM 방식 대신 RadixAttention을 도입하여 KV 캐시 관리 효율을 높였습니다. 특히 flashinfer_backend.py에서 Cross-attention을 위한 인덱싱 로직이 개선되었습니다.

# Before: 디코더 seq_lens_cpu를 그대로 사용하여 Cross-attention KV 길이 계산 오류 발생
# After: encoder_lens_cpu를 별도로 계산하여 정확한 KV 길이 참조
encoder_lens_cpu = encoder_lens.cpu() if encoder_lens is not None else None
# ...
else:
    # Cross-attention: attend to encoder tokens only
    paged_kernel_lens = encoder_lens
    kv_start_idx = torch.zeros_like(encoder_lens)
    kv_lens_cpu = encoder_lens_cpu

이 변경은 GPU와 CPU 간의 불필요한 동기화를 줄이고, Cross-attention 연산 시 올바른 인코더 길이를 참조하도록 하여 CUDA Graph 캡처 성공률을 높였습니다.

2. CUDA Graph 캡처를 위한 설정 최적화

CUDA Graph를 사용하기 위해서는 캡처 시점에 모든 커널이 그래프에 포함되어야 합니다. Whisper의 인코더 길이를 고정값으로 설정하여 캡처 과정에서 Cross-attention 커널이 누락되는 문제를 해결했습니다.

# 캡처 시 encoder_len_fill_value를 0에서 max_source_positions(1500)으로 변경
# 이를 통해 Cross-attention 커널이 그래프 캡처 시점에 포함됨
set_encoder_len_fill_value(max_source_positions)

3. OpenAI 호환성 강화 (Verbose JSON)

API 서버단에서는 verbose_json 포맷을 지원하도록 확장하여, 타임스탬프 정보를 포함한 세밀한 전사(Transcription) 결과를 반환할 수 있게 되었습니다.

# serving_transcription.py 내 신규 파싱 로직
for token_id in output_ids:
    if token_id >= TIMESTAMP_BASE_TOKEN_ID:
        timestamp = (token_id - TIMESTAMP_BASE_TOKEN_ID) * TIMESTAMP_BASE_OFFSET
        # 타임스탬프 토큰을 기준으로 세그먼트 분리

왜 이게 좋은가

이번 최적화의 핵심은 **'정적 그래프(Static Graph)로의 전환'**입니다.

1. 성능 향상: 벤치마크 결과, 처리량이 2.40 req/s에서 3.26 req/s로 36% 증가했습니다. 이는 CUDA Graph를 통해 커널 실행 오버헤드를 획기적으로 줄였기 때문입니다.
2. 정확도 유지: 연산 로직을 변경했음에도 불구하고 WER(Word Error Rate)은 12.77%로 동일하게 유지되었습니다.
3. 교훈: 추론 엔진에서 CUDA Graph를 도입할 때 가장 큰 걸림돌은 '동적 길이(Dynamic Length)'입니다. encoder_len_fill_value를 최대치로 고정하여 정적 그래프를 생성하고, 실제 연산 시에는 마스킹을 통해 올바른 길이를 제어하는 패턴은 고성능 추론 엔진 설계의 정석입니다.

리뷰어 피드백 반영

리뷰 과정에서 encoder_out_cache_loc를 trtllm_mha 백엔드에 추가하여 향후 확장성을 확보했고, position_ids의 OOB(Out-of-Bounds) 문제를 방지하기 위해 클램핑 로직을 추가하는 등 안정성 측면에서도 꼼꼼한 검토가 이루어졌습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.cuda.make_graphed_callables.html
• https://docs.sglang.ai/index.html

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.