dropbox/gemlite

해결하는 문제

GemLite는 LLM 추론에서 저비트 행렬 곱셈의 성능 저하 문제를 해결합니다. 기본 Torch AO 카ーネル과 비교해 프리필 및 디코딩 단계의 실행 속도를 크게 향상시키는 고도로 최적화된 카ーネル 세트를 제공하여, 양자화 모델의 지연 시간을 줄이고 처리량을 증가시킵니다.

작동 방식

GemLite는 Triton을 사용하여 다양한 워크로드에 맞춘 전용 행렬 곱셈 카르네이를 구현합니다:

• GEMM: 일반적인 행렬 곱셈을 위한 표준 Tensor Core 기반 카르네이.
• GEMM Split-K: 배치 크기 2~32인 배치 처리 LLM 디코딩을 최적화하기 위해 K 차원을 여러 작업으로 나누어 부분 합을 계산합니다.
• GEMV: 소규모 배치(배치 크기 M=1)를 위해 활성화 값을 1차원 청크로 나누어 설계되었습니다.
• GEMV RevSplit-K: 배치 크기 1 디코딩을 최적화하기 위해 프로그램당 작업량을 두 배로 늘려 스케일 및 제로 로딩 오버헤드를 줄이는 새로운 알고리즘입니다.

FP16, BF16, FP8, INT8, 그리고 8, 4, 2, 1비트 저비트 가중치를 포함한 다양한 정밀도를 지원하며, Blackwell과 같은 최신 하드웨어용 MXFP 및 NVFP 포맷도 지원합니다.

대상 사용자

GemLite는 특히 NVIDIA GPU에서 높은 성능을 유지하면서 모델 크기와 메모리 대역폭 요구를 줄이기 위해 양자화를 사용하는 LLM 추론 최적화에 종사하는 개발자 및 연구자를 대상으로 합니다.

주요 특징

• 매우 빠른 성능 향상: 기본 Torch AO 카르네ル 대비 프리필은 최대 78배, 디코딩은 최대 36배 빠릅니다.
• 광범위한 정밀도 지원: FP16, BF16, FP8, INT8, MXFP/NVFP4와 호환됩니다.
• 통합 가능성: vLLM(HQQ 또는 TorchAO 경유), SGLang와 호환되며 torch.compile과도 호환됩니다.
• 오토튜닝 캐싱: 최적의 카르네이 설정을 저장하고 로드할 수 있어 시작 시 반복적인 오토튜닝 시간을 제거할 수 있습니다.
• 유연한 패킹: HQQ 스타일의 가중치 패킹과 가중치 및 활성화의 채널별 스케일링을 지원합니다.