[sglang] SGLang 성능 최적화: PDL 도입과 안전한 CUDA 동기화로 DSV3.2/GLM-5 가속하기

PR 링크: sgl-project/sglang#23965 상태: Merged | 변경: +0 / -0

들어가며

대규모 언어 모델(LLM) 추론 엔진인 SGLang에서 최근 DeepSeek-V3.2(DSV3.2) 및 GLM-5 모델의 저지연(Low-latency) 성능을 대폭 개선한 업데이트가 있었습니다. 이번 PR의 핵심은 NVIDIA Hopper 아키텍처(H100 등)에서 지원하는 PDL(Programmatic Dependency Launch) 기능을 다양한 커널에 적용하고, 기존에 잠재적 위험 요소였던 인라인 어셈블리 기반의 의존성 제어 코드를 공식 CUDA API로 교체하여 안정성을 높인 것입니다.

이 변경을 통해 추론 속도가 69.04 TPS에서 71.39 TPS로 약 3.4% 향상되었으며, 특히 멀티 스트림 환경에서의 커널 실행 효율이 최적화되었습니다.

코드 분석: 핵심 변경 사항

1. PyTorch 2.0 torch.compile을 활용한 연산 융합

nsa_indexer.py에서는 단순하지만 반복적으로 발생하는 텐서 스케일링 연산을 torch.compile로 래핑하여 오버헤드를 줄였습니다.

Before:

# 매번 Python 인터프리터를 통해 실행됨
weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale

After:

@torch.compile(dynamic=True)
def _apply_q_scale_and_softmax_scale(
    self, weights: torch.Tensor, q_scale: torch.Tensor
):
    return weights.unsqueeze(-1) * q_scale * self.softmax_scale

# ... 중략 ...
weights = self._apply_q_scale_and_softmax_scale(weights, q_scale)

unsqueeze와 곱셈 연산이 하나의 커널로 융합(Fusion)되어 실행되므로, GPU 메모리 대역폭 소모와 커널 런칭 오버헤드가 감소합니다.

2. Triton 커널에 PDL(Programmatic Dependency Launch) 적용

set_mla_kv_buffer_kernel과 같은 KV 캐시 관리 커널에 PDL을 적용했습니다. PDL은 커널 간의 의존성을 하드웨어 수준에서 관리하여 CPU의 개입 없이 다음 커널을 트리거할 수 있게 합니다.

After (Triton):

def set_mla_kv_buffer_kernel(
    ...,
    USE_GDC: tl.constexpr = False,
):
    # ...
    if USE_GDC:
        tl.extra.cuda.gdc_wait() # 의존성 대기

    # 커널 로직 수행

    if USE_GDC:
        tl.extra.cuda.gdc_launch_dependents() # 후속 커널 트리거

이 방식은 특히 Hopper 아키텍처에서 launch_pdl=True 옵션과 함께 사용되어, 커널 간 공백 시간(Gap)을 최소화합니다.

3. 안전하지 않은 인라인 어셈블리 제거 및 표준 API 교체

기존 코드에서는 griddepcontrol을 직접 호출하기 위해 인라인 어셈블리를 사용했으나, 이는 메모리 배리어(Memory Barrier)를 보장하지 않아 데이터 오염의 위험이 있었습니다. 이를 NVIDIA가 권장하는 공식 CUDA API로 교체했습니다.

Before:

#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
  asm volatile("griddepcontrol.wait;");
#endif
// ...
  asm volatile("griddepcontrol.launch_dependents;");

After:

#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
  cudaGridDependencySynchronize();
#endif
// ...
  cudaTriggerProgrammaticLaunchCompletion();

cudaGridDependencySynchronize는 단순 대기를 넘어 필요한 메모리 가시성을 보장하므로, 병렬 실행 환경에서 훨씬 안전합니다.

왜 이게 좋은 최적화인가?

1. 하드웨어 가속 활용: NVIDIA Hopper 아키텍처의 전용 하드웨어 기능을 활용하여 소프트웨어 스케줄링 오버헤드를 하드웨어 수준으로 내렸습니다.
2. 안정성 확보: asm volatile은 컴파일러 최적화에 의해 순서가 바뀌거나 메모리 동기화가 누락될 위험이 큽니다. 이를 공식 API로 교체함으로써 디버깅하기 어려운 레이스 컨디션(Race Condition)을 원천 차단했습니다.
3. 실질적인 성능 향상: 벤치마크 결과 71.39 TPS를 달성하며 실질적인 처리량 향상을 입증했습니다. 이는 특히 토큰 생성 속도가 중요한 실시간 서비스에서 큰 차이를 만듭니다.

결론

이번 PR은 최신 GPU 아키텍처의 기능을 깊이 있게 이해하고 이를 안전하게 적용하는 시니어 엔지니어링의 정석을 보여줍니다. 단순히 "빠르게" 만드는 것을 넘어, "안전하고 견고하게" 최적화하는 것이 고성능 추론 엔진 개발의 핵심임을 다시 한번 확인시켜 줍니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__PROG__LAUNCH.html

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