LizTail improves parallel computing for neural network tasks

LizTail: Unifying Shape Adaptation and Spatial Parallelism with Segment-Reconfigurable Systolic Arrays

Hardware Architecture

Summary

Computers use special grids called systolic arrays to speed up tasks like understanding language or recognizing images, but these grids often don't match well with the different shapes and sizes of tasks. The authors created LizTail, a flexible grid that can break into smaller parts or bend to fit these tasks better, making it faster and more efficient. This design helps computers run certain AI tasks up to 25 times faster while using less energy. The system also comes with software tools to manage how tasks are spread across the grid parts.

What this means in practice

  • For hardware accelerators designers: Design neural-network chips that efficiently handle diverse model sizes by physically reconfiguring computational arrays for better speed and energy use.
  • For edge ai device engineers: Implement AI inference engines that adapt to varying task sizes and improve decoding speeds in language models on resource-constrained devices.

Authors

Gang Liu, Chen Yang, Junyan Hu, Jianguo Chen, Lixin Duan

Abstract

Systolic arrays are widely used to accelerate neural-network inference, yet fixed two-dimensional arrays struggle with the substantial diversity in both matrix geometry and available parallelism across CNNs and Transformer-based large language models. This mismatch is particularly severe for skinny matrices, small-batch autoregressive decoding, and operators containing multiple independent partitions, leading to low processing-element (PE) utilization and repeated spatial folding. Existing reconfigurable arrays typically improve either shape adaptation or distributed execution, whereas supporting both through fine-grained routing incurs considerable interconnect, buffering, and control overheads. We present LizTail (named after lizard tail autotomy, where predefined fracture planes allow segments to detach and remain independently active, analogous to LizTail's predefined reconfiguration boundaries), a multi-mode systolic array that enables constrained segment-level reconfiguration by partitioning a $128\times128$ array into 32 regular $128\times4$ segments and adding connectivity only at predefined endpoints. The resulting substrate supports trimming, long-path extension through $180^\circ$ turn-back connections, and distributed execution in which each subarray can itself be trimmed or extended. Differentiated PE datapaths and segment-aware CEB/SIB buffers provide the required compute and data-delivery flexibility without unrestricted PE-level routing. A hardware-constrained mapper selects the execution organization, dataflow, task partition, and tiling. Across representative workloads, LizTail achieves geometric-mean speedups of $5.02\times$ on CNN inference and $25.67\times$ on LLM decoding over a TPU-like array, while reducing EDP by up to $38.23\times$. LizTail is open-sourced at https://github.com/kmjbwxa/LizTail.