[sglang] SGLang에서 GLM-5 모델 성능 최적화: Aiter 백엔드 활용 및 텐서 패딩 전략

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

들어가며

최신 대규모 언어 모델(LLM)을 특정 하드웨어 아키텍처, 특히 AMD GPU 환경에서 구동할 때 하드웨어의 연산 효율을 극대화하는 것은 매우 중요합니다. 이번 SGLang의 PR은 GLM-5 모델을 AMD의 Aiter 백엔드에서 실행할 때 발생하는 성능 병목을 해결하기 위해 텐서 패딩(Tensor Padding) 전략과 커널 파라미터 최적화를 도입했습니다. 특히 모델의 헤드(head) 수가 하드웨어 커널의 요구사항보다 적을 때 발생하는 비효율을 어떻게 해결했는지 살펴봅니다.

코드 분석

1. 텐서 패딩을 통한 Aiter Sparse Attention 최적화

Aiter의 mla_decode_fwd 커널은 효율적인 연산을 위해 헤드 수가 16의 배수여야 한다는 제약이 있습니다. GLM-5와 같이 TP(Tensor Parallelism) 8 환경에서 헤드 수가 8인 경우, 이 제약 조건을 만족하지 못해 성능 저하가 발생합니다.

이를 해결하기 위해 nsa_backend.py에서 need_pad_heads 로직을 도입했습니다.

# Before
# (패딩 로직 없음)

# After
self.need_pad_heads = self.num_q_heads < 16
self.head_repeat_factor = (16 // self.num_q_heads if self.num_q_heads < 16 else 1)

# Forward 시 패딩 적용
if self.need_pad_heads:
    q_kernel = q.view(-1, layer.tp_q_head_num, layer.head_dim).repeat_interleave(self.head_repeat_factor, dim=1)
    o_kernel = q.new_empty((q.shape[0], layer.tp_q_head_num * self.head_repeat_factor, layer.v_head_dim))

repeat_interleave를 사용하여 헤드 수를 16으로 맞춤으로써 커널이 최적의 성능을 낼 수 있는 환경을 강제로 조성했습니다.

2. 하드코딩된 파라미터 제거

nsa_indexer.py에서는 불필요하게 하드코딩되었던 파라미터들을 제거하여 커널이 내부적으로 최적의 설정을 선택하도록 개선했습니다.

-                ChunkK=128,
-                TotalCuCount=256,
-                WavePerEU=5,

이러한 변경은 커널의 유연성을 높이고, 특정 하드웨어 세대나 상황에 맞춰 커널이 스스로 최적화된 구성을 찾게 함으로써 유지보수성을 향상시킵니다.

왜 이게 좋은가

이번 최적화의 핵심은 **'하드웨어 가속기의 요구사항에 맞춘 데이터 정렬(Data Alignment)'**입니다.

1. 성능 향상: 헤드 수를 16의 배수로 패딩함으로써, GPU의 연산 유닛이 데이터를 처리할 때 메모리 정렬 및 벡터 연산 효율이 극대화됩니다. 이는 특히 GLM-5와 같이 헤드 수가 적은 모델에서 큰 성능 차이를 만듭니다.
2. 유연성: 하드코딩된 파라미터를 제거함으로써, 향후 Aiter 커널이 업데이트되어도 SGLang 측에서 코드를 수정할 필요 없이 자동으로 최적화된 파라미터를 사용하게 됩니다.

일반적으로 LLM 추론 최적화 시, 커널의 요구사항(Alignment)을 맞추기 위해 텐서를 패딩하는 것은 메모리 사용량이 약간 늘어나는 단점이 있지만, 연산 속도 측면에서는 압도적인 이득을 가져오는 경우가 많습니다.

결론

이번 PR은 특정 모델(GLM-5)의 특성을 고려하여 하드웨어 백엔드(Aiter)의 제약 조건을 우회하는 영리한 전략을 보여줍니다. 성능 프로파일링 결과가 이를 뒷받침하며, 하드웨어 가속기 최적화의 핵심은 결국 '하드웨어가 가장 좋아하는 데이터 형태를 제공하는 것'임을 다시 한번 확인시켜 줍니다.

참고 자료

• https://github.com/ROCm/aiter
• https://docs.sglang.ai/

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