July 8, 2026 · TUTORIALS · GPU
Learning FlashAttention the Hard Way — Part 2
Efficient CUDA Kernel Generation
This is Part 2/2. Part 1 covers the math. Attention becomes an associative reduction with the help of the small carrier state triple
(m, d, o). This part covers the efficient GPU lowering.
Congratulations on clearing the Algebra Boss in Part 1! However, if you took the shortcut and went straight to this article, don't worry, this one will be just as painful!
In the first part we explored how to recognize "secretly associative" operations that allow efficient parallelization. Now, we're going to take those to the metal and generate efficient GPU code. In the end we will reproduce the SOTA handwritten FlashAttention-2 kernel on RTX 5090 (FA-3 and FA-4, unfortunately, will not run on consumer GPUs) with each optimization explained and quantified. This is primarily educational piece, however, it comes with an interesting result as well: given no attention-specific knowledge, the autotuner lands on the exact geometry FlashAttention-2's author picked by hand and matches its latency to the microsecond. I also keep the negative results: transferable optimizations from FA-3 and FA-4 applied, measured and refused.
Kernels in this article are not handwritten: Emmy compiler generates them. Emmy can apply a wide range of optimizations like register-tiling, smem operand staging, etc., each controlled by a set of knobs. This approach is much more scalable as optimizations are not limited to attention. Also, the tool might be useful for you if you're interested in learning how ML compilers work.
You can reproduce every kernel in this article by invoking the emmy compile CLI command after you install it in a Python virtual environment. A quick bootstrap script:
git clone https://github.com/cloudrift-ai/emmy.git
cd emmy && make setup && source venv/bin/activate
Throughout the article, I will be using two attention kernels: SDPA_P for the real performance benchmarks and a smaller SDPA_L for code listings. Define them in your shell if you wish to reproduce kernels on the fly.
# attention kernel for performance benchmark
SDPA_P='q=torch.randn(1,8,4096,64,dtype=torch.float16);\
k=torch.randn(1,8,4096,64,dtype=torch.float16);\
v=torch.randn(1,8,4096,64,dtype=torch.float16);\
F.scaled_dot_product_attention(q,k,v)'
# attention kernel for code listing
SDPA_L='q=torch.randn(1,4,128,64,dtype=torch.float16);\
k=torch.randn(1,4,128,64,dtype=torch.float16);\
v=torch.randn(1,4,128,64,dtype=torch.float16);\
F.scaled_dot_product_attention(q,k,v)'
The Scoreboard
| rung | what changed | µs | speedup |
|---|---|---|---|
| Move 0 — scalar streaming kernel | the fused shared-score chain, one FMA at a time | 25959 | — |
| Move 1 — tensor streaming kernel | both matmuls on the tensor cores | 612 | baseline |
| Move 2 — smem operand staging | K/V tiles through shared memory | 511 | 1.20× |
| — the bank-conflict fix (padding) | +16 B slab rows; 132M conflict replays → ~0 | 315 | 1.62× |
| — the C→A register repack | P feeds P@V without touching smem | 271 | 1.16× |
| — the Q hoist | loop-invariant Q fragments load once | 233 | 1.16× |
| Move 3 — pipelining | prefetch KV tile i+1 under tile i's math | 222 | 1.05× |
— the x4 drain-pairing pass | two ldmatrix.x2 fuse into one x4 | 219 | 1.01× |
| — TMA transport | rank-4 box copies, hardware swizzle | 211 | 1.04× |
| Move 4 — register tiling | two independent softmax chains per warp | 206 | 1.02× |
| FlashAttention-2 (PyTorch) | the hand-written reference | 206 | 1.00× |
Four optimizations were refused: a warp-scoped barrier, the fast-exp path, deeper rings, and warp specialization.
The Move Stack
Optimizations are controlled by the EMMY_KNOBS environment variable. Let's take a look at the codec it uses to describe them by unpacking the following example:
export EMMY_KNOBS='TILE=a:mma_m16n8k16_f16/w4x1/f2x8/k4,STAGE=d2/tma/ring,REDUCE=g4k'
-
TILE=a:mma_m16n8k16_f16/w4x1/f2x8/k4controls the tiling of contraction (matmul) kernelsameans atom (tensor) to use, e.g.mma_m16n8k16_f16;a:noneora:scalarforces Emmy to emit the scalar kernelw4x1defines dimensions of the warp grid needed for tensor cores,4x1in this casef2x8defines each warp's register sub-tilek4defines the K-chunk, four 16-wide atom-K steps = a 64-key block
-
STAGE=d2/tma/ringcontrols the operand staging pipelined2defines a depth-2 gmem→smem pipeline, i.e. double-buffering,d3would force triple-bufferingtmais the transport to use (the Blackwell box-copy engine), other options arecp.asyncand a legacy register transportringinstructs the kernel to prefetch tile i+1's copy under tile i's math, i.e. leverage the asynchrony available undercp.asyncortmatransport
REDUCE=g4kis used for reduction kernelsg4splits the reduction across 4 CTAs, so each CTA is reducing a respective portion of the input arraykcombines the four partial results with a separate kernel;awould combine them with atomics
Move 0 — Scalar Streaming Kernel
First, let's generate the naive kernel for contrast. To do that, we use the knob PLACE@fold=cut that prevents the fusion of the online-softmax kernel, and TILE=a:scalar that forces scalar matmuls.
The kernels are exact listings from emmy compile commands up to variable names and comments. Listings were post-processed for readability, but the structure was kept intact.
EMMY_KNOBS="PLACE@fold=cut,PLACE@tuple=cut,TILE=a:scalar,REDUCE=" emmy compile -c "$SDPA_P"
// ── Kernel 1 of 2: the score producer ──
// S = (Q·Kᵀ)·scale for every (query, key) pair, written out to HBM in full.
// One thread computes one score cell of the [1, 8, 4096, 4096] matrix.
extern "C" __global__ __launch_bounds__(256)
void k_sdpa_reduce_3cd3bc(
const __half* q,
const __half* k,
const __half* sdpa_scale,
__half* sdpa_scaled)
{
int _gid = blockIdx.x * blockDim.x + threadIdx.x;
// Decode the flat thread id → (batch·head a0, query row a1, key col a2).
int a0 = _gid / 16777216;
int a1_b = (_gid / 4096) % 4096;
int a2_b = (_gid) % 4096;
int a1_u = (_gid) % 1; // register sub-tile index — always 0 at scalar tier
int a2_u = (_gid) % 1;
// Dot product over the head dim (D=64): the raw QKᵀ score for this row/col.
float acc0__c0_0 = 0.0f;
#pragma unroll
for (int a3 = 0; a3 < 64; a3++) {
__half in2__ar0 = q[a0 * 262144LL + (a1_b + a1_u) * 64 + a3];
__half in1__bc0 = k[a0 * 262144LL + (a2_b + a2_u) * 64 + a3];
__half acc0__v__c0_0 = in1__bc0 * in2__ar0;
acc0__c0_0 += __half2float(acc0__v__c0_0);
}
__half in0__c0_0 = sdpa_scale[0];
float v1__c0_0 = acc0__c0_0 * __half2float(in0__c0_0); // apply the 1/√D scale
// The villain: [.., 4096, 4096] score write to HBM
sdpa_scaled[a0 * 16777216LL + (a1_b + a1_u) * 4096 + (a2_b + a2_u)] = __float2half(v1__c0_0);
}
// ── Kernel 2 of 2: two-pass softmax, then P@V ──
// Reads the score matrix back from HBM three times: a row-max pass, a Σexp
// pass, then a normalize-and-accumulate pass against V. One thread per row.
extern "C" __global__ __launch_bounds__(256)
void k_sdpa_reduce_3ce0d7(
const __half* sdpa_scaled,
const __half* v,
__half* sdpa)
{
int _gid = blockIdx.x * blockDim.x + threadIdx.x;
int a0 = _gid / 4096; // batch·head
int a1 = (_gid) % 4096; // query row
// Pass 1 — row max m over the 4096 keys (for softmax numerical stability).
__half acc0 = __float2half(-1e+30f);
for (int a2 = 0; a2 < 4096; a2++) {
__half in0 = sdpa_scaled[a0 * 16777216LL + a1 * 4096 + a2];
acc0 = __hmax(acc0, in0);
}
// Pass 2 — denominator d = Σ e^{sⱼ − m}.
float acc1 = 0.0f;
for (int a2 = 0; a2 < 4096; a2++) {
__half in1 = sdpa_scaled[a0 * 16777216LL + a1 * 4096 + a2];
__half v1 = __float2half(expf((in1 - acc0)));
acc1 += __half2float(v1);
}
float v2 = 1.0f / acc1; // reciprocal denominator, applied per element
// Pass 3 — O = Σ softmax(s)ⱼ · vⱼ, one accumulator per head-dim channel.
for (int a3 = 0; a3 < 64; a3++) {
float acc2 = 0.0f;
for (int a4 = 0; a4 < 4096; a4++) {
__half in2 = sdpa_scaled[a0 * 16777216LL + a1 * 4096 + a4];
__half in3 = v[a0 * 262144LL + a4 * 64 + a3];
float v6 = __half2float(in3) * (v2 * expf((in2 - acc0)));
acc2 += v6;
}
sdpa[a0 * 262144LL + a1 * 64 + a3] = __float2half(acc2);
}
}
Two kernels. The first computes every score and writes the full [1, 8, 4096, 4096] tensor to HBM; the second reads it straight back for the softmax and the P@V.
Now, let's enable fusion but pin the scalar tier explicitly. TILE=a:scalar keeps the matmuls off the tensor cores; TILE@pj=f64 puts the whole 64-channel output row into one thread's registers, so the score is computed ONCE per streamed key and feeds all 64 channels. This is FA-1.
The math behind this transformation is covered in Part 1. TLDR version: we don't need to compute the softmax maximum as a separate pass: with the help of a small state triple
(m, d, o)we can update the maximum and the desired output value in the same loop body, eliminating both the separate pass and the costly intermediate matrix materialization.
EMMY_KNOBS="PLACE=fuse,TILE=a:scalar,TILE@pj=f64,REDUCE=" emmy compile -c "$SDPA_P"
// FA-1: one fused kernel, one QUERY ROW per thread. Each score is consumed the
// instant it is produced and never written to HBM — the softmax runs online,
// streaming over the KV axis. Carrier (m_i, l_i, O_i_0..63) = running
// (row-max, denominator, output row). The 64 unrolled channel lines are elided
// to the first two throughout.
extern "C" __global__ __launch_bounds__(256)
void k_sdpa(
const __half* k,
const __half* q,
// you may wonder why do I pass single constant as a gmem array
// all constants are passed via gmem to avoid CUDA-graph recompilation
// the worst offender is the seq_len constant that changes depending on input
const float* _flash_scale,
const __half* v,
__half* sdpa)
{
int _gid = blockIdx.x * blockDim.x + threadIdx.x;
int b0 = _gid / 32768; // batch
int b1 = (_gid / 4096) % 8; // head
int m = (_gid) % 4096; // query row this thread owns
float l_i = 0.0f; // denominator d
float O_i_0 = 0.0f; // output row o — one register per channel
float O_i_1 = 0.0f;
// … O_i_2 … O_i_63 …
float m_i = -1e+30f; // running max m (identity −∞)
// the streaming reduce
for (int kv = 0; kv < 4096; kv++) {
// score sⱼ = q · kⱼ over the head dim — computed once, shared by all channels.
float sacc = 0.0f;
for (int dd = 0; dd < 64; dd++) {
__half k_e = k[b0 * 2097152LL + b1 * 262144LL + kv * 64 + dd];
__half q_e = q[b0 * 2097152LL + b1 * 262144LL + m * 64 + dd];
__half sacc__v = k_e * q_e;
sacc += __half2float(sacc__v);
}
float scale_c = _flash_scale[0];
float s = sacc * scale_c; // scaled score sⱼ
// Online-softmax update — the twist that keeps the reduce associative:
float m_i__t0 = fmaxf(m_i, s); // m_new = max(m_i, sⱼ)
float m_i__t2 = expf((m_i - m_i__t0)); // α = e^{m_i − m_new}
float m_i__t3 = l_i * m_i__t2; // rescale old denominator
float m_i__t5 = expf((s - m_i__t0)); // p = e^{sⱼ − m_new}
float m_i__t6_0 = O_i_0 * m_i__t2; // rescale old output, per channel
float m_i__t6_1 = O_i_1 * m_i__t2;
// … m_i__t6_2 … m_i__t6_63 …
l_i = m_i__t3 + m_i__t5; // d ← d·α + p
// P@V: fold vⱼ weighted by p into the row (V read as uint4 = 8 halves/load).
float O_i__pv_0 = 0.0f;
float O_i__pv_1 = 0.0f;
// … O_i__pv_2 … O_i__pv_63 …
for (int pj = 0; pj < 1; pj++) {
float O_i__p = m_i__t5;
uint4 _v_v_e_0 = *reinterpret_cast<const uint4*>(
&v[b0 * 2097152LL + b1 * 262144LL + kv * 64]);
const __half* _v_v_e_0_h = reinterpret_cast<const __half*>(&_v_v_e_0);
__half v_e_0 = _v_v_e_0_h[0];
__half v_e_1 = _v_v_e_0_h[1];
float O_i__pv__v_0 = __half2float(v_e_0) * O_i__p;
float O_i__pv__v_1 = __half2float(v_e_1) * O_i__p;
// … the other seven uint4 V loads and 62 channel MACs …
O_i__pv_0 += O_i__pv__v_0;
O_i__pv_1 += O_i__pv__v_1;
}
O_i_0 = m_i__t6_0 + O_i__pv_0; // o ← o·α + p·vⱼ, per channel
O_i_1 = m_i__t6_1 + O_i__pv_1;
// … O_i_2 … O_i_63 …
m_i = fmaxf(m_i, s); // commit the new running max
}
// deferred normalize (last moment), stored back as uint4 = 8 halves/store
float O_i__proj_0 = O_i_0 / l_i;
float O_i__proj_1 = O_i_1 / l_i;
// … O_i__proj_2 … O_i__proj_63 …
__half _vs[8] = { __float2half(O_i__proj_0), __float2half(O_i__proj_1), /* … */ };
*reinterpret_cast<uint4*>(&sdpa[b1 * 262144LL + m * 64])
= *reinterpret_cast<const uint4*>(_vs);
// … the other seven uint4 stores …
}
This is a pretty terrible kernel that achieves only about 2.5% FMA-pipe utilization on RTX 5090. Better scheduling could claw some of that back, but not much: every scalar MAC spends ~5 instructions (two loads, a multiply, a convert, an add) to retire 2 flops, so even a perfectly scheduled scalar kernel tops out roughly 7× above this one. So let's jump straight to the tensor-core version and optimize that.
Move 1 — Tensor Cores
The biggest optimization lever we can pull is the tensor cores. First, let's introduce a few helpers that will be used in the code listing below to make it more readable.
Even though CUDA comes with high-level APIs for tensor-core instructions, I am not using them in Emmy. They are stateful, and thus are more difficult to use in codegen, and require additional heavy dependencies. Instead, Emmy emits a preamble for each kernel that mimics
nvcuda::wmmafunctions. Operands are still stored in registers, so we don't have those nicewmma::fragmentstructures in the code. However, since some optimizations require register repacking, thewmma::fragmentAPI doesn't cut it anyway.
// one mma.sync.m16n8k16 tensor-core op: d = a·b + c, fp32 accumulate,
// fp16 operands
static __device__ __forceinline__ void emmy_mma_m16n8k16_f16(
float* d, const unsigned* a, const unsigned* b, const float* c) {
asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
"{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
: "=f"(d[0]),"=f"(d[1]),"=f"(d[2]),"=f"(d[3])
: "r"(a[0]),"r"(a[1]),"r"(a[2]),"r"(a[3]),"r"(b[0]),"r"(b[1]),
"f"(c[0]),"f"(c[1]),"f"(c[2]),"f"(c[3]));
}
// read one A-fragment straight from global memory — each lane loads its own 4
// registers at the PTX m16n8k16 lane→element map (the B loads
// emmy_mma_load_b_gmem[_trans] are the same idea, N-major)
template <typename T>
static __device__ __forceinline__ void emmy_mma_load_a_gmem(
unsigned* r, const T* g, int ldm) {
int lane = threadIdx.x & 31, grp = lane >> 2, tig = lane & 3;
#pragma unroll
for (int i = 0; i < 4; ++i) {
int row = grp + ((i & 1) ? 8 : 0); // M: groupID, +8 for the 2nd row block
int col = (tig << 1) + ((i & 2) ? 8 : 0); // K: 2*tid_in_group, +8 for k16
const T* p = g + row * ldm + col;
unsigned packed;
((T*)&packed)[0] = p[0]; ((T*)&packed)[1] = p[1]; // .f16x2: (col, col+1)
r[i] = packed;
}
}
// the C→A register repack: two k-adjacent fp32 C-fragments → one fp16 A-operand
// fragment, four cvt.rn.f16x2 — no shuffle, no smem round-trip (feeds P@V)
static __device__ __forceinline__ void emmy_c_to_a_f16(
unsigned* a, const float* c0, const float* c1) {
asm("cvt.rn.f16x2.f32 %0, %1, %2;\n" : "=r"(a[0]) : "f"(c0[1]), "f"(c0[0]));
asm("cvt.rn.f16x2.f32 %0, %1, %2;\n" : "=r"(a[1]) : "f"(c0[3]), "f"(c0[2]));
asm("cvt.rn.f16x2.f32 %0, %1, %2;\n" : "=r"(a[2]) : "f"(c1[1]), "f"(c1[0]));
asm("cvt.rn.f16x2.f32 %0, %1, %2;\n" : "=r"(a[3]) : "f"(c1[3]), "f"(c1[2]));
}
Now, let's force Emmy to emit flash attention kernels leveraging tensor cores as follows:
EMMY_KNOBS="PLACE=fuse,TILE=a:mma_m16n8k16_f16/w1x1/f1x2/k4,STAGE=" \
emmy compile -c "$SDPA_L" --target sm_120 --ir cuda
Both matmuls are now using mma.sync. The softmax lives between the C-fragments, P@V fed by the C→A register repack:
extern "C" __global__ __launch_bounds__(32)
void k_sdpa(
const __half* q,
const float* _flash_scale,
const __half* k,
const __half* v,
__half* sdpa
) {
// thread-id decode: global id → batch b0, head b1, query-row tile m, warp lane
int _gid = blockIdx.x * blockDim.x + threadIdx.x;
int b0 = _gid / 1024; // batch index
int b1 = (_gid / 256) % 4; // head index
int m = (_gid / 32) % 8; // query-row (M16) tile
int _lane = (_gid) % 32; // lane within the warp
// carrier: m_i = running row-max, l_i = running denominator (Σ p),
// O_i_f = the f32 output fragments accumulated across the KV stream
float m_i0 = -1e+30f;
float l_i0 = 0.0f;
float m_i1 = -1e+30f;
float l_i1 = 0.0f;
float O_i_f[8][4] = {};
// A fragments: load this warp's Q row-tile straight from gmem (m16n8k16 A)
unsigned _sacc_a[4][4];
#pragma unroll
for (int _r2 = 0; _r2 < 4; _r2 += 1) {
emmy_mma_load_a_gmem(
_sacc_a[_r2],
&q[b0 * 32768 + b1 * 8192 + m * 16 * 64 + _r2 * 16],
64);
}
float scale_c = _flash_scale[0]; // softmax scale 1/√d
// KV streaming loop: one 16-key block per step (flash online softmax)
#pragma unroll
for (int kv0 = 0; kv0 < 128; kv0 += 16) {
float sacc_f[2][4] = {};
unsigned _sacc_b[2][2];
// QKᵀ score mma: S = Q @ Kᵀ on mma.sync (Kᵀ via transposed-B load)
#pragma unroll
for (int _r1 = 0; _r1 < 4; _r1 += 1) {
#pragma unroll
for (int _r0 = 0; _r0 < 2; _r0 += 1) {
emmy_mma_load_b_gmem_trans(_sacc_b[_r0], &k[b0 * 32768 + b1 * 8192 + (kv0 + _r0 * 8) * 64 + _r1 * 16], 64);
emmy_mma_m16n8k16_f16(sacc_f[_r0], _sacc_a[_r1], _sacc_b[_r0], sacc_f[_r0]);
}
}
// score scale: sⱼ ← S · scale_c
#pragma unroll
for (int _r0 = 0; _r0 < 2; _r0 += 1) {
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
sacc_f[_r0][_e] = sacc_f[_r0][_e] * scale_c;
}
}
// online-softmax row-max: reduce the 4 score regs, then a __shfl_xor
// butterfly (lanes 2,1) folds the row across the group → block row-max
float _rmx0 = fmaxf(fmaxf(fmaxf(sacc_f[0][0], sacc_f[0][1]), sacc_f[1][0]), sacc_f[1][1]);
float _rmx1 = fmaxf(fmaxf(fmaxf(sacc_f[0][2], sacc_f[0][3]), sacc_f[1][2]), sacc_f[1][3]);
_rmx0 = fmaxf(_rmx0, __shfl_xor_sync(0xffffffff, _rmx0, 2));
_rmx1 = fmaxf(_rmx1, __shfl_xor_sync(0xffffffff, _rmx1, 2));
_rmx0 = fmaxf(_rmx0, __shfl_xor_sync(0xffffffff, _rmx0, 1));
_rmx1 = fmaxf(_rmx1, __shfl_xor_sync(0xffffffff, _rmx1, 1));
// twist: m_new = max(m_i, block-max), α = e^{m_i − m_new}
float _mn0 = fmaxf(m_i0, _rmx0); // m_new (rows 0-7)
float _mn1 = fmaxf(m_i1, _rmx1); // m_new (rows 8-15)
float _al0 = expf((m_i0 - _mn0)); // α = e^{m_i − m_new}
float _al1 = expf((m_i1 - _mn1)); // α (rows 8-15)
// p = e^{sⱼ − m_new}: rescale each score reg to the new row-max
float _p_f0[4];
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
_p_f0[_e] = expf(sacc_f[0][_e] - ((_e < 2) ? (_mn0) : (_mn1)));
}
float _p_f1[4];
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
_p_f1[_e] = expf(sacc_f[1][_e] - ((_e < 2) ? (_mn0) : (_mn1)));
}
// denominator: reduce Σ p over the 4 regs, __shfl butterfly across group
float _rsm0 = _p_f0[0] + _p_f0[1] + _p_f1[0] + _p_f1[1];
float _rsm1 = _p_f0[2] + _p_f0[3] + _p_f1[2] + _p_f1[3];
_rsm0 = _rsm0 + __shfl_xor_sync(0xffffffff, _rsm0, 2);
_rsm1 = _rsm1 + __shfl_xor_sync(0xffffffff, _rsm1, 2);
_rsm0 = _rsm0 + __shfl_xor_sync(0xffffffff, _rsm0, 1);
_rsm1 = _rsm1 + __shfl_xor_sync(0xffffffff, _rsm1, 1);
float l_i__n0 = (l_i0 * _al0) + _rsm0; // d ← d·α + p
float l_i__n1 = (l_i1 * _al1) + _rsm1;
l_i0 = l_i__n0;
l_i1 = l_i__n1;
// output rescale: o ← o·α before this block's P@V accumulates in
#pragma unroll
for (int _r0 = 0; _r0 < 8; _r0 += 1) {
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
O_i_f[_r0][_e] = O_i_f[_r0][_e] * ((_e < 2) ? (_al0) : (_al1));
}
}
// C→A repack: the f32 P fragment converts in-place to the f16 A
// operand for P@V — lane maps align, no shuffle / smem round-trip
unsigned _pa[4];
emmy_c_to_a_f16(_pa, _p_f0, _p_f1);
// P@V mma: O_i += P @ V on mma.sync, streaming V's 8 N-tiles from gmem
unsigned _O_i__pv_b[8][2];
#pragma unroll
for (int _r0 = 0; _r0 < 8; _r0 += 1) {
emmy_mma_load_b_gmem(_O_i__pv_b[_r0], &v[b0 * 32768 + b1 * 8192 + kv0 * 64 + _r0 * 8], 64);
emmy_mma_m16n8k16_f16(O_i_f[_r0], _pa, _O_i__pv_b[_r0], O_i_f[_r0]);
}
m_i0 = _mn0; // commit the new row-max
m_i1 = _mn1;
}
// final normalize: o ← o / l (the accumulated denominator)
#pragma unroll
for (int _r0 = 0; _r0 < 8; _r0 += 1) {
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
O_i_f[_r0][_e] = O_i_f[_r0][_e] / ((_e < 2) ? (l_i0) : (l_i1));
}
}
// store: repack each C fragment as __half2 pairs into the output rows
#pragma unroll
for (int _r0 = 0; _r0 < 8; _r0 += 1) {
const int _g = (threadIdx.x & 31) >> 2; const int _t = (threadIdx.x & 31) & 3;
*reinterpret_cast<__half2*>(&sdpa[b1 * 8192 + m * 16 * 64 + _r0 * 8 + _g * 64 + _t * 2])
= __floats2half2_rn(O_i_f[_r0][0], O_i_f[_r0][1]);
*reinterpret_cast<__half2*>(&sdpa[b1 * 8192 + m * 16 * 64 + _r0 * 8 + (_g + 8) * 64 + _t * 2])
= __floats2half2_rn(O_i_f[_r0][2], O_i_f[_r0][3]);
}
}
This is a much more reasonable kernel that takes 612µs on RTX 5090. It will serve as a baseline.
Move 2 — SMEM Staging (1.20x)
Once KV is tiled, each tile's K and V rows are read by every query row in the block. The STAGE@<kv> codec leverages shared memory to avoid re-reading arguments from global memory.
Staging brings two new helper families into the preamble: the cp.async fill and the ldmatrix drain:
// one 16-byte cp.async: async copy gmem→smem that bypasses the register file
// (.cg = cache-global, bypass L1 — the streaming form). commit closes a batch;
// wait<N> blocks until ≤ N batches are in flight.
static __device__ __forceinline__ void emmy_cp_async_cg(void* smem, const void* gmem) {
asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n"
:: "r"(__cvta_generic_to_shared(smem)), "l"(gmem) : "memory");
}
static __device__ __forceinline__ void emmy_cp_async_commit() {
asm volatile("cp.async.commit_group;\n" ::: "memory");
}
template <int N>
static __device__ __forceinline__ void emmy_cp_async_wait() {
asm volatile("cp.async.wait_group %0;\n" :: "n"(N) : "memory");
}
// one ldmatrix.x4: load a warp-distributed 16-bit fragment from shared memory
// into 4 registers (the drain that feeds the mma from the staged slab;
// emmy_ldmatrix_x2[_trans] is the 2-register B form)
static __device__ __forceinline__ void emmy_ldmatrix_x4(unsigned* r, const void* smem) {
asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
: "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
: "r"(__cvta_generic_to_shared(smem)));
}
Compile command:
EMMY_KNOBS="PLACE=fuse,TILE=a:mma_m16n8k16_f16/w1x1/f1x2/k4,STAGE=d1/cp" \
emmy compile -c "$SDPA_L" --target sm_120 --ir cuda # helper preamble elided
extern "C" __global__ __launch_bounds__(32)
void k_sdpa(
const __half* q,
const float* _flash_scale,
const __half* k,
const __half* v,
__half* sdpa
) {
// ---- thread-id decode: 32 lanes = one warp, which owns one (batch,head,m-tile)
int _gid = blockIdx.x * blockDim.x + threadIdx.x;
int b0 = _gid / 1024; // batch (4 heads × 8 m-tiles × 32 lanes = 1024)
int b1 = (_gid / 256) % 4; // head (0..3)
int m = (_gid / 32) % 8; // query m-tile of 16 rows within the 128-long seq
int _lane = (_gid) % 32; // lane within the warp
// K / V staging slabs: 16 rows × 72 halves; 72 = 64 head-dim + 8 pad, which
// dodges shared-memory bank conflicts on the ldmatrix drain.
__shared__ __half _k_smem[1152];
__shared__ __half _v_smem[1152];
// online-softmax carriers: running max m_i and denominator l_i (two row-groups
// per lane), plus the f32 output accumulator O_i (8 fragments × 4 elems).
float m_i0 = -1e+30f;
float l_i0 = 0.0f;
float m_i1 = -1e+30f;
float l_i1 = 0.0f;
float O_i_f[8][4] = {};
// load Q into the A fragments once — Q is reused across every KV tile.
unsigned _sacc_a[4][4];
#pragma unroll
for (int _r2 = 0; _r2 < 4; _r2 += 1) {
emmy_mma_load_a_gmem(_sacc_a[_r2], &q[b0 * 32768 + b1 * 8192 + m * 16 * 64 + _r2 * 16], 64);
}
float scale_c = _flash_scale[0];
// ---- KV loop: stream 16 key/value rows per step (flash-attention schedule).
for (int kv0 = 0; kv0 < 128; kv0 += 16) {
// cooperative cp.async gmem→smem fill of the K tile into the padded slab
// (16 rows × 8 vec-halves per lane-sweep; single-buffered, d1).
for (int _fk = threadIdx.x; _fk < 128; _fk += 32) {
emmy_cp_async_cg(&_k_smem[_fk / 8 * 72 + 8 * (_fk % 8)], &k[b0 * 32768 + b1 * 8192 + (kv0 + _fk / 8) * 64 + 8 * (_fk % 8)]);
}
// same cooperative cp.async fill for the V tile.
for (int _fv = threadIdx.x; _fv < 128; _fv += 32) {
emmy_cp_async_cg(&_v_smem[_fv / 8 * 72 + 8 * (_fv % 8)], &v[b0 * 32768 + b1 * 8192 + (kv0 + _fv / 8) * 64 + 8 * (_fv % 8)]);
}
emmy_cp_async_commit(); // publish the cp.async group
emmy_cp_async_wait<0>(); // wait for all in-flight copies
__syncthreads(); // slabs now visible to the whole warp
// QKᵀ: drain K from the padded slab with ldmatrix, feed the mmas.
float sacc_f[2][4] = {};
unsigned _sacc_b[2][2];
#pragma unroll
for (int _r1 = 0; _r1 < 4; _r1 += 1) {
// paired ldmatrix → two col-adjacent K (B) fragments at once.
emmy_ldmatrix_x4_pair(
_sacc_b[0],
_sacc_b[1],
&_k_smem[_r1 * 16 + (((threadIdx.x & 31) % 8) + ((threadIdx.x & 31) / 16) * 8) * 72 + (((threadIdx.x & 31) / 8) % 2) * 8]
);
#pragma unroll
for (int _r0 = 0; _r0 < 2; _r0 += 1) {
emmy_mma_m16n8k16_f16(sacc_f[_r0], _sacc_a[_r1], _sacc_b[_r0], sacc_f[_r0]);
}
}
// apply the softmax scale to the QKᵀ scores.
#pragma unroll
for (int _r0 = 0; _r0 < 2; _r0 += 1) {
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
sacc_f[_r0][_e] = sacc_f[_r0][_e] * scale_c;
}
}
// row max over this tile's scores, then reduce across the 4 lanes that
// share a row group (xor 2, xor 1).
float _rmx0 = fmaxf(fmaxf(fmaxf(sacc_f[0][0], sacc_f[0][1]), sacc_f[1][0]), sacc_f[1][1]);
float _rmx1 = fmaxf(fmaxf(fmaxf(sacc_f[0][2], sacc_f[0][3]), sacc_f[1][2]), sacc_f[1][3]);
_rmx0 = fmaxf(_rmx0, __shfl_xor_sync(0xffffffff, _rmx0, 2));
_rmx1 = fmaxf(_rmx1, __shfl_xor_sync(0xffffffff, _rmx1, 2));
_rmx0 = fmaxf(_rmx0, __shfl_xor_sync(0xffffffff, _rmx0, 1));
_rmx1 = fmaxf(_rmx1, __shfl_xor_sync(0xffffffff, _rmx1, 1));
// online-softmax update: new running max, then the rescale factor α.
float _mn0 = fmaxf(m_i0, _rmx0); // m_new
float _mn1 = fmaxf(m_i1, _rmx1);
float _al0 = expf((m_i0 - _mn0)); // α = e^{m_i − m_new}
float _al1 = expf((m_i1 - _mn1));
// p = e^{sⱼ − m_new} for each score element.
float _p_f0[4];
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
_p_f0[_e] = expf(sacc_f[0][_e] - ((_e < 2) ? (_mn0) : (_mn1)));
}
float _p_f1[4];
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
_p_f1[_e] = expf(sacc_f[1][_e] - ((_e < 2) ? (_mn0) : (_mn1)));
}
// row sums of p, reduced across the row-group lanes.
float _rsm0 = _p_f0[0] + _p_f0[1] + _p_f1[0] + _p_f1[1];
float _rsm1 = _p_f0[2] + _p_f0[3] + _p_f1[2] + _p_f1[3];
_rsm0 = _rsm0 + __shfl_xor_sync(0xffffffff, _rsm0, 2);
_rsm1 = _rsm1 + __shfl_xor_sync(0xffffffff, _rsm1, 2);
_rsm0 = _rsm0 + __shfl_xor_sync(0xffffffff, _rsm0, 1);
_rsm1 = _rsm1 + __shfl_xor_sync(0xffffffff, _rsm1, 1);
float l_i__n0 = (l_i0 * _al0) + _rsm0; // d ← d·α + p
float l_i__n1 = (l_i1 * _al1) + _rsm1;
l_i0 = l_i__n0;
l_i1 = l_i__n1;
// rescale the running output by α before this tile: o ← o·α.
#pragma unroll
for (int _r0 = 0; _r0 < 8; _r0 += 1) {
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
O_i_f[_r0][_e] = O_i_f[_r0][_e] * ((_e < 2) ? (_al0) : (_al1));
}
}
// repack the P scores from the C-fragment layout into an A operand.
unsigned _pa[4];
emmy_c_to_a_f16(_pa, _p_f0, _p_f1);
// P@V: drain V from the padded slab (transposed ldmatrix) and accumulate
// into the output fragments, two n-tiles per ldmatrix pair.
unsigned _O_i__pv_b[8][2];
emmy_ldmatrix_x4_trans_pair(
_O_i__pv_b[0],
_O_i__pv_b[1],
&_v_smem[0 + ((threadIdx.x & 31) % 16) * 72 + ((threadIdx.x & 31) / 16) * 8]
);
#pragma unroll
for (int _r0 = 0; _r0 < 2; _r0 += 1) {
emmy_mma_m16n8k16_f16(O_i_f[_r0], _pa, _O_i__pv_b[_r0], O_i_f[_r0]);
}
// … the other three V ldmatrix pairs (n-tiles 2..7), elided …
m_i0 = _mn0; // commit the new running max
m_i1 = _mn1;
__syncthreads(); // slabs free to be refilled next tile
}
// normalize by the softmax denominator: o ← o / d.
#pragma unroll
for (int _r0 = 0; _r0 < 8; _r0 += 1) {
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
O_i_f[_r0][_e] = O_i_f[_r0][_e] / ((_e < 2) ? (l_i0) : (l_i1));
}
}
// store the output fragments back to gmem as __half2 pairs.
#pragma unroll
for (int _r0 = 0; _r0 < 8; _r0 += 1) {
const int _g = (threadIdx.x & 31) >> 2; const int _t = (threadIdx.x & 31) & 3;
*reinterpret_cast<__half2*>(&sdpa[b1 * 8192 + m * 16 * 64 + _r0 * 8 + _g * 64 + _t * 2])
= __floats2half2_rn(O_i_f[_r0][0], O_i_f[_r0][1]);
*reinterpret_cast<__half2*>(&sdpa[b1 * 8192 + m * 16 * 64 + _r0 * 8 + (_g + 8) * 64 + _t * 2])
= __floats2half2_rn(O_i_f[_r0][2], O_i_f[_r0][3]);
}
}
Codegen Fixes (1.16x each)
The C→A register repack (271µs). P is the first mma's output (a C-fragment) and the second mma's A operand. The first version of this kernel bounced P through a shared-memory slab and ldmatrix-ed it back. However, the m16n8 C-fragment is element-for-element lane-aligned with the m16k16 A-fragment's k-halves. Two adjacent score fragments convert into one P@V operand fragment with four cvt.rn.f16x2.f32 instructions.
The Q hoist (233µs). Q's operand index carries the query and head-dim axes and never the stream axis, so its fragments are loop-invariant. Kernels in this article hoist it and keep in resident registers. FA-2 has this as a compile-time option, disabled by default, re-reading Q from shared memory every block. On this geometry, hoisting it is helpful
Move 3 — Pipelining (1.05x)
Pipelining allows us to overlap the operand loading with computation. The d2/cp/ring knob enables double-buffering of the K/V slabs: prefetch tile i+1's copies under tile i's math.
One more helper: fuse two slab-adjacent ldmatrix.x2 drains into one ldmatrix.x4:
// two slab-adjacent B fragments in one ldmatrix.x4: the four registers land
// straight in b0[0..1] and b1[0..1] (emmy_ldmatrix_x4_trans_pair is the x4.trans
// form, for the canonical-B slab)
static __device__ __forceinline__ void emmy_ldmatrix_x4_pair(
unsigned* b0, unsigned* b1, const void* smem) {
asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
: "=r"(b0[0]), "=r"(b0[1]), "=r"(b1[0]), "=r"(b1[1])
: "r"(__cvta_generic_to_shared(smem)));
}
The full pipelined kernel (d2/cp/ring) — tile i+1's cp.async fill issued before the wait_group 1 on tile i, two slabs in flight:
EMMY_KNOBS="PLACE=fuse,TILE=a:mma_m16n8k16_f16/w1x1/f1x2/k4,STAGE=d2/cp/ring" \
emmy compile -c "$SDPA_L" --target sm_120 --ir cuda
extern "C" __global__ __launch_bounds__(32)
void k_sdpa(
const __half* q,
const float* _flash_scale,
const __half* k,
const __half* v, __half* sdpa)
{
// ---- thread-id decode: this warp owns one (batch b0, head b1, query-tile m).
int _gid = blockIdx.x * blockDim.x + threadIdx.x;
int b0 = _gid / 1024; // batch index
int b1 = (_gid / 256) % 4; // head index (4 heads)
int m = (_gid / 32) % 8; // query row-tile (8 tiles of 16 rows = 128)
int _lane = (_gid) % 32; // lane within the warp
// ---- depth-2 ring buffers: each K/V slab holds TWO 16-row tiles (16*72 each,
// 72 = 64 + 8 pad to dodge bank conflicts); the (kv0/16 % 2) slot alternates.
__shared__ __half _k_smem[2304]; // 2 * 16 * 72
__shared__ __half _v_smem[2304];
// ---- online-softmax carrier: running max m_i and denom l_i (per row-half
// 0/1), plus the f32 output accumulator O_i (8 n-tiles x 4 elems).
float m_i0 = -1e+30f;
float l_i0 = 0.0f;
float m_i1 = -1e+30f;
float l_i1 = 0.0f;
float O_i_f[8][4] = {};
// ---- load Q into the A-fragment registers once (reused every KV tile).
unsigned _sacc_a[4][4];
#pragma unroll
for (int _r2 = 0; _r2 < 4; _r2 += 1) {
emmy_mma_load_a_gmem(_sacc_a[_r2], &q[b0 * 32768 + b1 * 8192 + m * 16 * 64 + _r2 * 16], 64);
}
float scale_c = _flash_scale[0];
// ---- PROLOGUE: issue the cp.async fill of KV tile 0 into ring slot 0 before
// entering the loop, then commit it (the loop waits on it at wait<1>).
for (int _fk = threadIdx.x; _fk < 128; _fk += 32) {
emmy_cp_async_cg(
&_k_smem[_fk / 8 * 72 + 8 * (_fk % 8)],
&k[b0 * 32768 + b1 * 8192 + _fk / 8 * 64 + 8 * (_fk % 8)]
);
}
for (int _fv = threadIdx.x; _fv < 128; _fv += 32) {
emmy_cp_async_cg(
&_v_smem[_fv / 8 * 72 + 8 * (_fv % 8)],
&v[b0 * 32768 + b1 * 8192 + _fv / 8 * 64 + 8 * (_fv % 8)]
);
}
emmy_cp_async_commit();
// ---- KV loop: stream 128 keys/values in 16-row tiles, flash-attention style.
for (int kv0 = 0; kv0 < 128; kv0 += 16) {
// PREFETCH tile i+1 into the OTHER ring slot ((kv0/16 + 1) % 2) while the
// current tile is still in flight; the source row clamps to 112 on the
// last iteration so the out-of-range prefetch reads in-bounds (unused).
for (int _fk = threadIdx.x; _fk < 128; _fk += 32) {
emmy_cp_async_cg(
&_k_smem[((kv0 / 16 + 1) % 2 * 16 + _fk / 8) * 72 + 8 * (_fk % 8)],
&k[b0 * 32768 + b1 * 8192 + (((kv0 + 16 < 128) ? (kv0 + 16) : (112)) + _fk / 8) * 64 + 8 * (_fk % 8)]
);
}
for (int _fv = threadIdx.x; _fv < 128; _fv += 32) {
emmy_cp_async_cg(
&_v_smem[((kv0 / 16 + 1) % 2 * 16 + _fv / 8) * 72 + 8 * (_fv % 8)],
&v[b0 * 32768 + b1 * 8192 + (((kv0 + 16 < 128) ? (kv0 + 16) : (112)) + _fv / 8) * 64 + 8 * (_fv % 8)]
);
}
emmy_cp_async_commit();
// wait<1>: keep exactly ONE copy group in flight (the prefetch above),
// draining only the current tile's fill; sync before reading smem.
emmy_cp_async_wait<1>();
__syncthreads();
// ---- Q@K^T: ldmatrix the current slot's K tile and accumulate the 16x16
// score fragment sacc_f over the 4 k-steps.
float sacc_f[2][4] = {};
unsigned _sacc_b[2][2];
#pragma unroll
for (int _r1 = 0; _r1 < 4; _r1 += 1) {
emmy_ldmatrix_x4_pair(
_sacc_b[0],
_sacc_b[1],
&_k_smem[kv0 / 16 % 2 * 16 * 72 + _r1 * 16 + (((threadIdx.x & 31) % 8) + ((threadIdx.x & 31) / 16) * 8) * 72 + (((threadIdx.x & 31) / 8) % 2) * 8]
);
#pragma unroll
for (int _r0 = 0; _r0 < 2; _r0 += 1) {
emmy_mma_m16n8k16_f16(sacc_f[_r0], _sacc_a[_r1], _sacc_b[_r0], sacc_f[_r0]);
}
}
// ---- scale scores by 1/sqrt(d).
#pragma unroll
for (int _r0 = 0; _r0 < 2; _r0 += 1) {
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
sacc_f[_r0][_e] = sacc_f[_r0][_e] * scale_c;
}
}
// ---- tile row-max, reduced across the 4 lanes sharing a row (xor 2 then
// xor 1); fold into the running max to get the new max _mn.
float _rmx0 = fmaxf(fmaxf(fmaxf(sacc_f[0][0], sacc_f[0][1]), sacc_f[1][0]), sacc_f[1][1]);
float _rmx1 = fmaxf(fmaxf(fmaxf(sacc_f[0][2], sacc_f[0][3]), sacc_f[1][2]), sacc_f[1][3]);
_rmx0 = fmaxf(_rmx0, __shfl_xor_sync(0xffffffff, _rmx0, 2));
_rmx1 = fmaxf(_rmx1, __shfl_xor_sync(0xffffffff, _rmx1, 2));
_rmx0 = fmaxf(_rmx0, __shfl_xor_sync(0xffffffff, _rmx0, 1));
_rmx1 = fmaxf(_rmx1, __shfl_xor_sync(0xffffffff, _rmx1, 1));
float _mn0 = fmaxf(m_i0, _rmx0);
float _mn1 = fmaxf(m_i1, _rmx1);
float _al0 = expf((m_i0 - _mn0)); // α = e^{m_i − m_new}
float _al1 = expf((m_i1 - _mn1)); // α = e^{m_i − m_new}
// ---- p = e^{s_j − m_new} for every score element (row-half picks _mn).
float _p_f0[4];
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
_p_f0[_e] = expf(sacc_f[0][_e] - ((_e < 2) ? (_mn0) : (_mn1)));
}
float _p_f1[4];
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
_p_f1[_e] = expf(sacc_f[1][_e] - ((_e < 2) ? (_mn0) : (_mn1)));
}
// ---- row-sum of p (lane-reduced), then d ← d·α + p.
float _rsm0 = _p_f0[0] + _p_f0[1] + _p_f1[0] + _p_f1[1];
float _rsm1 = _p_f0[2] + _p_f0[3] + _p_f1[2] + _p_f1[3];
_rsm0 = _rsm0 + __shfl_xor_sync(0xffffffff, _rsm0, 2);
_rsm1 = _rsm1 + __shfl_xor_sync(0xffffffff, _rsm1, 2);
_rsm0 = _rsm0 + __shfl_xor_sync(0xffffffff, _rsm0, 1);
_rsm1 = _rsm1 + __shfl_xor_sync(0xffffffff, _rsm1, 1);
float l_i__n0 = (l_i0 * _al0) + _rsm0; // d ← d·α + p
float l_i__n1 = (l_i1 * _al1) + _rsm1;
l_i0 = l_i__n0;
l_i1 = l_i__n1;
// ---- rescale the running output by α before adding this tile: o ← o·α.
#pragma unroll
for (int _r0 = 0; _r0 < 8; _r0 += 1) {
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
O_i_f[_r0][_e] = O_i_f[_r0][_e] * ((_e < 2) ? (_al0) : (_al1));
}
}
// ---- repack the P scores from C-fragment layout into an A operand.
unsigned _pa[4];
emmy_c_to_a_f16(_pa, _p_f0, _p_f1);
// ---- P@V: ldmatrix.trans the current slot's V tile in n-tile pairs and
// accumulate into O_i (each drained pair feeds two mmas).
unsigned _O_i__pv_b[8][2];
emmy_ldmatrix_x4_trans_pair(
_O_i__pv_b[0],
_O_i__pv_b[1],
&_v_smem[kv0 / 16 % 2 * 16 * 72 + ((threadIdx.x & 31) % 16) * 72 + ((threadIdx.x & 31) / 16) * 8]
);
#pragma unroll
for (int _r0 = 0; _r0 < 2; _r0 += 1) {
emmy_mma_m16n8k16_f16(O_i_f[_r0], _pa, _O_i__pv_b[_r0], O_i_f[_r0]);
}
emmy_ldmatrix_x4_trans_pair(
_O_i__pv_b[2],
_O_i__pv_b[3],
&_v_smem[kv0 / 16 % 2 * 16 * 72 + 16 + ((threadIdx.x & 31) % 16) * 72 + ((threadIdx.x & 31) / 16) * 8]
);
emmy_mma_m16n8k16_f16(O_i_f[2], _pa, _O_i__pv_b[2], O_i_f[2]);
emmy_mma_m16n8k16_f16(O_i_f[3], _pa, _O_i__pv_b[3], O_i_f[3]);
emmy_ldmatrix_x4_trans_pair(
_O_i__pv_b[4],
_O_i__pv_b[5],
&_v_smem[kv0 / 16 % 2 * 16 * 72 + 32 + ((threadIdx.x & 31) % 16) * 72 + ((threadIdx.x & 31) / 16) * 8]
);
emmy_mma_m16n8k16_f16(O_i_f[4], _pa, _O_i__pv_b[4], O_i_f[4]);
emmy_mma_m16n8k16_f16(O_i_f[5], _pa, _O_i__pv_b[5], O_i_f[5]);
emmy_ldmatrix_x4_trans_pair(
_O_i__pv_b[6],
_O_i__pv_b[7],
&_v_smem[kv0 / 16 % 2 * 16 * 72 + 48 + ((threadIdx.x & 31) % 16) * 72 + ((threadIdx.x & 31) / 16) * 8]
);
emmy_mma_m16n8k16_f16(O_i_f[6], _pa, _O_i__pv_b[6], O_i_f[6]);
emmy_mma_m16n8k16_f16(O_i_f[7], _pa, _O_i__pv_b[7], O_i_f[7]);
// ---- commit the new running max; sync before the next tile reuses smem.
m_i0 = _mn0;
m_i1 = _mn1;
__syncthreads();
}
// ---- normalize by the softmax denominator: o ← o / d.
#pragma unroll
for (int _r0 = 0; _r0 < 8; _r0 += 1) {
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
O_i_f[_r0][_e] = O_i_f[_r0][_e] / ((_e < 2) ? (l_i0) : (l_i1));
}
}
// ---- store O to gmem as __half2 pairs (two row-halves per n-tile).
#pragma unroll
for (int _r0 = 0; _r0 < 8; _r0 += 1) {
const int _g = (threadIdx.x & 31) >> 2;
const int _t = (threadIdx.x & 31) & 3;
*reinterpret_cast<__half2*>(&sdpa[b1 * 8192 + m * 16 * 64 + _r0 * 8 + _g * 64 + _t * 2])
= __floats2half2_rn(O_i_f[_r0][0], O_i_f[_r0][1]);
*reinterpret_cast<__half2*>(&sdpa[b1 * 8192 + m * 16 * 64 + _r0 * 8 + (_g + 8) * 64 + _t * 2])
= __floats2half2_rn(O_i_f[_r0][2], O_i_f[_r0][3]);
}
}
New Hardware: TMA (1.04x)
The last transport upgrade leverages the TMA transport. This transport was primarily introduced to feed lightning fast WGMMA tensor core instructions on Hopper/Blackwell datacenter GPUs. Alas, WGMMA instructions are not available on consumer GPUs. However, TMA uses fewer instructions than cp.async and gives us free swizzle, reducing the LSU load.
The TMA path adds its own preamble helpers: one for the box copy, the mbarrier the consumer waits on for the copy
to land, and the swizzle XOR every ldmatrix drain reads its slab address through:
// one TMA box copy: a rank-4 tile gmem→smem in a single instruction, from the
// __grid_constant__ descriptor; the transfer's completion is signaled on the
// mbarrier (transaction-byte count).
static __device__ __forceinline__ void cp_async_bulk_tensor_4d(
void* smem, const CUtensorMap* desc,
int c0, int c1, int c2, int c3, unsigned long long* mbar) {
unsigned saddr = __cvta_generic_to_shared(smem);
unsigned maddr = __cvta_generic_to_shared(mbar);
asm volatile("cp.async.bulk.tensor.4d.shared::cta.global.mbarrier::complete_tx::bytes "
"[%0], [%1, {%2, %3, %4, %5}], [%6];\n"
:: "r"(saddr), "l"(desc), "r"(c0), "r"(c1), "r"(c2), "r"(c3), "r"(maddr) : "memory");
}
// block until the TMA transfer parity flips (try_wait suspends the warp rather
// than hot-spinning while the copy drains); mbarrier_init / _arrive_expect_tx
// set the slot up on the producer side.
static __device__ __forceinline__ void mbarrier_wait_parity(
unsigned long long* mbar, int phase) {
unsigned addr = __cvta_generic_to_shared(mbar);
asm volatile("{.reg .pred P; bw: mbarrier.try_wait.parity.shared.b64 P, [%0], %1; @!P bra bw;}\n"
:: "r"(addr), "r"(phase) : "memory");
}
// B128 slab swizzle: XOR a b16 smem element index the way the TMA copy engine
// permuted the 16-byte chunks on fill — every ldmatrix drain reads through it.
static __device__ __forceinline__ int emmy_swizzle_b128(int e) {
return e ^ (((e >> 6) & 7) << 3);
}
The full kernel on the TMA ring (d2/tma/ring).
EMMY_KNOBS="PLACE=fuse,TILE=a:mma_m16n8k16_f16/w1x1/f1x2/k4,STAGE=d2/tma/ring,WSPEC=" \
emmy compile -c "$SDPA_L" --target sm_120 --ir cuda # helper preamble elided
extern "C" __global__ __launch_bounds__(32)
void k_sdpa(
const __half* k,
const __half* v,
const __half* q,
const float* _flash_scale,
__half* sdpa,
const CUtensorMap* __restrict__ _desc_k,
const CUtensorMap* __restrict__ _desc_v)
{
// Thread-id decode: 32 lanes/warp, one warp owns one 16-row query tile.
int _gid = blockIdx.x * blockDim.x + threadIdx.x; // global thread id
int b0 = _gid / 1024; // batch index
int b1 = (_gid / 256) % 4; // head index
int m = (_gid / 32) % 8; // query-tile index (16 rows) within the head
int _lane = (_gid) % 32; // lane within the warp
// Carrier init. K/V ride a 2-slot smem ring; the dense 16×64 slabs use the
// hardware swizzle (no pad). Softmax carriers m_i/l_i and the o accumulator
// are split into rows 0..7 (…0) and rows 8..15 (…1) of the query tile.
__shared__ __align__(1024) __half _k_smem[2048]; // K ring: 2 × 16×64
__shared__ __align__(1024) __half _v_smem[2048]; // V ring: 2 × 16×64
__shared__ unsigned long long _mbar[2]; // one mbarrier per slot
float m_i0 = -1e+30f; // running row-max m_i (rows 0..7)
float l_i0 = 0.0f; // running denom d (rows 0..7)
float m_i1 = -1e+30f; // running row-max m_i (rows 8..15)
float l_i1 = 0.0f; // running denom d (rows 8..15)
float O_i_f[8][4] = {}; // output accumulator o (8 n-tiles × 4 elems)
unsigned _sacc_a[4][4]; // Q fragment (A operand): 4 k-slices × 4 regs
// A loads: pull the stationary Q tile into registers once — it is reused
// as the mma A operand across every K/V tile of the ring.
#pragma unroll
for (int _r2 = 0; _r2 < 4; _r2 += 1) {
emmy_mma_load_a_gmem(
_sacc_a[_r2],
&q[b0 * 32768 + b1 * 8192 + m * 16 * 64 + _r2 * 16], 64);
}
float scale_c = _flash_scale[0]; // 1/sqrt(head_dim) softmax scale
// One thread initializes both ring-slot mbarriers (1 arrival each: the TMA).
if (threadIdx.x == 0) {
mbarrier_init(&_mbar[0], 1);
mbarrier_init(&_mbar[1], 1);
}
__syncthreads();
// Prologue box-copy: the elected thread issues the first K/V TMA into slot 0.
// arrive.expect_tx tells the mbarrier how many bytes to await (4096 = two
// 16×64 __half slabs); the TMA engine signals _mbar[0] when they land.
if (threadIdx.x == 0) {
mbarrier_arrive_expect_tx(&_mbar[0], 4096);
cp_async_bulk_tensor_4d(&_k_smem[0], _desc_k, 0, 0, b1, b0, &_mbar[0]);
cp_async_bulk_tensor_4d(&_v_smem[0], _desc_v, 0, 0, b1, b0, &_mbar[0]);
}
for (int kv0 = 0; kv0 < 128; kv0 += 16) {
// Prefetch the NEXT ring slot while this one is consumed (double-buffer).
// The elected thread issues the TMA box-copy against the tensor-map
// descriptor; c1 is the K/V row offset (clamped to 112 on the last tile).
if (threadIdx.x == 0) {
mbarrier_arrive_expect_tx(&_mbar[(kv0 / 16 + 1) % 2], 4096);
cp_async_bulk_tensor_4d(
&_k_smem[(kv0 / 16 + 1) % 2 * 16 * 64], _desc_k, 0,
((kv0 + 16 < 128) ? (kv0 + 16) : (112)), b1, b0,
&_mbar[(kv0 / 16 + 1) % 2]);
cp_async_bulk_tensor_4d(
&_v_smem[(kv0 / 16 + 1) % 2 * 16 * 64], _desc_v, 0,
((kv0 + 16 < 128) ? (kv0 + 16) : (112)), b1, b0,
&_mbar[(kv0 / 16 + 1) % 2]);
}
// Parity-wait on THIS slot: suspend until its TMA tx has fully landed.
mbarrier_wait_parity(&_mbar[kv0 / 16 % 2], kv0 / 16 / 2 % 2);
float sacc_f[2][4] = {}; // QKᵀ scores S (2 n-tiles × 4 elems)
unsigned _sacc_b[2][2]; // K fragment (B operand): 2 n-tiles
#pragma unroll
for (int _r1 = 0; _r1 < 4; _r1 += 1) {
// ldmatrix drain: paired load of the K slab from THIS ring slot.
// The smem index passes through the hardware swizzle's XOR
// (emmy_swizzle_b128) — ring slot picked by (kv0/16 % 2), so the
// read never needs padding.
emmy_ldmatrix_x4_pair(
_sacc_b[0],
_sacc_b[1],
&_k_smem[emmy_swizzle_b128(kv0 / 16 % 2 * 16 * 64 + _r1 * 16
+ (((threadIdx.x & 31) % 8) + ((threadIdx.x & 31) / 16) * 8) * 64
+ (((threadIdx.x & 31) / 8) % 2) * 8)]
);
#pragma unroll
for (int _r0 = 0; _r0 < 2; _r0 += 1) {
// QKᵀ mma: S += Q · Kᵀ over this k-slice.
emmy_mma_m16n8k16_f16(
sacc_f[_r0], _sacc_a[_r1], _sacc_b[_r0], sacc_f[_r0]);
}
}
// Apply the softmax scale to the raw scores.
#pragma unroll
for (int _r0 = 0; _r0 < 2; _r0 += 1) {
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
sacc_f[_r0][_e] = sacc_f[_r0][_e] * scale_c;
}
}
// Online-softmax update. Tile row-max, then a 2-step butterfly reduce
// across the 4 lanes of each row group.
float _rmx0 = fmaxf(fmaxf(fmaxf(sacc_f[0][0], sacc_f[0][1]), sacc_f[1][0]), sacc_f[1][1]);
float _rmx1 = fmaxf(fmaxf(fmaxf(sacc_f[0][2], sacc_f[0][3]), sacc_f[1][2]), sacc_f[1][3]);
_rmx0 = fmaxf(_rmx0, __shfl_xor_sync(0xffffffff, _rmx0, 2));
_rmx1 = fmaxf(_rmx1, __shfl_xor_sync(0xffffffff, _rmx1, 2));
_rmx0 = fmaxf(_rmx0, __shfl_xor_sync(0xffffffff, _rmx0, 1));
_rmx1 = fmaxf(_rmx1, __shfl_xor_sync(0xffffffff, _rmx1, 1));
float _mn0 = fmaxf(m_i0, _rmx0); // m_new = max(m_i, tile-max), rows 0..7
float _mn1 = fmaxf(m_i1, _rmx1); // m_new, rows 8..15
float _al0 = expf((m_i0 - _mn0)); // α = e^{m_i − m_new}
float _al1 = expf((m_i1 - _mn1)); // α = e^{m_i − m_new}
float _p_f0[4];
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
// p = e^{sⱼ − m_new} (n-tile 0; _e<2 → rows 0..7 else 8..15)
_p_f0[_e] = expf(sacc_f[0][_e] - ((_e < 2) ? (_mn0) : (_mn1)));
}
float _p_f1[4];
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
// p = e^{sⱼ − m_new} (n-tile 1)
_p_f1[_e] = expf(sacc_f[1][_e] - ((_e < 2) ? (_mn0) : (_mn1)));
}
// Row-sum of p, reduced across the same 4 lanes → this tile's Σp.
float _rsm0 = _p_f0[0] + _p_f0[1] + _p_f1[0] + _p_f1[1];
float _rsm1 = _p_f0[2] + _p_f0[3] + _p_f1[2] + _p_f1[3];
_rsm0 = _rsm0 + __shfl_xor_sync(0xffffffff, _rsm0, 2);
_rsm1 = _rsm1 + __shfl_xor_sync(0xffffffff, _rsm1, 2);
_rsm0 = _rsm0 + __shfl_xor_sync(0xffffffff, _rsm0, 1);
_rsm1 = _rsm1 + __shfl_xor_sync(0xffffffff, _rsm1, 1);
float l_i__n0 = (l_i0 * _al0) + _rsm0; // d ← d·α + p (rows 0..7)
float l_i__n1 = (l_i1 * _al1) + _rsm1; // d ← d·α + p (rows 8..15)
l_i0 = l_i__n0;
l_i1 = l_i__n1;
// Rescale the running output by α before this tile's P@V: o ← o·α.
#pragma unroll
for (int _r0 = 0; _r0 < 8; _r0 += 1) {
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
O_i_f[_r0][_e] = O_i_f[_r0][_e] * ((_e < 2) ? (_al0) : (_al1));
}
}
// C→A repack: the P scores live in the C-fragment layout; convert them
// in-register into the A operand for the P@V mma (no smem round-trip).
unsigned _pa[4];
emmy_c_to_a_f16(_pa, _p_f0, _p_f1);
unsigned _O_i__pv_b[8][2]; // V fragment (B operand): 8 n-tiles
// P@V mma: drain V from the ring slot with a transposed paired ldmatrix
// (swizzled index) and accumulate o += P · V, 16 head-dim cols at a time.
emmy_ldmatrix_x4_trans_pair(
_O_i__pv_b[0], _O_i__pv_b[1],
&_v_smem[emmy_swizzle_b128(kv0 / 16 % 2 * 16 * 64
+ ((threadIdx.x & 31) % 16) * 64
+ ((threadIdx.x & 31) / 16) * 8)]);
#pragma unroll
for (int _r0 = 0; _r0 < 2; _r0 += 1) {
emmy_mma_m16n8k16_f16(O_i_f[_r0], _pa, _O_i__pv_b[_r0], O_i_f[_r0]);
}
emmy_ldmatrix_x4_trans_pair(
_O_i__pv_b[2], _O_i__pv_b[3],
&_v_smem[emmy_swizzle_b128(kv0 / 16 % 2 * 16 * 64 + 16
+ ((threadIdx.x & 31) % 16) * 64
+ ((threadIdx.x & 31) / 16) * 8)]);
emmy_mma_m16n8k16_f16(O_i_f[2], _pa, _O_i__pv_b[2], O_i_f[2]);
emmy_mma_m16n8k16_f16(O_i_f[3], _pa, _O_i__pv_b[3], O_i_f[3]);
emmy_ldmatrix_x4_trans_pair(
_O_i__pv_b[4], _O_i__pv_b[5],
&_v_smem[emmy_swizzle_b128(kv0 / 16 % 2 * 16 * 64 + 32
+ ((threadIdx.x & 31) % 16) * 64
+ ((threadIdx.x & 31) / 16) * 8)]);
emmy_mma_m16n8k16_f16(O_i_f[4], _pa, _O_i__pv_b[4], O_i_f[4]);
emmy_mma_m16n8k16_f16(O_i_f[5], _pa, _O_i__pv_b[5], O_i_f[5]);
emmy_ldmatrix_x4_trans_pair(
_O_i__pv_b[6], _O_i__pv_b[7],
&_v_smem[emmy_swizzle_b128(kv0 / 16 % 2 * 16 * 64 + 48
+ ((threadIdx.x & 31) % 16) * 64
+ ((threadIdx.x & 31) / 16) * 8)]);
emmy_mma_m16n8k16_f16(O_i_f[6], _pa, _O_i__pv_b[6], O_i_f[6]);
emmy_mma_m16n8k16_f16(O_i_f[7], _pa, _O_i__pv_b[7], O_i_f[7]);
m_i0 = _mn0; // commit m_i ← m_new for the next tile
m_i1 = _mn1;
__syncthreads(); // slot fully consumed before the ring reuses it
}
// Normalize by the softmax denominator: o ← o / d.
#pragma unroll
for (int _r0 = 0; _r0 < 8; _r0 += 1) {
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
O_i_f[_r0][_e] = O_i_f[_r0][_e] / ((_e < 2) ? (l_i0) : (l_i1));
}
}
// Store the output tile to gmem as __half2 pairs (two row blocks: _g, _g+8).
#pragma unroll
for (int _r0 = 0; _r0 < 8; _r0 += 1) {
const int _g = (threadIdx.x & 31) >> 2; const int _t = (threadIdx.x & 31) & 3;
*reinterpret_cast<__half2*>(&sdpa[b1 * 8192 + m * 16 * 64 + _r0 * 8 + _g * 64 + _t * 2])
= __floats2half2_rn(O_i_f[_r0][0], O_i_f[_r0][1]);
*reinterpret_cast<__half2*>(&sdpa[b1 * 8192 + m * 16 * 64 + _r0 * 8 + (_g + 8) * 64 + _t * 2])
= __floats2half2_rn(O_i_f[_r0][2], O_i_f[_r0][3]);
}
}
The memory counters already beat the handwritten kernel: 3.31M LSU instructions vs FA-2's 4.36M, 4K conflicts vs 22K.
Move 4 — Register Tiling (1.02x)
Traditionally, register tiling is an FMA-based optimization (Volkov's better-performance-at-lower-occupancy argument). However, it can also be used to increase the tensor-pipe utilization. Bump FM to 2 (f1x8 → f2x8) and each warp takes two query tiles: two independent (m, l, O) carriers run against shared K/V fragments, every B fragment ldmatrix-ed once and fanned across both chains' mmas.
EMMY_KNOBS="PLACE=fuse,TILE=a:mma_m16n8k16_f16/w1x1/f2x2/k4,STAGE=d2/tma/ring,WSPEC=" \
emmy compile -c "$SDPA_L" --target sm_120 --ir cuda # helper preamble elided
extern "C" __global__ __launch_bounds__(32)
void k_sdpa(
const __half* k,
const __half* v,
const __half* q,
const float* _flash_scale,
__half* sdpa,
const CUtensorMap* __restrict__ _desc_k,
const CUtensorMap* __restrict__ _desc_v)
{
// One warp (32 threads) per CTA. thread-id decode: flat row id → batch b0,
// head b1, query-tile-pair m (each warp owns 32 query rows = TWO 16-row
// query tiles), and the lane within the warp.
int _gid = blockIdx.x * blockDim.x + threadIdx.x;
int b0 = _gid / 512; // batch
int b1 = (_gid / 128) % 4; // head
int m = (_gid / 32) % 4; // query-tile-pair (32 rows / warp)
int _lane = (_gid) % 32;
// TMA ring transport: double-buffered K/V slabs in smem + their mbarriers.
__shared__ __align__(1024) __half _k_smem[2048];
__shared__ __align__(1024) __half _v_smem[2048];
__shared__ unsigned long long _mbar[2];
// Carrier of chain 0 (_q0) — the first 16-row query tile: its own online-
// softmax running max m_i and denom l_i (one per 8-col N-half, ..0 / ..1),
// its O accumulator, and the persistent A (query) mma fragments in registers.
float m_i_q00 = -1e+30f;
float l_i_q00 = 0.0f;
float m_i_q01 = -1e+30f;
float l_i_q01 = 0.0f;
float O_i_q0_f[8][4] = {};
unsigned _sacc_q0_a[4][4];
#pragma unroll
for (int _r2 = 0; _r2 < 4; _r2 += 1) {
emmy_mma_load_a_gmem(_sacc_q0_a[_r2], &q[b0 * 32768 + b1 * 8192 + m * 32 * 64 + _r2 * 16], 64);
}
// Carrier of chain 1 (_q1) — the second 16-row query tile (rows m*32 + 16):
// identical structure (its own m_i / l_i / O / A), run INDEPENDENTLY of
// chain 0 but against the SAME shared K/V fragments (the f2x2 win).
float m_i_q10 = -1e+30f;
float l_i_q10 = 0.0f;
float m_i_q11 = -1e+30f;
float l_i_q11 = 0.0f;
float O_i_q1_f[8][4] = {};
unsigned _sacc_q1_a[4][4];
#pragma unroll
for (int _r2 = 0; _r2 < 4; _r2 += 1) {
emmy_mma_load_a_gmem(_sacc_q1_a[_r2], &q[b0 * 32768 + b1 * 8192 + (m * 32 + 16) * 64 + _r2 * 16], 64);
}
float scale_c = _flash_scale[0];
// Prime the TMA ring: init both mbarriers, then kick the first K/V bulk copy.
if (threadIdx.x == 0) {
mbarrier_init(&_mbar[0], 1);
mbarrier_init(&_mbar[1], 1);
}
__syncthreads();
if (threadIdx.x == 0) {
mbarrier_arrive_expect_tx(&_mbar[0], 4096);
cp_async_bulk_tensor_4d(&_k_smem[0], _desc_k, 0, 0, b1, b0, &_mbar[0]);
cp_async_bulk_tensor_4d(&_v_smem[0], _desc_v, 0, 0, b1, b0, &_mbar[0]);
}
// Streaming KV loop: one 16-key tile per step over the 128-key sequence.
for (int kv0 = 0; kv0 < 128; kv0 += 16) {
// Prefetch the NEXT tile into the other ring slot while computing this one.
if (threadIdx.x == 0) {
mbarrier_arrive_expect_tx(&_mbar[(kv0 / 16 + 1) % 2], 4096);
cp_async_bulk_tensor_4d(
&_k_smem[(kv0 / 16 + 1) % 2 * 16 * 64], _desc_k, 0,
((kv0 + 16 < 128) ? (kv0 + 16) : (112)), b1, b0,
&_mbar[(kv0 / 16 + 1) % 2]);
cp_async_bulk_tensor_4d(
&_v_smem[(kv0 / 16 + 1) % 2 * 16 * 64], _desc_v, 0,
((kv0 + 16 < 128) ? (kv0 + 16) : (112)), b1, b0,
&_mbar[(kv0 / 16 + 1) % 2]);
}
// Wait for THIS tile's TMA fill to land in smem.
mbarrier_wait_parity(&_mbar[kv0 / 16 % 2], kv0 / 16 / 2 % 2);
// S = Q @ Kᵀ — two score accumulators, one per query chain.
float sacc_q0_f[2][4] = {};
float sacc_q1_f[2][4] = {};
unsigned _sacc_b[2][2];
#pragma unroll
for (int _r1 = 0; _r1 < 4; _r1 += 1) {
// ONE ldmatrix loads the shared K fragment (swizzled smem address)…
emmy_ldmatrix_x4_pair(
_sacc_b[0],
_sacc_b[1],
&_k_smem[emmy_swizzle_b128(kv0 / 16 % 2 * 16 * 64 + _r1 * 16
+ (((threadIdx.x & 31) % 8) + ((threadIdx.x & 31) / 16) * 8) * 64
+ (((threadIdx.x & 31) / 8) % 2) * 8)]
);
#pragma unroll
for (int _r0 = 0; _r0 < 2; _r0 += 1) {
// …and it is FANNED into both chains' mmas: one K load, two chains.
emmy_mma_m16n8k16_f16(sacc_q0_f[_r0], _sacc_q0_a[_r1], _sacc_b[_r0], sacc_q0_f[_r0]);
emmy_mma_m16n8k16_f16(sacc_q1_f[_r0], _sacc_q1_a[_r1], _sacc_b[_r0], sacc_q1_f[_r0]);
}
}
// Apply the softmax scale to both chains' scores.
#pragma unroll
for (int _r0 = 0; _r0 < 2; _r0 += 1) {
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
sacc_q0_f[_r0][_e] = sacc_q0_f[_r0][_e] * scale_c;
}
}
#pragma unroll
for (int _r0 = 0; _r0 < 2; _r0 += 1) {
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
sacc_q1_f[_r0][_e] = sacc_q1_f[_r0][_e] * scale_c;
}
}
// --- Chain 0 online softmax (its own m/l/O carrier) ---
float _rmx_q00 = fmaxf(fmaxf(fmaxf(sacc_q0_f[0][0], sacc_q0_f[0][1]), sacc_q0_f[1][0]), sacc_q0_f[1][1]);
float _rmx_q01 = fmaxf(fmaxf(fmaxf(sacc_q0_f[0][2], sacc_q0_f[0][3]), sacc_q0_f[1][2]), sacc_q0_f[1][3]);
// Warp-reduce the row max across the 4 lanes sharing a query row.
_rmx_q00 = fmaxf(_rmx_q00, __shfl_xor_sync(0xffffffff, _rmx_q00, 2));
_rmx_q01 = fmaxf(_rmx_q01, __shfl_xor_sync(0xffffffff, _rmx_q01, 2));
_rmx_q00 = fmaxf(_rmx_q00, __shfl_xor_sync(0xffffffff, _rmx_q00, 1));
_rmx_q01 = fmaxf(_rmx_q01, __shfl_xor_sync(0xffffffff, _rmx_q01, 1));
float _mn_q00 = fmaxf(m_i_q00, _rmx_q00); // m_new = max(m_i, rowmax)
float _mn_q01 = fmaxf(m_i_q01, _rmx_q01);
float _al_q00 = expf((m_i_q00 - _mn_q00)); // α = e^{m_i − m_new}
float _al_q01 = expf((m_i_q01 - _mn_q01));
float _p_q0_f0[4];
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
// p = e^{sⱼ − m_new}
_p_q0_f0[_e] = expf(sacc_q0_f[0][_e] - ((_e < 2) ? (_mn_q00) : (_mn_q01)));
}
float _p_q0_f1[4];
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
_p_q0_f1[_e] = expf(sacc_q0_f[1][_e] - ((_e < 2) ? (_mn_q00) : (_mn_q01)));
}
float _rsm_q00 = _p_q0_f0[0] + _p_q0_f0[1] + _p_q0_f1[0] + _p_q0_f1[1];
float _rsm_q01 = _p_q0_f0[2] + _p_q0_f0[3] + _p_q0_f1[2] + _p_q0_f1[3];
_rsm_q00 = _rsm_q00 + __shfl_xor_sync(0xffffffff, _rsm_q00, 2);
_rsm_q01 = _rsm_q01 + __shfl_xor_sync(0xffffffff, _rsm_q01, 2);
_rsm_q00 = _rsm_q00 + __shfl_xor_sync(0xffffffff, _rsm_q00, 1);
_rsm_q01 = _rsm_q01 + __shfl_xor_sync(0xffffffff, _rsm_q01, 1);
float l_i_q0__n0 = (l_i_q00 * _al_q00) + _rsm_q00; // d ← d·α + p
float l_i_q0__n1 = (l_i_q01 * _al_q01) + _rsm_q01;
l_i_q00 = l_i_q0__n0;
l_i_q01 = l_i_q0__n1;
#pragma unroll
for (int _r0 = 0; _r0 < 8; _r0 += 1) {
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
// o ← o·α
O_i_q0_f[_r0][_e] = O_i_q0_f[_r0][_e] * ((_e < 2) ? (_al_q00) : (_al_q01));
}
}
// Repack chain 0's P (f32 C fragment) → f16 A fragment for the P@V mma.
unsigned _pa_q0[4];
emmy_c_to_a_f16(_pa_q0, _p_q0_f0, _p_q0_f1);
// --- Chain 1 online softmax (same moves, own carrier) ---
float _rmx_q10 = fmaxf(fmaxf(fmaxf(sacc_q1_f[0][0], sacc_q1_f[0][1]), sacc_q1_f[1][0]), sacc_q1_f[1][1]);
float _rmx_q11 = fmaxf(fmaxf(fmaxf(sacc_q1_f[0][2], sacc_q1_f[0][3]), sacc_q1_f[1][2]), sacc_q1_f[1][3]);
_rmx_q10 = fmaxf(_rmx_q10, __shfl_xor_sync(0xffffffff, _rmx_q10, 2));
_rmx_q11 = fmaxf(_rmx_q11, __shfl_xor_sync(0xffffffff, _rmx_q11, 2));
_rmx_q10 = fmaxf(_rmx_q10, __shfl_xor_sync(0xffffffff, _rmx_q10, 1));
_rmx_q11 = fmaxf(_rmx_q11, __shfl_xor_sync(0xffffffff, _rmx_q11, 1));
float _mn_q10 = fmaxf(m_i_q10, _rmx_q10); // m_new = max(m_i, rowmax)
float _mn_q11 = fmaxf(m_i_q11, _rmx_q11);
float _al_q10 = expf((m_i_q10 - _mn_q10)); // α = e^{m_i − m_new}
float _al_q11 = expf((m_i_q11 - _mn_q11));
float _p_q1_f0[4];
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
// p = e^{sⱼ − m_new}
_p_q1_f0[_e] = expf(sacc_q1_f[0][_e] - ((_e < 2) ? (_mn_q10) : (_mn_q11)));
}
float _p_q1_f1[4];
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
_p_q1_f1[_e] = expf(sacc_q1_f[1][_e] - ((_e < 2) ? (_mn_q10) : (_mn_q11)));
}
float _rsm_q10 = _p_q1_f0[0] + _p_q1_f0[1] + _p_q1_f1[0] + _p_q1_f1[1];
float _rsm_q11 = _p_q1_f0[2] + _p_q1_f0[3] + _p_q1_f1[2] + _p_q1_f1[3];
_rsm_q10 = _rsm_q10 + __shfl_xor_sync(0xffffffff, _rsm_q10, 2);
_rsm_q11 = _rsm_q11 + __shfl_xor_sync(0xffffffff, _rsm_q11, 2);
_rsm_q10 = _rsm_q10 + __shfl_xor_sync(0xffffffff, _rsm_q10, 1);
_rsm_q11 = _rsm_q11 + __shfl_xor_sync(0xffffffff, _rsm_q11, 1);
float l_i_q1__n0 = (l_i_q10 * _al_q10) + _rsm_q10; // d ← d·α + p
float l_i_q1__n1 = (l_i_q11 * _al_q11) + _rsm_q11;
l_i_q10 = l_i_q1__n0;
l_i_q11 = l_i_q1__n1;
#pragma unroll
for (int _r0 = 0; _r0 < 8; _r0 += 1) {
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
// o ← o·α
O_i_q1_f[_r0][_e] = O_i_q1_f[_r0][_e] * ((_e < 2) ? (_al_q10) : (_al_q11));
}
}
// Repack chain 1's P → f16 A fragment for its P@V mma.
unsigned _pa_q1[4];
emmy_c_to_a_f16(_pa_q1, _p_q1_f0, _p_q1_f1);
// O += P @ V. Load the shared V fragment ONCE (transposed, swizzled) and
// fan it into both chains' output mmas — again one V load, two chains.
unsigned _O_i__pv_b[8][2];
emmy_ldmatrix_x4_trans_pair(
_O_i__pv_b[0],
_O_i__pv_b[1],
&_v_smem[emmy_swizzle_b128(kv0 / 16 % 2 * 16 * 64
+ ((threadIdx.x & 31) % 16) * 64
+ ((threadIdx.x & 31) / 16) * 8)]
);
#pragma unroll
for (int _r0 = 0; _r0 < 2; _r0 += 1) {
emmy_mma_m16n8k16_f16(O_i_q0_f[_r0], _pa_q0, _O_i__pv_b[_r0], O_i_q0_f[_r0]);
emmy_mma_m16n8k16_f16(O_i_q1_f[_r0], _pa_q1, _O_i__pv_b[_r0], O_i_q1_f[_r0]);
}
emmy_ldmatrix_x4_trans_pair(
_O_i__pv_b[2],
_O_i__pv_b[3],
&_v_smem[emmy_swizzle_b128(kv0 / 16 % 2 * 16 * 64 + 16
+ ((threadIdx.x & 31) % 16) * 64
+ ((threadIdx.x & 31) / 16) * 8)]
);
emmy_mma_m16n8k16_f16(O_i_q0_f[2], _pa_q0, _O_i__pv_b[2], O_i_q0_f[2]);
emmy_mma_m16n8k16_f16(O_i_q1_f[2], _pa_q1, _O_i__pv_b[2], O_i_q1_f[2]);
emmy_mma_m16n8k16_f16(O_i_q0_f[3], _pa_q0, _O_i__pv_b[3], O_i_q0_f[3]);
emmy_mma_m16n8k16_f16(O_i_q1_f[3], _pa_q1, _O_i__pv_b[3], O_i_q1_f[3]);
emmy_ldmatrix_x4_trans_pair(
_O_i__pv_b[4],
_O_i__pv_b[5],
&_v_smem[emmy_swizzle_b128(kv0 / 16 % 2 * 16 * 64 + 32
+ ((threadIdx.x & 31) % 16) * 64
+ ((threadIdx.x & 31) / 16) * 8)]
);
emmy_mma_m16n8k16_f16(O_i_q0_f[4], _pa_q0, _O_i__pv_b[4], O_i_q0_f[4]);
emmy_mma_m16n8k16_f16(O_i_q1_f[4], _pa_q1, _O_i__pv_b[4], O_i_q1_f[4]);
emmy_mma_m16n8k16_f16(O_i_q0_f[5], _pa_q0, _O_i__pv_b[5], O_i_q0_f[5]);
emmy_mma_m16n8k16_f16(O_i_q1_f[5], _pa_q1, _O_i__pv_b[5], O_i_q1_f[5]);
emmy_ldmatrix_x4_trans_pair(
_O_i__pv_b[6],
_O_i__pv_b[7],
&_v_smem[emmy_swizzle_b128(kv0 / 16 % 2 * 16 * 64 + 48
+ ((threadIdx.x & 31) % 16) * 64
+ ((threadIdx.x & 31) / 16) * 8)]
);
emmy_mma_m16n8k16_f16(O_i_q0_f[6], _pa_q0, _O_i__pv_b[6], O_i_q0_f[6]);
emmy_mma_m16n8k16_f16(O_i_q1_f[6], _pa_q1, _O_i__pv_b[6], O_i_q1_f[6]);
emmy_mma_m16n8k16_f16(O_i_q0_f[7], _pa_q0, _O_i__pv_b[7], O_i_q0_f[7]);
emmy_mma_m16n8k16_f16(O_i_q1_f[7], _pa_q1, _O_i__pv_b[7], O_i_q1_f[7]);
// Commit both chains' running maxes for the next KV tile.
m_i_q00 = _mn_q00;
m_i_q01 = _mn_q01;
m_i_q10 = _mn_q10;
m_i_q11 = _mn_q11;
__syncthreads();
}
// Normalize chain 0 by its softmax denominator and store the output tile.
#pragma unroll
for (int _r0 = 0; _r0 < 8; _r0 += 1) {
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
O_i_q0_f[_r0][_e] = O_i_q0_f[_r0][_e] / ((_e < 2) ? (l_i_q00) : (l_i_q01));
}
}
#pragma unroll
for (int _r0 = 0; _r0 < 8; _r0 += 1) {
const int _g = (threadIdx.x & 31) >> 2; const int _t = (threadIdx.x & 31) & 3;
*reinterpret_cast<__half2*>(&sdpa[b1 * 8192 + m * 32 * 64 + _r0 * 8 + _g * 64 + _t * 2])
= __floats2half2_rn(O_i_q0_f[_r0][0], O_i_q0_f[_r0][1]);
*reinterpret_cast<__half2*>(&sdpa[b1 * 8192 + m * 32 * 64 + _r0 * 8 + (_g + 8) * 64 + _t * 2])
= __floats2half2_rn(O_i_q0_f[_r0][2], O_i_q0_f[_r0][3]);
}
// Normalize chain 1 by its own denominator and store the second tile.
#pragma unroll
for (int _r0 = 0; _r0 < 8; _r0 += 1) {
#pragma unroll
for (int _e = 0; _e < 4; _e++) {
O_i_q1_f[_r0][_e] = O_i_q1_f[_r0][_e] / ((_e < 2) ? (l_i_q10) : (l_i_q11));
}
}
#pragma unroll
for (int _r0 = 0; _r0 < 8; _r0 += 1) {
const int _g = (threadIdx.x & 31) >> 2; const int _t = (threadIdx.x & 31) & 3;
*reinterpret_cast<__half2*>(&sdpa[b1 * 8192 + (m * 32 + 16) * 64 + _r0 * 8 + _g * 64 + _t * 2])
= __floats2half2_rn(O_i_q1_f[_r0][0], O_i_q1_f[_r0][1]);
*reinterpret_cast<__half2*>(&sdpa[b1 * 8192 + (m * 32 + 16) * 64 + _r0 * 8 + (_g + 8) * 64 + _t * 2])
= __floats2half2_rn(O_i_q1_f[_r0][2], O_i_q1_f[_r0][3]);
}
}
This last move achieves parity with the SOTA handwritten FA-2 kernel. Both kernels now sit at ~167 TFLOPS: about 80% of the RTX 5090's fp16/fp32-accumulate tensor roofline. Asynchronous tensor-core instructions available in datacenter Hopper/Blackwell GPUs (FA-3 and FA-4 optimizations) would allow us to utilize the leftover 20% of the tensor pipe, but they're not available on consumer Blackwell. Most likely, this is the performance limit we can achieve without sacrificing accuracy.
Refused Optimizations
A few optimizations were attempted but were left out:
- Warp-scoping the handoff barrier (~0µs). Narrowing the C→A smem handoff's
__syncthreadsto__syncwarpdid nothing. - Fast
exp(~0µs). Libmexpfis a ~15-instruction sequence;__expfis one FMUL plusMUFU.EX2. Swapping ~70 exponentials moved nothing across two different kernel structures. Thus, foldingscale·log₂einto the score multiply and running the loop in base 2 — which FA-2 does — was rejected as well. - Deeper rings (
d3,d4): occupancy loss beats extra overlap at every geometry. - Warp specialization (213 vs 206µs). The FA-3 move is not relevant for
sm_120. Dedicating a producer warp slows down the overall pipeline.
Relationship to FA-3 and FA-4
FlashAttention-3 is a Hopper kernel. Its three contributions are wgmma (warpgroup MMA — an asynchronous tensor-core instruction), TMA, and a warp-specialized producer/consumer pipeline deep enough to exploit both. Consumer Blackwell (sm_120) has TMA, but no wgmma: its tensor core is the synchronous mma.sync. Warp specialization without async MMA is a pessimization for flash attention.
FlashAttention-4 is a B200 kernel. It targets tcgen05 MMA — asynchronous and accumulating into TMEM, the 256 KB tensor memory that exists only on datacenter Blackwell (sm_100) — plus 2-CTA MMA across paired SMs. None of that silicon is present on sm_120. And FA-4's portable-sounding ideas fail here too: the FMA-emulated exp answers a datacenter-Blackwell imbalance (tensor throughput doubled while the SFU stood still); on consumer Blackwell the tensor pipe is still the limit.
So the honest frontier statement for consumer Blackwell is FA-2 parity. TMA is the one post-FA-2 optimization that can be leveraged (it shows up as the LSU lead); the FA-3/4 optimizations stay locked until the necessary hardware reaches consumer GPUs.
Comparison with FA-2
First, the final bench end to end against PyTorch (whose SDPA dispatches to FlashAttention-2 on this shape):
EMMY_KNOBS="PLACE=fuse,TILE=a:mma_m16n8k16_f16/w4x1/f2x8/k4,STAGE=d2/tma/ring,WSPEC=" \
emmy run -c "$SDPA_P" --bench --target sm_120
Backend Latency (us) vs Eager
-------------------------------------
Eager PyTorch 205 1.00x
Emmy 205 1.00x
Kernel us % grid block smem regs occ PLACE@fold TILE@dd TILE@pj REDUCE STAGE
k_sdpa 204.7 100.0% 256 128 32.0K 254 17% fuse a:mma_m16n8k16_f16/w4x1/f2x8/k4 a:mma_m16n8k16_f16/w4x1/f2x8/k4 d2/tma/ring
TOTAL 204.7
The exact instantiation PyTorch ran (from NCU's demangled name): flash_fwd_kernel<Flash_fwd_kernel_traits<64, 128, 128, 4, Is_Q_in_regs=false, Share_Q_K_smem=false, half_t>>. Here is how the generated winner stacks up against the Dao-AILab source:
| Dimension | Emmy (generated, tuned) | FlashAttention-2 (hand-written) |
|---|---|---|
| CTA query tile | 128 rows = 4 warps × 2 reg tiles × 16 | identical |
| per-warp ILP | 2 independent (m,l,O) chains | identical |
| KV block | 64 keys | 128 keys |
| Q residency | register fragments, loaded once | smem, re-read per KV block (Is_Q_in_regs=false) |
| K/V transport | rank-4 TMA boxes, hw swizzle, 1-thread issue | cooperative cp.async, Swizzle<3,3,3> layout |
| P → A operand | fragment repack | identical |
| exp path | libm expf (fast path measured flat) | exp2f with scale·log₂e folded |
| LSU instructions | 3.31M | 4.36M |
| latency | 206µs | 206µs |
TLDR: The geometry converged. The autotuner's winning point, 128 query rows, 4 warps, 2 m-atoms per warp, is Dao's hand-picked constant set. The FA-2 source comments even document the author's manual sweep ("Using 8 warps is 18% slower… block size (256 x 64) is 85% slower, because of register spilling"), and they read line-for-line like the autotuner's fork table. TMA saves about 1M LSU instructions compared to FA-2 but measures flat: transport is not the bottleneck for this kernel. The fast-exp optimization is not needed on RTX 5090.
Wrap-Up
That's it. Bosses are slain, mobs are fleeing or cowering in fear. You're the hero of the Attentionium. The quest reward are these three knowledge bits that will help you in your next optimization journey:
- The schedule needs only associativity. Every optimization — stream the KV axis, carry
(m, d, o)in registers, run the matmuls on tensor cores, stage and pipeline the tiles, register-tile the queries — is licensed by associativity, not by attention. - Attention belongs to a class of secretly associative operations. The transport-of-structure principle helps to prove it. The Third Homomorphism Theorem helps to discover secretly associative loops.
- All FlashAttention optimizations are automatically discoverable. The generic set of scheduling and optimization moves that the Emmy compiler applies automatically discovers all FlashAttention optimizations and reaches parity with the SOTA handwritten kernel that took engineers years to optimize.
Bonfire lit
References
Series
- Learning FlashAttention the Hard Way — Part 1. The algebra: the twisted monoid, transport of structure, the numerical analysis.
- GPU Matmul Optimization. Register tiling, padding, swizzle, split-K — the same moves, on the kernel that taught them.
FlashAttention
- Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022.
- Tri Dao. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. 2023. The kernel this post ties; source for the face-off section (v2.7.4,
flash_fwd_kernel.h/kernel_traits.h/softmax.h). - Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision. 2024.
wgmma, TMA, warp specialization, FP8 — the Hopper rungs. - Ted Zadouri, Markus Hoehnerbach, Jay Shah, Timmy Liu, Vijay Thakkar, Tri Dao. FlashAttention-4. 2026.
tcgen05, tensor memory, the FMA-emulatedexp, the conditional rescale — the B200 rungs. - Tri Dao, et al. Flash-Decoding for long-context inference. 2023. Split-KV — the monoid's split-K.
GPU programming primers
- Simon Boehm. How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance. The why behind tiling, staging, and pipelining.
- Vasily Volkov. Better Performance at Lower Occupancy. GTC 2010. The register-tiling argument — the one FA-2's 255-register design and our reg_m rung both descend from.
- NVIDIA. PTX ISA —
ldmatrix,mma.sync,cp.async.bulk.tensor. The lane maps behind the C→A repack and the paired-x4drains.