[sglang] SGLang의 AMD AITER AllReduce 최적화: 하드코딩된 제약 제거 및 성능 개선

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

들어가며

LLM 추론 성능을 극대화하기 위해 커널 융합(Kernel Fusion)은 필수적입니다. 특히 분산 환경에서 AllReduce와 RMSNorm을 결합하는 것은 통신 오버헤드를 줄이는 핵심 전략입니다. 하지만 기존 SGLang의 AITER(AMD Iterative Tensor) 융합 경로에는 모델별로 hidden_dim을 일일이 등록해야 하는 하드코딩된 허용 목록(allowlist)이 존재하여 유지보수 부담이 컸습니다. 본 PR은 이 제약을 제거하고, AITER C++ 백엔드의 유연한 디스패치 로직을 활용하도록 구조를 개선했습니다.

코드 분석

1. python/sglang/srt/layers/communicator.py

기존에는 total_bytes 비교 연산에서 <를 사용하여 경계 조건에 문제가 있었습니다. AITER 내부 로직은 <=를 기준으로 동작하므로, 이를 일치시켜 정확한 커널 활성화를 보장합니다.

-        and total_bytes < 8 * 1024 * 8192
+        and total_bytes <= 8 * 1024 * 8192

2. python/sglang/srt/distributed/parallel_state.py

가장 핵심적인 변경사항입니다. 기존에는 hidden_dim이 {512, 1024, 2048, 4096} 중 하나여야만 1-stage 커널을 사용할 수 있었습니다. 이제는 AITER C++ 디스패치가 지원 여부를 결정하도록 위임하고, Python 측에서는 128KB라는 메모리 임계값만 체크합니다.

-            use_1stage_ar = total_bytes <= 128 * 1024 and hidden_dim in {
-                512,
-                1024,
-                2048,
-                2880,
-                4096,
-            }
+            use_1stage_ar = total_bytes <= 128 * 1024

왜 이게 좋은가

1. 유지보수성 향상: 새로운 모델이 추가될 때마다 Python 코드를 수정할 필요가 없습니다. AITER C++ 백엔드에서 지원하는 모든 hidden_dim을 자동으로 수용합니다.
2. 성능 개선: 벤치마크 결과, hidden_dim=2880인 모델에서 디코드 성능이 최대 2.10배 향상되었습니다. 전체 토큰 처리량(throughput) 또한 약 3.7% 증가했습니다.
3. 정확성 보장: ROCm/aiter#2586에서 수정된 수치 정확성 문제를 통합하여, 융합 커널 사용 시 GSM8K 벤치마크 점수가 베이스라인 대비 +2.3pp 향상되는 등 수치적 안정성을 확인했습니다.

교훈

• 백엔드 위임: 비즈니스 로직(Python)에서 하드웨어 특화 제약(C++)을 관리하려 하지 마십시오. 백엔드 디스패처가 충분히 똑똑하다면, 상위 계층은 최소한의 가드레일(예: 메모리 임계값)만 설정하는 것이 훨씬 유지보수에 유리합니다.
• 경계 조건 확인: 분산 시스템에서 커널을 융합할 때는 통신 라이브러리와 Python 레이어 간의 경계 조건(> vs >=)을 엄격히 일치시켜야 합니다.

References

• SGLang GitHub Repository
• ROCm/aiter

참고 자료

• https://github.com/sgl-project/sglang
• https://github.com/ROCm/aiter

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