JADE: Joint Architecture-Dataflow Exploration for LLM Inference on Heterogeneous In- and Near-Memory Computing Systems

Large Language Model (LLM) inference challenges memory/computing organization and dataflow optimization on traditional hardware stacks due to its various attention mechanisms and autoregressive structure. In-memory computing (IMC) and near-memory computing (NMC) architectures have emerged as potential solutions to reduce the memory burden in LLMs. However, efficient LLM orchestration on IMC/NMC, as well as dedicated IMC/NMC architecture optimization, are interrelated, encompassing a hardware-software co-design space of a myriad of parameters, which still remains largely underexplored. This work provides a joint architecture-dataflow exploration (JADE) framework that covers both architectural factors on computing/memory/communication, and algorithmic factors on LLM dataflow. We identify key configurations in four categories: computing and memory organization, interconnection topology, matrix multiplication (MM) unrolling and mapping, and attention mechanism. A unified MM representation enables flexible unrolling and mapping on various architectural organizations. Dataflow is routed based on the collective communication patterns inherent in the distributed LLMs and the features of common topologies. Finally, case studies show the design tradeoffs between key design options and provide insights into both homogeneous and heterogeneous IMC/NMC design for LLMs.