<figure><img src="image_1.png" /></figure>

HAL
open science

# ACCELERATING FUSED WMMA OPERATIONS
VIA IN-REGISTER DATA PERMUTATION

Tommy Morales

► To cite this version:

Tommy Morales. ACCELERATING FUSED WMMA OPERATIONS VIA IN-REGISTER DATA PERMUTATION. 2025. hal-05230428

HAL Id: hal-05230428

https://hal.science/hal-05230428v1

Preprint submitted on 29 Aug 2025

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

Public Domain
ACCELERATING FUSED WMMA OPERATIONS VIA IN-REGISTER
DATA PERMUTATION

Tommy Morales

ENSIIE

Évry, France

tommy.morales@ensiie.eu

ABSTRACT

The performance of GPU kernels for AI, such as FlashAttention, is increasingly limited by on-chip data movement after achieving I/O-awareness. A critical bottleneck in fused operations using the WMMA API is the costly round-trip to shared memory required for post-matmul reductions, which breaks data locality and introduces synchronization overhead. To address this, we propose a novel technique that performs reductions entirely within the GPU’s register file. By reverse-engineering the swizzled data layout of WMMA accumulator fragments, we devised a set of in-register permutations that realign data for efficient intra-warp reduction using shfl_sync instructions. A proof-of-concept implementation on a fused matrix-multiply-max operation demonstrates up to a 7x speedup across modern NVIDIA architectures (Blackwell and Hopper) over the standard shared memory-based approach. This result indicates that direct manipulation of fragment data layouts is a powerful new vector for micro-architectural optimization beyond the current state-of-the-art.

# 1 Introduction

Attention-based Transformer architectures remain the dominant choice for Large Language Models (LLMs). While variations like Grouped-Query Attention (GQA) have emerged, most designs stem from the Multi-Head Attention (MHA) mechanism introduced in the seminal "Attention Is All You Need" paper. Optimizing the computational cost of the attention forward pass is critical for both training and inference. To address this, I/O-aware algorithms have become the standard since the introduction of FlashAttention, which minimizes expensive data transfers between GPU high-bandwidth memory (HBM) and on-chip SRAM.

With state-of-the-art implementations like FlashAttention achieving near-optimal occupancy and effectively eliminating the I/O bottleneck, performance improvements now depend on on-chip compute efficiency. This shifts focus to optimizing calculations performed by Tensor Cores via the WMMA API. However, a common pattern in fused kernels-performing a matrix multiplication followed by a reduction or element-wise operation (e.g., softmax)-introduces a new bottleneck: the data round-trip through shared memory. Current practice requires storing the accumulator fragment from registers to shared memory using `wmma::store_matrix_sync`, synchronizing via `__syncthreads()`, and then reloading the data into registers for further processing. This approach incurs significant latency (see Section 4.3) and breaks data locality, forcing subsequent operations to re-fetch data from the slower shared memory cache instead of keeping it within the register file.

In this paper, we propose a technique to bypass this shared memory round-trip by performing post-matmul operations directly in registers. We first detail our analysis of the WMMA fragment’s internal data layout, revealing the mapping between matrix elements and specific thread registers. We then exploit this knowledge to develop an in-register data permutation algorithm. This allows us to realign data for efficient parallel reduction using shfl_sync instructions. Our benchmarks, conducted on a fused matrix-multiply-max kernel, a core operation in algorithms like FlashAttention, demonstrate an 7x speedup compared to the standard shared memory-based approach.
## 1.1 Our Contributions

The contributions of this paper are fourfold:

1. We present an experimental methodology for reverse-engineering the undocumented, architecture-specific data layout of WMMA fragments.

2. We introduce a novel in-register permutation algorithm that eliminates the costly shared memory round-trip for post-matmul reduction operations.

3. We demonstrate the effectiveness of this technique through a rigorous, comparative analysis on two state-of-the-art GPU architectures (NVIDIA Hopper and Blackwell), achieving up to an **7x speedup** and an **7x improvement in energy efficiency** on a key micro-benchmark.

4. We provide a complete algorithm (Section 5) and a corresponding open-source implementation detailing how our technique can replace the standard softmax computation within a full FlashAttention kernel, proving its applicability to state-of-the-art models.

The remainder of this paper is structured as follows. Section 2 reviews the state-of-the-art in attention optimization. Section 3 details our reverse-engineering methodology for the WMMA fragment layout and presents our novel in-register permutation algorithm. Section 4 provides a comprehensive experimental validation of our technique across two GPU architectures. Finally, Section 5 demonstrates how our method can be integrated as a core component within a full FlashAttention kernel.

## 2 Related Work

The paradigm for sequence modeling was fundamentally altered by [1] with the introduction of the Transformer architecture. By replacing recurrent neural networks (RNNs) with a purely attention-based mechanism, their work overcame the limitations of sequential processing and enabled massive parallelization on modern hardware like GPUs. The introduction of Multi-Head Attention (MHA) allowed the model to jointly attend to information from different representation subspaces at different positions, forming the computational core of most modern Large Language Models (LLMs).

While the Transformer architecture is highly parallelizable, its performance on physical hardware is overwhelmingly dominated by data movement, not arithmetic complexity. This reflects a broader shift in high-performance computing, where the cost of fetching data from off-chip DRAM is orders of magnitude higher than a floating-point operation, rendering classical complexity analysis (e.g., [2]) less relevant for practical performance. This has led to the development of **I/O-aware algorithms**. The seminal work in this area for Transformers is FlashAttention by [3], which redesigns the attention algorithm to minimize data transfers between the GPU's High-Bandwidth Memory (HBM) and its on-chip SRAM. By using techniques such as kernel fusion and tiling, FlashAttention computes the exact attention output with significantly fewer memory reads/writes. Subsequent versions [4] [5] have further refined this approach.

With the I/O bottleneck largely addressed by these state-of-the-art methods, the new performance frontier lies in optimizing computation and data movement *on-chip*. Modern NVIDIA GPUs provide specialized hardware units, Tensor Cores, to accelerate matrix multiplications. These are typically accessed via high-level abstractions like the WMMA (Warp-Level Matrix-Multiply-Accumulate) API [6]. However, while powerful, the WMMA API treats the physical layout of data within the GPU's register file as an opaque "black box". This abstraction, designed for ease of use, can introduce subtle but significant on-chip bottlenecks, such as the costly round-trip to shared memory for post-matmul operations. This paper addresses this specific on-chip bottleneck, exploring optimizations that lie beyond the current SOTA by manipulating the fragment data layout directly.

It is important to note that I/O-aware attention is not a monolithic field. While FlashAttention and its derivatives optimize the standard Multi-Head/Grouped-Query Attention (MHA/GQA) mechanism, a parallel and architecturally distinct approach has emerged with Multi-Head Latent Attention (MLA), introduced by [7]. Instead of optimizing the computation order, MLA fundamentally alters the architecture by compressing the Key-Value cache itself into a low-rank latent representation. Recent work such as TransMLA [8] has even demonstrated methods for converting pre-trained GQA models to the MLA format. While our direct application focuses on the FlashAttention-style softmax, the underlying technique of in-register data manipulation is a kernel-level optimization that addresses a bottleneck fundamental to any algorithm leveraging the WMMA API, including potential future MLA implementations.
# 3 WMMA Fragments Manipulation

Fragments are the form in which the WMMA API stores matrices to perform operations - (mainly Matrix Multiplications via `mma_synch()`) - on tensor cores. In order to manipulate and reorganize the data layout, reverse engineering this black box was an essential first step.

## 3.1 Extracting the data layout

To manipulate data within a WMMA fragment, it is imperative to understand how the 256 elements of a 16x16 matrix are distributed across the 8 registers of each of the 32 threads in a warp. As this mapping is undocumented, we developed a micro-kernel to chart it experimentally. The principle is to create a "tracer matrix" where the value of each element encodes its own (row, col) coordinate. This matrix is then passed through the WMMA pipeline via an identity multiplication, preserving the tracer values. The resulting fragment is dumped to global memory, allowing a host-side script to reconstruct the mapping. The overall process is described in Algorithm 1.

<caption>**Algorithm 1 WMMA Fragment Layout Reverse-Engineering**</caption>

1: **Input:** None
2: **Output:** A 16x16 table Mapping(row)[col] = (lane, frag_idx)
3: **procedure** DEVICE KERNEL(output_buffer)
4:     tracer_matrix ← 16x16 shared memory array
5:     lane ← thread's lane index (0-31)
6:     // Each thread in the block populates one element of the tracer matrix
7:     tracer_matrix[threadId.x][threadId.y][threadId.x] ← threadId.x * 100 + threadId.x
8:     __syncthreads()
9:     **if** lane < 32 **then**
10:         a_frag ← wmma::load_matrix_sync(tracer_matrix)
11:         id_frag ← Load Identity Matrix into a fragment
12:         acc_frag ← wmma::mma_sync(a_frag, id_frag)
13:         **for** i ← 0 to 7 **do**
14:             output_buffer[lane * 8 + i] ← acc_frag.x[i]
15:             **end for**
16:         **end if**
17:     **end procedure**
18: **procedure** HOST ANALYSIS(output_buffer)
19:     **for** lane ← 0 to 31 **do**
20:         **for** frag_idx ← 0 to 7 **do**
21:             value ← output_buffer[lane * 8 + frag_idx]
22:             row ← [value/100]
23:             col ← value%100
24:             Mapping(row)[col] ← (lane, frag_idx)
25:         **end for**
26:     **end for**
27:     **return** Mapping
28: **end procedure**

The mapping we get is detailed in Figure 1.
<table><thead><tr><th></th><th>C00</th><th>C01</th><th>C02</th><th>C03</th><th>C04</th><th>C05</th><th>C06</th><th>C07</th><th>C08</th><th>C09</th><th>C10</th><th>C11</th><th>C12</th><th>C13</th><th>C14</th><th>C15</th></tr></thead><tbody><tr><td>R00</td><td>0</td><td>0</td><td>1</td><td>1</td><td>2</td><td>2</td><td>3</td><td>3</td><td>0</td><td>0</td><td>1</td><td>1</td><td>2</td><td>2</td><td>3</td><td>3</td></tr><tr><td>R01</td><td>4</td><td>4</td><td>5</td><td>5</td><td>6</td><td>6</td><td>7</td><td>7</td><td>4</td><td>4</td><td>5</td><td>5</td><td>6</td><td>6</td><td>7</td><td>7</td></tr><tr><td>R02</td><td>8</td><td>8</td><td>9</td><td>9</td><td>10</td><td>10</td><td>11</td><td>11</td><td>8</td><td>8</td><td>9</td><td>9</td><td>10</td><td>10</td><td>11</td><td>11</td></tr><tr><td>R03</td><td>12</td><td>12</td><td>13</td><td>13</td><td>14</td><td>14</td><td>15</td><td>15</td><td>12</td><td>12</td><td>13</td><td>13</td><td>14</td><td>14</td><td>15</td><td>15</td></tr><tr><td>R04</td><td>16</td><td>16</td><td>17</td><td>17</td><td>18</td><td>18</td><td>19</td><td>19</td><td>16</td><td>16</td><td>17</td><td>17</td><td>18</td><td>18</td><td>19</td><td>19</td></tr><tr><td>R05</td><td>20</td><td>20</td><td>21</td><td>21</td><td>22</td><td>22</td><td>23</td><td>23</td><td>20</td><td>20</td><td>21</td><td>21</td><td>22</td><td>22</td><td>23</td><td>23</td></tr><tr><td>R06</td><td>24</td><td>24</td><td>25</td><td>25</td><td>26</td><td>26</td><td>27</td><td>27</td><td>24</td><td>24</td><td>25</td><td>25</td><td>26</td><td>26</td><td>27</td><td>27</td></tr><tr><td>R07</td><td>28</td><td>28</td><td>29</td><td>29</td><td>30</td><td>30</td><td>31</td><td>31</td><td>28</td><td>28</td><td>29</td><td>29</td><td>30</td><td>30</td><td>31</td><td>31</td></tr><tr><td>R08</td><td>0</td><td>0</td><td>1</td><td>1</td><td>2</td><td>2</td><td>3</td><td>3</td><td>0</td><td>0</td><td>1</td><td>1</td><td>2</td><td>2</td><td>3</td><td>3</td></tr><tr><td>R09</td><td>4</td><td>4</td><td>5</td><td>5</td><td>6</td><td>6</td><td>7</td><td>7</td><td>4</td><td>4</td><td>5</td><td>5</td><td>6</td><td>6</td><td>7</td><td>7</td></tr><tr><td>...</td><td>...</td><td>...</td><td>...</td><td>...</td><td>...</td><td>...</td><td>...</td><td>...</td><td>...</td><td>...</td><td>...</td><td>...</td><td>...</td><td>...</td><td>...</td></tr></tbody></table>

Figure 1: Experimentally discovered data layout on an RTX 5080 GPU. The mapping of *lane* to matrix coordinates exhibits a repeating pattern every 8 rows and 8 columns.

## 3.2 Reorganising the data

While a full de-swizzle of the data might seem intuitive, our analysis shows it is suboptimal. A more cycle-efficient solution consists of a partial in-register permutation combined with carefully constructed masks for the reduction phase. As conceptually illustrated in Figure 2, a minimal swap operation de-interleaves the data within each thread's registers. This consolidates all elements belonging to a single logical matrix row into contiguous registers within a thread sub-group, making the data ready for an efficient parallel reduction. This permutation, requiring only a few clock cycles, is detailed in Section 3.3.

However, in a complete kernel implementation like FlashAttention, it is often necessary to store thread-local accumulators (e.g., the Oaccum matrix) back to shared memory and later on HBM. For this store operation to be correct and efficient, we must use the forward mapping formulas derived from our reverse-engineering analysis. These formulas allow us to calculate the correct global memory address for each element stored in a thread's registers.

<table><caption>LANE 0 REGISTERS: BEFORE vs AFTER PERMUTATION</caption><thead><tr><th></th><th>BEFORE<br>(Data for a single row is scattered)</th><th>AFTER<br>(Data is consolidated)</th></tr></thead><tbody><tr><td>x[0,1]</td><td>ROW 0</td><td>ROW 0</td></tr><tr><td>x[2,3]</td><td>ROW 8</td><td>SWAP</td></tr><tr><td>x[4,5]</td><td>ROW 0</td><td>SWAP</td></tr><tr><td>x[6,7]</td><td>ROW 8</td><td>SWAP</td></tr></tbody></table>

Figure 2: Conceptual view of the in-register permutation for a single thread (e.g., lane=0). The swap operation reconsolidates scattered data segments from a single logical matrix row (e.g., ROW 0) into contiguous registers, making subsequent parallel reduction feasible.

The mapping from a thread's lane ID (*lane*, where *lane* ∈ [0, 31]) and its register index (*i*, where *i* ∈ [0, 7]) to a matrix coordinate (*row*, *col*) can be described by the following formulas:

$$ \text{row} = (\text{lane}/4) + (((i \gg 1)&1) \times 8) \quad (1) $$

$$ \text{col} = ((\text{lane}\%4) \times 2) + (i\%2) + ((i \gg 2) \times 8) \quad (2) $$

Algorithm 2 describes how these formulas are used to write a thread-local array (*local_data*) to its correct, swizzled position in a 2D global output matrix.
**Algorithm 2** Storing Thread-Local Data using Forward Mapping

1: **Input:** A thread-local array `local_data[0..7]`, global output matrix Out

2: **Context:** Executed by each thread *lane* in a warp

3: `block_row_start` ← Starting row index of the output tile in Out

4: `block_col_start` ← Starting col index of the output tile in Out

5: **for** *i* ← 0 to 7 **do**

6:   **//** Calculate target coordinate using forward mapping formulas

7:   `target_row` ← `block_row_start` + (`lane`/4) + (((`i` ≫ 1)&1) × 8)

8:   `target_col` ← `block_col_start` + ((`lane`%4) × 2) + (`i`%2) + ((`i` ≫ 2) × 8)

9:   **//** Write the local data to the correct global memory address

10: `Out[target_row][target_col]` ← `local_data[i]`

11: **end for**

### 3.3 In-Register Parallel Reduction

Following the in-register permutation, the data is logically aligned for efficient parallel processing. All data elements for two distinct matrix rows are now held within a single 4-thread sub-group (a "quad"). This enables any associative row-wise reduction (e.g., max, sum) to be performed entirely in registers using `shfl_xor_sync` instructions.

The process involves two stages, as described in Algorithm 3. First, each thread computes a local reduction on the data held in its registers. For example, after permutation, a thread's registers `x[0]-x[3]` contain all its data for the first set of 8 rows, and `x[4]-x[7]` for the second. The second stage is an intra-quad parallel reduction, where the 4 threads of a sub-group cooperatively reduce their local values. This is achieved using `shfl_xor_sync` with a dynamic mask to constrain the communication within the quad.

**Algorithm 3** Generic In-Register Parallel Reduction

1: **Input:** A permuted accumulator fragment `acc_frag`, an associative binary operator `⊕`

2: **Output:** For each quad, two reduced values `val_A` and `val_B`

3: **Context:** Executed by each thread *lane* in a warp

4: **//** Stage 1: Local reduction within each thread's registers

5: `local_A` ← `identityElement`

6: **for** *k* ← 0 to 3 **do**

7:   `local_A` ← `local_A` ⊕ `acc_frag.x[k]`

8: **end for**

9: `local_B` ← `identityElement`

10: **for** *k* ← 4 to 7 **do**

11:   `local_B` ← `local_B` ⊕ `acc_frag.x[k]`

12: **end for**

13: **//** Stage 2: Parallel reduction within each 4-thread quad

14: `mask` ← `0xF` << (((`lane`/4) * 4) ▶ e.g., `0xF` for quad 0, `0xF0` for quad 1)

15: `partner_A` ← `__shfl__xor_sync(`mask`, `local_A`, 1)

16: `val_A` ← `local_A` ⊕ `partner_A`

17: `partner_A` ← `__shfl__xor_sync(`mask`, `val_A`, 2)

18: `val_A` ← `val_A` ⊕ `partner_A`

19: **//** Repeat for B

20: `partner_B` ← `__shfl__xor_sync(`mask`, `local_B`, 1)

21: `val_B` ← `local_B` ⊕ `partner_B`

22: `partner_B` ← `__shfl__xor_sync(`mask`, `val_B`, 2)

23: `val_B` ← `val_B` ⊕ `partner_B`

24:

25: **if** `lane`%4 = 0 **then** ▶ Lead thread of the quad now holds the final reduced values

26: **return** `val_A`, `val_B`

27: **end if**

The mask `0xF` << (((`lane`/4) * 4) dynamically creates a 4-bit window corresponding to the thread's sub-group (e.g., 0x00000000F for threads 0-3, 0x000000F0 for threads 4-7). This allows all 8 sub-groups to perform their reductions in parallel without interfering with one another, achieving maximum data-level parallelism.
# 4 Experimental Validation

To quantify the performance benefits of our in-register technique, we conducted a micro-benchmark comparing it against the standard shared memory-based approach.

## 4.1 Setup

All experiments were performed on an NVIDIA RTX 5080 GPU (Blackwell architecture) and an NVIDIA H100 (Hopper architecture). The key technical specifications and software versions are detailed in Table 1.

<table><caption>Table 1: Experimental Setup Details.</caption><thead><tr><td>Component</td><td>Configuration</td></tr></thead><tbody><tr><td>Operating System</td><td>Ubuntu 24.04.3 LTS</td></tr><tr><td>CUDA Toolkit</td><td>12.9</td></tr><tr><td>NVIDIA Driver</td><td>575.57.08</td></tr></tbody></table>

Initial attempts looping over kernel launches and using an internal loop within the CUDA kernel were found to be unreliable due to launch overhead and aggressive compiler optimizations (Dead Code Elimination), respectively.

To mitigate these issues, we adopted a script-driven methodology. For latency measurements, an external bash script launches the executable, which performs a single kernel execution hundreds of times. The script discards the initial warm-up runs and averages the subsequent results, providing a stable measurement free from both compiler interference and cold-start anomalies. A similar approach was used for power measurement, where a sustained load was generated by the script to allow for stable readings from nvidia-smi. All benchmarking scripts are provided in the public repository to ensure full reproducibility.

## 4.2 Kernels Under Test

* **Baseline Kernel (Shared Memory):** This kernel represents the standard, textbook approach. It performs the matrix multiplication using wmma::mma_sync, stores the full accumulator fragment to a _shared_ memory array via wmma::store_matrix_sync, synchronizes the warp with _syncthreads(), and finally has each thread in the warp read its assigned row from shared memory to compute the maximum value.

* **Proposed Kernel (In-Register):** This kernel implements our technique. After the wmma::mma_sync, it performs the in-register permutation and subsequent intra-quad parallel reduction as described in Algorithm 3. The final result is written directly to global memory, completely bypassing the shared memory round-trip.

## 4.3 Results

Our experiments reveal a significant and robust advantage for the proposed in-register technique, demonstrating consistency across two distinct state-of-the-art GPU architectures.

The latency measurements, presented in Table 2, show a stable and substantial speedup of approximately 7x on both NVIDIA's Hopper and Blackwell architectures. This consistency validates the fundamental nature of our optimization, proving its effectiveness beyond a single hardware generation. The shared memory round-trip constitutes a major bottleneck on both platforms, and its elimination yields a massive performance gain in latency-bound scenarios.

<table><caption>Table 2: Performance Comparison Across GPU Architectures.</caption><thead><tr><td>Architecture</td><td>Kernel</td><td>Latency (ms)</td><td>Speedup</td></tr></thead><tbody><tr><td rowspan="2">NVIDIA H100 (Hopper)</td><td>Baseline (Shared Mem)</td><td>0.216</td><td>1.0x</td></tr><tr><td><strong>In-Register (Ours)</strong></td><td><strong>0.03</strong></td><td><strong>7x</strong></td></tr><tr><td rowspan="2">NVIDIA RTX 5080 (Blackwell)</td><td>Baseline (Shared Mem)</td><td>0.0825</td><td>1.0x</td></tr><tr><td><strong>In-Register (Ours)</strong></td><td><strong>0.0127</strong></td><td><strong>6.5x</strong></td></tr></tbody></table>

To understand the source of this acceleration, a micro-architectural analysis was conducted under a high-occupancy workload to simulate a realistic, throughput-bound scenario (Table 3). This analysis confirms the speedup is not due to a
change in occupancy—which is nearly identical for both kernels as it is limited by register count—but is instead rooted in the complete elimination of on-chip data movement and synchronization overhead. Our kernel reduces the shared memory footprint and the number of synchronization barriers to zero. It is noteworthy that while latency gains are dramatic, throughput under high occupancy is comparable. This illustrates a core principle of GPU architecture: latency hiding. While the GPU’s scheduler can effectively hide the baseline kernel’s shared memory latency by switching to other active warps, our technique fundamentally removes this latency source. This provides greater scheduling flexibility and reduces on-chip power draw, which directly translates to the significant performance improvements observed in our latency-bound benchmarks.

<table><caption>Table 3: Micro-architectural Analysis under High Occupancy. Occupancy data from Nsight Compute, other metrics from code analysis and `ptxas` compiler output.</caption><thead><tr><th>Metric</th><th>Architecture</th><th>Baseline Kernel</th><th>In-Register Kernel</th></tr></thead><tbody><tr><td rowspan="2">Achieved Occupancy (%)</td><td>H100</td><td>24.50</td><td>25.07</td></tr><tr><td>RTX 5080</td><td>44.10</td><td>43.77</td></tr><tr><td>Registers per Thread</td><td>H100 & 5080</td><td>22</td><td><b>22</b></td></tr><tr><td>Shared Memory Footprint</td><td>H100 & 5080</td><td>1024 bytes</td><td><b>0 bytes (-100%)</b></td></tr><tr><td>Synchronization Barriers</td><td>H100 & 5080</td><td>1</td><td><b>0 (-100%)</b></td></tr></tbody></table>

Finally, the energy efficiency gains detailed in Table 4 directly mirror the performance improvements. By completing the operation significantly faster at a nearly identical power draw, our in-register method consumes approximately 7x less energy per operation. This demonstrates that the technique not only accelerates computation but does so with a massive improvement in performance-per-watt, a critical consideration for large-scale AI workloads.

<table><caption>Table 4: Energy Efficiency Comparison. Energy/Op = Avg. Power (W) × Latency (s).</caption><thead><tr><th>Architecture</th><th>Kernel</th><th>Energy/Op (µJ)</th><th>Efficiency Gain</th></tr></thead><tbody><tr><td rowspan="2">NVIDIA H100</td><td>Baseline (Shared Mem)</td><td>18.45</td><td>1.0x</td></tr><tr><td><b>In-Register (Ours)</b></td><td><b>2.57</b></td><td><b>7.0x</b></td></tr><tr><td rowspan="2">NVIDIA RTX 5080</td><td>Baseline (Shared Mem)</td><td>5.04</td><td>1.0x</td></tr><tr><td><b>In-Register (Ours)</b></td><td><b>0.78</b></td><td><b>6.5x</b></td></tr></tbody></table>

# 5 Application to a Full FlashAttention Kernel

The true potential of our technique is realized when it is integrated as the core softmax engine within a full FlashAttention kernel. The standard implementation of the softmax step within FlashAttention’s main loop relies on the same shared memory round-trip that our micro-benchmark identified as a bottleneck. Our in-register method can replace this entire sequence.

Algorithm 4 details the complete, fused softmax calculation performed directly on the accumulator fragment `work_frag` resulting from the $S_{ij} = Q_i K_j^T$ matrix multiplication. The algorithm first applies the necessary scaling and permutation, then calculates the row-wise max, subtracts it, computes the exponentials, calculates the row-wise sum, and finally normalizes to get the probability matrix $P_{ij}$. All intermediate steps, including the update of running statistics, are performed entirely within registers, maximizing data locality and eliminating all `__syncthreads()` stalls from the softmax computation.
**Algorithm 4** In-Register Fused Softmax for FlashAttention Tile

1: **Input:** Accumulator fragment work_frag (containing $S_{ij}$), previous stats $m_{old}, l_{old}$

2: **Output:** Probability fragment p_frag (containing $P_{ij}$), updated stats $m_{new}, l_{new}$

3: **Context:** Executed by each thread lane in a warp after the $Q_i K_j^T$ mma_sync

4: // 1. Scale and Permute

5: work_frag ← work_frag / √$d_k$

6: Apply in-register permutation to work_frag (as per Section 2.2)

7: // 2. Compute current tile's row-wise max ($m_{ij}$)

8: $m_{local_A}, m_{local_B} \leftarrow$ Perform in-register Max reduction (Algorithm 3) on work_frag

9: // 3. Update global row-wise max

10: $m_{new} \leftarrow \max(m_{old}, m_{ij})$

11: $l_{old\_rescaled} \leftarrow l_{old} \times \exp(m_{old} - m_{new})$

12: // 4. Compute exponentials and current tile's row-wise sum ($l_{ij}$)

13: work_frag ← exp(work_frag - $m_{new}$)

14: $l_{local_A}, l_{local_B} \leftarrow$ Perform in-register Sum reduction (Algorithm 3) on work_frag

15: // 5. Update global sum normalizer

16: $l_{new} \leftarrow l_{old\_rescaled} + l_{ij}$

17: // 6. Final normalization to get $P_{ij}$

18: p_frag ← work_frag / $l_{new}$

19: Apply inverse permutation to p_frag to prepare for $P_{ij}V_j$ multiplication.

# 6 Conclusion

In this work, we identified and addressed a critical on-chip bottleneck in fused WMMA operations: the costly data round-trip through shared memory for post-matmul reductions. We have demonstrated that by reverse-engineering the opaque data layout of WMMA fragments, it is possible to design a novel in-register permutation algorithm that completely bypasses this bottleneck.

Our experimental results are conclusive. On a key micro-benchmark, our technique yields a speedup of up to **7x** and improves energy efficiency by up to **7x** compared to the standard shared memory-based approach on modern NVIDIA GPUs. We have provided a detailed algorithm showing how this technique can be directly integrated into the softmax computation of state-of-the-art models like FlashAttention, along with an open-source implementation for full reproducibility<fnref n="1" />.

The primary limitation of this method is its inherent dependency on architecture-specific data layouts, sacrificing portability for maximum performance. Nevertheless, this work establishes that direct, in-register manipulation of fragment data is a powerful and largely unexplored vector for optimization, offering substantial gains beyond the current I/O-aware paradigm.

# References

[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In *Advances in neural information processing systems*, pages 5998–6008, 2017.

[2] Volker Strassen. Gaussian elimination is not optimal. *Numerische mathematik*, 13(4):354–356, 1969.

[3] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In *Advances in Neural Information Processing Systems*, 2022.

[4] Tri Dao. FlashAttention-2: Faster attention with better parallelism and work partitioning. *arXiv preprint arXiv:2307.08691*, 2023.

[5] Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. FlashAttention-3: Fast and accurate attention with asynchrony and low-precision. *arXiv preprint arXiv:2407.08608*, 2024.

[6] NVIDIA Corporation. CUDA C++ Programming Guide: Warp-Level Matrix Operations. https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html.

[7] DeepSeek-AI. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model. *arXiv preprint arXiv:2405.04434*, 2024.

<footnote n="1">¹https://github.com/Z6BQw2/in-register-wmma-poc</footnote>
[8] Fanxu Meng, Pingzhi Tang, Xiaojuan Tang, Zengwei Yao, Xing Sun, and Muhan Zhang. TransMLA: MLA Is All You Need. *arXiv preprint arXiv:2502.07864*, 2025.