RSH-SpMM: A Row-Structured Hybrid Kernel for Sparse Matrix-Matrix Multiplication on GPUs

이 논문은 GPU 기반 희소 행렬 - 행렬 곱셈 (SpMM) 의 성능 저하를 해결하기 위해 적응형 행 분할, RS-Tile 표현, 그리고 부하 균형 하이브리드 커널을 도입한 RSH-SpMM 을 제안하며, 다양한 희소 작업에서 기존 최첨단 방법보다 1.27 배에서 6.13 배까지의 가속화를 달성함을 보여줍니다.

핵심

📚 배경: 왜 이 기술이 필요한가요?

1. 문제 상황: 꽉 찬 도서관 vs 빈 도서관
컴퓨터가 데이터를 계산할 때, 보통은 모든 책장 (데이터) 이 꽉 차 있는 경우를 가정합니다. 이때는 일렬로 책을 정리하고 빠르게 꺼내면 됩니다 (이걸 '텐서 코어'라고 하는 강력한 계산기가 잘합니다).

하지만 현실의 데이터 (예: SNS 친구 관계, 추천 시스템 등) 는 **대부분의 책장이 비어있고, 책이 몇 권만 놓여있는 '빈 도서관'**과 같습니다.

• 기존 방법의 한계:
  • CUDA 코어 (일반 작업자): 빈 책장 사이를 뛰어다니며 책을 찾아야 해서 너무 느립니다.
  • 텐서 코어 (초고속 로봇): 빈 책장을 채우기 위해 가상의 책 (패딩) 을 넣어서라도 무리하게 일렬로 정리하려다 보니, 실제 작업보다 가짜 책 정리하는 데 시간이 더 걸려 비효율적입니다.

2. RSH-SpMM 의 등장: 똑똑한 도서관 사서
이 논문은 "빈 도서관의 특성을 파악해서, 빈 부분은 일반 사서가, 꽉 찬 부분은 로봇이 처리하게 하는" 새로운 시스템을 제안합니다.

---

🚀 RSH-SpMM 의 핵심 아이디어 (3 단계)

이 시스템은 도서관을 효율적으로 운영하기 위해 세 가지 전략을 사용합니다.

1. 🧩 똑똑한 분류 (RS-Tile 형식)

• 비유: 도서관 사서가 들어오는 책들을 보고 "이 책들은 옆에 모여서 정리하면 좋겠다 (로봇이 처리)", "이 책들은 너무 멀거나 적어서 로봇이 오면 시간만 낭비되니 내가 따로 처리하겠다 (일반 사서)"라고 분류합니다.
• 기술적 의미: 행렬 (데이터) 을 분석해서, 빈틈이 적은 부분은 로봇 (텐서 코어) 이 처리할 수 있도록 잘게 나누고, 너무 희박하거나 불규칙한 부분은 일반 작업자 (CUDA 코어) 가 처리하도록 따로 떼어냅니다.

2. 🔄 책장 재배치 (국소성 인식 재순서)

• 비유: 원래 책장 순서가 엉망이라서, 같은 주제 책들이 도서관 구석구석에 흩어져 있었습니다. 사서는 "주제가 비슷한 책들을 서로 옆으로 모아놓는" 작업을 먼저 합니다.
• 기술적 의미: 데이터의 구조를 분석하여, 비슷한 패턴을 가진 행 (Row) 들을 서로 붙여놓습니다. 이렇게 하면 로봇이 처리할 때 빈 책장을 채우는 일이 줄어들고, 한 번에 많은 일을 처리할 수 있게 됩니다.

3. ⚖️ 부하 균형 잡기 (적응형 로드 밸런싱)

• 비유: 어떤 책장은 책이 너무 많아서 로봇이 한 번에 처리하기엔 무겁고, 어떤 책장은 너무 가벼워서 로봇이 오면 오히려 비효율적입니다. 사서는 **"무거운 책장은 로봇이 여러 번 나눠서 처리하고, 가벼운 책은 일반 사서가 빠르게 처리"**하도록 동적으로 배분합니다.
• 기술적 의미: 데이터의 밀도가 급격히 변하는 상황에서도, 로봇과 일반 작업자의 작업을 적절히 나누어 어느 한쪽이 쉬지 않고 계속 일할 수 있도록 균형을 맞춥니다.

---

🏆 결과: 얼마나 빨라졌나요?

이 새로운 시스템을 적용한 결과, 기존 방식들보다 1.27 배에서 최대 6.13 배까지 속도가 빨라졌습니다.

• 기존 방식: 빈 도서관을 로봇이 처리하려다 지치거나, 일반 사서가 처리하느라 너무 느림.
• RSH-SpMM: 빈 부분은 사서가, 꽉 찬 부분은 로봇이 처리하되, 서로의 능력을 최대한 활용하도록 유연하게 배분함.

💡 한 줄 요약

**"데이터가 꽉 찬 곳은 강력한 로봇 (텐서 코어) 이, 빈 곳은 유연한 사서 (CUDA 코어) 가 처리하도록, 두 팀이 완벽하게 협력하게 만든 똑똑한 도서관 관리 시스템"**입니다.

이 기술은 인공지능 (AI) 이 친구 추천을 하거나, 복잡한 과학 시뮬레이션을 할 때 훨씬 더 빠르고 안정적으로 작동하게 해줍니다.

기술 요약

1. 문제 정의 (Problem)

희소 행렬 - 행렬 곱셈 (SpMM) 은 그래프 신경망 (GNN), 과학적 시뮬레이션, 희소 딥러닝 등 다양한 분야에서 핵심 연산입니다. 그러나 현대 GPU 의 성능을 극대화하는 데는 다음과 같은 근본적인 어려움이 존재합니다.

• 극심한 불규칙성 (Extreme Irregularity): 실제 세계의 희소 행렬은 행 길이 분포가 치우쳐 있고 (heavy-tailed), 지역적 밀도가 급격히 변하며, 비영 (nonzero) 패턴이 파편화되어 있습니다.
• Tensor Core 활용의 한계: Tensor Core 는 밀집된 (dense) 데이터 타일에서 높은 처리량을 제공하지만, SpMM 의 불규칙한 구조는 Tensor Core 가 요구하는 정렬된 타일 형식과 맞지 않습니다. 이로 인해 패딩 (padding) 이 발생하고 유효한 연산 주기가 줄어들어 Tensor Core 활용도가 낮아집니다.
• 기존 방법의 부족:
  • CUDA Core 만 사용하는 방법: 불규칙성을 잘 처리하지만, Tensor Core 의 높은 연산 성능을 활용하지 못해 처리량 (throughput) 이 낮습니다.
  • Tensor Core 전용 방법: 고정된 윈도우 크기와 경직된 타일 규칙을 사용하여 밀도가 낮은 행을 처리할 때 성능이 급격히 저하됩니다.
  • 기존 하이브리드 방법: 행 단위의 분할이 거칠어 (coarse-grained) 미세한 구조적 불일치를 해결하지 못하며, 여전히 Tensor Core 활용도가 낮거나 CUDA Core 에 불균형한 부하를 줍니다.

2. 제안 방법론 (Methodology: RSH-SpMM)

저자들은 RSH-SpMM을 제안했습니다. 이는 행 (Row) 구조를 기반으로 한 미세 입자 (fine-grained) 하이브리드 프레임워크로, GPU 의 이질적인 연산 유닛 (Tensor Core 와 CUDA Core) 을 SpMM 의 불규칙한 구조에 맞춰 최적화합니다.

핵심 구성 요소:

1. RS-Tile (Row-Structured Tile) 표현:

   • 희소 행렬을 두 부분으로 분해합니다.
   • TC Block Part: Tensor Core 에 적합한 밀집된 타일을 형성할 수 있는 구조적으로 일관된 행들을 모아 8x8 크기의 타일로 압축합니다. 비트맵 (Bitmap) 을 사용하여 인덱싱 오버헤드를 줄입니다.
   • CUDA Residual Part: 비영 개수가 매우 적거나 이웃 행과 구조가 너무 다른 행들은 Tensor Core 경로에서 제외하여, 오버헤드가 적은 CUDA Core 전용 경로로 보냅니다.
   • 이 방식은 메타데이터 오버헤드를 줄이고 타일 밀도를 높입니다.

2. 미세 입자 하이브리드 실행 전략 (Fine-grained Hybrid Execution):

   • 적응형 행 분할 (Adaptive Row Partitioning): 행의 비영 개수와 이웃 행과의 구조적 유사성을 기반으로 행을 TC 경로 또는 CUDA 경로로 동적으로 할당합니다.
   • 파이프라인화된 Tensor Core 커널: RS-Tile 형식을 활용하여 공유 메모리에서 레지스터로 데이터를 이동하고, 비트맵 기반 디코딩을 통해 밀집된 MMA (Matrix Multiply-Accumulate) 연산을 수행합니다. 더블 버퍼링을 통해 메모리 지연을 숨기고 Tensor Core 를 지속적으로 활용합니다.
   • CUDA 커널: 분리된 잔여 행 (residual rows) 을 처리하며, 타일 재구성 오버헤드 없이 경량화된 연산 경로를 제공합니다.

3. 로컬리티 인식 재순서화 (Locality-aware Reordering):

   • 행의 구조적 유사성을 높이기 위해 행 순서를 재배열합니다.
   • 가중 Jaccard 유사도 (Weighted Jaccard Similarity): 빈번한 열 (high-frequency columns) 보다는 정보량이 많은 희소 열을 더 중요하게 평가하는 가중치를 적용합니다.
   • kNN 그래프 및 MST: 행 간의 유사성 그래프를 구성하고 최소 신장 트리 (MST) 를 통해 전역 순서를 결정하며, 2-opt 스왑 등을 통해 지역적 정렬을 최적화합니다. 이는 인접 행 간의 구조적 파편화를 줄여 타일 형성을 용이하게 합니다.

4. 부하 균형 (Load Balancing):

   • 행 밀도의 급격한 변화로 인한 워프 (warp) 간 부하 불균형을 해결하기 위해, "초장" 행은 선택적으로 분할하고 구조적으로 불일치하는 행은 초기에 CUDA 경로로 분리하여 Tensor Core 의 점유율 (occupancy) 을 유지합니다.

3. 주요 기여 (Key Contributions)

1. RS-Tile 표현: Tensor Core 정렬 타일을 노출하면서도 불규칙 행을 CUDA 경로로 라우팅하는 컴팩트한 행 구조 표현을 제안하여 메타데이터 오버헤드를 줄이고 타일 밀도를 향상시켰습니다.
2. 미세 입자 하이브리드 커널: 적응형 행 분할, 부하 균형이 잡힌 파이프라인화된 Tensor Core 커널, 경량 CUDA 커널을 통합하여 경직된 타일 규칙을 제거하고 이중 희소 인덱스 구조 유지 필요성을 없앴습니다.
3. 로컬리티 인식 재순서화: 구조적으로 유사한 행을 그룹화하여 타일 파편화를 줄이고 MMA 활용도를 높이는 새로운 재순서화 기법을 제안했습니다.
4. 성능 검증: 다양한 실제 워크로드에서 기존 최첨단 (SOTA) SpMM 설계 (CUDA Core, Tensor Core, 하이브리드 기반) 를 능가하는 성능을 입증했습니다.

4. 실험 결과 (Results)

• 실험 환경: NVIDIA RTX 4090 (Ada Lovelace) 및 RTX 3090 (Ampere) GPU.
• 데이터셋: SuiteSparse 컬렉션의 512 개 행렬 및 GNN/스파스 학습에 사용되는 9 개의 대표적 실제 행렬.
• 성능 향상:
  • cuSPARSE 대비: RTX 4090 에서 평균 2.35 배, RTX 3090 에서 2.86 배의 속도 향상.
  • 최적화 대상별 비교:
    • CUDA Core 기반 (Sputnik, RoDe) 대비 2.06~2.61 배 향상.
    • Tensor Core 기반 (TC-GNN, DTC-SpMM, Acc-SpMM) 대비 최대 6.13 배까지 향상 (불규칙한 행 밀도에서 TC-GNN 이 실패하거나 느린 경우 제외).
    • 하이브리드 방법 (HC-SpMM) 대비 평균 2.10 배 향상.
  • 범용성: SuiteSparse 컬렉션 전체에서 cuSPARSE 대비 평균 3.21 배 속도 향상 (80% 이상의 행렬에서 1.24~8.2 배 범위).
• 엔드 투 엔드 성능 (GCN 훈련): 6 층 GCN 모델 훈련 시, 기존 방법 (PyG, cuSPARSE, TC-GNN 등) 대비 평균 1.06~1.49 배 빠른 훈련 시간을 기록했습니다.
• 하드웨어 활용도: SM 처리량 (Throughput) 이 28% 에서 33% 로 증가했고, Tensor Core 활용도는 Acc-SpMM 의 5.6% 에서 RSH-SpMM 의 **8.8%**로 크게 개선되었습니다.

5. 의의 및 결론 (Significance)

RSH-SpMM 은 현대 GPU 의 이질적인 아키텍처 (CUDA Core 와 Tensor Core) 와 실제 희소 행렬의 극심한 구조적 불규칙성 사이의 간극을 성공적으로 메웠습니다.

• 구조 인식형 실행: 행 단위에서 발생하는 불규칙성을 미세하게 감지하여 적절한 연산 유닛에 작업을 할당함으로써, 고정된 윈도우 기반 방법들의 한계를 극복했습니다.
• 안정적인 성능: 행 밀도가 급격히 변하거나 구조가 파편화된 경우에도 Tensor Core 의 높은 처리량을 안정적으로 유지하며, CUDA Core 를 효율적인 백업 경로로 활용합니다.
• 실용적 가치: GNN 및 대규모 언어 모델 (LLM) 추론 등 희소 연산이 핵심인 분야에서 GPU 가속 효율을 획기적으로 높여, 실제 애플리케이션의 실행 시간을 단축할 수 있음을 입증했습니다.

결론적으로, RSH-SpMM 은 희소 행렬 곱셈을 위한 새로운 표준으로 자리 잡을 수 있는 강력한 하이브리드 커널 설계입니다.