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

AI Analysis

Abstract

Key Highlights

• Initial FlashAttention-2 kernel on a RISC-V EPAC chip achieves only 16.2% of peak floating-point utilization.
• Identified performance bottlenecks include suboptimal vector instruction lengths in the value product and register spilling in the softmax's software exponential.
• Proposed reorientation of the recurrence transposes score and output matrices, aligning vector lengths of all kernel regions with the query block size.
• Moving exponential polynomial coefficients to scalar registers reduces softmax vector memory instructions by 78.7% and moved bytes by a factor of five.
• The combined techniques yield a 1.84x to 2.08x speed-up in whole-kernel cycles on EPAC silicon, increasing FP utilization to 29.7% of peak.
• The reoriented softmax transforms per-row reductions into elementwise vector accumulations, eliminating reduction instructions.

Technical Details

• The target platform is the EPAC RISC-V CPU, featuring 16,384-bit (2KB) vector registers capable of holding 512 FP32 elements (vlmax).
• The stock PyTorch CPU FlashAttention-2 kernel's value product (gemm_pv) vectorizes along the head dimension (d=64), utilizing only 64 of 512 available elements.
• The software exponential in the softmax region uses a degree-6 polynomial whose coefficients, when stored in vector registers, cause spilling due to the limited 32 architectural vector registers.
• Reorienting the vector axis involves transposing the score tile (S_ij) and the running output (O_i), making the query axis (B_q) the contiguous axis for all operations.
• This reorientation allows the vector length of every kernel region to be controlled by a single parameter, B_q, which can be set to vlmax (512).
• Storing exponential coefficients in scalar registers, rather than vector registers, eliminates spill traffic, reducing vector memory instructions in the softmax by 78.7%.
• A fused traversal in the softmax, enabled by a row fitting in a single vector register, further reduces memory instructions by 30.9% by processing maximum, exponential, and sum in one pass.
• The score product (gemm_qk) in the natural layout is limited by strided memory access patterns when reading the transposed key block (K_j), a problem addressed by packing (a pre-existing technique).

Implications

• The study highlights that the choice of vector axis orientation is a first-order architectural decision for FlashAttention performance on long-vector processors, unlike on short-vector or GPU architectures.
• Existing FlashAttention implementations, optimized for GPUs or short-vector CPUs, are suboptimal for long-vector RISC-V processors and require significant restructuring.
• The proposed software optimizations provide a clear path to significantly improve FlashAttention performance on current and future long-vector RISC-V architectures.
• Kernel designers for long-vector machines should prioritize aligning data access patterns and vector lengths with the architecture's wide register capabilities.
• The findings are particularly relevant for prefill and training phases of Transformer models, where the query block size (B_q) can be set to match the maximum vector length (vlmax).

Limitations

• The study does not compare performance against GPUs or short-vector CPUs; the reported gains are conditional on the use of a long-vector machine and a query block longer than the head dimension.
• The fused traversal optimization in the softmax is only effective when the key block size (B_kv) is less than or equal to the maximum vector length (vlmax).
• The reorientation technique is most advantageous when the query block (B_q) is longer than the head dimension (d) and is set to the register width (vlmax).
• The reoriented kernel is not suitable for autoregressive decode scenarios, where the query block typically consists of a single row or query group.
• A factor of three in peak performance remains unclaimed, with the score product (gemm_qk) being the largest remaining bottleneck, requiring further optimization within the matrix-multiplication routine itself (e.g., register blocking, tuning for decoupled units without chaining).
• The evaluation uses RVV 0.7.1, though the authors state that the instructions used have 1.0 counterparts and the identified limits are not affected by this difference.
• Instruction latency values are documented by the vendor only at vl=512 FP32 elements, which may simplify the analysis for other vector lengths.