[논문] PRISM: 메모리 내 처리 기반 희소 MTTKRP를 통한 텐서 분해 가속화

출처: arXiv - 2605.29728v1

개요

본 논문은 PRISM이라는 새로운 Processing‑In‑Memory (PIM) 솔루션을 소개한다. PRISM은 CP‑ALS 텐서 분해에서 핵심 커널인 희소 Matricized Tensor Times Khatri‑Rao Product (spMTTKRP)를 가속화한다. UPMEM의 DRAM 기반 PIM 칩 내부에서 데이터를 가까이 두고 연산을 수행함으로써, CPU와 GPU 구현에서 발생하는 메모리 대역폭 병목을 해소하고 실제 희소 텐서에서 최대 2.6배의 속도 향상을 달성한다.

주요 기여

• 첫 번째 PIM 기반 spMTTKRP 엔진: 메모리 집약적인 spMTTKRP 커널을 UPMEM PIM DRAM 내부에서 직접 실행하는 완전한 파이프라인을 설계하였다.
• 포괄적인 설계 공간 탐색: 텐서 파티셔닝 전략, 수치 형식(FP32, FP16, 정수) 및 커널 마이크로‑옵티마이제이션을 평가하여 UPMEM 아키텍처의 제약에 맞추었다.
• 이기종 CPU‑PIM 협업: CPU가 제어 중심 작업을 담당하고 PIM 코어가 대부분의 연산을 수행하는 경량 오프로드 방식을 도입해 양쪽의 장점을 모두 활용한다.
• 성능‑효율성 향상: 최고 CPU‑전용 CP‑ALS 구현 대비 최대 2.37× 속도 향상, CPU‑PIM 하이브리드 사용 시 2.64× 향상을 보이며 이론적 피크 성능에 훨씬 가까운 비율을 달성한다.
• 오픈‑소스 평가 프레임워크: UPMEM 하드웨어에서 다른 희소 텐서 워크로드를 벤치마크할 수 있도록 스크립트와 커널을 제공한다.

방법론

1. 문제 분해 – 저자들은 spMTTKRP를 세 단계로 나눈다: (a) 희소 텐서 슬라이싱, (b) Khatri‑Rao 곱 조회, (c) 출력 행렬에 누적.
2. PIM 매핑 – 각 UPMEM DRAM 칩에는 다수의 작은 RISC‑V 코어가 존재한다. PRISM은 텐서 슬라이스를 개별 코어에 할당해 비어 있지 않은 원소를 CPU로 옮기지 않고 DRAM에서 직접 스트리밍한다.
3. 파티셔닝 전략 – 부하 균형과 칩 간 통신 최소화를 위해 행 기반, 열 기반, 하이브리드 파티셔닝을 실험한다.
4. 수치 형식 튜닝 – 32‑bit 부동소수점, 16‑bit 부동소수점, 8‑bit 정수 표현을 사용해 정밀도 감소가 속도와 분해 품질에 미치는 영향을 조사한다.
5. 커널 최적화 – 루프‑퓨전, 소프트웨어 프리패치, 제한된 ALU를 활용한 벡터‑유사 SIMD 에뮬레이션 등을 포함한다.
6. 이기종 오케스트레이션 – 호스트 CPU의 경량 런타임이 작업 청크를 스케줄하고 부분 결과를 수집·최종 축소를 수행하면서 PIM 코어는 순수 연산 루프에 전념한다.

결과 및 분석

• 메모리 트래픽 감소: 데이터를 DRAM 내부에 유지함으로써 PRISM은 CPU‑전용 실행 대비 외부 메모리 접근을 약 45 % 절감한다.
• 자원 효율성: PIM 구현은 UPMEM 시스템의 이론적 피크 FLOP/byte 비율의 약 70 %에 도달하지만, GPU 기준은 약 30 %에 머문다.
• 정밀도 트레이드‑오프: FP16 사용 시 대부분의 데이터셋에서 분해 오차 손실이 거의 없으면서 약 10 % 추가 속도 향상이 나타난다. 반면 8‑bit 정수 양자화는 수렴에 악영향을 주어 권장되지 않는다.
• 확장성 한계: 텐서가 다수의 PIM 칩에 분산될 경우, 특히 불규칙한 희소성 패턴을 가진 워크로드에서는 칩 간 동기화 오버헤드가 지배적으로 작용한다.

실용적 함의

• 모델 학습 가속: CP‑ALS를 활용하는 머신러닝 파이프라인(예: 추천 시스템, 토픽 모델링, 신호 처리)은 정확도 손실 없이 학습 시간을 크게 단축할 수 있다.
• 에너지 효율적인 추론: PIM 코어는 CPU/GPU 조합에 비해 연산당 전력 소모가 현저히 낮아, 엣지 환경에서의 분석(예: 디바이스 내 추천) 구현이 보다 현실적이다.
• 단순화된 데이터 파이프라인: 거대한 희소 텐서를 호스트로 스트리밍할 필요가 없어, 저장과 연산이 공동 위치하는 데이터 중심 아키텍처 설계가 가능해진다.
• 툴링 통합: 오픈‑소스 커널은 얇은 C‑API를 통해 TensorLy, SPLATT 등 인기 텐서 라이브러리에 래핑될 수 있어 기존 코드베이스가 투명하게 spMTTKRP를 오프로드하도록 지원한다.
• 향후 하드웨어 로드맵: 이번 결과는 차세대 PIM 지원 서버에 대한 설득력 있는 사용 사례를 제시하며, 하드웨어 벤더가 네이티브 Khatri‑Rao 프리미티브와 더 큰 온‑칩 메모리를 제공하도록 유도한다.

제한점 및 향후 연구

• 분산 메모리 오버헤드: 현재 UPMEM 플랫폼은 빠른 칩 간 통신을 지원하지 않아 매우 큰 텐서로 확장하는 데 제약이 있다. 차세대 고대역폭 메쉬 네트워크가 이 문제를 완화할 수 있다.
• 정밀도 민감도: FP16이 많은 데이터셋에서 잘 동작하지만, 일부 과학 응용에서는 여전히 FP32 정확도가 요구되어 정밀도 감소에 따른 속도 이득이 제한된다.
• 알고리즘 범위: PRISM은 CP‑ALS 내 spMTTKRP에만 초점을 맞추었으며, Tucker 분해, 텐서‑트레인 등 다른 텐서 연산으로 확장하는 연구가 남아 있다.
• 소프트웨어 생태계: 주류 ML 프레임워크와의 통합이 아직 수동 커널 호출 수준에 머물러 있다. 보다 높은 수준의 추상화 레이어가 도입되면 채택 장벽이 낮아질 것이다.

전반적으로 PRISM은 Processing‑in‑Memory가 단순한 이론적 호기심이 아니라, 실제 현업 분석 워크로드의 핵심 머신러닝 원시 연산에 구체적인 성능·효율성 이점을 제공할 수 있음을 보여준다.

저자

• Daniel Pacheco
• Leonel Sousa
• Aleksandar Ilic

논문 정보

• arXiv ID: 2605.29728v1
• 분류: cs.DC
• 발표일: 2026년 5월 28일
• PDF: PDF 다운로드