Writing Speed-of-Light Attention Backward Kernel for 5090 in CUDA

August 23, 2026

Overview

In this post, I’ll iteratively optimize an implementation of the FlashAttention backward kernel written in CUDA for the 5090. My goal is to deepen my understanding of GPU performance and learn how to use it effectively, including shared-memory swizzling, register pressure, MMA, asynchronous data movement with TMA and mbarrier, and double buffering. One disclaimer: I did this outside of my work at NVIDIA, using only public information and materials.The code was done together with codex. At the time of writing, it was still not autonomous enough and needed constant steering and correction, but it made iteration much faster, e.g: editing, trying out my ideas, and so on.

There are already many excellent blog posts on writing high-performance GPU kernels I recommend reading these and getting familiar with CUDA C++ first.. However, there is much less on the backward kernel of attention, which is well known to be much harder to implement and optimize than the forward pass, so this feels like the natural next step.The RTX 5090 is also the GPU I have at home, and its SM120 is a comparatively approachable place to start. It uses warp-level mma instructions rather than exposing Hopper’s wgmma or datacenter Blackwell’s tcgen05 and Tensor Memory programming model, so there are fewer architecture-specific mechanisms to manage.

Setup

All benchmarks use BF16 multi-head causal attention on an RTX 5090. The query, key, and value tensors have shape [B, H, N, d], with B=8, H=16, N=4096, and d=128This is the same setup as the FlashAttention paper..

For reproducibility, I set the 5090’s power limit to 575 W and lock GPU clock to 2.407 GHz this is the boost clock NVIDIA uses for its rated BF16 Tensor Core peak of 209.5 TFLOPS here.. The benchmark command is:

sudo nvidia-smi -i 0 -pl 575
sudo nvidia-smi -i 0 -lgc 2407,2407
python kernels/attention/main.py \
    --shape 8_16_4096_128 \
    --kernel cudnn \
    --kernel fa \
    --kernel attention_v1 \
    --kernel attention_v2 \
    --kernel attention_v3 \
    --direction backward \
    --causal
sudo nvidia-smi -i 0 -rgc

The fused attention backward pass recomputes QK^T, then computes the four gradient matrix multiplications for dV, dP, dQ, and dK Softmax, masking, scaling, and other non-matmul operations execute as part of the kernels but are not included in this conventional FLOP count. See Appendix C for the full compute and memory lower-bound calculation..

FLOPs = 10 * B * H * N^2 * d * 0.5
TFLOPS = FLOPs / latency_ms / 1e9
Kernel Latency Effective TFLOPS % of theoretical peak
F.sdpa() (cuDNN) 9.5079 ms 144.55 69.00%
F.sdpa() (Flash Attention) 8.5584 ms 160.59 76.65%
v1 10.3287 ms 133.07 63.52%
v2 9.5374 ms 144.10 68.78%
v3 8.5866 ms 160.06 76.40%

We begin by rederiving the backward pass, and then map it onto the hardware with a corresponding pseudocode.

Attention

Mathematical Formulation

For clarity, the derivation and pseudocode below describe one attention head for one batch element, so Q,K,V,O have shape [N,d]. Hence, the forward pass is

$$ \begin{aligned} S &= \frac{QK^T}{\sqrt{d}} + M \\ L_i &= \log \sum_j \exp(S_{ij}) \\ P_{ij} &= \exp(S_{ij} - L_i) \\ O &= PV \end{aligned} $$

where $M$ is the attention mask and $L$ is the row-wise log-sum-exp. The forward pass saves $L$ so we can recompute $P$ without storing $S$ in the backward pass.

Let $\mathcal{L}$ be the scalar loss and use $dX = \partial \mathcal{L} / \partial X$ for the gradient of any tensor $X$. Given the upstream gradient $dO = \partial \mathcal{L} / \partial O \in \mathbb{R}^{N \times d}$, backpropagation gives Applying chain rule to $V$ we have $\frac{\partial \mathcal{L}}{\partial V_{kj}}=\sum_i \frac{\partial \mathcal{L}}{\partial O_{ij}}\frac{\partial O_{ij}}{\partial V_{kj}}$. Since $O_{ij}=\sum_k P_{ik}V_{kj}$, we have $\frac{\partial O_{ij}}{\partial V_{kj}}=P_{ik}$. Therefore, $\frac{\partial \mathcal{L}}{\partial V_{kj}}=\sum_i P_{ik}dO_{ij}$. Hence, $dV=P^T dO$. The expressions for $dP$, $dQ$, and $dK$ follow from the similar argument.

$$ \begin{aligned} dV &= P^T dO \\ dP &= dO V^T \\ dS &= \operatorname{dsoftmax}(dP) \\ dQ &= \frac{dS K}{\sqrt{d}} \\ dK &= \frac{dS^T Q}{\sqrt{d}}. \end{aligned} $$

One can derive the gradient of the softmax applied row-wise for one row as below, with complete derivation in the appendix.

$$ ds =p\odot\left(dp-\left(p^Tdp\right)\mathbf{1}\right). $$

Stacking the rows gives we have:

$$ dS=P\odot\left(dP-D\mathbf{1}^T\right), \qquad D=\operatorname{rowsum}(P\odot dP) =\operatorname{rowsum}(dO\odot O). $$

where $D\in\mathbb{R}^{N\times 1}$ contains one softmax correction term per row. The complete backward pass is therefore

$$ \begin{aligned} dV &= P^T dO \\ dP &= dO V^T \\ D &= \operatorname{rowsum}(dO\odot O) \\ dS &= P\odot\left(dP-D\mathbf{1}^T\right) \\ dQ &= \frac{dS K}{\sqrt{d}} \\ dK &= \frac{dS^T Q}{\sqrt{d}}. \end{aligned} $$

In the kernel implementation, we will first split $Q$, $O$, and $dO$ into $T_r$ row tiles $Q_i, O_i, dO_i \in \mathbb{R}^{B_r \times d}$ and $K$, $V$ into $T_c$ tiles $K_j, V_j \in \mathbb{R}^{B_c \times d}$ (from here on, subscripts will index tiles rather than elements).

Each tile pair $(i,j)$ gives us a $B_r \times B_c$ block of $S_{ij} = Q_i K_j^T / \sqrt{d}$, from which together with $L_i$ (saved in the forward pass), we can recompute $P_{ij} = \exp(S_{ij} - L_i)$ on the fly.

$D$ depends only on $dO$ and $O$, so we precompute it once before the main loop.

Tiling the attention backward pass and recomputing one score/probability tile.

The same tiling applies to the three gradient matmuls.

Start with $dV=P^T dO$. Each row $j$ of $dV$ is a weighted sum of every query row of $dO$ with weights from row $j$ of $P.T$ (or column $j$ of $P$) :

Row $j$ of $P^T$ (column $j$ of $P$) weights the query rows of $dO$ into row $j$ of $dV$.

$dK=dS^T Q/\sqrt{d}$ follows the same pattern.

With $dQ=dS K/\sqrt{d}$: each query row of $dQ$ is a weighted sum of every key row of $K$, with weights from row of $dS$.

Hence, we have the following::

$$ dV_j = \sum_i (P_{ij})^T\, dO_i, \qquad dK_j = \frac{1}{\sqrt{d}} \sum_i (dS_{ij})^T\, Q_i, \qquad dQ_i = \frac{1}{\sqrt{d}} \sum_j dS_{ij}\, K_j. $$

The remaining product $dP_{ij}=dO_i V_j^T$ lives inside a single tile, so it is not a reduction.

Notice that a tile pair $(i,j)$ will write to both a $j$-indexed accumulator $dK_j$ and $dV_j$ and an $i$-indexed one $dQ_i$, so one CUDA block cannot own all three. We can either parallelize over key/value tiles $j$ and combine the $dQ_i$ partials, or parallelize over query tiles $i$ and combine $dK_j$ and $dV_j$.

I choose to parallelize over $j$: one block owns $dK_j$ and $dV_j$, loops over query tiles, and atomically adds each partial $dQ_i$ in HBM.

Traversal of one fixed K/V column tile in the causal backward pass. Tiles above the diagonal are skipped, the diagonal tile is partially masked, and the remaining tiles are fully valid.

At a high level, the kernel will follow the pseudocode below. The high-level loop structure also follows Algorithm 2 of FlashAttention-2. We specialize the outline to the v1 kernel by making placement in registers, shared memory, and HBM explicit, also adding causal tile skipping and elementwise diagonal masking, and showing the FP32 accumulator-to-BF16 output conversion.

#----------------------------------------------------------------------------
# Tiled attention backward.
#
# Inputs:
#   Q, K, V       N x d
#   O, dO         N x d
#   L             N x 1, row-wise logsumexp of S
#
# Outputs:
#   dQ, dK, dV    N x d
#
# Tile shapes:
#   Q_i, O_i, dO_i    B_r x d
#   K_j, V_j          B_c x d
#   S_ij, P_ij        B_r x B_c
#----------------------------------------------------------------------------

alpha = 1 / sqrt(d)

# D in the equations; one scalar correction per query row.
delta = row_sum(dO * O, keepdim=true)             # N x 1

# Column-tile blocks share an FP32 dQ accumulator. The final dQ is BF16.
dQ_accum = zeros(N, d, dtype=FP32)

parallel for j = 0 ... T_c - 1:                   # One CUDA block per K/V tile.
    load K_j, V_j from HBM to SRAM

    dK_j = zeros(B_c, d)                          # FP32 register accumulators.
    dV_j = zeros(B_c, d)                          # FP32 register accumulators.

    # Tiles i < j are fully masked. Starting at j prevents those tiles from
    # loading Q/dO or launching any MMA work.
    first_i = j if causal else 0

    for i = first_i ... T_r - 1:                  # Sequential within block j.
        load Q_i, dO_i from HBM to SRAM

        # Recompute this tile instead of materializing S and P in HBM.
        S_ij = alpha * Q_i @ transpose(K_j)
        dP_ij = dO_i @ transpose(V_j)

        # With aligned square tiles, only i == j crosses the causal diagonal.
        diagonal = causal and i == j
        for query, key in tile(i, j):
            if diagonal and key > query:
                P_ij[query, key] = 0
                dS_ij[query, key] = 0
            else:
                P_ij[query, key] = exp(
                    S_ij[query, key] - L_i[query]
                )
                dS_ij[query, key] = (
                    P_ij[query, key] *
                    (dP_ij[query, key] - delta_i[query])
                )

        # This block owns dV_j and dK_j.
        dV_j += transpose(P_ij) @ dO_i
        dK_j += alpha * transpose(dS_ij) @ Q_i

        # Every column tile contributes to dQ_i, so preserve the partial sums
        # in FP32 and combine them atomically.
        atomic_add(dQ_accum_i, alpha * dS_ij @ K_j)

    store BF16(dK_j), BF16(dV_j) from registers to HBM

# A second kernel runs after every column-tile block has finished. The kernel
# boundary provides the grid-wide synchronization needed before conversion.
parallel for element in dQ:
    dQ[element] = BF16(dQ_accum[element])

Now, with the math and the pseudocode clear, we can write our backward kernels.

Version 1

Let’s start with a baseline implementation of the pseudocode above. I will optimize for correctness first, but this version still uses the standard practices from the start, e.g: using TMA, and Tensor Core. As an example, algorithm from the legendary Simon’s blog only use CUDA cores to optimize the kernel.

Precomputing D

First, we will implement the row-wise reduction D = rowsum(dO * O).

We compute it beforehand because the same $D_i$ is reused by every K/V column block that processes query row $i$ Computing it inside the main kernel would make every column block recompute the same dot product. Precomputing this also means that in the main kernel, we only need the small vector $D$ and does not have to load $O$..

delta = row_sum(dO * O, keepdim=true)             # N x 1

We assign one warp to process one [batch, head, query row] at a time. The launch will create fewer warps than rows, and a grid-stride loop lets each warp process several rows sequentially. Each of the 32 lanes in 1 warp compute 4 of the 128 products O * dO and sums those four. Then we will use __shfl_down_sync to reduce the 32 lane sums so lane 0 can write $D_i$. A simplified eight-lane __shfl_down_sync reduction. A simplified eight-lane __shfl_down_sync tree reduction.

//----------------------------------------------------------------------------
const int rows = batch * heads * sequence;

const int d_blocks = (rows + D_THREADS - 1) / D_THREADS;
compute_D_v1_kernel<<<d_blocks, D_THREADS>>>(O, dO, D, rows);

// One warp computes one row of D = rowsum(dO * O) at a time.
__global__ void compute_D_v1_kernel(
    const nv_bfloat16 *O,
    const nv_bfloat16 *dO,
    float *D,
    int rows) {
    const int lane = threadIdx.x % 32;
    const int warp = threadIdx.x / 32;

    for (int row = blockIdx.x * D_WARPS_PER_BLOCK + warp;
         row < rows;
         row += gridDim.x * D_WARPS_PER_BLOCK) {
            
        const size_t offset =
            static_cast<size_t>(row) * HEAD_DIM +
            lane * 4;
        const nv_bfloat162 *output =
            reinterpret_cast<const nv_bfloat162 *>(
                O + offset);
        const nv_bfloat162 *grad =
            reinterpret_cast<const nv_bfloat162 *>(
                dO + offset);
        const float2 output_0 =
            __bfloat1622float2(output[0]);
        const float2 output_1 =
            __bfloat1622float2(output[1]);
        const float2 grad_0 =
            __bfloat1622float2(grad[0]);
        const float2 grad_1 =
            __bfloat1622float2(grad[1]);

        float value =
            output_0.x * grad_0.x +
            output_0.y * grad_0.y +
            output_1.x * grad_1.x +
            output_1.y * grad_1.y;

        // Add the 32 partial sums; the full rowsum lands in lane 0.
        #pragma unroll
        for (int delta = 16; delta > 0; delta /= 2)
            value += __shfl_down_sync(
                0xffffffff,
                value,
                delta);

        if (lane == 0)
            D[row] = value;
    }
}

Main backward kernel

With D precomputed, we can now write the main tiled backward kernel.

As in the previous section, one CUDA block owns one K/V tile $K_j, V_j$ for a (batch, head) pair, keeps it fixed, and walks sequentially over query tiles.

We will set $B_r = B_c = 64$ There are a few reasons why 64 is a reasonable starting choice. 4 $64\times128$ BF16 tiles $K$, $V$, $Q$, and $dO$ and the 2 $64\times64$ $P$ and $dS$ take total 80 KiB, which is the largest power-of-two that still fits the 99 KiB opt-in shared-memory limit. 64 is also a multiple of the m16n8k16 MMA shape in both dimensions, so no fragment is partial., so each of Q, K, V, and dO tile is $64 \times 128$ and the score tile $S_{ij}$ is $64 \times 64$. Each block uses 16 warps.

The three grid dimensions select the K/V tile, attention head, and batch element, with blockIdx.x choosing the K/V rows beginning at key_start = 64 * blockIdx.x.

constexpr int BLOCK_SIZE = 64;  // B_r = B_c = 64 in v1.

// BLOCK_Q is B_r, BLOCK_KV is B_c. Both equal BLOCK_SIZE here.
using Config = AttentionV1Config<BLOCK_SIZE>;

// Grid dimensions: (K/V tile, head, batch).
const dim3 blocks(
    sequence / Config::BLOCK_KV,
    heads,
    batch);

// True enables causal masking. For now, don't worry about WARPS_PER_ROW_GROUP,
// NUM_WARPS, or the Q_tmap, K_tmap, V_tmap, and dO_tmap arguments. I explain
// the warp layout and tensor maps later.
attention_v1_bwd_kernel<
    true,
    BLOCK_SIZE,
    WARPS_PER_ROW_GROUP,
    NUM_WARPS>
    <<<blocks,                   // Grid dimensions.
       NUM_WARPS * 32,           // Threads per block.
       Config::SMEM_BYTES>>>(    // Dynamic shared-memory bytes per block.
        Q,
        K,
        V,
        dO,
        L,
        D,
        dQ_accum,
        dK,
        dV,
        Q_tmap,
        K_tmap,
        V_tmap,
        dO_tmap,
        heads,
        sequence);

The next step is to move the data to the Tensor Cores to perform MMA operations.

We will follow the typical flow:

Each query-tile iteration applies that flow as In v1, there is only one Q/dO buffer in shared memory; I am not using double buffering just yet.:

// Once per block: TMA load K_j, V_j; ldmatrix into K_frag, V_frag.

for each query tile i:
    // Wait for TMA Q_i, dO_i.

    // 1. S_ij = Q_i K_j^T
    //    ldmatrix Q;  mma with K_frag
    // 2. dP_ij = dO_i V_j^T
    //    ldmatrix dO; mma with V_frag
    // 3. P_ij, dS_ij from S and dP; store to smem
    // 4. dV_j += P^T dO,  dK_j += dS^T Q
    //    ldmatrix P, dS, Q, dO;  mma
    //    TMA issue next Q/dO
    // 5. dQ_i += dS K
    //    ldmatrix dS, K;  mma;  atomicAdd

Global to shared memory with TMA

A TMA load usually has two parts.

First, the host encodes an opaque CUtensorMap descriptor describing the tensor’s global-memory layout and the tile to transfer. Then, one device thread passes that descriptor and a set of starting coordinates to cp.async.bulk.tensor.

We create a tensor map for each of Q, K, V, and dO, which are flattened to [batch * heads, N, d]. TMA lists dimensions from fastest-first, so the descriptor uses the order {d, N, batch * heads}.cuTensorMapEncodeTiled also takes a CUtensorMapSwizzle. TMA applies the permutation as it writes the tile: it can keep the same row-major order it reads from HBM (NONE), or rearrange each destination row into 32-byte, 64-byte, or 128-byte swizzled panels. Those layouts let later ldmatrix loads avoid bank conflicts. We use the largest, CU_TENSOR_MAP_SWIZZLE_128B.

// Host-side setup: one CUtensorMap per source tensor (Q, K, V, dO)
// so a later TMA load knows how to move a panel.
template <int BLOCK_SIZE>
static void init_tensor_map(
    CUtensorMap *tensor_map,
    const nv_bfloat16 *data,
    int tensors,
    int sequence,
    int head_dim) {
    constexpr uint32_t RANK = 3;
    constexpr uint32_t PANEL_ELEMENTS =
        128 / sizeof(nv_bfloat16);

    // TMA dimension order: {d column, sequence row, batch-head}.
    const uint64_t global_dims[RANK] = {
        static_cast<uint64_t>(head_dim),
        static_cast<uint64_t>(sequence),
        static_cast<uint64_t>(tensors),
    };

    // Byte strides for the sequence and batch-head dimensions. The
    // d-column stride is implicitly sizeof(nv_bfloat16).
    const uint64_t global_strides[RANK - 1] = {
        static_cast<uint64_t>(head_dim) *
            sizeof(nv_bfloat16),
        static_cast<uint64_t>(sequence) *
            head_dim * sizeof(nv_bfloat16),
    };

    // One TMA instruction transfers one 64-column (128-byte) panel of
    // BLOCK_SIZE rows. A d=128 BF16 row is 256 bytes, so two panels
    // make the 64x128 tile.
    const uint32_t box_dims[RANK] = {
        PANEL_ELEMENTS,
        BLOCK_SIZE,
        1,
    };
    const uint32_t element_strides[RANK] = {1, 1, 1};

    const CUresult status = cuTensorMapEncodeTiled(
        tensor_map,
        CU_TENSOR_MAP_DATA_TYPE_BFLOAT16,
        RANK,
        const_cast<nv_bfloat16 *>(data),
        global_dims,
        global_strides,
        box_dims,
        element_strides,
        CU_TENSOR_MAP_INTERLEAVE_NONE,
        CU_TENSOR_MAP_SWIZZLE_128B,
        CU_TENSOR_MAP_L2_PROMOTION_NONE,
        CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);

    if (status != CUDA_SUCCESS)
        throw std::runtime_error("cuTensorMapEncodeTiled failed");
}

Second, we need to issue that copy on the device side. One thread passes a descriptor and starting coordinates to cp.async.bulk.tensor.3d that starts an asynchronous copy of one box_dims panel from global memory to shared memory. We wait for it with mbarrier.try_wait.parity: it asks whether a given phase has finished (even phases have parity 0, odd phases have parity 1) and retries until the answer is yes.

Those PTX instructions are wrapped in tma_3d_g2s() g2s means global-to-shared memory. and the mbarrier_*() helpers. See Appendix A.2 for further details.

The kernel then loads $K_j$ and $V_j$ once at the start of the block:

// Two mbarrier objects sit in shared memory just after the tile buffers.
const uint32_t smem_addr = cvta_shared(smem); // 32-bit shared address
const uint32_t kv_barrier_addr =
    smem_addr + DATA_BYTES; // DATA_BYTES = K, V, Q, dO, P, dS
const uint32_t query_barrier_addr =
    kv_barrier_addr + sizeof(uint64_t); // 8-byte Q/dO barrier after K/V

if (tid == 0) {
   mbarrier_init(kv_barrier_addr, 1);
   mbarrier_init(query_barrier_addr, 1);
   mbarrier_fence_init();
}
__syncthreads();

// One elected thread in warp 0 issues four TMA loads: two 64-column
// panels each for K_j and V_j.
if (warp_id == 0 && elect_one_sync()) {
   for (int panel = 0; panel < HEAD_DIM / PANEL_ELEMENTS; panel++) {
      const int d_start = panel * PANEL_ELEMENTS;
      const int panel_offset = panel * PANEL_BYTES;

      // Load K_j tile.
      tma_3d_g2s(
         cvta_shared(K_tile) + panel_offset, // smem dest
         &K_tmap,                            // tensor map
         d_start,                            // d
         key_start,                          // N
         batch_head_idx,                     // batch * head
         kv_barrier_addr);

      // Load V_j tile.
      tma_3d_g2s(
         cvta_shared(V_tile) + panel_offset,
         &V_tmap,
         d_start,
         key_start,
         batch_head_idx,
         kv_barrier_addr);
   }
   mbarrier_arrive_expect_tx(
      kv_barrier_addr,
      2 * TILE_BYTES); // One complete K tile plus one complete V tile.
}

if (warp_id == 0)
   mbarrier_wait(kv_barrier_addr, 0); // Phase 0.
__syncthreads(); // Other warps must not read K/V until that wait completes.

load_query_tiles() uses the same pattern for $Q_i$ and $dO_i$ on query_barrier_addr, once per query-tile iteration:

auto load_query_tiles = [&](int batch_head_idx, int start) {
    if (warp_id == 0 && elect_one_sync()) {
        for (int panel = 0;
             panel < HEAD_DIM / PANEL_ELEMENTS;
             panel++) {
            const int d_start = panel * PANEL_ELEMENTS;
            const int panel_offset = panel * PANEL_BYTES;

            // Load Q_i tile.
            tma_3d_g2s(
                cvta_shared(Q_tile) + panel_offset, // smem dest
                &Q_tmap,                            // tensor map
                d_start,                            // d
                start,                              // N
                batch_head_idx,                     // batch * head
                query_barrier_addr);

            // Load dO_i tile.
            tma_3d_g2s(
                cvta_shared(dO_tile) + panel_offset,
                &dO_tmap,
                d_start,
                start,
                batch_head_idx,
                query_barrier_addr);
        }
        mbarrier_arrive_expect_tx(
            query_barrier_addr,
            2 * TILE_BYTES); // One complete Q tile plus one complete dO tile.
    }
};

After $Q_i$ and $dO_i$ are no longer needed for the current $dK_j,dV_j$ update, we issue the next asynchronous $Q_i$/$dO_i$ transfer before starting the $dQ_i$ phase:

// All warps have finished reading the current Q_i and dO_i buffers.
__syncthreads();

const int next_query_start = query_start + BLOCK_Q;
if (next_query_start < sequence)
   load_query_tiles(batch_head_idx, next_query_start);

Shared memory to registers

Once TMA phase is complete, we use ldmatrix to move fragments from shared memory to registers.

TMA stored the tiles in a 128-byte-swizzled layout, so the warps cannot use ordinary row-major offsets when loading their fragments. We made a helper function below that maps a logical (row, 16-byte chunk) to its physical offset.

//----------------------------------------------------------------------------
// Split rows wider than 128 bytes into independent 128-byte panels, then XOR
// the row within each panel's 16-byte chunk index. Global memory remains
// row-major. TMA creates this layout in shared memory, and ldmatrix loads use
// the same physical offset.
template <int ROW_BYTES>
__device__ inline int swizzle_128b_panel_offset(int row, int chunk) {
    constexpr int CHUNK_BYTES = 16; // one 16-byte chunk
    constexpr int PANEL_BYTES = 128;
    constexpr int CHUNKS_PER_PANEL = PANEL_BYTES / CHUNK_BYTES;

    static_assert(ROW_BYTES % PANEL_BYTES == 0);

    // `chunk` is a 16-byte chunk index, not an element column.
    const int panel = chunk / CHUNKS_PER_PANEL;
    const int chunk_in_panel = chunk % CHUNKS_PER_PANEL;
    const int swizzled_chunk = chunk_in_panel ^ (row % CHUNKS_PER_PANEL);
    return row * ROW_BYTES + panel * PANEL_BYTES + swizzled_chunk * CHUNK_BYTES;
}
Logical-to-physical mapping for 128-byte shared-memory panels, with each cell being a 16-byte chunk.

tile_addr() applies this mapping to the 128-column Q, K, V, and dO tiles. The 64-column P and dS tiles use the same mapping through matrix_addr() and matrix_offset(). See Appendix A.1 for further details.

Each lane uses those helpers to get the physical address of one 16-byte chunk. The 32 lanes then execute ldmatrix_x4 together: their lane-supplied addresses identify four $8\times8$ matrices, whose elements are distributed into the register layout expected by mma.m16n8k16. See Appendix A.3 for further details.

The ldmatrix.x4 instruction loads four $8\times8$ matrices from shared memory into warp-distributed registers.

Putting the loop together

We now expand that query-tile outline, starting with steps 1 and 2.

Steps 1 and 2: Recompute $S_{ij}$ and $dP_{ij}$

The kernel issues mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32.m16n8k16 is the largest warp-level MMA shape for BF16. Other input types have different shapes, for example, FP8 can use m16n8k32. mma.sync is warp-wide: all 32 lanes issue it together and collectively do one $16\times8\times16$ multiply-add.Each lane passes only its slice of the fragment: A[4] holds 8 BF16s of A, B[2] holds 4 BF16s of B, and C[4] holds 4 FP32 accumulators of C / D. See Appendix A.3 for further details.

$$ S_{ij}=\alpha Q_iK_j^T, \qquad [64\times128][128\times64]=[64\times64]. $$

The MMA instructions compute $Q_iK_j^T$, and the kernel applies $\alpha$ when it uses the accumulated scores to reconstruct $P_{ij}$.

Let’s divide the $64\times64$ score output tile into $64/16=4$ fragments along its rows and $64/8=8$ fragments along its columns. This gives us $4\times8=32$ output fragments, each with the shape of $16\times8$. Along the reduction dimension, each fragment requires $128/16=8$ MMA steps.

The MMA m16n8k16 C / D fragment layout.

For v1, we use 16 warps as a $4\times4$ grid to process the score output fragments: four row groups and four column partitions.This choice is not really free. Four warps share one row group, so all four will load the same $Q_i$ and $dO_i$ A fragments from shared memory. We will revisit that redundancy in the v2 kernel. Each warp owns two adjacent $16\times8$ fragments.

// 16-row region of the score tile.
const int row_group = warp_id / WARPS_PER_ROW_GROUP;
// Matching 16-column region.
const int column_partition = warp_id % WARPS_PER_ROW_GROUP;

Each warp will then own a $16\times16$ region of the $64\times64$ score tile. For every 16-element slice of the reduction dimension, it issues two MMA instructions, one per $16\times8$ output fragment, so eight reduction steps take $2\times8=16$ MMAs per warp for this matmul.

How one $16\times64$ score row group is split into eight $16\times8$ MMA output tiles, and the resulting $4\times4$ warp layout. Each column partition assigns two adjacent output tiles to one warp.

Before the query loop, each warp loads its $K_j$ and $V_j$ fragments into K_frag and V_frag with ldmatrix_x4. Those registers are reused for every query tile:

constexpr int SCORE_FRAGMENTS_PER_WARP =
    BLOCK_KV / MMA_N / WARPS_PER_ROW_GROUP; // 2

uint32_t K_frag[HEAD_DIM / MMA_K][SCORE_FRAGMENTS_PER_WARP][2];
uint32_t V_frag[HEAD_DIM / MMA_K][SCORE_FRAGMENTS_PER_WARP][2];

#pragma unroll
for (int d_offset = 0; d_offset < HEAD_DIM; d_offset += MMA_K) {
    // Each ldmatrix_x4 fills two fragments, so this steps by 2.
    // With SCORE_FRAGMENTS_PER_WARP == 2 this runs only once.
    #pragma unroll
    for (int local_score_fragment = 0;
         local_score_fragment < SCORE_FRAGMENTS_PER_WARP;
         local_score_fragment += 2) {
        const int score_fragment =
            column_partition * SCORE_FRAGMENTS_PER_WARP +
            local_score_fragment;
        // K_tile is 64 by 128.
        // key_row = which key
        // key_col = which 16-byte chunk along that key.
        const int key_row =
            score_fragment * MMA_N +
            lane % 8 +
            (lane / 16) * 8;
        const int key_col =
            d_offset + ((lane / 8) % 2) * 8;
        uint32_t pair[4];
        ldmatrix_x4(
            pair,
            tile_addr(K_tile, key_row, key_col));
        // Each B fragment is 16 of d by 8 keys. One ldmatrix_x4 holds
        // both 8-wide halves of this 16-wide d slice, for two key eights:
        //              d 0..7         d 8..15
        //   keys 0..7  pair[0]        pair[1]         fragment local
        //   keys 8..15 pair[2]        pair[3]         fragment local+1
        K_frag[d_offset / MMA_K][local_score_fragment][0] = pair[0];
        K_frag[d_offset / MMA_K][local_score_fragment][1] = pair[1];
        K_frag[d_offset / MMA_K][local_score_fragment + 1][0] = pair[2];
        K_frag[d_offset / MMA_K][local_score_fragment + 1][1] = pair[3];

        // V_frag is filled identically from V_tile at the same coordinates.
        // ...
    }
}

Once these fragments are in registers, we enter the loop:

// Each warp owns two 16x8 score fragments and two 16x8 dP fragments.
// The trailing [4] holds this lane's four FP32 accumulator registers for
// one fragment.
float score_acc[SCORE_FRAGMENTS_PER_WARP][4] = {};
float dP_acc[SCORE_FRAGMENTS_PER_WARP][4] = {};

// HEAD_DIM / MMA_K = 128 / 16 = 8 reduction slices along d.
#pragma unroll
for (int d_offset = 0; d_offset < HEAD_DIM; d_offset += MMA_K) {
    uint32_t Q_reg[4];
    uint32_t dO_reg[4];
    // Q_tile is 64 by 128.
    // query_row = which query
    // query_col = which 16-byte chunk along that query.
    const int query_row =
        row_group * MMA_M + lane % MMA_M;
    const int query_col =
        d_offset + (lane / MMA_M) * 8;

    // Q_i and dO_i have the same 64x128 layout and are both row-major A
    // operands, so they can use the same coordinates.
    // Warps sharing the same row_group load the same A fragments, then pair them
    // with different K/V B fragments selected by column_partition.
    ldmatrix_x4(
        Q_reg,
        tile_addr(Q_tile, query_row, query_col));
    ldmatrix_x4(
        dO_reg,
        tile_addr(dO_tile, query_row, query_col));

    // Two score MMAs and two dP MMAs per reduction slice.
    #pragma unroll
    for (int local_score_fragment = 0;
         local_score_fragment < SCORE_FRAGMENTS_PER_WARP;
         local_score_fragment++) {
        mma_m16n8k16(
            Q_reg,
            K_frag[d_offset / MMA_K][local_score_fragment],
            score_acc[local_score_fragment]);
        mma_m16n8k16(
            dO_reg,
            V_frag[d_offset / MMA_K][local_score_fragment],
            dP_acc[local_score_fragment]);
    }
}

Step 3: Reconstruct $P_{ij}$ and form $dS_{ij}$

score_acc and dP_acc hold $S_{ij}$ and $dP_{ij}$ in FP32 registers. We now load the saved log-sum-exp $L_i$ and the precomputed row reduction $D_i$ to compute $P_{ij}$ and $dS_{ij}$.On the diagonal tile, masked scores are replaced with -inf first, so the exponential makes both $P_{ij}$ and $dS_{ij}$ zero at those positions.

// L and D for this lane's two query rows, 8 apart in the C fragment.
const float L_row[2] = {L[stat_row], L[stat_row + 8]};
const float D_row[2] = {D[stat_row], D[stat_row + 8]};

// Fully masked 64x64 tiles were skipped before entering the loop.
// On the diagonal tile, we compute dense MMA fragments and then
// zeros invalid entries through score=-inf.
if constexpr (CAUSAL) {
    if (query_start == key_start &&
        key_in_tile > query_in_tile)
        score_acc[local_score_fragment][value] = -INFINITY;
}


float P[2];
float ds[2];
#pragma unroll
for (int value = 0; value < 2; value++) {
    // __expf is CUDA's fast approximate FP32 exponential. 
    // It is faster but less accurate than expf.
    P[value] = __expf(
        score_acc[local_score_fragment][half * 2 + value] *
            softmax_scale -
        L_row[half]);
    ds[value] = P[value] *
        (dP_acc[local_score_fragment][half * 2 + value] -
         D_row[half]);
}

// __floats2bfloat162_rn converts two FP32 inputs to BF16 using
// round-to-nearest-even and packs them into one nv_bfloat162, allowing each
// lane to store two adjacent columns with one 32-bit write. 
// Across all lanes, row halves, score fragments, and warps, these writes
// fill the 64x64 P_ij and dS_ij tiles in shared memory. Later MMAs load
// those tiles with ldmatrix, which can only read from shared memory, not
// from other warps' registers.
const int index = matrix_offset(query_in_tile, key_in_tile);
reinterpret_cast<nv_bfloat162 *>(P_tile + index)[0] =
    __floats2bfloat162_rn(P[0], P[1]);
reinterpret_cast<nv_bfloat162 *>(dS_tile + index)[0] =
    __floats2bfloat162_rn(ds[0], ds[1]);

Step 4: Accumulate $dK_j$ and $dV_j$

The three remaining products all share one shape:

$$ dV_j = P_{ij}^T dO_i, \qquad dK_j = \alpha\, dS_{ij}^T Q_i, \qquad dQ_i = \alpha\, dS_{ij} K_j, $$
$$ [64\times64][64\times128]=[64\times128]. $$

As in the steps above, we split the $64\times128$ output into $16\times8$ fragments: $64/16=4$ along the rows and $128/8=16$ along the columns, on the same $4\times4$ warp grid.

// The block owns this K/V column tile. dK_j and dV_j start at zero and
// stay in FP32 registers until every contributing query tile is done.
// OUTPUT_MMAS_PER_WARP = HEAD_DIM / MMA_N / WARPS_PER_ROW_GROUP = 4
// fragments per warp.
// The trailing [4] is this lane's four FP32 C registers for one fragment.
float dK_acc[OUTPUT_MMAS_PER_WARP][4] = {}; 
float dV_acc[OUTPUT_MMAS_PER_WARP][4] = {};

// matrix_row, matrix_col, query_row, and output_col come from the warp
// assignment and the current reduction slice.
for (int reduce_k = 0; reduce_k < BLOCK_Q; reduce_k += MMA_K) {
    uint32_t dST_reg[4];
    uint32_t PT_reg[4];
    uint32_t Q_reg[4];
    uint32_t dO_reg[4];
    // P_tile is [query, key]. dV_j = P^T dO is [key, d], so the A
    // operand must be [key, query] = P^T. ldmatrix.trans reads the tile
    // that way without moving it, same story for dS and dK.
    ldmatrix_x4_trans(
        dST_reg, matrix_addr(dS_tile, matrix_row, matrix_col));
    ldmatrix_x4_trans(
        PT_reg, matrix_addr(P_tile, matrix_row, matrix_col));

    // Q_reg and dO_reg hold the corresponding column-major B fragments.
    ldmatrix_x4_trans(
        Q_reg, tile_addr(Q_tile, query_row, output_col));
    ldmatrix_x4_trans(
        dO_reg, tile_addr(dO_tile, query_row, output_col));

    // Pseudocode updates:
    //   dK_j += dS_ij^T @ Q_i
    //   dV_j += P_ij^T @ dO_i
    mma_m16n8k16(dST_reg, Q_reg, dK_acc[local_output_n]);
    mma_m16n8k16(dST_reg, Q_reg + 2, dK_acc[local_output_n + 1]);
    mma_m16n8k16(PT_reg, dO_reg, dV_acc[local_output_n]);
    mma_m16n8k16(PT_reg, dO_reg + 2, dV_acc[local_output_n + 1]);
}

// dK_acc remains unscaled here. The kernel will apply softmax_scale when it
// writes the completed dK_j after the query loop.

Step 5: Compute and atomically accumulate partial $dQ_i$

After step 4, the current $Q_i$ and $dO_i$ values are no longer needed. We can issue the next TMA transfer now because the remaining calculation reads only $dS_{ij}$ and $K_j$.

Each warp then computes its slice of $dQ_i=\alpha\,dS_{ij}K_j$. Unlike $dK_j$ and $dV_j$, every K/V column-tile block writes the same $dQ_i$, so this partial has to be added into the global FP32 accumulator with atomicAdd:

// ds_row, ds_col, key_row, output_col, and offset come from the warp
// assignment and the current reduction slice.
float dQ_acc[OUTPUT_MMAS_PER_WARP][4] = {};
for (int reduce_k = 0; reduce_k < BLOCK_KV; reduce_k += MMA_K) {
    uint32_t dS_reg[4];
    uint32_t K_reg[4];
    ldmatrix_x4(
        dS_reg, matrix_addr(dS_tile, ds_row, ds_col));
    ldmatrix_x4_trans(
        K_reg, tile_addr(K_tile, key_row, output_col));
    mma_m16n8k16(dS_reg, K_reg, dQ_acc[local_output_n]);
    mma_m16n8k16(dS_reg, K_reg + 2, dQ_acc[local_output_n + 1]);
}

float *acc = dQ_acc[local_output_n];
atomicAdd(
    reinterpret_cast<float2 *>(dQ_accum + offset),
    make_float2(
        acc[0] * softmax_scale,
        acc[1] * softmax_scale));
atomicAdd(
    reinterpret_cast<float2 *>(
        dQ_accum + offset + 8 * HEAD_DIM),
    make_float2(
        acc[2] * softmax_scale,
        acc[3] * softmax_scale));

After the query loop, the block converts dK_acc and dV_acc from FP32 registers to BF16 and stores them to global dK and dV. A second kernel then converts dQ_accum from FP32 to BF16.The kernel is BF16 in and BF16 out, matching $Q$, $K$, $V$, and $dO$. Tensor Cores accumulate in FP32 and we then cast the results back to BF16.

// This block owns dK_j and dV_j, so no other block writes these elements.
// Convert each lane's two FP32 pairs to BF16 and store them without atomics.
#pragma unroll
for (int local_output_n = 0;
     local_output_n < OUTPUT_MMAS_PER_WARP;
     local_output_n++) {
    // Calculate the global-memory offset for this lane's owned key row and
    // pair of adjacent output columns.
    float *dk = dK_acc[local_output_n];
    float *dv = dV_acc[local_output_n];

    // Convert the FP32 pair to BF16, then store it to global dK.
    reinterpret_cast<nv_bfloat162 *>(dK + offset)[0] =
        __floats2bfloat162_rn(
            dk[0] * softmax_scale,
            dk[1] * softmax_scale);
    reinterpret_cast<nv_bfloat162 *>(
        dK + offset + 8 * HEAD_DIM)[0] =
        __floats2bfloat162_rn(
            dk[2] * softmax_scale,
            dk[3] * softmax_scale);
    // Convert the FP32 pair to BF16, then store it to global dV.
    reinterpret_cast<nv_bfloat162 *>(dV + offset)[0] =
        __floats2bfloat162_rn(dv[0], dv[1]);
    reinterpret_cast<nv_bfloat162 *>(
        dV + offset + 8 * HEAD_DIM)[0] =
        __floats2bfloat162_rn(dv[2], dv[3]);
}

// This kernel is launched after the main backward kernel. The kernel boundary
// guarantees that every atomic contribution to dQ_accum is complete first.
__global__ void convert_dQ_v1_kernel(
    const float *dQ_accum,
    nv_bfloat16 *dQ,
    size_t elements) {
    for (size_t index =
             blockIdx.x * blockDim.x + threadIdx.x;
         index < elements;
         index += blockDim.x * gridDim.x)
        // Convert one FP32 to BF16, then store it to global dQ.
        dQ[index] =
            __float2bfloat16_rn(dQ_accum[index]);
}

That’s it for the full backward kernel for v1. We can now start benchmarking and think about how to improve it : ).

Benchmarking

With the command from the setup section, v1 reaches 10.3287 ms, or 133.07 effective TFLOPS. This is already close to cuDNN, but still leaves a decent gap to the compute lower bound of 6.56 ms.

Using -Xptxas=-v, ptxas reports this for v1:

ptxas info : Function properties for attention_v1_bwd_kernel
    56 bytes stack frame, 60 bytes spill stores, 60 bytes spill loads
ptxas info : Used 128 registers

That means we have spills. 512 threads per block times 128 registers is already the whole 64K:

512 threads * 128 registers = 65,536 registers/block

These spill counts are static. Using Nsight Compute, we can see if local-memory traffic is actually present at runtime:

ncu \
    --profile-from-start off \
    --kernel-name 'regex:attention_v1_bwd_kernel' \
    --launch-count 1 \
    --set full \
    --import-source yes \
    --force-overwrite \
    -o profile_attention_v1_bwd \
    .venv/bin/python kernels/attention/main.py \
      --shape 8_16_4096_128 \
      --profile attention_v1 \
      --direction backward \
      --causal
Nsight Compute memory chart for the v1 main backward kernel.

The memory chart agrees with the compiler report: v1 executes 52.04 million local-memory instructions.Local here means CUDA local memory: a thread-private address space backed by device memory and cached through L1 and L2, not fast on-chip scratchpad memory. The compiler commonly uses it for stack storage and register spills.

Let’s look at the workload analysis in Nsight.

Nsight Compute workload analysis for the v1 main backward kernel.

The compute-workload analysis shows that v1 only achieves 35.5% Tensor Core pipe utilization on sm__pipe_tensor_cycles_active.*.pct_of_peak_sustained_*. However, on consumer cards, this metric is kinda maxxed out at 50% for BF16/FP16 MMA with FP32 accumulate, so we can use rule of thumb to scale it by 2x to about 71%Consumer GeForce SMs have a 50% limit on this pipe metric for HMMA (FP16, BF16, TF32) with FP32 accumulate. It is a counter defect, not a real hardware cap. Look here. Later Tensor utilization numbers in this post use the same 2x correction.. Nice, but still far from making full use of the Tensor Cores.

Version 2

To reduce register pressure and avoid spills, we can first try to put two warps on each row group instead of four. A block will then shrink from 16 warps (512 threads) to 8 (256 threads). Each warp now owns twice as many column fragments, so each thread holds twice as many accumulators and needs more registers, but halving the thread count raises our register budget from 128 to 255. Even at 255, that would still spill if everything doubled. However, the Q/dO fragments shared by a row group and the other per-thread state stay the same size, so the extra accumulators fit and the spills disappear. Using -Xptxas=-v, ptxas reports this for v2:

ptxas info : Function properties for attention_v2_bwd_kernel
    0 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
ptxas info : Used 255 registers

The same change also cuts redundant operand movement. This is the same $Q$/$dO$ redundancy I mentioned in v1.

For $S=QK^T$ on query rows 0-15 and one 16-wide slice of $d$, every warp in the row group loads the same $16\times16$ Q fragment, but only the K columns differ. In v1, four warps each ldmatrix that Q, but in v2, two warps each ldmatrix the same Q, so those ldmatrixs are cut in half.

Both v1 and v2 perform the same eight score MMAs, but v2 uses two warps per row group instead of four, so each loaded Q fragment is reused across four MMAs instead of two.

$dP$, $dK$, $dV$, and $dQ$ also do the same thing with their row-shared operand ($dO$, $P$, or $dS$). However, $K$ and $V$ are not shared that way, so total ldmatrix instructions / traffic drops, but not by half. Moreover, $L$ and $D$ were also loaded by every warp in the row group, so those global loads cut in half.

Seeing the profile measurement in Nsight, we can confirm both effects. Local traffic goes to zero, matching the compiler report. ldmatrix falls 26.7% and global loads fall 50% with the $L$/$D$ warps. Together that is 21.2% fewer instructions:

Metric v1 v2 Change
Latency 10.3287 ms 9.5374 ms -7.7%
Effective TFLOPS 133.07 144.10 +8.3%
Local loads 34.3M 0 -100%
Local stores 17.7M 0 -100%
ldmatrix instructions 223.6M 164.0M -26.7%
Global loads 17.0M 8.5M -50.0%
Executed instructions 1.938B 1.528B -21.2%
Nsight Compute memory chart for the v2 main backward kernel.

Version 3

v2 has one Q/dO tile in shared memory. $Q_i$ is still live through $dK$ and $dO_i$ is still live through $dV$, so we cannot overwrite that buffer until those products finish.

So we can now try double bufferingThis is the N=2 case of N-stage pipelining: two buffers, and we ping-pong between them., which essentially means prefetching the next iteration while we compute on this one. TMA fills the next query tile while we finish the current iteration.

How double buffering works in practice is illustrated in the figured below.

Schematic Q/dO pipeline in v2 and v3. (Note that the box widths illustrate ordering and possible overlap, not measured execution time.)

The same schedule in pseudocode is:

v2 (one stage)

# K_j, V_j already in smem.
async_load(Q, dO, tile=0)

for i in query_tiles:
    wait(Q, dO)

    compute S[i,j] and dP[i,j]
    form P[i,j] and dS[i,j]
    accumulate dV[j] and dK[j]

    synchronize()
    # The only Q/dO slot is now free.

    if tile i+1 exists:
        async_load(Q, dO, tile=i+1)

    # That load overlaps only dQ.
    accumulate dQ[i]

v3 (two stages)

# K_j, V_j, P, dS still one copy.
async_load(Q[0], dO[0], tile=0)
stage = 0

for i in query_tiles:
    wait(Q[stage], dO[stage])

    if tile i+1 exists:
        async_load(Q[stage^1], dO[stage^1], tile=i+1)

    # That Q/dO load overlaps the rest of i.
    compute S[i,j] and dP[i,j]
    form P[i,j] and dS[i,j]
    accumulate dV[j] and dK[j]

    synchronize()
    # Q[stage] and dO[stage] are now free.

    accumulate dQ[i]
    stage ^= 1

Fitting the second stage

Double buffering does not come for free. Since now we have two Q/dO tiles which means we need more shared memory, so just using the same layout as v2 naively will not fitA $64\times128$ BF16 tile takes up 16,384 bytes. Adding a second Q/dO stage naively to v2 would add $2\times16{,}384=32{,}768$ bytes, increasing the shared-memory allocation from 81,936 to 114,704, which is more than the RTX 5090’s limit of 101,376 bytes (99 KiB) of opt-in shared memory per block..

So in v3, we instead halves BLOCK_Q from 64 to 32 while keeping BLOCK_KV=64. Two 32-row Q/dO stages occupy the same space as the single 64-row stage in v2. The P and dS would also shrink from $64\times64$ to $32\times64$, reducing the total allocation to 73,744 bytes (72.0 KiB).

How it works in code is something like this:

load_query_tiles(tensor, first_query_start, 0);
int query_stage = 0;

for (int query_start = first_query_start;
     query_start < sequence;
     query_start += BLOCK_Q) {
    if (warp_id == 0)
        mbarrier_wait(query_barrier_addr, query_phase);
    __syncthreads();

    const int next_query_start = query_start + BLOCK_Q;
    if (next_query_start < sequence)
        load_query_tiles(tensor, next_query_start, query_stage ^ 1);

    // Compute S, dP, dV, dK, and dQ from query_stage.

    query_phase ^= 1;
    query_stage ^= 1;
}

Halving the query tile doubles the iteration count, so v3 pays for the overlap with 5.5% more executed instructions than v2, 1.528B against 1.612B.The added work lands where the tile count predicts. TMA transfers double exactly, 1.10M to 2.16M, as do the per-iteration $L$/$D$ global loads, 8.52M to 17.04M. ldmatrix rises only 25.3%, from 164.0M to 205.5M, because $K$ and $V$ are still loaded once per block and reused across every query tile. However, we change how often those warps can issue.

See metrics below. Prefetching into the second stage removes most of the long scoreboard stallA long scoreboard stall is a warp waiting on a global-memory dependency. In v2 that is mostly the wait for the current $Q_i$ and $dO_i$ to arrive., so the schedulers find something to issue on a larger fraction of cycles, and we utilize the Tensor Cores more.

Metric v2 v3
Latency 9.5374 ms 8.5866 ms
Effective TFLOPS 144.10 160.06
Long scoreboard stall 0.93 0.21
Issue slots busy 10.83% 12.79%
Tensor utilization 75.64% 84.62%

v3 is 76.40% of the theoretical peak.

End Notes

I already knew most of these, including swizzled shared memory, ldmatrix and mma fragment layouts, register pressure, TMA with mbarrier, and double buffering, but knowing them and using them effectively to make kernels fast are essentially two different things. It was really fun tore-derive all the maths (which I think is non-trivial), and to think about how to map it efficiently onto the hardware.

I currently stopped at v3 for now. There is certainly more performance left, but chasing it would mean diminishing returns for my long-term goal. For me, the point of this project was to act as a stepping stone toward designing new algorithms that are not only elegant, but also map efficiently to the hardware and run extremely fast in practice.

Nsight Compute produced most of the measurements in this post, while Excalidraw was where I wrote most of the diagrams. All kernel versions, correctness tests, and benchmarks are in the GitHub repository.

Appendix A: Kernel Helpers

The kernel excerpts above use thin wrappers around the PTX instructions for lane election, TMA transfers, mbarrier, ldmatrix, and Tensor Core MMA, plus the swizzled shared-memory address helpers.

A.1 Shared-memory addresses

cvta_shared() converts generic C++ pointers to 32-bit shared-memory addresses. tile_addr() is the ldmatrix address for 128-column Q, K, V, or dO. matrix_addr() is the same for 64-column P and dS; matrix_offset() is the element index used when storing P and dS.

__device__ inline uint32_t cvta_shared(const void *ptr) {
    return static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
}

// Byte address for ldmatrix on 128-col Q/K/V/dO (two 128-byte panels).
auto tile_addr = [&](const nv_bfloat16 *base, int row, int col) {
    constexpr int VECTOR_ELEMENTS = 16 / sizeof(nv_bfloat16); // 16-byte chunk
    const int panel = col / PANEL_ELEMENTS; // [0, 64) -> 0, [64, 128) -> 1
    const int chunk = (col % PANEL_ELEMENTS) / VECTOR_ELEMENTS; // 16-byte chunk index 0..7
    return cvta_shared(base) +
        panel * PANEL_BYTES +
        swizzle_128b_panel_offset<128>(row, chunk);
};

// Element index for C++ stores into 64-col P/dS.
auto matrix_offset = [&](int row, int col) {
    constexpr int VECTOR_ELEMENTS = 16 / sizeof(nv_bfloat16); // 16-byte chunk
    const int chunk = col / VECTOR_ELEMENTS; // 16-byte chunk index
    return
        swizzle_128b_panel_offset<MATRIX_ROW_BYTES>(row, chunk) /
            sizeof(nv_bfloat16) +
        col % VECTOR_ELEMENTS;
};

// Byte address for ldmatrix on 64-col P/dS (one 128-byte panel).
auto matrix_addr = [&](const nv_bfloat16 *base, int row, int col) {
    constexpr int VECTOR_ELEMENTS = 16 / sizeof(nv_bfloat16); // 16-byte chunk
    const int chunk = col / VECTOR_ELEMENTS; // 16-byte chunk index
    return cvta_shared(base) +
        swizzle_128b_panel_offset<MATRIX_ROW_BYTES>(row, chunk);
};

A.2 TMA and synchronization

__device__ inline bool elect_one_sync() {
    int elected = 0;
    asm volatile(
        "{\n"
        ".reg .pred p;\n"
        "elect.sync _|p, %1;\n"
        "@p mov.s32 %0, 1;\n"
        "}"
        : "+r"(elected)
        : "r"(0xffffffff));
    return elected;
}

__device__ inline void mbarrier_init(uint32_t addr, int count) {
    asm volatile(
        "mbarrier.init.shared::cta.b64 [%0], %1;"
        :
        : "r"(addr), "r"(count));
}

__device__ inline void mbarrier_fence_init() {
    asm volatile("fence.mbarrier_init.release.cluster;");
}

__device__ inline void mbarrier_arrive_expect_tx(
    uint32_t addr,
    int bytes) {
    asm volatile(
        "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 "
        "_, [%0], %1;"
        :
        : "r"(addr), "r"(bytes)
        : "memory");
}

__device__ inline void mbarrier_wait(uint32_t addr, int phase) {
    asm volatile(
        "{\n"
        ".reg .pred done;\n"
        "wait:\n"
        "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 "
        "done, [%0], %1;\n"
        "@!done bra.uni wait;\n"
        "}"
        :
        : "r"(addr), "r"(phase)
        : "memory");
}

__device__ inline void tma_3d_g2s(
    uint32_t dst,
    const void *tensor_map,
    int x,
    int y,
    int z,
    uint32_t mbarrier) {
    asm volatile(
        "cp.async.bulk.tensor.3d.shared::cta.global."
        "mbarrier::complete_tx::bytes "
        "[%0], [%1, {%2, %3, %4}], [%5];"
        :
        : "r"(dst), "l"(tensor_map), "r"(x), "r"(y), "r"(z),
          "r"(mbarrier)
        : "memory");
}

A.3 Register fragments and Tensor Cores

__device__ inline void ldmatrix_x4(uint32_t reg[4], uint32_t addr) {
    asm volatile(
        "ldmatrix.sync.aligned.m8n8.x4.shared.b16 "
        "{%0, %1, %2, %3}, [%4];"
        : "=r"(reg[0]), "=r"(reg[1]), "=r"(reg[2]), "=r"(reg[3])
        : "r"(addr));
}

__device__ inline void ldmatrix_x4_trans(
    uint32_t reg[4],
    uint32_t addr) {
    asm volatile(
        "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 "
        "{%0, %1, %2, %3}, [%4];"
        : "=r"(reg[0]), "=r"(reg[1]), "=r"(reg[2]), "=r"(reg[3])
        : "r"(addr));
}

__device__ inline void mma_m16n8k16(
    const uint32_t A[4],
    const uint32_t B[2],
    float C[4]) {
    asm volatile(
        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
        "{%0, %1, %2, %3}, "
        "{%4, %5, %6, %7}, "
        "{%8, %9}, "
        "{%0, %1, %2, %3};"
        : "+f"(C[0]), "+f"(C[1]), "+f"(C[2]), "+f"(C[3])
        : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]),
          "r"(B[0]), "r"(B[1]));
}

Appendix B: Softmax Backward Derivation

For a fixed row of $S$ and $P$, with $s=S_{r,:}^T$ and $p=P_{r,:}^T$, and similarly for $ds$ and $dp$. Then

$$ p_i=\frac{e^{s_i}}{\sum_k e^{s_k}}. $$

Define $Z=\sum_k e^{s_k}$, so that $p_i=e^{s_i}/Z$. Then we have:

$$ \frac{\partial e^{s_i}}{\partial s_j}=\delta_{ij}e^{s_i}, \qquad \frac{\partial Z}{\partial s_j}=e^{s_j}. $$

Using the quotient rule gives

$$ \begin{aligned} \frac{\partial p_i}{\partial s_j} &=\frac{\delta_{ij}e^{s_i}Z-e^{s_i}e^{s_j}}{Z^2} \\ &=\delta_{ij}p_i-p_ip_j \\ &=p_i\left(\delta_{ij}-p_j\right). \end{aligned} $$

Applying the chain rule,

$$ ds_j =\frac{\partial \mathcal{L}}{\partial s_j} =\sum_i\frac{\partial \mathcal{L}}{\partial p_i} \frac{\partial p_i}{\partial s_j} =\sum_i dp_i p_i\left(\delta_{ij}-p_j\right). $$

Expanding the sum,

$$ \begin{aligned} ds_j &=\sum_i dp_i p_i(\delta_{ij}-p_j) \\ &=p_j dp_j-p_j\sum_i p_i dp_i \\ &=p_j\bigl(dp_j-p^T dp\bigr). \end{aligned} $$

In vector form this is

$$ ds =p\odot\bigl(dp-(p^T dp)\mathbf{1}\bigr). $$

Appendix C: Estimating the Speed of Light

Using the same shape as the benchmarks: B=8, H=16, N=4096, d=128.

Five matmuls ($QK^T$, $dV$, $dP$, $dQ$, $dK$). Causal attention skips the upper triangle, so we count about half the query-key pairsThe exact causal triangle contains $N(N+1)/2$ valid query-key pairs. I use the $N^2/2$ approximation for simplicity.Softmax, masking, scaling, and other non-matmul operations are not in this FLOP count.:

$$ \text{FLOPs} =5\cdot 2BH\Bigl(\frac{N^2}{2}\Bigr)d =5BHN^2d \approx 1.374\ \text{TFLOPs}. $$

The memory lower bound is to read $Q$, $K$, $V$, $O$, $dO$ (BF16) and $L$ (FP32) once each, and write $dQ$, $dK$, $dV$ (BF16) onceA column-parallel schedule reloads the same $Q$, $dO$, $L$, and $D$ in every column-tile block, so the kernel reads more than this bound.:

$$ \begin{aligned} \text{read} &=5(BHNd)\cdot 2+(BHN)\cdot 4 &&\approx 673\ \text{MB}, \\ \text{write} &=3(BHNd)\cdot 2 &&\approx 403\ \text{MB}. \end{aligned} $$

About $1.076$ GB if every tensor moves once and intermediates stay on chip. Locked clocks: $209.5$ TFLOP/s dense BF16 Tensor (FP32 acc) and $1.792$ TB/s.


Roofline model

    Peak Bound
Compute 1.374 TFLOPs 209.5 TFLOP/s 6.56 ms
Memory 1.076 GB 1.792 TB/s 0.60 ms

So under the roofline model, the kernel is compute-bound.

← Back to all posts