[sglang] SGLang 토크나이저 매니저: O(n²) 복사 비용 제거를 통한 스트리밍 성능 최적화
PR 링크: sgl-project/sglang#22567 상태: Merged | 변경: +None / -None
들어가며
대규모 언어 모델(LLM) 기반 애플리케이션에서 사용자 경험을 좌우하는 중요한 요소 중 하나는 응답 속도입니다. 특히 스트리밍 방식으로 출력을 제공할 때, 중간 단계에서의 비효율적인 연산은 전체 시스템의 성능을 크게 저하시킬 수 있습니다. 이번 PR은 sgl-project/sglang 레포지토리의 토크나이저 매니저에서 발생하는 심각한 성능 병목 현상, 즉 비증분 스트리밍(non-incremental streaming) 시 O(n²) 복사 비용 문제를 해결하여 스트리밍 성능을 획기적으로 개선하는 내용을 담고 있습니다.
기존 sglang의 tokenizer_manager는 비증분 스트리밍 모드(기본 설정)에서 state.output_ids.copy()와 state.get_text() 연산이 각 스텝마다 전체 시퀀스 길이에 비례하는 비용을 발생시켰습니다. 이는 총 O(n²)의 복사 및 문자열 재구축 비용으로 이어져, 16k 출력 토큰에서 def-stream 방식이 iso-stream/nostream 방식보다 2.3배 느려지는 병목 현상을 초래했습니다. 이 PR은 이러한 비효율성을 제거하고, 불필요한 복사와 문자열 생성을 지연시킴으로써 스트리밍 성능을 대폭 향상시키는 것을 목표로 합니다.
코드 분석: O(n²) 병목 제거 전략
이번 PR의 핵심은 python/sglang/srt/managers/tokenizer_manager.py 파일의 _handle_batch_output 함수와 _wait_one_response 함수에 대한 변경입니다. 주요 변경 사항은 비증분 스트리밍 시 중간 결과 처리 방식을 최적화하는 데 있습니다.
1. _handle_batch_output 함수 변경: 불필요한 복사 및 문자열 생성 지연
_handle_batch_output 함수는 배치 출력을 처리하고, 각 스텝마다 out_dict를 생성하여 다음 단계로 전달합니다. 기존에는 비증분 스트리밍의 중간 단계에서도 state.output_ids.copy()와 state.get_text()를 호출하여 output_ids와 text를 out_dict에 포함시켰습니다. 하지만 프로파일링 결과, 스트리밍 소비자는 중간 청크에서 output_ids를 읽지 않으며, _wait_one_response는 비증분 스트리밍의 경우 마지막 out_dict만 사용하고 중간 out_dict는 버린다는 사실이 밝혀졌습니다. 이로 인해 불필요한 O(n) 연산이 매 스텝마다 발생하여 O(n²)의 총 비용이 발생했습니다.
Before:
else:
output_token_ids = state.output_ids.copy()
output_text = state.get_text()
out_dict = {
"text": output_text,
"output_ids": output_token_ids,
"meta_info": meta_info,
}
After:
elif state.finished:
out_dict = {
"text": state.get_text(),
"output_ids": state.output_ids.copy(),
"meta_info": meta_info,
}
else:
# Non-incremental intermediate: pass reference (no
# copy) and defer text to _wait_one_response to avoid
# O(n) per-step cost that compounds to O(n^2).
out_dict = {
"text": None,
"output_ids": state.output_ids,
"meta_info": meta_info,
}
변경 후에는 state.finished 상태가 아닐 때, 즉 중간 단계에서는 output_ids를 .copy() 없이 레퍼런스(참조)로 전달하고, text는 None으로 설정하여 생성을 지연시킵니다. 이는 asyncio 이벤트 루프가 단일 스레드이므로 안전합니다. 소비자가 참조를 직렬화하기 전에 state.output_ids가 extend()될 수 있는 await 지점이 없기 때문입니다.
BatchTokenIDOutput 경로에서도 동일한 최적화가 적용되었습니다.
Before (BatchTokenIDOutput):
else:
output_token_ids = state.output_ids.copy()
out_dict = {
"output_ids": output_token_ids,
"meta_info": meta_info,
}
After (BatchTokenIDOutput):
elif state.finished:
out_dict = {
"output_ids": state.output_ids.copy(),
"meta_info": meta_info,
}
else:
out_dict = {
"output_ids": state.output_ids,
"meta_info": meta_info,
}
2. _wait_one_response 함수 변경: 지연된 텍스트 해결
_handle_batch_output에서 text 생성을 지연시켰으므로, 실제로 text가 필요한 시점에 이를 해결해야 합니다. _wait_one_response 함수는 _handle_batch_output에서 생성된 out_dict를 받아 최종 응답을 준비하는 역할을 합니다. 비증분 스트리밍의 경우, 이 함수는 마지막 out_dict만 사용하므로, 해당 시점에 state.get_text()를 호출하여 text를 생성하도록 변경되었습니다.
Before:
else:
out = out_list[-1]
if finished:
# For non-streaming cases, response has not been sent yet (`response_sent_to_client_time` has not been set yet).
# Record response sent time right before we log finished results and metrics.
After:
else:
out = out_list[-1]
# Resolve deferred text for non-incremental streaming.
# _handle_batch_output sets "text": None on intermediate chunks
# to avoid O(n) string rebuild per step (O(n^2) total).
if (
is_stream
and not incremental_stream
and "text" in out
and out["text"] is None
):
out["text"] = state.get_text()
if finished:
# For non-streaming cases, response has not been sent yet (`response_sent_to_client_time` has not been set yet).
# Record response sent time right before we log finished results and metrics.
이 변경으로 인해 _handle_batch_output에서 매 스텝마다 발생하던 O(n) 문자열 재구축 비용이 _wait_one_response에서 최종적으로 한 번만 발생하게 되어, 전체 O(n²) 비용이 O(n)으로 줄어듭니다.
왜 이 최적화가 좋은가?
이 최적화는 불필요한 연산을 제거하고 필요한 시점에만 자원을 할당하는 지연 평가(Lazy Evaluation) 원칙을 효과적으로 적용한 사례입니다. 특히, output_ids의 불필요한 복사를 제거하고 text 생성을 지연시킴으로써, 다음과 같은 성능 향상을 달성했습니다.
벤치마킹 결과 (Qwen/Qwen3-1.7B | 4x H200 | 400 prompts | input=128 | output=16384):
| Metric | Before (main) | After (fix) | Change |
|---|---|---|---|
| Output tok/s | 8,004 | 14,618 | +82.6% |
| Duration (s) | 402 | 220 | -45% |
| Mean TTFT (ms) | 69,959 | 1,159 | -98.3% |
| Mean TPOT (ms) | 11.83 | 6.26 | -47% |
| Mean ITL (ms) | 11.96 | 6.15 | -49% |
가장 병목이 심했던 def-stream 16384 출력 토큰 시나리오에서 출력 토큰/초가 82.6% 증가하고, 전체 실행 시간이 45% 단축되었습니다. 특히 Mean TTFT (Time To First Token)는 98.3% 감소하여, 첫 토큰 응답 시간이 극적으로 개선되었음을 보여줍니다. 이는 사용자 경험에 직접적인 영향을 미치는 매우 중요한 지표입니다.
이 최적화는 nostream 모드에는 영향을 미치지 않아 회귀(regression)가 없음을 확인했습니다. 여전히 def-stream이 iso-stream/nostream보다 약간 느리지만, 이는 토크나이저 매니저가 아닌 SSE 직렬화 및 HTTP 청크 인코딩과 같은 다른 오버헤드 때문이며, O(n²) 병목은 완전히 제거되었습니다.
일반적 교훈:
- 프로파일링의 중요성: 이 최적화는 철저한 프로파일링을 통해 O(n²) 병목의 정확한 원인을 파악한 덕분에 가능했습니다. 성능 문제가 의심될 때는 항상 프로파일링 도구를 사용하여 실제 병목 지점을 찾는 것이 중요합니다.
- 불필요한 복사 피하기: 파이썬에서 리스트나 문자열과 같은 큰 객체를 불필요하게 복사하는 것은 성능에 치명적일 수 있습니다. 특히 루프 내에서 반복적으로 복사가 일어나는 경우 O(n²) 또는 그 이상의 복잡도를 야기할 수 있으므로, 참조 전달(pass by reference)을 적극적으로 고려해야 합니다.
- 지연 평가(Lazy Evaluation): 데이터가 실제로 필요할 때까지 연산을 지연시키는 것은 리소스 사용을 최적화하는 강력한 방법입니다. 중간 결과가 버려지거나 최종 단계에서만 사용되는 경우, 지연 평가는 상당한 성능 이득을 가져올 수 있습니다.
- 스트리밍 시스템 설계: 스트리밍 시스템에서는 각 중간 단계의 오버헤드가 전체 성능에 큰 영향을 미칩니다. 각 단계에서 발생하는 연산의 복잡도와 필요성을 신중하게 검토하여 불필요한 작업을 최소화해야 합니다.
이번 PR은 SGLang의 스트리밍 성능을 크게 향상시켰으며, 대규모 언어 모델 서비스의 효율성을 높이는 데 기여할 것입니다.
References
- Python
list.copy()documentation - Python
str.join()(related toget_text()performance) - asyncio event loop
참고 자료
- https://docs.python.org/3/library/stdtypes.html#list.copy
- https://docs.python.org/3/library/stdtypes.html#str.join
- https://docs.python.org/3/library/asyncio-eventloop.html
⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.
관련 포스트
- [sglang] SGLang Ngram 추측 디코딩: 외부 코퍼스 기반 Suffix Automaton 통합으로 성능 최적화
- [sglang] SGLang: Piecewise CUDA Graph와 Sliding Window Attention의 효율적인 공존
- [sglang] Dumper 디버그 유틸리티 리팩토링: 설정 구조 개선과 Non-intrusive 모드 도입
- [sglang] 미사용 BatchMultimodalOutput/DecodeReq 제거로 코드베이스 정리
- [Ray Serve] SGLang 서버의 순차 배치 처리를 동시 실행으로 전환
PR Analysis 의 다른글
- 이전글 [sglang] Whisper 모델 추론 성능 극대화: 동시 Prefill 요청을 위한 배치 인코더 최적화
- 현재글 : [sglang] SGLang 토크나이저 매니저: O(n²) 복사 비용 제거를 통한 스트리밍 성능 최적화
- 다음글 [sglang] Intel GPU 가속을 위한 SGLang MoE 커널 최적화: GPT-OSS bf16 지원 분석
댓글