[sglang] MiniMax-H3 모델을 24GB GPU에서 가속화하는 INT8 양자화 및 플러그형 어텐션 최적화

PR 링크: sgl-project/sglang#34581 상태: Merged | 변경: +616 / -45

들어가며

MiniMax-H3와 같은 대규모 DiT(Diffusion Transformer) 모델은 BF16 기준 60GB 이상의 가중치를 요구하여, 24GB VRAM을 가진 소비자용 GPU(예: RTX 4090)에서는 실행이 매우 어렵습니다. 기존에는 레이어별 오프로딩(layerwise offload)을 사용하더라도 연산 병목과 메모리 대역폭 문제로 인해 추론 속도가 매우 느렸습니다. 본 PR은 kitchen_int8 양자화와 플러그형 어텐션 백엔드를 도입하여, 24GB GPU 환경에서 추론 속도를 최대 2.48배 개선하는 솔루션을 제시합니다.

코드 분석

1. kitchen_int8 온라인 양자화

가장 핵심적인 변경은 DiT 블록 내 4개의 GEMM 연산을 INT8로 처리하는 것입니다. 기존의 오프라인 체크포인트 방식 대신, 로드 시점에 BF16 가중치를 즉석에서 양자화하여 메모리 오버헤드를 최소화했습니다.

# Before: BF16 연산 수행 (기본값)
# After: kitchen_int8을 통한 융합 커널 호출
quantized 209 linear layers (37.38 GiB BF16 -> 18.69 GiB INT8)

이 방식은 H3의 qkv 그룹화 구조를 보존하면서도, 레이어별로 CPU-GPU 간 데이터를 주고받는 방식으로 전체 60GB 가중치를 한 번에 로드하지 않아 OOM(Out of Memory)을 방지합니다.

2. 대규모 M 차원 GEMM을 위한 행 분할(Row-splitting)

특정 커널 설정에서 M=32700과 같은 큰 행렬 연산 시 CUTLASS Heuristic이 최적의 설정을 선택하지 못하는 문제가 있었습니다. 이를 해결하기 위해 행을 분할하여 처리하도록 변경했습니다.

# Row-splitting 적용 전후 비교
qkv_proj 26.51 ms -> 17.94 ms (1.48x 향상)

이 최적화는 SGLANG_KITCHEN_INT8_MAX_ROWS 환경 변수를 통해 제어 가능하며, 불필요한 Stream-K 오버헤드를 제거하여 연산 효율을 극대화합니다.

3. 플러그형 어텐션 백엔드

FlashAttention이 연산 병목의 41.8%를 차지함에 따라, 사용자가 상황에 맞춰 어텐션 알고리즘을 교체할 수 있도록 설계했습니다.

# 예시: SageAttention과 SolAttention 하이브리드 모드 사용
--attention-backend sol_attn \
--attention-backend-config dense_backend=sage_attn,dense_steps=10

왜 이게 좋은가

성능 측정 결과, BF16 + FlashAttention 대비 kitchen_int8 + FA 조합만으로도 1.34배의 속도 향상을 보였으며, Sage→Sol 하이브리드 어텐션 사용 시 2.48배까지 가속되었습니다.

핵심 교훈:

1. 데이터 레이아웃의 중요성: 사전 양자화된 체크포인트를 그대로 사용하면 qkv 재정렬 과정에서 정확도 오류가 발생할 수 있습니다. 온라인 양자화는 이러한 레이아웃 문제를 원천 차단합니다.
2. Heuristic의 한계: 대규모 행렬 연산 시 범용 커널 설정이 항상 최적은 아닙니다. 문제 크기에 따른 행 분할 전략이 성능에 큰 영향을 미칩니다.
3. 병목 지점 파악: 프로파일링을 통해 Attention과 GEMM이 전체 연산의 76%를 차지함을 확인하고, 이 두 지점을 집중 공략한 것이 주효했습니다.

리뷰어 피드백 반영

리뷰어들은 문서화의 중요성을 강조했습니다. 이에 따라 cookbook 및 환경 변수 설정 페이지를 업데이트하여 사용자가 최적화된 설정을 쉽게 적용할 수 있도록 가이드를 보강했습니다.

참고 자료

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

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