[sglang] SGLang의 AMD GPU 성능 최적화: Aiter CK 커널을 활용한 LayerNorm 오버헤드 제거

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

들어가며

대규모 언어 모델(LLM) 추론 엔진인 SGLang은 다양한 하드웨어 백엔드를 지원합니다. 특히 AMD GPU(ROCm 환경)에서 LayerNorm 연산은 기존에 PyTorch 네이티브 구현에 의존하고 있었습니다. 하지만 이 방식은 입력과 출력 시점에 불필요한 데이터 타입 캐스팅(dtype cast)이 발생하여, 결과적으로 cast -> layernorm -> cast라는 3단계 커널 호출을 유발했습니다. 이는 특히 NSA(Non-local Sparse Attention) 인덱서와 같은 고성능이 요구되는 연산에서 병목 현상을 일으켰습니다. 본 글에서는 aiter 라이브러리의 CK(Composable Kernel)를 활용하여 이 문제를 어떻게 해결했는지 살펴봅니다.

코드 분석

1. nsa_indexer.py: 데이터 타입 조정

먼저 NSA 인덱서에서 k_norm의 데이터 타입을 float32에서 bfloat16으로 변경하여, 최적화된 커널 경로를 탈 수 있도록 준비했습니다.

# Before
self.k_norm = LayerNorm(self.head_dim, dtype=torch.float32)

# After
self.k_norm = LayerNorm(
    self.head_dim, dtype=torch.bfloat16 if _use_aiter else torch.float32
)

2. layernorm.py: Aiter CK 커널 도입

forward_hip 메서드에서 aiter 라이브러리의 layernorm2d_fwd를 직접 호출하도록 로직을 추가했습니다. 이를 통해 불필요한 캐스팅 없이 연산을 수행합니다.

# After
if (
    _has_aiter_layer_norm
    and x.dtype in (torch.bfloat16, torch.float16)
    and x.dtype == self.dtype
):
    orig_shape = x.shape
    x = x.reshape(-1, self.hidden_size)
    return layer_norm(x, self.weight, self.bias, self.variance_epsilon).view(
        orig_shape
    )
else:
    return self.forward_native(x)

왜 이게 좋은가

이번 최적화의 핵심은 커널 퓨전(Kernel Fusion)과 불필요한 오버헤드 제거입니다. 기존 3개의 커널 호출이 1개로 줄어들면서, 레이어당 처리 시간이 약 12us에서 4us로 3배가량 단축되었습니다.

일반적 교훈

1. 커널 호출 횟수 최소화: GPU 연산에서 커널 실행(Launch) 자체는 오버헤드가 큽니다. 여러 단계의 연산을 하나의 커널로 통합하는 것은 성능 최적화의 기본입니다.
2. 타입 캐스팅 방지: 연산 중간에 발생하는 불필요한 dtype 변환은 메모리 대역폭을 낭비하고 추가 커널을 유발합니다. 가능한 한 연산 파이프라인 전체에서 일관된 데이터 타입을 유지하는 것이 중요합니다.
3. 하드웨어 특화 라이브러리 활용: PyTorch의 범용 구현체보다 특정 하드웨어(AMD ROCm 등)에 최적화된 커널 라이브러리(aiter, CK 등)를 사용하는 것이 성능 극대화에 유리합니다.

이러한 최적화는 GLM-5-FP8 모델 벤치마크에서 처리량(Throughput)을 최대 1.4% 향상시키고, TPOT(Time Per Output Token)를 2.8%까지 개선하는 성과를 거두었습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html
• https://github.com/sgl-project/sglang

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