PaperlayerView exact revision on arXiv

Source-linked AI summary

FLAT: An Optimized Dataflow for Mitigating Attention Bottlenecks

Sheng-Chun Kao, Suvinay Subramanian, Gaurav Agrawal, Amir Yazdanbakhsh, Tushar Krishna

arXiv:2107.06419v7cs.LGcs.AR

TL;DR

Attention models face low data reuse and quadratic intermediate-memory growth, which make long sequences challenging. FLAT uses fused attention dataflow and tiling; it improves performance and energy use over existing dataflow optimizers and scales to sequences up to 64K tokens.

  • Problem

    Attention layers face low reuse and quadratic intermediate-memory growth, creating memory-bandwidth demands that challenge long-sequence scalability.

    Source: §1 · p. 1, §10 · p. 15

  • Method

    FLAT combines inter-operator fusion, interleaved execution, and efficient tiling to increase data reuse in attention layers.

    Source: §1 · p. 2, §10 · p. 15

  • Results

    Against state-of-the-art dataflow optimizers, FLAT delivers 1.75× speedup and 44% energy savings on Edge, and 1.65× speedup and 55% energy savings on Cloud.

    Source: §1 · p. 2

  • Takeaways & Limitations

    FLAT unblocks transformer-model scalability up to 64K tokens, while conventional DNN dataflows for attention reach their efficiency limit above 512 tokens.

    Source: §1 · p. 2

Abstract

from arXiv · show

Attention mechanisms, primarily designed to capture pairwise correlations between words, have become the backbone of machine learning, expanding beyond natural language processing into other domains. This growth in adaptation comes at the cost of prohibitively large memory requirements and computational complexity, especially at higher number of input elements. This limitation is due to inherently limited data reuse opportunities and quadratic growth in memory footprints, leading to severe memory-boundedness and limited scalability of input elements. This work addresses these challenges by devising a tailored dataflow optimization, called FLAT, for attention mechanisms without altering their functionality. This dataflow processes costly attention operations through a unique fusion mechanism, transforming the memory footprint quadratic growth to merely a linear one. To realize the full potential of this bespoke mechanism, we propose a tiling approach to enhance the data reuse across attention operations. Our method both mitigates the off-chip bandwidth bottleneck as well as reduces the on-chip memory requirement. FLAT delivers 1.94x (1.76x) speedup and 49% and (42%) of energy savings compared to the state-of-the-art Edge (Cloud) accelerators with no customized dataflow optimization. When on-chip resources are scarce (20 KB-200 KB), FLAT yields, on average, 1.5x end-to-end latency reduction across a diverse range of conventional attention-based models with input sequence lengths ranging from 512-token to 64K-token. Our evaluations demonstrate that state-of-the-art DNN dataflow applied to attention operations reach the efficiency limit for inputs above 512 elements. In contrast, FLAT unblocks transformer models for inputs with up to 64K elements

1. Introduction

Longer sequences are important for tasks that depend on global context, but attention layers pose memory and bandwidth challenges. FLAT combines inter-operator fusion and tiling to improve data reuse and support longer sequences.

1. Introduction

  • Long sequences matter when input properties emerge from global context, with demand spanning language, protein folding, summarization, and audio generation.The passage gives examples of protein function and text meaning differing when considered in full context rather than as fragments.Source: §1 · p. 1
  • Attention presents low-reuse activation-activation operators and many-to-many dependencies, while intermediate tensors grow quadratically with sequence length.These properties make the operators memory-bound, constrain fusion, and pressure on-chip memory capacity.Source: §1 · pp. 1–2
  • FLAT, or Fused Logit Attention Tiling, fuses many-to-many operators while preserving dependencies, then tiles their execution to stage intermediate tensors on limited on-chip memory.The approach targets off-chip bandwidth pressure and scalability up to 64K inputs.Source: §1 · p. 2
  • The paper also develops a map-space exploration framework to search loop orders and tile sizes under processing-element and on-chip-memory constraints.The evaluation spans multiple attention-based models, including BERT, TrXL, FlauBERT, and T5.Source: §1 · p. 2
  • 1.75× and 1.65× speedups, with 44% and 55% energy savings, are reported for Edge and Cloud accelerators, respectively, versus state-of-the-art dataflow optimizers.FLAT also yields 1.5× end-to-end latency reduction and 1.4× end-to-end energy savings for 512-token inputs with scarce on-chip resources.Source: §1 · p. 2

2. Background on Attention-Based Models

An attention model is organized into blocks, layers, and operators, with attention computing pairwise token relationships through logits, normalization, and weighted values. Accelerator performance depends on operational intensity and dataflow, which determine data-reuse opportunities and realized performance.

2. Background on Attention-Based Models

  • An attention-based model comprises blocks containing attention and normalization layers followed by fully connected layers; layers comprise operators.A model may also include embedding, positional encoding, masking, or task-specific layers.Source: §2 · p. 2
  • Attention projects inputs into Query, Key, and Value tensors, computes pairwise Logit scores, normalizes each row with Softmax, and weights Value vectors to produce outputs.For input [B,N,D], the Logit tensor has dimension [B,H,N,N], and the attention output has dimension [B,H,N,d].Source: §2 · p. 3
  • Q, K, V, and O are activation-weight operators, whereas Logit and Attend are activation-activation operators that often dominate latency and power, especially at long sequence lengths.The operators perform GEMM computations on their respective tensors.Source: §2 · p. 3

2.1. DNN Accelerators - Performance Considerations

  • Operational intensity is arithmetic operations divided by memory accesses, serving as a proxy for an operator's maximum possible performance on given hardware.Lower operational intensity implies fewer data-reuse opportunities and a greater likelihood of memory-bandwidth limitation.Source: §2.1 · p. 3
  • Dataflow stages data through the memory hierarchy to processing elements and determines realized performance through tiling, compute order, and parallelism.A dataflow together with specific tile sizes is called a mapping.Source: §2.1 · p. 3

3. Challenges with Running Attention Layers

Attention layers challenge conventional dataflow with memory-bound L/A operators, complex many-to-many fusion dependencies, and intermediate tensors that grow quadratically with sequence length.

3.1. Challenge 1: Low Operational Intensity of L/A

  • L/A operators remain memory-bound at low performance, and increasing batch size does not help, unlike for some FC operators.The roofline analysis compares attention operators with CONV and FC operators at sequence length 512.Source: §3.1 · p. 4, Figure 2 · p. 4

3.2. Challenge 2: Complexity of Op Fusion for L/A

  • Tensor-tensor fusion in attention is difficult because many-to-many operators require complex dependency management and intermediate-data staging.FLAT addresses these dependencies with a specialized inter-operator dataflow.Source: Figure 3 · p. 4, §3.2 · p. 5

3.3. Challenge 3: Tensor Footprint of L/A

  • The L–A intermediate tensor is O(BHN^2), exceeds 8 GB beyond 16K sequence length, and makes whole-tensor on-chip staging unscalable.The proposed tiling controls the fused operator’s active memory footprint according to on-chip capacity.Source: §3.3 · p. 5, Figure 4 · p. 5

4. FLAT Dataflow Concept

FLAT fuses Logit (L) and Attend (A) operators in a specialized dataflow. It targets their low operational intensity while accounting for fusion dependencies and limited on-chip memory.

4. FLAT Dataflow Concept

  • FLAT is a specialized dataflow strategy combining intra-operator dataflow with inter-operator execution of L and A in concert.The method targets the two memory-bandwidth-bound operators in attention layers.Source: §4 · p. 5

4.1. Identifying Tensor Fusion Opportunity

  • FLAT fuses L and A because their fused operation has higher effective operational intensity, while other operator pairs offer less benefit or risk performance losses.Fusing additional operators also requires more on-chip staging and potentially finer-grained execution.Source: §4.1 · pp. 5–6

4.2. Challenges with Tensor Fusion Implementation

  • L–A fusion must preserve Softmax’s reduction dependency, while tiling must address the O(BHN^2) intermediate tensor that can exceed on-chip capacity.The implementation uses FLAT to enable L–A fusion despite these constraints.Source: §4.2 · p. 6

5. FLAT Dataflow Implementation

FLAT interleaves fused L and A operations through shared outer loops and operator-specific inner loops. Tiling granularity balances Softmax dependencies, on-chip capacity, data reuse, and compute utilization.

5. FLAT Dataflow Implementation

  • The fused operators share outer loops, while their distinct inner loops execute one after another in an interleaved schedule.This loop organization applies to consecutive tensor operators, with tile sizes set to manage dependencies and on-chip memory.Source: §5 · p. 6

5.1. FLAT-tile and Execution Granularity

  • FLAT-tile specifies how many slices of the partial intermediate tensor are computed in one fused pass, with granularity progressing from rows to heads and batches.The full intermediate tensor held on-chip is the M-Gran baseline.Source: §5.1 · p. 6

5.2. Managing Constraints from Data Dependency

  • Row-granularity is the minimum tile size because Softmax requires an [1, N] input across the key dimension.FLAT restricts tile sizes to multiples of this dependency-respecting unit.Source: §5.2 · p. 6

5.3. Managing Constraints from On-Chip Memory

  • Finer tiling reduces on-chip footprint, but row tiling can reduce key reuse and PE utilization, so FLAT co-explores inter- and intra-operator dataflow.For L–A, finer M-, B-, and H-Gran tiling does not forgo algorithmic reuse opportunities because these activation-activation operators have none.Source: §5.3 · p. 7

6. Evaluation Methodology

The evaluation uses an analytical accelerator model and exhaustive map-space exploration to compare attention dataflows across hardware configurations. It models memory hierarchy, data movement, and energy, and evaluates workloads spanning several transformer models and sequence lengths.

6.1. Accelerator Modeling Methodology

  • The analytical cost model estimates performance and energy across accelerator configurations, modeling compute arrays, distribution and reduction networks, scratchpads, global buffers, spilling, shared bandwidth limits, and softmax runtime.Energy estimates use Accelergy and account for compute, on-chip memory, and off-chip communication.Source: §6.1 · p. 8
  • Exhaustive map-space exploration searches intra- and inter-operator dataflows, including cross-layer fusion, and single-layer simulation results are within 1% of MAESTRO’s results.The framework compares dataflows with and without FLAT under the same exploration workflow.Source: §6.1 · p. 8
  • Evaluations model Edge and Cloud platforms, vary on-chip memory capacity, and include BERT, TransformerXL, FlauBERT, T5, and XLM at sequence lengths from 512 to 64K, plus a 256K case.All models use batch size 64; the paper states this choice is immaterial to the dataflow optimization.Source: §6.1 · pp. 8–9, Table 2 · p. 9
  • At 200K buffer capacity, FLAT improves over FLEX by 1.7× on L/A layers and 1.1× end-to-end.The reported table also compares FLAT against Naïve across 2GB, 20M, and 200K buffer sizes.Source: §6.1 · p. 9, Table 4 · p. 9
  • Utilization is reported as ideal arithmetic runtime divided by achieved dataflow runtime, indicating distance from the compute bound and compute-resource utilization.Ideal runtime is total model computations divided by target accelerator peak FLOPs.Source: §6.1 · p. 9

7. Evaluation I: FLAT Dataflow Efficacy

Across Edge and Cloud models, FLAT improves speed and energy use relative to FLEX, with larger gains at longer sequence lengths and constrained buffers. Its evaluation also examines compatibility with GPU and sparse-attention settings and a protein-sequencing case study.

7. Evaluation I: FLAT Dataflow Efficacy

  • At 200KB, FLAT-Opt improves over FLEX-Opt by 1.7× on L/A operations and 1.1× end-to-end; L/A computation rises from 10% at 512 tokens to 47% at 4K and 78% at 16K.Across the reported energy comparisons, FLAT-Opt reduces consumption by around 1.5×–2.0× relative to FLEX-Opt.Source: §7 · pp. 9–10, §1.0 · p. 11

1.0 Naïve FLEX FLAT

  • FLAT-Opt reaches nearly 1.0 compute utilization with 10×–100× less on-chip buffer in Edge and 100×–1000× less in Cloud than FLEX-Opt.The evaluation sweeps buffer sizes, sequence lengths, and model granularities.Source: Figure 9 · p. 9, §1.0 · p. 10, Figure 11 · p. 10
  • On the TrEMBL protein-sequencing workload, FLAT enables larger transformer models and/or longer sequence lengths in the 8K- and 16K-sequence cases.The case study uses a Tesla T4 GPU with 16GB memory and follows the methodology used in Performer.Source: §1.0 · p. 12, Table 6 · p. 13, Table 7 · p. 13

8.1. Cloud and Edge Accelerators

  • At 64K tokens, FLAT-Opt runs 2.8× faster than FLEX-Opt-Edge and 3.07× faster than FLEX-Opt-Cloud.The comparison uses the selected Edge and Cloud accelerator design points.Source: §8.1 · p. 13
  • FLAT-Opt delivers 1.75× speedup and 44% energy savings over FLEX-Opt-Edge, and 1.65× speedup and 55% energy savings over FLEX-Opt-Cloud across transformer models.These are the end-to-end comparisons reported for the Edge and Cloud platforms.Source: Table 5 · p. 12, §8.1 · p. 13
  • FLAT reduces off-chip bandwidth requirements against FLEX-Opt by 82% on average in Cloud and 71% on average in Edge.For increasing sequence length, the reported on-chip memory requirement grows quadratically for FLEX and linearly for FLAT.Source: §8.1 · p. 13

8.2. FLAT Compatibility with other Accelerators

  • On a Tesla T4 GPU, FLAT runs faster than baseline and supports sequence lengths up to 64K; sparse L/A operators remain memory-bound after 50% pruning, and FLAT raises their compute-to-memory time ratio toward or above 1.The sparse-attention comparison evaluates BERT and XLM under TPU-v3 resources.Source: §8.2 · pp. 13–14, Figure 16 · p. 14

9. Related Works

Prior work applies algorithmic changes and limited fusion to attention, while FLAT targets dependent MatMul-Softmax-MatMul fusion. The section distinguishes FLAT from GNN dataflows, approximation-based attention accelerators, and compiler fusion.

9. Related Works

  • FLAT targets dependent MatMul-Softmax-MatMul fusion, which prior head-fusion work does not explore.The cited head-fusion approaches add a loop over independent heads, a pattern already captured by FLEX.Source: §9 · p. 14
  • Algorithmic attention optimizations include compression, sparsity, low-rank and kernel methods, and can affect model quality; they are orthogonal to FLAT.The passage says FLAT can be used alongside these techniques on DNN accelerators.Source: §9 · p. 14
  • GNN fusion dataflows optimize matrix-matrix operations, whereas attention fusion handles matrix-Softmax-matrix operations and a quadratic intermediate activation matrix.The passage also distinguishes the GNN activation-weight and activation-activation computations from attention’s two activation-activation computations.Source: §9 · p. 14
  • A3 and ELSA use approximate computation, and Sanger predicts and rearranges sparse attention matrices; these techniques trade performance against model quality.FLAT is described instead as a generic dataflow technique that does not impact model quality and can be leveraged on most existing accelerators.Source: §9 · p. 14
  • ML compilers use fusion in a limited fashion, typically combining matrix operators with element-wise operators.The passage contrasts this with broader fusion targets discussed for attention.Source: §9 · p. 14

10. Conclusion

FLAT addresses long-sequence attention’s low data reuse and quadratic intermediate-memory growth with fusion, interleaved execution, and tiling. The dataflow is designed to reduce bandwidth demands and scale to long sequences.

10. Conclusion

  • FLAT combines inter-operator fusion, interleaved execution, and efficient tiling to optimize attention layers for long sequences.The approach targets low reuse in certain attention operators and quadratic growth of the intermediate memory footprint.Source: §10 · p. 15
  • The dataflow is designed to increase operational intensity and compute utilization while reducing off-chip bandwidth requirements.Source: §10 · p. 15
  • FLAT is presented as a way to scale attention-based models to long sequence lengths.Source: §10 · p. 15

Related papers

Loading 2107.06419v7…