TranSAGE: Automating Spatial Accelerator Generation for Large-Scale Transformer Models

1 Introduction

• Large-scale Transformer inference faces rising computational and memory demands, while optimizing spatial accelerators requires substantial hardware and software expertise.The introduction also notes that inference-time scaling techniques amplify these challenges.Source: §1 · p. 1
• Existing mapping frameworks largely target predefined, fixed or limited-configurability architectures rather than jointly exploring mapping strategies and architectural designs.Earlier generators also target systolic arrays or fixed fusion patterns, limiting opportunities for cross-operator optimization.Source: §1 · pp. 1–2
• TranSAGE elevates operation clusters to execution units, fusing parallel-suited operations into hardware pipelines that execute sequentially across clusters.Clusters share and reuse on-chip compute and memory resources.Source: §1 · p. 2
• Its DSE engine uses dynamic programming for fusion under a topological order and decomposed optimization to select each cluster’s dataflow parameters.Spatial architecture primitives and lowering passes bridge high-level operations to RTL implementations.Source: §1 · p. 2
• Up to 2.42× higher throughput and 13.04× greater energy efficiency than commercial GPU/NPU baselines; against state-of-the-art accelerator generators, up to 11.33× higher throughput and 6.07× higher energy efficiency.These are the reported evaluation results for TranSAGE.Source: §1 · p. 3

2.1 Spatial Accelerators

• Spatial-accelerator design considers PE-array organization and PU interconnect, including centralized global SRAM or specialized topologies with restricted connectivity.The stated scope focuses on a single compute core with on-chip resources and hierarchical memory.Source: §2.1 · p. 3

2.2 Dataflow and Representation

• Dataflow specifies tensor operations’ temporal and spatial mapping, data movement across hierarchy levels, and parallel or sequential execution.Inter-operation representations capture transfers and scheduling across operations, including pipelined execution.Source: §2.2 · pp. 3–4

2.3 Issues and Motivation

• Existing generators often lack efficient multi-operation dataflows, instead supporting single operations or fixed fusion patterns that can leave operations sequential and increase PE underutilization and memory access.Grouping operations into pipelined clusters is proposed to improve throughput and reduce external-memory accesses.Source: §2.3 · p. 4
• Insufficient or excessive fusion can both hurt performance: pipelined operations may compete for compute and SRAM resources, reducing PE utilization and data reuse.The relative merits of fusion depend on workload; decoding may favor fusing memory-bound operations when prefill does not.Source: §2.3 · p. 4
• Dataflow optimization remains difficult because per-operation tiling, unrolling, and loop-order choices combine into an exponentially large search space.For a six-operation cluster in a sequence-length-1024 attention layer, the search space can reach 10^21 configurations.Source: §2.3 · pp. 4–5
• Diverse cluster topologies and varying PE and SRAM allocations complicate RTL generation and mapping clusters onto one architecture with shared resources.Source: §2.3 · p. 5

3 Framework Overview

• TranSAGE takes PyTorch code, schedules operation clusters with optimized dataflows, and generates RTL for hardware implementation.Its two-stage workflow uses a generation stage followed by an inference stage.Source: §3 · p. 5
• The generation frontend canonicalizes Linalg IR, fuses memory-oriented operations, and uses DSE to select operation fusion and dataflow parameters.Backend lowering constructs a primitive graph that maps supported clusters onto shared hardware and generates binary and RTL code.Source: §3 · p. 5
• At inference, a host runtime selects a binary for each request using the input shape range associated with that binary.The runtime fills in loop bounds and memory addresses; shapes within the range are supported without architectural modification.Source: §3 · p. 5

4 Multi-Operation Pipeline Scheduling

• The scheduling problem jointly chooses operation clusters and each cluster’s tiling, unrolling, and loop-order parameters to optimize energy delay product (EDP).The graph partitioning stage searches fusion schemes, while dataflow optimization selects parameters within each cluster.Source: §4 · p. 5

4.1 Operation Fusion

• Dynamic programming evaluates connected clusters over topologically ordered operations and selects partitions minimizing overall EDP.It stores optimal EDP and cut points, then backtracks to reconstruct the partition and attach each cluster’s dataflow parameters.Source: §4.1 · pp. 5–6

4.2 Intra-Cluster Dataflow Optimization

• TranSAGE decomposes dataflow search by enumerating coupled L2 tiling and outer-order parameters before solving for L1 unroll factors and inner orders.The L1 parameters are optimized per operation under shared resource constraints; the solver formulation can be handled by linear programming tools such as GLPK.Source: §4.2 · p. 6, §4.2.2 · p. 8
• A dimension-relation network propagates tiling sizes across dependent operation dimensions, reducing enumeration to representative base dimensions.Equal, collapse, and constant relations capture producer-consumer dependencies, reshapes, and concat or slice constraints; collapse tiling also follows low-to-high dimension allocation for contiguous access.Source: §4.2.1 · pp. 6–7
• Valid outer orders require reuse distance one for directly forwarded tiles and matching producer-consumer access order.Tiling can change whether output accesses are consecutive, so valid outer-order vectors are generated for each tiling candidate.Source: §4.2.1 · p. 7
• An equivalence-based ordering cache reuses candidate outer-loop orders for tiling configurations sharing a dimension mask.The mask records whether each dimension is fully tiled; configurations in the same class have identical candidate order sets, avoiding repeated DFS enumeration.Source: §4.2.1 · p. 8
• Unrolling is limited to at most two dimensions to match the 2D PE array, prioritizing dimensions that preserve contiguous operand accesses.Selected unroll dimensions are placed at outer L1 loop levels to maximize data reuse across iterations.Source: §4.2.2 · p. 8

5.1 Spatial Architecture Primitives

• Spatial architecture primitives model an accelerator as a directed cyclic graph whose nodes are storage, computation, or routing components and whose edges represent dataflow dependencies.The primitive classes include DataPath connections, FIFO- and SRAM-based storage, routing operations, and arithmetic components.Source: §5.1 · p. 9

5.2 SAP Graph Construction

• The backend constructs a shared SAP graph for scheduled clusters through resource allocation, datapath construction, and datapath reorganization.Resources are assigned using unrolling factors and tiling sizes; reorganization removes redundant primitives and inserts configurable Muxes to select among dataflow sources at runtime.Source: §5.2 · pp. 9–10
• Datapath construction routes data between SRAM, DataNodes, arithmetic primitives, and output SRAM to implement each operation’s dependencies.Routing handles incompatible storage and DataNode shapes, while arithmetic primitives are connected according to dependencies in the Linalg body.Source: §5.2 · pp. 9–10

5.3 Binary and RTL Code Generation

• TranSAGE generates a binary per cluster and RTL by replacing each SAP graph primitive with its RTL template or hardware macro.Binaries encode tiling, unrolling, loop orders, and Mux controls; runtime fields supply loop bounds and memory addresses.Source: §5.3 · p. 10

6.1 Experimental Setup

• The evaluation uses five representative LLM workloads and compares TranSAGE with Orin NX, A100, Ascend 910B3, Gemmini, and LEGO.TranSAGE configurations run at 1 GHz using a 45 nm technology model, with energy estimates scaled to 8 nm for Orin NX and 7 nm for A100 and Ascend 910B3.Source: §6.1 · p. 10

6.2 Comparison with GPU and NPU

• For the memory-bound decode workload, TranSAGE achieves average throughput improvements of 2.41×, 1.62×, and 2.12× over the Orin NX, A100, and Ascend 910B3, respectively, demonstrating benefits of operation fusion.Figure 8 compares throughput and energy efficiency across GPUs and commercial accelerators.Source: §6.2 · p. 11, Figure 8 · p. 11

6.3 Comparison with SOTA Accelerator Generators

• For prefill, average throughput improves 1.58× over GEMMINI-OS, 2.22× over GEMMINI-WS, and 1.82× over LEGO under the edge configuration.Under the server configuration, improvements are 4.14×, 3.28×, and 5.50× over the same baselines; Figure 9 compares throughput and energy efficiency.Source: §6.3 · p. 11, Figure 9 · p. 12
• 11.33× higher decode throughput than GEMMINI-OS under the server configuration; decode energy-efficiency improvement reaches 6.07× over the baselines.The generator comparison also reports up to 96.15% prefill utilization under the edge configuration, but only 3.16% decode utilization there because of limited memory bandwidth.Source: §6.3 · pp. 11–12, Figure 9 · p. 12

6.4 Ablation Study

• Full-search execution time grows exponentially with maximum cluster size and becomes impractical above size 4, whereas TranSAGE’s search time grows almost linearly and finishes within tens of minutes at size 6.Figure 10 presents the maximum-cluster-size and reconfiguration-area ablation.Source: §6.4 · p. 12, Figure 10 · p. 12
• Increasing maximum cluster size rapidly lowers EDP before gains saturate as operations compete for shared PE and SRAM resources; reconfiguration area overhead remains at most 0.11%.The reported contention reduces utilization and data reuse; multiplexers and the limited number of distinct dataflow configurations keep reconfiguration overhead low.Source: §6.4 · pp. 12–13, Figure 10 · p. 12

7.1 Analysis of Fused Clusters

• Fusion patterns vary with workload and resources: edge and long-sequence configurations generally produce more clusters than server and short-sequence configurations.Operations such as Softmax, RoPE, and Gate projection may be split across clusters or fused across conventional stage boundaries.Source: §7.1 · p. 13
• Fusing data-dependent operations that use complementary compute resources can reduce memory accesses and improve utilization, as in Softmax and RoPE elementwise pipelines.Fusion is beneficial for resource-competing operations only when they form a balanced pipeline: Q and K/V projections split during prefill, while gate and up projections fuse under server-prefill configurations because their output shapes and computational demands are balanced.Source: §7.1 · pp. 13–14

7.2 Limitations and Future Directions

• Generated physical compute, storage, and interconnect resources are fixed, so evolving model structures may require hardware regeneration; multi-chip DSE does not explore communication-aware parallelism beyond data parallelism.The authors identify greater generated-architecture flexibility and communication-aware multi-chip DSE as future work.Source: §7.2 · p. 14

• TranSAGE treats operation clusters as execution units to enable efficient pipelining and improve throughput and energy efficiency over GPUs, NPUs, and prior accelerator generators for prefill and decoding.Source: §8 · p. 14

• Spatial architecture primitives include point-to-point links, FIFO and SRAM storage, routing and broadcast, reductions, composition and decomposition, and floating-point arithmetic operations.Source: A Spatial Architecture Primitives · p. 15

B System Configurations

• Edge uses Orin NX, while Server uses A100/Ascend 910B3.Source: B System Configurations · p. 15
• All configurations operate at 1 GHz, and device-matched setups use the largest valid PE organization whose peak throughput does not exceed the target device.Each PE issues one operation per cycle; one multiplication and one addition count as two floating-point operations.Source: B System Configurations · p. 15

C LLaMA3-8B Results

• The LLaMA3-8B results distinguish Edge (batch 1) and Server (batch 8), with prefill and decode reported as throughput and energy-efficiency metrics.TP denotes throughput in GFLOP/s; EE denotes energy efficiency in GFLOP/J.Source: C LLaMA3-8B Results · p. 15

D.1 Abstract

• The artifact workflow starts from PyTorch code and generates Linalg MLIR, schedules kernels, and lowers computation to an SAP graph for performance modeling.It also provides intermediate MLIR and metrics including latency, energy, FLOPs, throughput, and energy efficiency.Source: §D.1 · p. 16

D.2 Artifact Checklist

• The artifact exposes generated and optimized Linalg MLIR, scheduled clusters with dataflow parameters, an SAP graph, and a modeled performance report.Source: §D.2 · p. 16

D.3 Availability and Installation

• The artifact is archived on Zenodo, with a matching GitHub release, and offers prebuilt-Docker and local-Dockerfile installation methods.The repository clone command uses the ae branch.Source: §D.3 · p. 16

D.4 Experimental Workflow

• The workflow has five dependent stages whose intermediate outputs can be inspected and reused; rapid mode evaluates a block 0 FFN workload, while full mode runs the complete suite.Full evaluation can take substantially longer, particularly during kernel scheduling.Source: §D.4 · p. 16
• Stage 1 generates MLIR for selected model, action, block, layer, batch-size, and sequence-length combinations, then writes the files to mlir/.Source: §D.4 · p. 16
• Stage 2 applies mixed-precision resolution, operation generalization, generic marking, dead-operation elimination, fusion, and tensor-operation folding, writing optimized MLIR to generic/.Source: §D.4 · p. 16
• Stage 3 searches tiling, loop, unrolling, and cluster choices for the selected edge or server configuration, with at most six operations per cluster by default.Scheduled MLIR files are written to scheduled/.Source: §D.4 · p. 16
• Stage 4 lowers scheduled workloads to the hardware-facing DAP dialect and constructs a SAP graph with SRAM, data-path, arithmetic, and routing operations.The resulting MLIR files are written to dap/.Source: §D.4 · pp. 16–17
• Stage 5 estimates latency, energy, and FLOPs, emitting annotated MLIR and a performance log containing workload headers and metric records.Source: §D.4 · p. 17