본문으로 건너뛰기

[sglang] LLM 서빙 최적화: Gumbel-max 트릭으로 CPU 병목 제거하기 (SGLang 사례)

PR 링크: sgl-project/sglang#38117 상태: Merged | 변경: +9 / -1

들어가며

LLM(Large Language Model) 서빙 프레임워크에서 성능 최적화는 보통 GPU 커널의 연산 속도를 높이는 데 집중됩니다. 하지만 모델의 크기가 작아지거나(Small models), 배치 사이즈가 작은 경우, 성능의 병목 지점은 GPU 연산이 아닌 CPU 디스패치(Dispatch)로 옮겨갑니다. 즉, GPU가 일을 마쳤음에도 불구하고 CPU가 다음 명령을 내리는 속도가 느려 전체 시스템이 기다리게 되는 'CPU-bound' 상태가 발생하는 것입니다.

최근 SGLang 프로젝트에 반영된 PR #34020339401은 이러한 CPU 병목을 해결하기 위해 샘플러(Sampler)의 핵심 로직을 수정했습니다. 기존의 torch.multinomial 함수가 가진 내부적인 CPU 오버헤드를 Gumbel-max 트릭으로 대체하여, 소형 모델에서 비약적인 성능 향상을 이끌어낸 사례를 분석해 보겠습니다.

문제의 핵심: torch.multinomial의 숨겨진 비용

LLM의 디코딩 과정에서 다음 토큰을 선택할 때, 확률 분포로부터 샘플링을 수행합니다. SGLang의 메인 샘플러는 이를 위해 PyTorch의 표준 함수인 torch.multinomial을 사용해 왔습니다. 하지만 프로파일링 결과, 이 함수가 매 스텝마다 상당한 CPU 시간을 점유하고 있음이 밝혀졌습니다.

그 이유는 크게 두 가지입니다:

  1. Philox Offset Bookkeeping: PyTorch의 난수 생성기 상태를 관리하기 위한 CPU 측의 계산 작업이 수반됩니다.
  2. Device-side Asserts: 입력된 확률값이 유효한지(음수가 아닌지 등)를 확인하기 위한 장치 측 검증 로직이 CPU 디스패치 경로를 방해하거나 불필요한 동기화를 유발할 수 있습니다.

GPU 연산 자체는 매우 짧지만, 이를 실행하기 위한 CPU의 준비 과정이 길어지면서 전체 디코딩 루프의 지연 시간(Latency)이 늘어나는 결과를 초래했습니다.

코드 분석: Gumbel-max 트릭 도입

이번 PR의 핵심 변경 사항은 python/sglang/srt/layers/sampler.py 파일에 집중되어 있습니다. 기존의 단순한 multinomial 호출을 수학적으로 동일한 결과를 내면서도 CPU 오버헤드가 적은 방식으로 변경했습니다.

Before

기존 코드에서는 확률 분포(probs)가 주어지면 직접 torch.multinomial을 호출했습니다.

# python/sglang/srt/layers/sampler.py (기존 로직)
def sampling_from_probs_torch(probs, sampling_seed=None):
    if sampling_seed is None:
        sampled_index = torch.multinomial(probs, num_samples=1)
    else:
        # ... (생략)

After

변경된 코드에서는 환경 변수(SGLANG_OPT_USE_GUMBEL_SAMPLE)에 따라 Gumbel-max 트릭을 사용하도록 분기합니다.

# python/sglang/srt/layers/sampler.py (변경 후 로직)
def sampling_from_probs_torch(probs, sampling_seed=None):
    if sampling_seed is None:
        if envs.SGLANG_OPT_USE_GUMBEL_SAMPLE.get():
            # Gumbel-max trick: torch.multinomial과 분포적으로 동일하지만,
            # multinomial의 CPU 측 philox offset 관리 및 유효성 검사로 인한
            # 디코딩 경로 정체를 방지합니다.
            q = torch.empty_like(probs, dtype=torch.float32).exponential_(1.0)
            q.clamp_min_(torch.finfo(torch.float32).tiny)
            sampled_index = (probs.float() / q).argmax(dim=-1, keepdim=True)
        else:
            sampled_index = torch.multinomial(probs, num_samples=1)
    else:
        # ... (생략)

구현 상세 설명

  • torch.empty_like(...).exponential_(1.0): 확률 분포와 동일한 크기의 텐서를 생성하고 지수 분포(Exponential distribution)를 따르는 난수로 채웁니다.
  • probs.float() / q: 원래의 확률값을 지수 난수로 나눕니다. 이는 수학적으로 Gumbel 분포를 활용한 샘플링과 동일한 효과를 냅니다.
  • argmax(dim=-1): 나눈 값 중 가장 큰 값의 인덱스를 선택합니다. 이 결과값의 분포는 원래 probs에서 multinomial 샘플링을 한 것과 통계적으로 일치합니다.

왜 이게 좋은 최적화인가?

1. 성능 수치 (Measured Impact)

이 최적화의 효과는 특히 소형 모델에서 극적으로 나타났습니다. RTX 5090 환경에서의 벤치마크 결과는 다음과 같습니다.

  • Qwen3.5-2B (Batch Size 1): 317.8 tok/s → 557.4 tok/s (+75% 향상)
  • Qwen3.5-2B (Batch Size 32): 4999 tok/s → 6162 tok/s (+23% 향상)
  • Qwen3.5-0.8B (Batch Size 32): 8593 tok/s → 9994 tok/s (+16% 향상)

대형 모델(Gemma-4-E4B 등)에서는 샘플링 연산이 전체 연산에서 차지하는 비중이 작아 성능 변화가 거의 없었지만, 이는 반대로 성능 저하(Regression)의 위험 없이 소형 모델의 성능만 비약적으로 끌어올렸음을 의미합니다.

2. CPU-GPU 비동기성 유지

torch.multinomial은 내부적으로 복잡한 상태 관리와 검증 로직을 포함하고 있어 CPU가 GPU에 커널을 던지는 속도를 늦춥니다. 반면, Gumbel-max 트릭은 exponential_, div, argmax와 같은 단순하고 최적화된 기본 연산들의 조합으로 이루어집니다. 이러한 연산들은 PyTorch에서 매우 가볍게 디스패치되므로, CPU 병목을 효과적으로 우회할 수 있습니다.

3. 수학적 동등성 (Correctness)

Gumbel-max 트릭은 확률론적으로 multinomial 샘플링과 완전히 동일한 분포를 가집니다. 따라서 모델의 출력 품질(Accuracy)을 전혀 희생하지 않으면서 속도만 높일 수 있는 '공짜 점심(Free Lunch)'에 가까운 최적화입니다. 이미 vLLM과 같은 타 프레임워크에서도 기본 샘플러로 채택하고 있는 검증된 방식입니다.

결론

이번 PR은 고성능 LLM 엔진을 개발할 때 GPU 커널 자체의 속도만큼이나 CPU 디스패치 경로의 효율성이 중요하다는 것을 잘 보여줍니다. 특히 모델이 작아지고 하드웨어가 강력해질수록(예: RTX 5090), 소프트웨어 스택의 오버헤드가 전체 성능의 발목을 잡게 됩니다.

시니어 엔지니어로서 배울 수 있는 교훈은 명확합니다. 표준 라이브러리 함수(torch.multinomial)가 항상 최선의 성능을 보장하는 것은 아니며, 병목이 의심될 때는 해당 함수의 내부 동작을 이해하고 더 가벼운 수학적 대안을 고민해야 한다는 것입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글