PaperlayerView exact revision on arXiv

Source-linked AI summary

Reorienting the Vector Axis: Restructuring FlashAttention for Long-Vector Processors

Emmanuel Irabor, Marc Casas

arXiv:2610.29489v1cs.ARcs.DCcs.PF

TL;DR

FlashAttention and most of its subsequent engineering target GPUs, leaving its traversal largely fixed despite the importance of vector-register utilization on long-vector processors. This paper reorients the recurrence and moves exponential coefficients into scalar registers; together, the techniques make the kernel 1.84×–2.08× faster on EPAC silicon at a head dimension of 64, while the same-design FPGA ranks the changes identically.

  • Problem

    FlashAttention, [19] and almost all subsequent engineering target GPUs, while long-vector processors make traversal orientation important to register utilization.

    Source: §I · p. 1

  • Method

    The paper transposes the score tile and running output to make queries the contiguous axis, and holds polynomial exponential coefficients in scalar registers.

    Source: §I · pp. 1–2

  • Results

    1.84×–2.08× faster on EPAC silicon at a head dimension of 64 with the reorientation and scalar-coefficient techniques; an FPGA of the same design ranks every change identically.

    Source: §X · pp. 9–10

  • Takeaways & Limitations

    On processors with 512-element registers, the recurrence’s orientation determines vector length across kernel regions and whether softmax reductions become elementwise accumulations.

    Source: B. Softmax becomes vertical · p. 6, §X · p. 9

    The matrix-multiplication routine still leaves gains on the table through register blocking and an inner kernel tuned for a decoupled unit without chaining.

    Source: §VIII · p. 9

Abstract

from arXiv · show

A Transformer is built around attention, whose cost grows with the square of the sequence length and dominates the layer on long sequences. FlashAttention made that cost manageable by tiling the attention matrix and updating the softmax as it goes, but almost all work since has targeted GPUs. On long-vector processors one vector register can hold up to 512 single-precision numbers. The savings of a vector instruction are proportional to how much of the register it uses, so the choice of which loop to vectorize is a first-order decision. We port PyTorch's CPU FlashAttention kernel to a RISC-V CPU with 2KB vector registers and find that it reaches 16.2% of peak floating-point throughput. The kernel walks each tile one query row at a time: the product against V is vectorized along the 64-element head dimension, the softmax's per-row maximum and sum are reductions that collapse a full register to a scalar, and the software exponential's polynomial coefficients compete with tile data for vector registers and cause spilling. We propose two techniques. The first transposes the score tile and the running output so that queries lie along the contiguous axis. The algorithm and its values are unchanged. A single parameter, the number of queries in a tile, now sets the vector length of every part of the kernel and can be matched to the machine. Both reductions also disappear. The second moves the polynomial's coefficients into scalar registers, removing the spills and cutting the softmax's vector memory instructions by 78.7%. On silicon the two together cut whole-kernel cycles by 1.84 to 2.08x at a head dimension of 64, raising floating-point utilization from 16.2% to 29.7% of peak. An FPGA of the same design ranks the changes identically and scores them higher, at 2.71x.

I. INTRODUCTION

FlashAttention, [19] and most subsequent engineering target GPUs, while long-vector processors make loop orientation a consequential choice because utilization depends on how much of a register each instruction uses. The paper transposes the recurrence’s score tile and running output so one query-block parameter controls vector lengths across the kernel, and moves exponential coefficients into scalar registers.

I. INTRODUCTION

  • FlashAttention, [19], the standard attention implementation, and almost all subsequent engineering target GPUs; on 512-element registers, a loop over a head dimension of 64 uses only one eighth of capacity.No instruction scheduler can recover the remaining throughput.Source: §I · p. 1
  • Transposing the score tile and running output makes the query axis contiguous, so one query-tile parameter sets every region’s vector length and turns per-query reductions into vectors.The paper also proposes an instruction-count model for when reorientation wins.Source: §I · p. 2
  • Holding the software exponential’s polynomial coefficients in scalar registers removes spill traffic that a 32-register file would otherwise make unavoidable.The paper also characterizes kernel regions on a 16,384-bit RISC-V processor, separating vector-length limits from access-pattern limits.Source: §I · p. 2

A. Attention and FlashAttention tiling

  • FlashAttention avoids explicitly materializing the Θ(N^2)-memory attention score matrix by tiling queries and keys and carrying a running maximum and sum.The tile shape is an implementation parameter, not part of the attention mathematics.Source: A. Attention and FlashAttention tiling · p. 2

B. Long-vector processors

Long-vector processors can issue instructions over hundreds of elements, but utilization falls when a kernel’s vector length uses only a fraction of register capacity. The stock-kernel characterization distinguishes an access-pattern limit in the score product from a vector-length limit in the value product.

B. Long-vector processors

  • A vector instruction’s gain is proportional to how much of the register it uses, making low active vector lengths costly on long-vector processors.In the studied FP32 setup, the maximum vector length is 512 elements.Source: B. Long-vector processors · p. 2

III. CHARACTERISING THE STOCK KERNEL

  • The score product is limited by its access pattern, whereas the value product is limited by vector length; each requires a different remedy.This region-level characterization separates limits that a whole-kernel measurement cannot distinguish.Source: §III · p. 2
  • The study ports PyTorch’s FlashAttention-2 CPU forward pass to RISC-V vector intrinsics, retaining its traversal, FP32 inner kernels, polynomial exponential, and vector reductions.The later FlashAttention variants discussed here rely on GPU-specific mechanisms absent from a vector unit.Source: §III · p. 2

A. Where the cycles go

The stock kernel’s cycle budget is distributed across softmax, value-product, and score-product work, and shortening the key block sharply increases whole-kernel cycles. The score product reaches full vector width but is constrained by strided operand access, which packing improves.

A. Where the cycles go

  • 5.17 FLOP/cycle against a peak of 32 is the stock kernel’s whole-kernel throughput.The reported cycle budget assigns 38% to softmax, 36% to the value product, and 21% to the score product.Source: A. Where the cycles go · pp. 2–3
  • 3.32× more whole-kernel cycles result when Bkv shrinks from 512 to 64 at Bq = 128, with softmax accounting for 93% of the increase.The passage attributes the increase to issuing the same work at shorter vector length, rather than doing more work.Source: A. Where the cycles go · p. 3

B. Wide vectors, strided operands: the score product

  • At Bq = 128 and Bkv = 512, gemm_qk averages vl = 512, but 88.9% of its vector memory instructions are strided because K is stored with d contiguous.At d = 64, the contraction-axis walk steps by a 256-byte row pitch; strided accesses cost 128–512 cycles versus 32–128 for unit-stride accesses.Source: B. Wide vectors, strided operands: the… · p. 3
  • Packing the key block leaves arithmetic and vector length unchanged while replacing almost every strided load with a unit-stride one, making gemm_qk 2.38× faster.The matrix-multiplication routine performs packing in every arm measured in the paper.Source: B. Wide vectors, strided operands: the… · p. 3

C. Contiguous operands, short vectors: the value product

The value product uses contiguous operands but vectors only across the 64-element head dimension, leaving 448 of 512 register elements unused. The softmax’s software exponential uses range reduction and a degree-6 polynomial, whose coefficients prove costly on this machine.

C. Contiguous operands, short vectors: the value product

  • gemm_pv requests 64 elements from 512-element registers, so its vectors use one-eighth of the architectural width regardless of block size.Both operands are contiguous along the product’s sweep axis, so packing does not address this limit.Source: C. Contiguous operands, short vectors:… · p. 3

IV. A LONG-VECTOR SOFTMAX

  • The softmax changes move polynomial coefficients into scalar registers and fold three passes over each score row into one.The software exponential’s coefficients, rather than its instructions, are costly on this machine.Source: §IV · p. 3, A. How the exponential is computed · p. 4

A. How the exponential is computed

  • Range reduction splits x into k ln 2 and a remainder, then a degree-6 polynomial approximates the remainder’s exponential through Horner’s multiply–add chain.The 2^k factor is constructed with integer work by placing k’s biased value in the FP32 exponent field.Source: A. How the exponential is computed · p. 3

B. The constants were vectors

Scalar coefficients and a fused score-row traversal reduce softmax-region traffic and cycles. The coefficient change removes spill-related memory instructions; the fused traversal also avoids reloading a row that fits in one vector register.

B. The constants were vectors

  • Seven vector-register coefficients compete with tile data and running state, causing spills; scalar coefficients preserve accuracy at one extra instruction per polynomial step.The scalar operand of the vector multiply–add cannot serve as Horner’s addend, so the scalar-file encoding uses an extra instruction.Source: B. The constants were vectors · p. 4
  • 78.7% fewer vector memory instructions and 1.54× faster cycles follow the scalar-coefficient change in the softmax region.Moved bytes fall by a factor of five; Table I covers Bq = 128, Bkv = 512, N = 1024, d = 64.Source: B. The constants were vectors · p. 4, Table I · p. 4

C. A fused traversal removes what remains

D. The ceiling

The ceiling analysis finds limited headroom in the softmax after the two changes, while the reoriented formulation changes vectorization without changing the algorithm’s values. Its performance depends on choosing a query block that matches the register width.

D. The ceiling

  • The two changes reduce softmax cycles by 2.49× and whole-kernel cycles by 1.27×, leaving at most 1.24× from eliminating softmax entirely.Softmax is then one fifth of the kernel budget.Source: D. The ceiling · p. 4
  • Fusing three passes requires Bkv ≤ VLMAX; the saving disappears above that boundary, while the transposed form has no such precondition.The natural layout’s measured optimum uses Bkv = 512, exactly VLMAX on this machine.Source: D. The ceiling · p. 4

V. REORIENTING THE VECTOR AXIS

  • The section presents a transposed recurrence, a structurally changed softmax, and an instruction-count model for when the transpose wins.The transposed form’s performance is analyzed separately from the softmax work.Source: §V · p. 4

A. The transposed formulation

  • Transposing the tile makes Bq set every region’s vector length, replaces the two reductions with elementwise operations, and leaves the algorithm’s values unchanged.The output transpose adds 9–13% to normalization but under 0.5% to the kernel.Source: A. The transposed formulation · p. 5
  • At Bq = 128 the reoriented kernel is slower than the fused form, but at Bq = 512 it is 1.43× faster.The register-width tile is the condition for the reported win.Source: A. The transposed formulation · p. 5

B. Softmax becomes vertical

The transposed softmax keeps per-query maximum and sum as vectors and accumulates elementwise down key rows. This removes reductions and scalar moves but reads scores twice.

B. Softmax becomes vertical

  • The transposed softmax makes two passes over key rows, keeping maximum and sum as query-axis vectors and replacing reductions with elementwise accumulation.One vector-length configuration covers the tile, with the state held in vector registers across key rows.Source: B. Softmax becomes vertical · p. 6
  • The reoriented form removes two reductions and their vector-to-scalar moves, while computing each ρi with one vector exponential per query block.The natural form instead computes separate scalar exponentials per query row.Source: B. Softmax becomes vertical · p. 6
  • The two-pass traversal reads each score twice, yielding 10% more memory instructions in the softmax region at each form’s best tile.Source: B. Softmax becomes vertical · p. 6

C. When the transpose wins

The transpose wins when its saved reduction and move costs outweigh extra unit-stride accesses, but the comparison depends on each region’s vector axis and assumptions about cache residency.

C. When the transpose wins

  • The transpose changes vector-instruction counts rather than arithmetic, with each region compared according to the number of rows it walks and instructions issued per row.Source: C. When the transpose wins · p. 6
  • For the value product, the transpose wins when Bq > d, and its gain falls as the head dimension grows.The comparison turns on whether the query block exceeds the head dimension.Source: C. When the transpose wins · p. 6
  • For the score product, the instruction ratio is Bq/Bkv, so the trade is even when the two block lengths are equal.Source: C. When the transpose wins · p. 6
  • Softmax trades reductions and scalar moves for two extra memory instructions per row, while both orientations perform the same elementwise work per row.The natural orientation reduces maximum and sum; the transposed orientation uses vector maximum and addition down key rows.Source: C. When the transpose wins · p. 6
  • The model prices the transpose’s row trade as at most 256 cycles given up versus at least 320 saved, assuming the second read hits in cache.The paper says its timing results are not derived from this model.Source: C. When the transpose wins · p. 6

VI. METHODOLOGY

The methodology measures FP32 kernel behavior on an EPAC chip and a same-design FPGA, with hardware counters and emulation serving distinct roles. Results are reported for one core and one attention head unless stated otherwise.

A. Target platform

  • EPAC is an in-order RISC-V scalar core with a decoupled vector unit and 16,384-bit registers; the kernel design depends on its listed parameters.Four of those parameters shape the analysis that follows.Source: A. Target platform · p. 6, Table III · p. 7
  • The vector unit uses RVV 0.7.1, but the identified loop-length limits do not depend on differences from ratified RVV 1.0.The processor has 32 architectural and 40 physical registers, and no hardware exponential.Source: A. Target platform · p. 7

B. Measurement targets and instrumentation

  • Measurements use the EPAC chip at 1 GHz and same-design FPGA at 50 MHz; each result is reported separately, not averaged across targets.Their sustained memory bandwidth differs, and every arm is vectorized rather than compared with scalar code.Source: B. Measurement targets and… · p. 7
  • The stock kernel reaches 16.2% of the machine’s 32 FP32 FLOP/cycle peak; measurements use one core, one head, d = 64, and N = 1024 unless stated.Source: B. Measurement targets and… · p. 7
  • RAVE emulation reports instruction mixes and effective vector lengths but models register width only, while hardware tracing supplies regional cycles and retired instructions.The emulation does not model caches, pipeline, or latencies.Source: B. Measurement targets and… · p. 7

C. Region attribution

The reoriented kernel improves end-to-end performance, with gains tied to vector length and softmax work. Region-level measurements distinguish those benefits from score-product behavior and show the same change ranking on both targets.

C. Region attribution

  • The region measurements attribute cycles to four named regions, with three running per tile and output_scaling once per query block.They use FP32 on one core at d = 64 and N = 1024 unless stated; cycles outside the four regions are under 1% in every arm.Source: C. Region attribution · p. 7

A. End-to-end performance

  • 1.84× over the tuned stock kernel and 1.43× over the fused kernel, the reoriented kernel raises utilization from 16.2% to 29.7% of peak.The reoriented arm is fastest at Bq = 512, Bkv = 128; the stock and scalar-coefficient arms are fastest at Bq = 128, Bkv = 512.Source: A. End-to-end performance · p. 7, Figure 3 · p. 8

B. The two targets rank the changes identically

  • Both targets rank the changes identically, while every improvement is larger on the FPGA than on the EPAC chip.Table V compares each arm at its own fastest tile; its speed-up factors, not cycle counts, can be compared across the targets.Source: B. The two targets rank the changes… · p. 8, Table V · p. 8
  • 2.01× faster against the fused kernel at their respective best tiles, the reoriented value product gains from vector length set by Bq rather than d.This region carries most of the whole-kernel gain.Source: C. Where the gain comes from · p. 8
  • The reoriented softmax is 1.37× faster at each arm’s best tile, while the score product stays within 1% of even at the optima.Reorientation removes softmax reductions and vector-to-scalar moves; the score product issues the same instructions per row in both orientations.Source: C. Where the gain comes from · p. 8
  • 3.42× versus 2.42×, the stock kernel’s worst-to-best tile ratio is higher than the reoriented kernel’s.The natural layout takes vector length from Bkv, while the reoriented layout takes it from Bq, reversing which block parameter favors performance.Source: C. Where the gain comes from · p. 9

D. How the gain varies

The section tests how the gain varies with sequence length and head dimension. The speed-up persists across measured sequence lengths, while a larger head dimension reduces the value-product advantage.

D. How the gain varies

  • 1.84 to 2.08×, the speed-up does not decay across the measured sequence-length range, and both kernels retain the same optimal tiles.The stated explanation is that the speed-up ratios do not contain N.Source: D. How the gain varies · p. 9
  • At d = 128, the reoriented value-product advantage falls from 2.01× to 1.26×.The head dimension affects this region through the Bq/d ratio.Source: D. How the gain varies · p. 9

VIII. DISCUSSION

The paper’s conclusion is conditional on long-vector machines and long query blocks, and it distinguishes that setting from autoregressive decode. It relates this approach to other vector-attention designs and identifies remaining headroom in the score product.

VIII. DISCUSSION

  • The benefit requires long vector registers and a long query block; the paper scopes it to prefill and training, not autoregressive decode.For decode, where a query block holds one row or one query group, the natural orientation is the right one.Source: §VIII · p. 9
  • The reoriented kernel reaches 29.7% of peak, leaving a factor of three unclaimed; the score product is the next optimization target.The softmax is a quarter of the budget, so removing it entirely would buy at most 1.33×.Source: §VIII · p. 9

IX. RELATED WORK

  • Unlike FlashAttention-V [23], which packs several heads to extend to decode, query reorientation needs no extra heads but requires a long query block.The paper presents the two approaches as applying to different cases.Source: §IX · p. 9

X. CONCLUSION

  • With exponential coefficients moved out of spilling vector registers, the reoriented kernel is 1.84×–2.08× faster on EPAC silicon, and the same-design FPGA ranks every change identically.Reorientation transposes the score tile and running output so queries are contiguous, putting every region under one query-tile parameter and replacing per-row reductions with elementwise accumulations.Source: §X · pp. 9–10

Related papers

Loading 2610.29489v1…