[sglang] AMD MI300에서의 GPT-OSS 120B 추론 성능 최적화: 새로운 KV 캐시 레이아웃과 커널 통합

PR 링크: sgl-project/sglang#27063 상태: Merged | 변경: +1874 / -54

들어가며

최근 LLM(거대 언어 모델)의 발전 속도는 눈부시지만, 모델의 크기가 커짐에 따라 추론 성능, 특히 GPU에서의 효율적인 연산은 여전히 중요한 과제입니다. 특히 AMD GPU 환경에서 대규모 모델의 성능을 극대화하는 것은 생태계 확장을 위해 필수적입니다. SGLang 레포지토리의 이 PR은 AMD ROCm 환경에서 gpt-oss-120b 모델의 추론 성능을 크게 향상시키는 것을 목표로 합니다. 핵심은 Atom의 pa_decode_gluon 커널을 SGLang의 AITER 백엔드에 통합하고, 이를 위해 새로운 5D 벡터화된 KV 캐시 레이아웃을 도입하는 것입니다.

이 글에서는 해당 PR의 코드 변경 사항을 상세히 분석하고, 이러한 변경이 왜 성능 향상으로 이어지는지, 그리고 어떤 기술적 교훈을 얻을 수 있는지 살펴보겠습니다.

코드 분석

이번 PR의 핵심은 AMD GPU(MI300/MI355)에서 gpt-oss-120b 모델의 디코딩 성능을 높이기 위해 새로운 KV 캐시 레이아웃과 최적화된 커널을 도입하는 것입니다. 변경 사항은 여러 파일에 걸쳐 있으며, 크게 환경 변수 설정, 메모리 풀 관리, 라이터/리더 커널, AITER 백엔드 로직, 그리고 기존 커널과의 통합으로 나눌 수 있습니다.

1. 환경 변수 및 설정 (environ.py, server_args.py)

새로운 KV 캐시 레이아웃을 활성화하기 위해 SGLANG_AITER_KV_CACHE_LAYOUT 환경 변수가 도입되었습니다. 이 변수는 `

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html

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