KAPLA: Pragmatic Representation and Fast Solving of Scalable NN Accelerator Dataflow

📝 Abstract

Dataflow scheduling decisions are of vital importance to neural network (NN) accelerators. Recent scalable NN accelerators support a rich set of advanced dataflow techniques. The problems of comprehensively representing and quickly finding optimized dataflow schemes thus become significantly more complicated and challenging. In this work, we first propose comprehensive and pragmatic dataflow representations for temporal and spatial scheduling on scalable multi-node NN architectures. An informal hierarchical taxonomy highlights the tight coupling across different levels of the dataflow space as the major difficulty for fast design exploration. A set of formal tensor-centric directives accurately express various inter-layer and intra-layer schemes, and allow for quickly determining their validity and efficiency. We then build a generic, optimized, and fast dataflow solver, KAPLA, which makes use of the pragmatic directives to explore the design space with effective validity check and efficiency estimation. KAPLA decouples the upper inter-layer level for fast pruning, and solves the lower intra-layer schemes with a novel bottom-up cost descending method. KAPLA achieves within only 2.2% and 7.7% energy overheads on the result dataflow for training and inference, respectively, compared to the exhaustively searched optimal schemes. It also outperforms random and machine-learning-based approaches, with more optimized results and orders of magnitude faster search speedup.

💡 Analysis

Dataflow scheduling decisions are of vital importance to neural network (NN) accelerators. Recent scalable NN accelerators support a rich set of advanced dataflow techniques. The problems of comprehensively representing and quickly finding optimized dataflow schemes thus become significantly more complicated and challenging. In this work, we first propose comprehensive and pragmatic dataflow representations for temporal and spatial scheduling on scalable multi-node NN architectures. An informal hierarchical taxonomy highlights the tight coupling across different levels of the dataflow space as the major difficulty for fast design exploration. A set of formal tensor-centric directives accurately express various inter-layer and intra-layer schemes, and allow for quickly determining their validity and efficiency. We then build a generic, optimized, and fast dataflow solver, KAPLA, which makes use of the pragmatic directives to explore the design space with effective validity check and efficiency estimation. KAPLA decouples the upper inter-layer level for fast pruning, and solves the lower intra-layer schemes with a novel bottom-up cost descending method. KAPLA achieves within only 2.2% and 7.7% energy overheads on the result dataflow for training and inference, respectively, compared to the exhaustively searched optimal schemes. It also outperforms random and machine-learning-based approaches, with more optimized results and orders of magnitude faster search speedup.

📄 Content

데이터플로우 스케줄링 결정은 신경망(NN) 가속기 설계에서 매우 중요한 역할을 차지합니다. 최근에 등장한 확장 가능한 NN 가속기들은 풍부하고 다양한 고급 데이터플로우 기법들을 지원하도록 설계되었으며, 이러한 기법들을 모두 포괄적으로 표현하고 최적화된 데이터플로우 스킴을 신속하게 찾아내는 문제는 그 복잡성이 크게 증가함에 따라 더욱 까다롭고 도전적인 과제로 부상하고 있습니다.

본 논문에서는 먼저 확장 가능한 다중 노드 NN 아키텍처 상에서 시간적(temporal) 및 공간적(spatial) 스케줄링을 위한 포괄적이면서도 실용적인 데이터플로우 표현 방식을 제안합니다. 여기서 제시하는 비공식적인 계층적 분류 체계는 데이터플로우 공간의 서로 다른 레벨들 간에 존재하는 긴밀한 결합(coupling)을 강조하는데, 이러한 결합이 바로 빠른 설계 탐색을 방해하는 주요 난점임을 지적합니다.

다음으로, 텐서 중심(tensor‑centric)으로 정의된 일련의 형식적인 지시어(directives)를 도입합니다. 이 지시어들은 레이어 간(inter‑layer) 및 레이어 내부(intra‑layer)에서 적용될 수 있는 다양한 스킴을 정확히 기술할 수 있도록 설계되었으며, 각 스킴의 타당성(validity)과 효율성(efficiency)을 빠르게 판단할 수 있는 메커니즘을 제공합니다. 구체적으로, 텐서의 차원, 데이터 이동 경로, 연산 순서 등을 명시적으로 기술함으로써 설계자가 복잡한 데이터 흐름을 직관적으로 이해하고, 자동화된 도구가 이를 기반으로 합리적인 탐색을 수행하도록 합니다.

이러한 실용적인 지시어들을 활용하여 우리는 KAPLA라는 이름의 범용적이고 최적화된 고속 데이터플로우 솔버를 구축했습니다. KAPLA는 설계 공간을 효율적으로 탐색하기 위해 두 단계의 전략을 채택합니다. 첫 번째 단계에서는 상위 레이어 간(inter‑layer) 관계를 분리(decouple)하여 빠른 가지치기(pruning)를 수행함으로써 전체 탐색 범위를 급격히 축소합니다. 두 번째 단계에서는 하위 레이어 내부(intra‑layer) 스킴을 해결하기 위해 새롭게 고안된 ‘하향식(bottom‑up) 비용 감소(cost descending)’ 방법을 적용합니다. 이 방법은 각 레이어의 비용 모델을 역순으로 평가하면서, 비용이 가장 크게 감소하는 방향으로 스킴을 선택해 나가므로 전통적인 탐색 방식보다 훨씬 적은 연산량으로 최적에 근접한 해를 도출할 수 있습니다.

실험 결과, KAPLA는 훈련(training)과 추론(inference) 각각에 대해 최적의 데이터플로우 스킴을 전수 탐색(exhaustive search)으로 찾았을 때와 비교했을 때, 에너지 오버헤드가 각각 단 2.2 %와 7.7 %에 불과함을 확인했습니다. 이는 KAPLA가 거의 최적에 가까운 에너지 효율성을 유지하면서도 탐색 시간을 크게 단축한다는 것을 의미합니다. 또한, 무작위(random) 탐색 및 기존의 머신러닝 기반 탐색 기법들과 비교했을 때, KAPLA는 더 높은 최적화 수준의 결과를 제공할 뿐만 아니라 탐색 속도 면에서도 수십 배에서 수백 배에 달하는 가속 효과를 보였습니다.

요약하면, 본 연구는 (1) 확장 가능한 다중 노드 신경망 가속기 환경에서 시간·공간 스케줄링을 포괄적으로 표현할 수 있는 새로운 데이터플로우 모델을 제시하고, (2) 텐서 중심의 형식적인 지시어 집합을 통해 레이어 간·레벨 내부 스킴의 타당성과 효율성을 빠르게 판단할 수 있는 기반을 마련했으며, (3) 이러한 기반 위에 구축된 KAPLA 솔버가 상위 레이어와 하위 레이어를 각각 특화된 방식으로 처리함으로써 설계 공간 탐색을 효율화하고, (4) 실제 에너지 소비와 탐색 속도 측면에서 기존 방법들을 크게 능가한다는 점을 입증했습니다.

이러한 성과는 차세대 신경망 가속기 설계에서 데이터플로우 스케줄링을 자동화하고 최적화하는 데 필요한 핵심 기술적 토대를 제공하며, 앞으로 더욱 복잡하고 대규모인 멀티노드 학습·추론 시스템에서도 높은 효율성을 유지하면서 설계 주기를 단축할 수 있는 가능성을 열어줍니다. 앞으로의 연구에서는 KAPLA의 비용 모델을 더욱 정교화하고, 다양한 하드웨어 백엔드(예: ASIC, FPGA, GPU)와의 연동을 통해 실제 제품 수준의 구현 가능성을 검증하는 작업이 진행될 예정입니다.

이와 같이, 데이터플로우 스케줄링에 대한 포괄적이고 실용적인 표현 체계와 고속 탐색 알고리즘의 결합은 NN 가속기 설계 분야에 새로운 패러다임을 제시하며, 향후 인공지능 시스템의 에너지 효율성과 성능을 동시에 끌어올리는 핵심 동력이 될 것입니다.