[sglang] SGLang 성능 최적화: FP8 모델을 위한 Inductor 컴파일러 경로 개선

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

들어가며

LLM 서빙 엔진인 SGLang에서 성능을 쥐어짜기 위해 가장 중요한 요소 중 하나는 GPU 커널 오버헤드를 줄이는 것입니다. 특히 FP8(8-bit Floating Point) 모델을 사용할 때, 수많은 작은 연산들이 개별 커널로 실행되면 오버헤드가 누적되어 전체 성능을 저하시킵니다.

최근 SGLang 프로젝트에 반영된 [perf: optimize PCG inductor path for FP8 models] PR은 Piecewise CUDA Graph (PCG) 환경에서 PyTorch Inductor 컴파일러를 사용할 때 발생하는 불필요한 커널 런치를 제거하는 데 집중했습니다. 이 최적화를 통해 FP8 임베딩 작업에서 GPU 시간은 20% 감소했고, 커널 호출 횟수는 24%나 줄어드는 놀라운 성과를 거두었습니다.

시니어 엔지니어의 관점에서 이 PR이 왜 훌륭한 최적화인지 코드와 함께 분석해 보겠습니다.

---

코드 분석: 무엇이 왜 바뀌었는가?

1. Stride를 보존하는 View 연산으로의 전환

models/utils.py의 apply_qk_norm 함수는 Query(Q)와 Key(K)에 정규화를 적용합니다. 기존 코드는 2D reshape을 사용했는데, 이는 Inductor 컴파일러가 메모리 레이아웃을 최적화하는 데 방해가 되었습니다.

Before:

# 2D reshape은 stride 정보를 깨뜨려 별도의 복사 커널을 유발함
q_by_head = q.reshape(-1, head_dim)

After:

# 3D view를 사용하여 기존 텐서의 stride를 유지
q_by_head = q.view(*q.shape[:-1], -1, head_dim)

왜 좋은가? 기존의 reshape(-1, head_dim)은 텐서의 차원을 강제로 2D로 뭉개버립니다. 이 경우 Inductor는 내부적으로 as_strided, clone, split_with_sizes와 같은 별도의 Triton 커널들을 생성하게 됩니다. 반면, view를 통해 차원 구조를 명시적으로 유지하면 Inductor가 이를 Stride-preserving 연산으로 인식하여, 뒤따르는 RMSNorm 커널과 하나로 합쳐버릴(Fuse) 수 있습니다. 실제로 이 변경만으로 요청당 약 81개의 불필요한 커널이 제거되었습니다.

2. Opaque 커널 대신 Pure PyTorch Ops 사용

fp8_utils.py에서는 FP8 양자화(Quantization) 방식을 변경했습니다. 기존에는 성능을 위해 직접 작성된 커스텀 커널(sgl_kernel.per_tensor_quant_fp8)을 호출했습니다.

Before:

qinput, x_scale = scaled_fp8_quant(
    input_2d,
    input_scale,
    num_token_padding=num_token_padding,
    use_per_token_if_dynamic=use_per_token_if_dynamic,
)

After:

if (
    input_scale is not None
    and input_scale.numel() == 1
    and get_global_server_args().piecewise_cuda_graph_compiler == "inductor"
):
    # Inductor가 주변 연산과 퓨전할 수 있도록 표준 연산으로 작성
    qinput = (
        (input_2d * input_scale.reciprocal())
        .clamp(min=fp8_min, max=fp8_max)
        .to(fp8_dtype)
    )
    x_scale = input_scale
else:
    # Eager 모드나 다른 컴파일러에서는 기존 커스텀 커널 사용
    qinput, x_scale = scaled_fp8_quant(...)

왜 좋은가? 역설적이게도 "가장 빠른 커스텀 커널"이 항상 최선은 아닙니다. Inductor 같은 컴파일러 입장에서는 C++/CUDA로 작성된 불투명한(Opaque) 커널 내부를 들여다볼 수 없습니다. 따라서 커스텀 커널 앞뒤로 반드시 메모리 동기화가 발생합니다.

하지만 위 코드처럼 reciprocal, clamp, cast 같은 표준 PyTorch 연산으로 작성하면, Inductor는 이 연산들을 앞의 RMSNorm이나 뒤의 GEMM 연산과 하나의 Triton 커널로 묶어버립니다. 결과적으로 per_tensor_quant_fp8_kernel 호출이 완전히 사라지고 전체 실행 시간은 단축됩니다.

---

왜 이게 좋은 최적화인가?

1. 명확한 성능 향상 수치

H200 GPU에서 Qwen3-0.6B 모델로 벤치마크한 결과는 압도적입니다.

• Throughput: 27.68 items/s → 34.38 items/s (+24%)
• Latency: 요청당 GPU 시간 9.8ms → 7.8ms (-20%)
• Efficiency: 커널 런치 횟수 581회 → 441회 (-24%)

2. 컴파일러 친화적 설계 (Compiler-Aware Design)

이 PR은 무조건적인 최적화가 아니라, piecewise_cuda_graph_compiler == "inductor"인 경우에만 전략을 수정합니다. Eager 모드에서는 여전히 수작업으로 최적화된 커스텀 커널을 사용하고, 컴파일러 모드에서는 컴파일러가 가장 잘 요리할 수 있는 "재료(Pure Ops)"를 제공하는 영리한 접근 방식을 취했습니다.

3. 리뷰 피드백의 반영

초기 논의 과정에서 전역 환경 변수를 수정하려는 시도가 있었으나, 리뷰어(ch-wan)의 제안에 따라 get_global_server_args()를 사용하여 부수 효과(Side effect)를 최소화하는 방향으로 수정되었습니다. 이는 대규모 시스템에서 안정성을 유지하며 최적화를 적용하는 모범 사례를 보여줍니다.

결론

현대적인 딥러닝 가속화의 핵심은 단순히 "빠른 커널"을 만드는 것이 아니라, "커널 간의 경계를 허무는 것"에 있습니다. 이번 SGLang의 업데이트는 컴파일러의 특성을 이해하고 그에 맞춰 코드를 작성하는 것이 얼마나 큰 성능 차이를 만드는지 잘 보여주는 사례입니다.

Inductor를 사용 중이라면, 여러분의 커스텀 커널이 오히려 컴파일러의 손발을 묶고 있지는 않은지 확인해 보시기 바랍니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.Tensor.view.html
• https://pytorch.org/docs/stable/torch.compiler.html
• https://pytorch.org/docs/stable/generated/torch.clamp.html

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