Accelerating GPU Inference of Large Language Models with Moderately Unstructured Sparse Weight Matrices

기술 요약: 중등도 비구조적 희소 가중치 행렬을 이용한 대규모 언어 모델의 GPU 추론 가속화

문제 정의
대규모 언어 모델(LLM)의 배포는 추론 비용이라는 결정적인 병목 현상에 직면해 있습니다. 모델 프루닝 기술(예: Wanda, RIA)은 가중치 행렬에 희소성(sparsity)을 도입하여 계산 부하를 줄일 수 있지만, 이는 모델 품질 요구 사항에 의해 제약을 받습니다. 소규모 모델과 달리, LLM은 높은 희소도 수준에 매우 민감합니다. 성능을 유지하기 위해 일반적으로 프루닝은 중등도 비구조적 희소성(약 50%) 수준으로 제한됩니다.

이 특정 희소도 수준에서, 기존의 희소 행렬 곱셈(SpMM) GPU 커널은 밀집 행렬 곱셈(dense matrix multiplication) 라이브러리인 cuBLAS보다 우수한 성능을 내지 못합니다. 이러한 성능 격차는 세 가지 근본적인 과제에서 비롯됩니다:

1. 텐서 코어 불일치: 현대 GPU는 고효율 희소 텐서 코어를 제공하지만, 이는 엄격하게 2:4 구조적 희소성만을 지원합니다. 따라서 불규칙한 비구조적 프루닝 패턴을 직접 가속화할 수 없습니다.
2. 메타데이터 오버헤드: 전통적인 저장 형식(예: CSR)은 모든 비제로(non-zero) 요소에 대해 전체 위치 메타데이터를 저장해야 합니다. 50% 희소도에서는 이 메타데이터가 값 자체만큼의 공간을 차지할 수 있어, 메모리 대역폭 이득을 상쇄해 버립니다.
3. 온칩 디코딩 병목 현상: 메타데이터를 줄이는 압축 기술(예: 비트맵 인코딩)은 워크로드를 온칩 디코딩으로 전이시킵니다. 고대역폭 메모리(HBM)를 갖춘 현대 GPU에서, 글로벌 메모리 액세스는 더 느린 CUDA 코어에서 수행되는 디코딩 작업과 완전히 겹쳐질(overlap) 수 없으므로 성능 정체(stall)를 유발합니다.

방법론
저자들은 중등도 비구조적 희소성을 위해 설계된 새로운 3계층 행렬 저장 형식과 최적화된 SpMM 커널을 포함하는 통합적인 솔루션을 제안합니다.

1. 3계층 저장 형식
제안된 형식은 하드웨어 호환성과 저장 효율성의 균형을 맞추기 위해 희소 행렬을 세 가지 별도의 계층으로 분해합니다:

• Sparse-TC 계층: 이 계층은 불규칙한 분포에서 2:4 구조적 희소성 패턴을 추출합니다. 4개의 연속된 가중치 중 2개의 비제로 값을 유지합니다. 패턴이 고정되어 있으므로 메타데이터로 값당 2비트만 필요하며, 디코딩 오버헤드 없이 텐서 코어와 직접 호환됩니다.
• Slot-Filling 계층: 2:4 패턴에 들어맞지 않는 나머지 비제로 요소들은 Sparse-TC 계층에 의해 남겨진 빈 슬롯으로 재배치됩니다. 이 약 17.7%의 요소들에 대한 메타데이터를 최소화하기 위해, 저자들은 병렬 차분 거리(Parallel Differential Distance, PDD) 인코딩을 도입합니다. PDD는 절대 인덱스를 저장하는 대신 재배치된 요소들 사이의 상대적 거리를 인코딩합니다. 대부분의 거리는 4비트로 인코딩될 만큼 작으며, 거리가 큰 경우 인코딩 한도를 초과하지 않도록 더미 포인트(dummy points)를 삽서하여 범위를 끊어줍니다.
• Residual 계층: 첫 두 계층에서 효율적으로 처리할 수 없는 나머지 비제로 요소들(통상 전체의 1% 미만)은 전통적인 CSR 형식으로 저장됩니다. 데이터의 양이 매우 적기 때문에 이곳의 메타데이터 오버헤드는 무시할 수 있는 수준입니다.

2. 공동 최적화된 SpMM 커널
이 형식을 기반으로, 저자들은 파이프라인 실행 모델 내에서 희소 텐서 코어와 CUDA 코어를 결합하여 사용하는 GPU 커널을 설계했습니다.

• 병렬 디코딩: CUDA 코어는 Slot-Filling 계층의 PDD 메타데이터를 병렬로 디코딩하는 작업을 담당합니다. 이 계층은 전체 데이터의 작은 부분만을 차지하므로 디코딩 오버헤드가 효과적으로 관리됩니다.
• 구조적 가속: 희소 텐서 코어는 wgmma 명령어(Hopper 아키텍처)를 사용하여 2:4 구조를 활용해 Sparse-TC 계층을 처리함으로써 높은 처리량을 달ка합니다.
• 잔여물 처리: CUDA 코어는 또한 희소한 Residual 계층을 처리하며, 결과를 별도의 온칩 버퍼에 기록한 뒤 파이프라인 끝에서 텐서 코어 출력과 합산(reduction)합니다.
• 파이프라인 중첩: 커널은 다음 세 가지 핵심 작업을 중첩시키는 미세한 파이프라인을 구현합니다:
  • 밀집 입력 행렬(L2에 캐싱됨) 및 희소 가중치 타일(Global Memory로부터) 로드.
  • CUDA 코어에서의 PDD 메타데이터 디코딩.
  • 희소 텐서 코어에서의 행렬 곱셈 실행.
    이 "희소하게 로드하고, 희소하게 계산하는(load-as-sparse, compute-as-sparse)" 전략은 온칩 계산이 글로벌 메모리 액세스와 겹치도록 하여 HBM 활용도를 극대화합니다.

주요 기여

• 새로운 저장 형식: 중등도 비구조적 희소성에 맞춤화된 3계층 형식을 통해 희소 텐서 코어 호환성을 확보하고, PDD 인코딩을 통해 메타데이터 오버헤드를 획기적으로 줄이며, 온칩 디코딩 비용을 최소화했습니다.
• 공동 최적화된 커널: 파이프라이닝을 사용하여 메모리 액세스와 계산을 중첩함으로써 희소 텐서 코어와 CUDA 코어를 효율적으로 조율하는 SpMM 커널을 설계했습니다.
• 밀집 연산을 능가한 최초의 사례: 제안된 방식은 HBM을 갖춘 현대 GPU에서 중등도 비구조적 희소 SpMM이 밀집 행렬 곱셈보다 빠를 수 있음을 입증한 최초의 사례입니다.

결과
NVIDIA H100 SXM5 GPU에서 OPT-30B 및 OPT-66B 모델을 사용하여 50%, 60%, 70% 희소도 수준에 대한 평가를 수행했습니다.

• 커널 수준 성능: 50% 희소도에서, 제안된 방법은 SpInfer(EuroSys'25 Best Paper) 대비 최대 1.64배의 속도 향상을 달성했으며, SparTA 및 FlashLLM과 같은 다른 베이스라인들을 크게 앞질렀습니다. 또한 밀집 cuBLAS 성능을 능가했습니다.
• 엔드 투 엔드 추론: 이 방법은 FlashLLM(VLDB'24) 대비 최대 1.41배의 엔드 투 엔드 속도 향상을 제공하며, 배치 크기 16에서 OPT-30B 기준 FlashLLM보다 최대 41.4% 높은 처리량을 보여주었습니다.
• 메모리 효율성: 50% 희소도에서 밀집 베이스라인 대비 피크 글로벌 메모리 사용량을 약 21.4% 감소시켰으나, 메타데이터에 대한 비트맵 인코딩을 사용하지 않기 때문에 SpInfer보다는 약간 높은 메모리 사용량을 기록했습니다.
• 확장성: 이 방법은 50%에서 70% 희소도 구간 전반에서 이점을 유지합니다.

의의 및 주장
본 논문은 기존 희소 커널이 중등도 비구조적 희소성을 효과적으로 활용하지 못하는 LLM 추론 최적화의 핵심적인 간극을 해결한다고 주장합니다. 불규칙한 프루닝 패턴과 구조적 하드웨어 가속 사이의 간극을 메우는 저장 형식을 도입함으로써, 저자들은 모델 품질을 희생하지 않고도 프루닝이 실제 추론 속도 향상을 가져올 수 있음을 입증했습니다.

이 연구는 높은 희소도(>70%)는 모델 품질을 저하시는 반면, 중등도 희소도(~50%)가 LLM의 "스위트 스팟(sweet spot)"임을 강조합니다. 제안된 솔루션은 기존 희소 커널의 비효율성으로 인해 접근 불가능했던 이 스위트 스팟의 성능 잠재력을 끌어올립니다. 저자들은 이 최적화가 특히 LLM 추론의 디코드 단계(배치 크기가 작고 메모리 제한적인 단계)를 겨냥하고 있어, 프리필(prefill)과 디코드 단계가 분리되는 경우가 많은 프로덕션 환경의 확장 가능한 배포에 매우 적합하다고 언급했습니다. 연구의 재현성과 확산을 위해 소스 코드는 공개되었습니다.