Evolving Memory-Aware Schedules for Transformer Inference on Systolic Array Accelerators

Transformer models impose severe computational and memory demands on modern hardware accelerators, making efficient workload scheduling a key bottleneck for energy- and latency-constrained deployment. Mapping multi-head self-attention (MHSA) onto systolic-array architectures with hierarchical on-chip memory requires joint optimization of compute ordering, array assignment, and memory traffic. Accurate evaluation of such mappings must therefore be memory-aware, as traditional compute-only analytical models overlook bandwidth limits, buffer contention, and data-reuse effects, leading to unrealistic performance estimates. In this work, we introduce a simulation-driven scheduling framework that couples black-box optimizers with a high-level, memory-aware simulator to explore the combinatorial space of mapping Transformer workloads on systolic-array accelerators. The scheduler iteratively refines candidate execution plans using performance feedback from TransInferSim, which provides cycle-level latency and energy estimates through detailed modeling of compute and memory components. Multiple black-box optimizers are compared within this loop, revealing how memory-aware evaluation influences convergence dynamics and the relative performance of candidate schedules compared to naive compute-only estimation. Our experiments showed a 8.49% improvement in the latency of MHSA computation when using our memory-aware scheduling framework compared to the compute-only optimized baseline. This demonstrates the critical importance of memory-aware evaluation for efficient real-world deployment of Transformer workloads on modern hardware accelerators.