본문으로 건너뛰기

[sglang] SGLang 성능 최적화: 과도한 Prefill CUDA Graph Padding 방지로 TTFT 개선

PR 링크: sgl-project/sglang#31487 상태: Merged | 변경: +62 / -4

들어가며

대규모 언어 모델(LLM) 서비스에서 응답 속도는 사용자 경험에 결정적인 영향을 미칩니다. 특히 첫 토큰 생성 시간(Time To First Token, TTFT)은 사용자가 모델의 응답을 체감하는 첫 순간이기에 매우 중요합니다. SGLang은 이러한 LLM 서빙의 성능을 극대화하기 위해 다양한 최적화 기법을 활용하며, 그중 하나가 바로 NVIDIA CUDA Graph입니다.

CUDA Graph는 GPU 커널 실행의 오버헤드를 줄여주는 강력한 도구입니다. 여러 GPU 커널 호출을 하나의 그래프로 캡처하여 한 번의 API 호출로 실행함으로써, 특히 반복적인 작업에서 CPU 오버헤드를 크게 줄일 수 있습니다. SGLang에서는 LLM의 'Prefill' 단계(프롬프트 처리)에서 이 CUDA Graph를 활용하여 성능을 개선해왔습니다.

하지만 모든 최적화 기법이 항상 최적의 결과를 가져오는 것은 아닙니다. 이번 PR(perf: avoid excessive prefill CUDA graph padding)은 기존 Prefill CUDA Graph 활용 방식의 한계를 명확히 인식하고, 이를 개선하여 TTFT를 획기적으로 단축한 사례입니다. 기존에는 CUDA Graph가 캡처된 토큰 버킷 크기에 맞춰 패딩(padding)을 과도하게 적용하는 경우가 있었고, 이는 불필요한 연산으로 이어져 오히려 eager execution보다 느려지는 역효과를 낳았습니다. 이 PR은 이러한 문제를 해결하고, CUDA Graph를 더욱 '스마트하게' 활용하는 방법을 제시합니다.

코드 분석

이번 최적화의 핵심은 PrefillCudaGraphRunner가 CUDA Graph를 실행할지 eager execution으로 폴백(fallback)할지 결정하는 로직에 있습니다. 과도한 패딩으로 인해 불필요한 연산량이 많아지는 경우를 방지하는 것이 목표입니다.

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

이 파일은 SGLang에서 Prefill 단계의 CUDA Graph 실행을 담당하는 핵심 컴포넌트입니다. can_run_graph 메서드는 주어진 ForwardBatch에 대해 CUDA Graph를 실행할 수 있는지 여부를 판단합니다.

Before: 기존 can_run_graph 로직

이전에는 can_run_graph 메서드에서 num_tokens가 max_num_tokens를 초과하지 않는지만 확인했습니다. load_batch() 함수가 내부적으로 다음 캡처 버킷 크기에 맞춰 토큰을 패딩하기 때문에, 정확한 shape 일치 여부보다는 단순히 최대 토큰 수를 넘지 않는지만 검사했습니다.

        if num_tokens > self.max_num_tokens:
            return False
        # No backend-level shape check here: load_batch bucket-pads
        # num_tokens up to the nearest captured shape, so eligibility is
        # bounded by num_tokens <= self.max_num_tokens (already
        # checked above), not by exact shape membership.

이 로직의 문제는 num_tokens가 max_num_tokens보다 작더라도, load_batch()가 패딩하는 과정에서 실제 필요한 토큰 수보다 훨씬 많은 토큰을 처리하게 될 수 있다는 점입니다. 예를 들어, 캡처 버킷이 [2816, 3072, 8192]로 설정되어 있을 때, 3208개의 토큰을 가진 요청은 8192개의 토큰을 처리하는 그래프를 재실행하게 되어 2.55배의 불필요한 연산을 수행하게 됩니다.

After: 개선된 can_run_graph 로직

PR은 이러한 과도한 패딩 문제를 해결하기 위해 새로운 상수를 도입하고, can_run_graph 메서드에 추가적인 검증 로직을 추가했습니다.

먼저, 과도한 패딩을 판단하는 기준이 되는 _MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR 상수가 정의되었습니다.

# A replay executes every padded token in its capture bucket. Sparse bucket
# lists can otherwise turn the lower launch overhead into substantially more
# model work than an exact-shape eager forward.
_MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR = 2

이 상수는 CUDA Graph 재실행 시 패딩된 토큰 수가 실제 토큰 수의 몇 배를 초과하면 안 되는지를 정의합니다. 현재는 2배로 설정되어 있습니다. 이는 리뷰어의 코멘트에서도 언급되었듯이, 성능 회귀가 관찰될 경우 조정될 수 있는 휴리스틱 값입니다.

그리고 can_run_graph 메서드 내부에 이 상수를 활용한 새로운 검증 로직이 추가되었습니다.

        padded_num_tokens = self._pad_to_bucket(num_tokens, self.capture_num_tokens)
        if padded_num_tokens > num_tokens * _MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR:
            return False
        # No exact-shape check here: load_batch bucket-pads to the nearest
        # captured shape. The factor above only rejects replays whose padded
        # model work is disproportionate to the useful token count.

새로운 로직은 다음과 같이 작동합니다:

  1. _pad_to_bucket 메서드를 호출하여 현재 num_tokens가 어떤 캡처 버킷으로 패딩될지 계산하고, 그 결과를 padded_num_tokens에 저장합니다.
  2. padded_num_tokens가 num_tokens의 _MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR (현재 2)배를 초과하는지 확인합니다.
  3. 만약 초과한다면, 이는 과도한 패딩으로 인해 불필요한 연산량이 너무 많아진다는 의미이므로, CUDA Graph 실행을 포기하고 False를 반환하여 eager execution으로 폴백하도록 합니다.

이 변경으로 인해, CUDA Graph의 낮은 런칭 오버헤드(launch overhead) 이점이 과도한 모델 연산량 증가로 상쇄되는 것을 방지할 수 있게 됩니다. 즉, CUDA Graph를 무조건 사용하는 것이 아니라, 효율적일 때만 사용하도록 하는 '스마트한' 전략을 적용한 것입니다.

test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py

새로운 로직의 정확성을 검증하기 위해 전용 유닛 테스트 파일이 추가되었습니다. 이는 _MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR가 2로 설정되었을 때의 동작을 명확히 테스트합니다.

test_rejects_more_than_two_x_token_padding

이 테스트는 패딩된 토큰 수가 실제 토큰 수의 2배를 초과할 때 CUDA Graph 실행이 거부되는지 확인합니다. capture_num_tokens가 [4, 16]이고 num_tokens가 5일 경우를 예로 듭니다. 5개의 토큰은 다음 버킷인 16으로 패딩됩니다. 이때 padded_num_tokens (16)은 num_tokens (5)의 2배(10)를 초과하므로, can_run_graph는 False를 반환해야 합니다.

    def test_rejects_more_than_two_x_token_padding(self):
        runner = self._make_runner()
        # num_tokens=5, padded_num_tokens=16. 16 > 5 * 2 (10) -> False
        self.assertFalse(runner.can_run_graph(self._make_forward_batch(5)))

test_accepts_two_x_token_padding

이 테스트는 패딩된 토큰 수가 실제 토큰 수의 2배와 같거나 작을 때 CUDA Graph 실행이 허용되는지 확인합니다. num_tokens가 8일 경우, 다음 버킷인 16으로 패딩됩니다. 이때 padded_num_tokens (16)은 num_tokens (8)의 2배(16)와 같으므로, can_run_graph는 True를 반환해야 합니다.

    def test_accepts_two_x_token_padding(self):
        runner = self._make_runner()
        # num_tokens=8, padded_num_tokens=16. 16 <= 8 * 2 (16) -> True
        self.assertTrue(runner.can_run_graph(self._make_forward_batch(8)))

이러한 유닛 테스트는 새로운 로직이 의도한 대로 정확하게 동작하며, _MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR라는 휴리스틱 값이 시스템에 안정적으로 통합되었음을 보장합니다.

왜 이게 좋은가

이 PR은 LLM 서빙의 핵심 성능 지표인 TTFT와 E2E(End-to-End) 지연 시간을 크게 개선하는 데 기여했습니다. 이는 단순히 CUDA Graph를 사용하는 것을 넘어, 그 사용 방식을 최적화함으로써 얻어진 결과입니다.

성능 개선 수치

PR 설명에 따르면, 8x NVIDIA H200, TP8, Kimi-K2.7-Code 모델 환경에서 bench_serving 벤치마크를 수행한 결과는 다음과 같습니다.

Metric Before After Delta
Mean TTFT 310.11 ms 283.35 ms -8.63%
P90 TTFT 462.46 ms 415.45 ms -10.17%
P99 TTFT 467.96 ms 454.51 ms -2.87%
Mean E2E 538.27 ms 510.13 ms -5.23%
Median E2E 513.29 ms 407.60 ms -20.59%

특히 Mean TTFT는 약 8.63%, P90 TTFT는 약 10.17% 감소하여 사용자 체감 성능이 크게 향상되었음을 보여줍니다. Median E2E 지연 시간은 무려 20.59% 감소했습니다. 이는 특정 요청들이 과도한 패딩으로 인해 심각한 지연을 겪고 있었음을 시사하며, 이 PR이 그러한 '꼬리 지연(tail latency)'을 효과적으로 줄였음을 의미합니다.

개별 요청의 Forward duration 변화를 보면 그 효과가 더욱 극명하게 드러납니다.

Raw prefill Before After Forward duration
3208 tokens graph 8192 eager 522.13 -> 316.29 ms (-39.42%)
3353 tokens graph 8192 eager 654.76 -> 459.02 ms (-29.89%)

3208 토큰과 3353 토큰 요청의 경우, 이전에는 8192 토큰 그래프를 사용하여 각각 2.55배, 2.44배의 불필요한 연산을 수행했습니다. 이 PR 적용 후에는 eager execution으로 폴백하면서 Forward duration이 각각 39.42%, 29.89% 감소했습니다. 이는 과도한 패딩으로 인한 비효율이 얼마나 컸는지를 명확히 보여줍니다.

일반적 교훈

이 최적화 사례는 다음과 같은 중요한 교훈을 제공합니다.

  1. 최적화는 '더 많이'가 아니라 '더 스마트하게': 단순히 강력한 최적화 기법(CUDA Graph)을 적용하는 것만이 능사가 아닙니다. 해당 기법의 작동 방식과 한계를 정확히 이해하고, 특정 상황에서 오히려 성능 저하를 일으킬 수 있음을 인지해야 합니다. 이 PR은 CUDA Graph의 런칭 오버헤드 감소 이점이 과도한 연산량 증가로 상쇄될 때에는 eager execution이 더 효율적일 수 있음을 보여줍니다.
  2. 휴리스틱의 중요성과 신중한 접근: _MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR와 같은 휴리스틱 값은 실제 시스템의 복잡성을 단순화하여 효율적인 결정을 내리는 데 필수적입니다. 하지만 이러한 값은 모델, 하드웨어, 워크로드에 따라 최적이 아닐 수 있으므로, mickqian 리뷰어의 코멘트처럼 지속적인 모니터링과 필요 시 조정이 중요합니다.
  3. 성능 병목 지점의 정확한 식별: 이 PR은 기존 can_run_graph 로직이 max_num_tokens만 확인하고 패딩으로 인한 실제 연산량 증가를 고려하지 않았다는 근본적인 문제를 정확히 식별했습니다. 성능 문제를 해결하기 위해서는 표면적인 현상 너머의 근본 원인을 찾아내는 것이 중요합니다.
  4. 테스트의 중요성: 새로운 로직이 추가되면서, test_prefill_cuda_graph_padding.py와 같은 전용 유닛 테스트가 함께 추가되어 변경 사항이 의도한 대로 동작함을 검증했습니다. 이는 코드 변경의 안정성과 신뢰성을 높이는 데 필수적인 과정입니다.

마치며

SGLang의 이번 Prefill CUDA Graph 최적화는 LLM 서빙 시스템의 성능을 한 단계 끌어올린 중요한 개선입니다. 과도한 패딩으로 인한 불필요한 연산을 줄이고, CUDA Graph를 더욱 지능적으로 활용함으로써 TTFT와 E2E 지연 시간을 크게 단축했습니다. 이는 단순히 기술을 적용하는 것을 넘어, 그 기술의 한계를 이해하고 상황에 맞게 최적의 전략을 선택하는 것이 얼마나 중요한지를 보여주는 모범적인 사례입니다. 앞으로도 SGLang이 이러한 지속적인 성능 개선 노력을 통해 더욱 빠르고 효율적인 LLM 서빙을 제공할 것으로 기대됩니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글