---
title: Warpgroup Matrix-Multiply-Accumulate (WGMMA)
url: https://www.emergentmind.com/topics/warpgroup-matrix-multiply-accumulate-wgmma
type: topic
---

# Warpgroup Matrix-Multiply-Accumulate (WGMMA)

Warpgroup Matrix-Multiply-Accumulate (WGMMA) is an advanced asynchronous General Matrix-Multiply (GEMM) primitive introduced on NVIDIA Hopper (SM90) Tensor Cores. It enables efficient, large-tile, warpgroup-cooperative matrix multiplications at high throughput, dramatically enhancing the performance of kernels that rely on dense or sparse GEMM, such as sparse matrix-matrix multiplication (SpMM) and attention mechanisms in large language models (LLMs) [2604.17834], [2312.11918].

## 1. Definition and Instruction Architecture

WGMMA comprises a family of asynchronous matrix-multiply-accumulate instructions executed at the warpgroup level, specifically leveraging 128 threads—grouped as four contiguous warps. Each instruction, for example `wgmma.mma_async.m<M>n<N>k<K>.f32.bf16.bf16`, computes a $M \times N \times K$ tensor-core GEMM where inputs are typically in BF16 or FP16 and the accumulator is FP32. Unlike previous generations (Ampere WMMA), which required chaining multiple warp-sized (32 thread) operations to cover larger tiles, WGMMA can perform a full $64\times64\times16$ operation in a single atomic invocation [2604.17834], [2312.11918].

Operands may be sourced directly from shared memory (SMEM), reducing register pressure and facilitating high throughput. Execution is explicitly asynchronous, requiring ordering primitives (`wgmma.fence`, `wgmma.commit_group`, `wgmma.wait_group`) to manage compute–memory concurrency. This deep pipelining enables WGMMA to overlap math with explicit data movement orchestrated by the Tensor Memory Accelerator (TMA), hiding latency and maximizing resource utilization [2604.17834].

## 2. Mathematical Tiling and Data Movement

The mathematical operation for SpMM with WGMMA is expressed as:

$$
C_{\text{block}} = \sum_{t=0}^{k/16-1} A^{(t)} B^{(t)},
$$

where $A^{(t)} \in \mathbb{R}^{64 \times 16}$ and $B^{(t)} \in \mathbb{R}^{16 \times BN}$ are tiles loaded per K-slice iteration (with $k_0=16$), $C$ is the output tile, and $BN$ is selected to divide the output columns evenly [2604.17834]. This formulation enables highly efficient partitioning and pipelining, a core requirement for maximizing GEMM throughput at scale.

WGMMA operations in attention (as in FlashAttention-2) generalize to

$$
C_{ij} \mathrel{+}{=} \sum_{p=0}^{bK-1} A_{ip} B_{jp} \quad (0 \leq i < bM,\, 0 \leq j < bN)
$$

with operand and tile shapes selectable at launch (e.g., $bM, bN \in \{64,128\}$, $bK=\text{head_dim}$), adapting to both register budgeting and softmax fusion strategies [2312.11918].

## 3. Producer–Consumer Pipeline and Asynchronous Tiling

To fully exploit asynchrony and hide global memory latency, high-performance kernels (notably AsyncSparse BCSR) implement a producer–consumer pipeline using WGMMA in conjunction with TMA. The control logic organizes computation around a circular buffer in shared memory (e.g., 3 stages for BCSR), with one warpgroup dedicated to TMA data loads (producer), while the remaining warpgroups concurrently consume tiles via WGMMA [2604.17834]. Lightweight synchronization objects (`mbarrier`, with phase bits) replace coarse grained `__syncthreads()` to ensure precise resource handover with minimal overhead.

This producer-consumer organization ensures that:

- TMA loads for tile $q$ are issued ahead of WGMMA computation for tile $q$.
- Consumers execute WGMMA on ready tiles while the producer loads the next K-slice.
- Synchronization between producer and consumers is managed using `mbarrier.wait()` and `mbarrier.arrive()`.
- The pipeline enables overlap: for each tile, as soon as WGMMA begins on tile $q$, TMA can begin loading tile $q+1$.

A representative pseudocode illustrates this flow, explicitly identifying producer/consumer warpgroups, shared memory tiling, register allocation, and synchronization points as used in AsyncSparse BCSR kernels [2604.17834]:

```cpp
__shared__ float A_buf[3][64][64];
__shared__ float B_buf[3][64][BN];
mbarrier_t empty[3], full[3];

if (wgId == 0) for (q=0..2) mbarrier.arrive(empty[q], init=1);

for (int blk = block_row_ptr[row]; blk < block_row_ptr[row+1]; ++blk) {
  int q = blk % 3;

  // Producer
  if (warpgroup_id == 0) {
    mbarrier.wait(empty[q]);
    if (threadId == 0) {
      cp.async.bulk.tensor(A_desc, ...);
      cp.async.bulk.tensor(B_desc, ...);
    }
  }

  // Consumers
  if (warpgroup_id >= 1) {
    mbarrier.wait(full[q]);
    wgmma.fence();
    wgmma.commit_group();
    for (int t=0; t<4; ++t) {
      wgmma.mma_async.m64nBNk16.f32.bf16.bf16(...);
    }
    wgmma.wait_group();
    mbarrier.arrive(empty[q]);
  }
}
```
[2604.17834]

## 4. Tile Parameterization, Data Layout, and Resource Management

Tile shapes and operand layouts are explicitly parameterized to maximize occupancy and minimize register spills. In SpMM, a pragmatic choice is $m=64$, $k_0=16$, with $n=BN$ selected for the target workload. In FlashAttention-2, $(bM, bN) \in \{64, 128\} \times \{64, 128\}$ is selected based on the head dimension to balance register pressure: larger tiles improve performance but can cause register spills that hurt throughput [2312.11918].

Operand layouts for SMEM are tightly coupled to GMMA atom layouts, with K-major 128-byte swizzling applied by TMA to eliminate bank conflicts [2312.11918]. Register allocation is dynamically managed via per-warpgroup `setmaxnreg.inc/dec.sync`, permitting multiple CTAs per SM when under the register file limit (e.g., 384 threads/CTA = 3 warpgroups). Producer warpgroups are register-light, while consumer warpgroups hold per-thread accumulators (e.g., for a $64\times BN$ tile, each thread holds 32 FP32 values) [2604.17834].

## 5. Kernel Fusion, CUTLASS Integration, and Practical Implementation

WGMMA is critical for fusing multiple GEMM operations within a single CUDA kernel, especially where intermediate data dependencies incur significant memory traffic on prior GPUs. In the FlashAttention-2 implementation, back-to-back WGMMA dispatches are used: GEMM-I computes $S_{ij}$, followed by in-register softmax, and GEMM-II computes $O_{ij}$, all within a kernel using the CUTLASS layout/tensor abstractions, removing all intermediate global memory writes for $S$ and $P$ matrices [2312.11918].

The pipeline is tightly controlled:

- TMA loads next operand (e.g., $V$) for GEMM-II while current GEMM-I executes.
- Operand partitioning, register/local memory reshaping, and layout alignment are handled by CUTLASS primitives (`partition_fragment_A/B/C`, `ReshapeTStoTP`).
- Only one warpgroup per CTA is used to avoid register file exhaustion.

Example invocation with CUTLASS:

```c++
auto smemLayoutQ = tile_to_shape(GMMA::Layout_K_SW128_Atom<MmaA>{}, tileShapeQ);
Tensor gQ = make_tensor(ptrQ, gmemLayoutQ);
auto tmaLoadQ = make_tma_copy(SM90_TMA_LOAD{}, gQ, smemLayoutQ, tileShapeQ, Int<1>{});

using TiledMma0 = /* ... */;
TiledMma0 tiledMma0;

__global__ void fused_fmha_kernel(/* … */) {
  cfk::copy(tQgQ(_,0), tQsQ(_,0), tmaLoadQ, tma_mbar[0]);
  // ...
  cfk::gemm_bar_wait(tiledMma0, tSrQ, tSrK, tSrS, tma_mbar[0]);
}
```
[2312.11918]

## 6. Performance Characteristics and System-Level Impact

Empirical ablation studies on SpMM (BF16, SuiteSparse, $N=1024$) quantify the performance contribution of WGMMA and its associated asynchrony:

| Optimization    | Throughput (TFLOPS) | Relative to Scalar | Description                        |
|-----------------|---------------------|--------------------|------------------------------------|
| opt0 (scalar)   | 0.12                | 0.08× cuSPARSE     | Scalar FMA on CUDA core            |
| opt1 (+WGMMA)   | 0.63                | 0.42×              | Add WGMMA only                     |
| opt2 (+TMA)     | 2.37                | 1.56×              | Add TMA overlapping                |
| opt3 (+special) | 6.44                | 4.31×              | Add warp specialization            |

In end-to-end SuiteSparse benchmarks, the WCSR kernel leveraging WGMMA achieves 23.5 TFLOPS (4.86× cuSPARSE, 2.40× FlashSparse) at density $\geq1\%$ and $N=1024$. In LLM inference, for Qwen2.5-7B with 90% block sparsity, BCSR kernels achieve 1.58–1.98× speedup on the FFN projection versus dense cuBLAS. End-to-end prefill, after integrating both sparse attention (MInference) and sparse FFN, produces up to 2.66× system-level speedup at 64K tokens [2604.17834].

FlashAttention-2 kernels that utilize WGMMA on Hopper observe a 20–50% FLOPs/s boost over legacy Ampere-optimized implementations. The threefold improvement in GEMM throughput—enabled by large tiles and the 128-thread warpgroup—increases memory and compute utilization, provided tile shapes are chosen to avoid register oversubscription [2312.11918].

## 7. Research Significance and Practical Considerations

WGMMA marks a critical shift in GPU programming for high-performance GEMM workloads:

- Asynchronous, warpgroup-level dispatch with large tiles enables deep overlap between compute and data movement.
- Bank-conflict-free layouts via hardware swizzling, and fine-grained mbarrier synchronization, facilitate higher SM occupancy.
- Register management and kernel fusion enabled by WGMMA (and supported by CUTLASS and CuTe) allow single-kernel implementations of operations that previously required costly intermediate storage.
- Performance scaling, even at low matrix density ($\geq 1\%$), establishes WGMMA as foundational to sparse and dense operations supporting contemporary LLM and scientific workloads.

A plausible implication is that, as workloads shift towards increasingly irregular or structured sparse data patterns and large-scale model inference, leveraging WGMMA-style primitives with explicit asynchrony, hardware-accelerated data movement (TMA), and cooperative threading will become standard in high-performance GPU kernels [2604.17834], [2312.11918].

Source: https://www.emergentmind.com/topics/warpgroup-matrix-multiply-accumulate-wgmma