Emre Albayrak

1.8x Faster Attention on an RTX 5080: an FP8 Flash Attention Kernel in CuTe-DSL

This blog’s explanation and structure will be different this time. Now I have the control, you are going to read this in my words. Minimal AI usage. Only for grammar fixes and diagrams, and some technical explanations.

I made a custom Flash Attention implementation on CuTe DSL that gets 123.6 ms on the 5080, passes FlashInfer’s 131.8, and cuDNN’s 226.2 (BF16). tbh, I did not expect that result either. How did I do it? Let me explain.

Where this even came from?

We were sitting with Ezgi in a cafe, and she showed me her blog post about profiling the newly released MiniMax H3 model on a single GB300. I read it and she suddenly asked “Want to take a shot at this on your 5080? Could turn into something interesting.” Heck yeah. Thanks Ezgi, I’m taking it from here.

Determining the target

I first tried to load the model onto the 5080 for profiling, but it did not fit in 16 GB. Then I came across this idea. I loaded a single DiT block of the model (1.3 GB of weights), and measured. H3 has 50 of them and they are sequential. So if I measure one block’s metrics and multiply by 50, I could get the whole DiT block’s metrics. So I did.

one H3 DiT block

one DiT block, RTX 5080, S=30272, msBF16shareFP8 GEMMsshare
SDPA attention241.748.4%241.758.7%
fc1 GEMM78.915.8%31.07.5%
qkv GEMM59.111.8%25.56.2%
fc2 GEMM + gate41.08.2%36.99.0%
out GEMM + gate21.64.3%19.04.6%
qknorm + rope33.26.6%33.28.1%
norms + modulate + silu_mul24.24.8%24.25.9%
block total499.9411.7

Turns out on a single block, attention is 48% of the whole process, the rest were ops like GEMM, residual etc. and I found my target. I’ll write an FP8(?) attention kernel for H3, and will contribute to SGL if it works as expected and fast enough. Challenge accepted.

Then I asked AI to search SGL for two things:

  1. Finding something about SM120 attention, like is there any supported backend for SM120, and if it does, how fast is it?
  2. Is there any active contribution or question about SM120 attention?

There were some questions about it but turns out SGL did not have any SM120-specific attention kernel. The only ones were for SM100 (Datacenter Blackwell), and SGLang did not even select cuDNN for SM120, it fell back to the default Torch SDPA. So we can’t benefit SM120’s capabilities on attention side. I really wanted to change that.

Okay, we have our target, and nobody had done anything yet. We are going to write an FA-style FP8 attention kernel. I decided to go with CuTe DSL this time. About why, I just went curious, wanted to create something with it, and SGL accepts DSL kernels.

I had never written an attention kernel at the Python level, with a DSL. So, I asked AI to search for any people that wrote DSL kernels. I came across Blake Ledden’s FP8 FA kernel (I call it Blake Attention) PR and decided to read it. Then I contacted him and asked if he could tell me which methods he applied, and for suggestions. He was very kind and told me everything he did, alongside with some suggestions. No worries, I’ll explain everything I applied from his kernel.

Learning CuTe-DSL

Okay. We have some methods, some learning paths to start. I have some C++ CUTLASS experience, and switching to CuTe DSL was not really a big deal for me. The difference is function names and some preparations. For start, I pair-coded 5 different DSL kernels with AI, I asked and it answered. I asked to analyze Blake’s attention kernel, and prepare me a learning series in just a few steps so in the end it led me to the information I needed to write this kernel. I also followed documentation examples, but they looked a little confusing.

I have a learning philosophy that first I get that example, and I try to reduce it to its simplest form, keeping only the necessary parts. So you get a very clean picture to learn from and focus only on that essential part. I’m not going to explain what I wrote and what they do in this blog, because I’m going to explain the whole attention kernel, so you learn the hardest but simplest way, unlike me.

The attention itself

Okay. Let’s start with understanding what we are going to do first. We have to perform an attention operation. When I started writing the kernel, in some parts of the code I could not understand some of the index or coordinate calculations, because they were difficult to conceptualize and visualize in theory.

One time while I was at the gym, I decided to watch Umar Jamil’s Flash Attention explanation video. And when I saw this diagram, everything started to click:

multi-head attention, after Umar Jamil's diagram

When I saw the Q, K and V pieces, I realized how we distribute those pieces across the GPU.

On H3, the shape of Q, K and V for one head is (30272, 128). We have 56 heads in total, so that makes (56, 30272, 128) for Q alone. And we have 3 matrices, Q, K and V, so the total becomes (3, 56, 30272, 128). Wow. But let’s not let this huge size intimidate us. Let’s start with Q, for one head.

One head owns (30272, 128), which means we have a total of 30272 tokens, each embedded with 128 values. We have to map this matrix onto the GPU. We have CTAs (blocks). However, the number of values that a CTA can fetch and process is, of course, limited. So we split this head across CTAs like this:

one head: Q, K, V, O in 64-row strips

Every CTA owns a small portion of this head, which is (64, 128) in our case, for Q, K and V alike. The Q tile is fetched once, but in every iteration we have to read a new K and V tile.

Now we have the Q, K and V pieces. One more question comes. How does one CTA process those tiles? With warps.

Threads inside a CTA work in groups of 32, called warps. Each CTA contains 128 threads, and that makes 4 warps per CTA. And we split the score tile like this:

score tile, 4 warp bands

With that way, each warp owns (16, 64) in a single iteration step, and each lane (single thread inside of warp) owns 32 values. But the warp does not compute this (16, 64) in one shot. The tensor core instruction gives 16 x 8 per call, so the warp calls it 8 times, one tile per call:

one warp's slab, 8 tiles

Let’s also take a look inside one tile. 16 x 8 = 128 scores across 32 lanes, so 4 scores per lane. Which lane holds which score? The PTX ISA decides that:

lane map of one 16 x 8 tile

The tensor core instruction is fixed at 16 rows x 8 keys per call, and a warp’s slab is 8 of those side by side like the diagram above.

Okay, we got Q, but there are 2 more, K and V. How are we going to process them? Simpler. The Q tile stays, we iterate over K and V strips. Each strip is (64, 128). All 128 threads load it into shared memory together. Then every warp reads the whole strip, not a slice like Q. Process, then load the next strip.

shared memory of one CTA

Okay, I said enough but there is one more thing. Softmax.

Softmax wants two things from a row before it can finish: the max of the row, and the sum of the row. But we never see a whole row. We see 64 keys, process them, then the next 64. The row is 30272 keys long.

So we cheat. For every query row we carry two numbers: m, the biggest score so far, and l, the sum of exps so far. A new strip comes in. If its max is bigger than m, every exp we summed before used the wrong max. The fix is one multiply: a = e^(m_old - m_new). Scale l by a, then add the exps of the new strip. O has the same problem, because P used the old max. Scale O by a too. We divide O by l only once, at the very end.

Small example. One row, 4 keys, 2 strips: [1, 2] then [4, 3].

If you had all 4 keys at once: max = 4, l = e^-3 + e^-2 + e^0 + e^-1 = 1.553. Same number. That is online softmax. Here is the same thing for the whole row, but for all 473 strips:

online softmax state per strip

Enough shrinking. I guess you got the idea about what we are going to do. If you came way down here, congrats! Now, we can inspect the code.

The kernel, “block by block”

Okay. Let’s scrub in and get our hands dirty. I’m going to explain the kernel block by block and link with diagrams above for you to ‘click’.

Who am I, what’s mine?

this CTA's Q strip and O strip

# thread 0..127 inside the CTA
tidx, _, _ = cute.arch.thread_idx()
# bidx: which 64-row strip, bidz: which head
bidx, _, bidz = cute.arch.block_idx()

# per-head FP8 descales, folded into two scalars
# S = qk_scale * (Q K^T)
qk_scale = mScales[0, bidz] * mScales[1, bidz] * softmax_scale
# undo the V descale and the x256 we put on P
pv_scale = mScales[2, bidz] / 256.0
# 30272, for the bounds check
sequence = mLSE.shape[0]

# online softmax state, two rows per thread
# first max(-inf, x) = x
m = cute.make_rmem_tensor((2,), cutlass.Float32)
m.fill(-cutlass.Float32.inf)
# first a * 0 + sum = sum
l = cute.make_rmem_tensor((2,), cutlass.Float32)
l.fill(0.0)

# views into global memory, no data moves here
# own Q strip (64, 128)
gQ = cute.local_tile(mQ[None, None, bidz], cta_tiler, (bidx, None, 0), proj=(1, None, 1))
# all K strips (64, 128, 473)
gK = cute.local_tile(mK[None, None, bidz], cta_tiler, (None, None, 0), proj=(None, 1, 1))
# all V strips (128, 64, 473), feature-major
gV = cute.local_tile(mV[None, None, bidz], (128, 64), (0, None))

Every CTA starts by learning its address. bidx is the strip, bidz is the head. Together they say: “I own Q rows bidx * 64 .. bidx * 64 + 63 of head bidz, and the same rows of O”. The filled boxes in the diagram are exactly that.

The three local_tile calls turn the address into views. gQ is my one strip, (64, 128). gK and gV are every strip of the head, (64, 128, 473) and (128, 64, 473). The last axis is the strip index, t0 .. t472 in diagram 1. No byte moves in this block. A view is an address, not a copy.

How to read local_tile: the tiler is the GEMM box (M, N, K) = (64, 64, 128). proj keeps the two axes each matrix uses. In the coordinate, an integer picks one box, None keeps all boxes as a new axis. The full walkthrough is folded below.

Reading local_tile, the full walkthrough

local_tile cuts a box out of a big tensor and hands you a view of one box, or of all boxes along an axis. No data moves. To read the call, we need to know what the tiler is and what 0, 1 and None mean.

gQ = cute.local_tile(mQ[None, None, bidz], cta_tiler, (bidx, None, 0), proj=(1, None, 1))

First we give mQ[None, None, bidz]. mQ is the input Q tensor. It has 3 axes: (token, feature, head). We pin the head axis to bidz, the head this CTA works on. The other two axes stay open. What is left is a (30272, 128) matrix: one head, all tokens, all features.

Next comes cta_tiler. The tiler is the GEMM’s box, not Q’s box. One iteration is a GEMM: S (64 x 64) = Q (64 x 128) x K^T (128 x 64). A GEMM has 3 axes: M = 64 query rows, N = 64 key rows, K = 128 features. So cta_tiler = (64, 64, 128) is (M, N, K), the whole GEMM box of one CTA. Each matrix touches only two of the three axes:

MatrixAxesBox
QM, K(64, 128)
KN, K(64, 128)
SM, N(64, 64)

One tiler, three matrices. proj picks the two axes each matrix uses. That is the whole reason it exists. In proj and in the coordinate, only three values appear:

ValueIn projIn the coordinate
1this tiler axis applies to my tensor, keep itnot used
Nonethis tiler axis does not apply to my tensor, drop itkeep all boxes on this axis as a new trailing axis
integer inot usedgive me box number i on this axis, the axis collapses

For Q, proj = (1, None, 1) says: keep M, drop N, keep K. The tiler shrinks from (64, 64, 128) to (64, 128). That is Q’s box.

The coordinate (bidx, None, 0) comes last. It has three slots, the same three as the tiler: M, N, K. Each slot answers one question: on this axis, which box do I want?

The first slot is M, the token axis. The tiler cuts it into boxes of 64 rows. 30272 rows make 473 boxes. bidx says: give me box number bidx. That is one strip, rows bidx * 64 to bidx * 64 + 63. Diagram 1 shows this as “CTA 1 owns rows 64..127”.

The second slot is N, the key axis. Q has no key axis. proj already dropped this slot, so the value here is never read. The None is filler. It exists only because the coordinate must have the same three slots as the tiler.

The third slot is K, the feature axis. The tiler cuts it into boxes of 128. The axis has 128 features, so there is exactly one box. 0 says: give me box 0, the only one. In practice this means the feature axis is not split at all.

After this, gQ is (64, 128): one strip of 64 rows, all 128 features. The CTA’s own strip.

Now K, with one change. proj = (None, 1, 1) drops M and keeps N and K, so the box is again (64, 128). The coordinate is (None, None, 0). The first None sits on the dropped M slot, filler again. The second None sits on N, a kept slot, and here it means: do not pick a box, give me all 473 of them as a new axis. So gK is (64, 128, 473). The third axis is the strip index, t0 .. t472 in diagram 1. The K loop walks that axis.

V has no proj. Its tiler (128, 64) is already in V’s own two axes, (feature, token), so there is nothing to drop. The coordinate (0, None) says: feature box 0, all token boxes. gV is (128, 64, 473).

m and l are the online softmax state from diagram 6, two query rows per thread. They start empty: max = -inf, sum = 0. qk_scale and pv_scale fold the FP8 descales of this head into two numbers. The QK^T block uses the first, the P.V block uses the second.

Watch out: V’s tiler is (128, 64), not cta_tiler. V is stored feature-major. One strip is [128 features, 64 tokens], so the P.V GEMM reads it without a transpose.

Make the room

global to shared, one CTA

# one CTA's shared memory: Q tile, K strip, V strip, P tile
@cute.struct
class SharedStorageQKVP:
    q: cute.struct.Align[cute.struct.MemRange[mQ.element_type, cute.cosize(sQ_layout)], 16]
    k: cute.struct.Align[cute.struct.MemRange[mK.element_type, cute.cosize(sK_layout)], 16]
    v: cute.struct.Align[cute.struct.MemRange[mV.element_type, cute.cosize(sV_layout)], 16]
    probabilities: cute.struct.Align[cute.struct.MemRange[mQ.element_type, cute.cosize(sP_layout)], 16]

smem = cutlass.utils.SmemAllocator()
storage = smem.allocate(SharedStorageQKVP.size_in_bytes(), byte_alignment=16)

# shared tensors: layout says the shape, swizzle says where each row lands
sQ = SharedStorageQKVP(storage).q.get_tensor(sQ_layout, swizzle=sQ_swizzle)
sK = SharedStorageQKVP(storage).k.get_tensor(sK_layout, swizzle=sK_swizzle)
sV = SharedStorageQKVP(storage).v.get_tensor(sV_layout, swizzle=sV_swizzle)
sP = SharedStorageQKVP(storage).probabilities.get_tensor(sP_layout, swizzle=sP_swizzle)
# V arrives as [128 features, 64 tokens], pick stage 0
sV_transposed = sV[None, None, 0]

# global -> shared copy: this thread's share of the source and the destination
thr_copy_Q = tiled_copy_Q.get_slice(tidx)
thr_copy_K = tiled_copy_K.get_slice(tidx)
thr_copy_V = tiled_copy_V.get_slice(tidx)

tQgQ = thr_copy_Q.partition_S(gQ)
tKgK = thr_copy_K.partition_S(gK)
tVgV = thr_copy_V.partition_S(gV)

tQsQ = thr_copy_Q.partition_D(sQ)
tKsK = thr_copy_K.partition_D(sK)
tVsV = thr_copy_V.partition_D(sV)

# 473 strips
k_tile_count = cute.size(tKgK, mode=[3])

# Q moves once, here. K and V move inside the loop.
cute.copy(tiled_copy_Q, tQgQ[None, None, None], tQsQ[None, None, None, 0])

We told the CTA where its data lives. Now we need to allocate some space in SMEM. For this, we define a struct for the Q, K, V and P tiles. Align forces the starting address of that field to be a multiple of 16 bytes. Nothing else. Let’s inspect a single Align line:

q: cute.struct.Align[cute.struct.MemRange[mQ.element_type, cute.cosize(sQ_layout)], 16]

Read it from the inside out. mQ.element_type is FP8, one byte. cute.cosize(sQ_layout) is how many elements the layout spans, 8192 for (64, 128, 1). MemRange[type, count] reserves that many elements as a raw byte range. No shape yet. Align[range, 16] starts the range on a 16-byte boundary, because a 128-bit cp.async needs an aligned destination. So the line says: 8 KB of raw shared memory, 16-byte aligned, called q. The shape and the swizzle come later, from get_tensor.

The struct only describes the memory, it does not own anything. That is the job of SmemAllocator(). It is a bump pointer that starts at the CTA’s SMEM base. allocate hands back a pointer to the next n bytes, rounded up to a 16-byte boundary, and moves the pointer forward. size_in_bytes() is 28 KB: q, k and v are 64 x 128 FP8, 8 KB each, and p is 64 x 64, 4 KB. Align padding is counted in, zero here.

SharedStorageQKVP(storage) lays the struct over that pointer. get_tensor(layout, swizzle) turns each raw field into a tensor. The layout gives the shape, (64, 128, 1). The trailing 1 is the pipeline stage axis. This version of the kernel has one stage, so that index is always 0. We will come back to this axis when we optimize. The swizzle decides where each row lands in the banks. ldmatrix reads 8 rows of 16 bytes at once. In a plain layout those 8 rows hit the same banks. The swizzle mixes the column bits, so the 8 rows spread over 8 different banks.

Then the global-to-shared copy. tiled_copy_Q is a cp.async, 16 bytes per thread per instruction. 128 threads move 2 KB per instruction, so one 8 KB tile takes 4 instructions per thread. get_slice(tidx) picks this thread’s share. partition_S slices the source, the global view. partition_D slices the destination, the shared tensor. tQgQ and tQsQ have the same shape, so cute.copy pairs them element for element. tKgK keeps the strip axis: mode 3 is 473, and that is k_tile_count.

The last line moves Q. Once, here, before the loop. Diagram 1 said “read once”, this is that line. K and V move inside the loop, one strip per iteration.

Watch out: nobody waits for the Q copy here. cp.async is asynchronous. The first loop iteration calls cp_async_wait_group(0), and that waits for everything in flight, Q included.

Who holds which numbers

one thread's values in tCrC and tCrO

# the MMA's view of the work: which 4 scores does this thread hold
thr_mma = tiled_mma.get_slice(tidx)

# this thread's share of each shared tile, in MMA operand layout
tCsQ = thr_mma.partition_A(sQ)
tCsK = thr_mma.partition_B(sK)
tCsV = thr_mma.partition_B(sV_transposed)
tCsP = thr_mma.partition_A(sP)
# P is written in C layout, read back in A layout
tCsP_C = thr_mma.partition_C(sP[None, None, 0])

# register fragments, same shape as the shared share, stage axis dropped
tCrQ = tiled_mma.make_fragment_A(tCsQ[None, None, None, 0])
tCrK = tiled_mma.make_fragment_B(tCsK[None, None, None, 0])
tCrV = tiled_mma.make_fragment_B(tCsV)

# S accumulator: 64 x 64 per CTA, 32 FP32 per thread, reset every strip
acc_shape = thr_mma.partition_shape_C((64, 64))
tCrC = cute.make_rmem_tensor(acc_shape, cutlass.Float32)

# O accumulator: 64 x 128 per CTA, 64 FP32 per thread, lives for all 473 strips
acc_shape_O = thr_mma.partition_shape_C((64, 128))
tCrO = cute.make_rmem_tensor(acc_shape_O, cutlass.Float32)
tCrO.fill(0.0)

# P: FP8 copy of the scores in C layout, and its A-layout fragment
tCrP_C = cute.make_fragment_like(tCrC, mQ.element_type)
tCrP = tiled_mma.make_fragment_A(tCsP[None, None, None, 0])
# 64 keys / 32 = 2
num_k_block_PV = cute.size(tCrP, mode=[2])

# shared -> register: ldmatrix, 4 matrices of 8 x 16 bytes per call
atom_copy_s2r_Q = cute.make_copy_atom(
    cute.nvgpu.warp.LdMatrix8x16x8bOp(transpose=False, num_matrices=4),
    mQ.element_type
)
atom_copy_s2r_K = cute.make_copy_atom(
    cute.nvgpu.warp.LdMatrix8x16x8bOp(transpose=False, num_matrices=4),
    mK.element_type
)
atom_copy_s2r_V = cute.make_copy_atom(
    cute.nvgpu.warp.LdMatrix8x16x8bOp(transpose=False, num_matrices=4),
    mV.element_type
)

# ldmatrix laid out to feed the MMA's A and B operands
tiled_copy_s2r_P = cute.make_tiled_copy_A(atom_copy_s2r_Q, tiled_mma)
ldmatrix_P = tiled_copy_s2r_P.get_slice(tidx)
tCsP_copy_view = ldmatrix_P.partition_S(sP)
tCrP_copy_view = ldmatrix_P.retile(tCrP)
tCsP_p = tCsP_copy_view[None, None, None, 0]

tiled_copy_s2r_Q = cute.make_tiled_copy_A(atom_copy_s2r_Q, tiled_mma)
tiled_copy_s2r_K = cute.make_tiled_copy_B(atom_copy_s2r_K, tiled_mma)
tiled_copy_s2r_V = cute.make_tiled_copy_B(atom_copy_s2r_V, tiled_mma)

ldmatrix_Q = tiled_copy_s2r_Q.get_slice(tidx)
ldmatrix_K = tiled_copy_s2r_K.get_slice(tidx)
ldmatrix_V = tiled_copy_s2r_V.get_slice(tidx)

# the same shared tile and the same registers, seen the way ldmatrix moves them
tCsQ_copy_view = ldmatrix_Q.partition_S(sQ)
tCrQ_copy_view = ldmatrix_Q.retile(tCrQ)
tCsK_copy_view = ldmatrix_K.partition_S(sK)
tCrK_copy_view = ldmatrix_K.retile(tCrK)
tCsV_copy_view = ldmatrix_V.partition_S(sV_transposed)
tCrV_copy_view = ldmatrix_V.retile(tCrV)

# 128 features / 32 = 4
num_k_block = cute.size(tCrQ, mode=[2])

# pick stage 0
tCsQ_p = tCsQ_copy_view[None, None, None, 0]
tCsK_p = tCsK_copy_view[None, None, None, 0]
tCsV_p = tCsV_copy_view

Block 2 gave us sQ, sK and sV in shared memory. But, that’s not enough. The tensor core reads its operands from registers only. So every operand needs two definitions. One tile in shared memory, one fragment in registers. They have different owners. The shared tile belongs to the whole CTA. All 4 warps read the same K strip, you can take a look at diagram 5. The register fragment belongs to one thread. Each lane holds its own 4 scores, diagram 4. Between the two sits ldmatrix.

shared tile to register fragment

tiled_mma.get_slice(tidx) picks this thread’s slice of the MMA. Block 2 did the same with tiled_copy_Q.get_slice(tidx), and the two slices differ. The copy slice says which 16 bytes this thread moves. The MMA slice says which 4 scores this thread holds. partition_A(sQ) cuts the shared tile in the MMA’s A-operand layout and hands back this thread’s share, tCsQ. partition_B does the same for K and V, partition_C for the output side. Still no bytes moved. tCsQ is a view.

make_fragment_A(tCsQ[..., 0]) opens registers of the same shape, stage axis dropped. That is tCrQ, empty for now. tCrC and tCrO are the accumulators. tCrC is 32 FP32 per thread for S: the 8 tiles of diagram 3, 4 values each. tCrO is 64 FP32 for O, twice as wide. tCrC resets every strip. tCrO is zeroed once here and lives through all 473 strips.

Now the bridge. LdMatrix8x16x8bOp is one ldmatrix call: 4 matrices of 8 rows x 16 bytes. Each lane hands in one row address. The hardware deals the bytes out in the MMA’s fragment order. make_tiled_copy_A(atom, tiled_mma) lays this atom out so its output lands where the MMA’s A operand expects it. ldmatrix_Q.partition_S(sQ) is the shared tile again, now cut the way ldmatrix reads it. Two partitions of one tile: tCsQ for the MMA, tCsQ_copy_view for the copy.

retile(tCrQ) is the register side of the same idea. ldmatrix writes registers in its own order. The MMA reads them in its own order. Both orders describe the same 16 registers. retile gives the copy a view of tCrQ in the copy’s terms. Nothing is copied. In the loop, cute.copy fills tCrQ_copy_view and cute.gemm reads tCrQ. Same registers, two names.

P gets three names, tCrP_C, tCsP_C and tCrP, because on FP8 the scores cannot stay in registers between the two GEMMs. You can see it from Block 6.

Watch out: nothing moves in this block. 60 lines, zero bytes. Every name is a view. Only cute.copy and cute.gemm inside the loop touch data.

The loop starts

one strip: global, shared, registers, mma

# 473 turns, one K strip and one V strip per turn
for k_tile in range(k_tile_count):
    # global -> shared, strip k_tile into stage 0
    cute.copy(tiled_copy_K, tKgK[None, None, None, k_tile], tKsK[None, None, None, 0])
    cute.copy(tiled_copy_V, tVgV[None, None, None, k_tile], tVsV[None, None, None, 0])

    # close the batch, wait for my copies, then wait for everyone's
    cute.arch.cp_async_commit_group()
    cute.arch.cp_async_wait_group(0)
    cute.arch.sync_threads()

    # fresh S for this strip. O keeps accumulating.
    tCrC.fill(0.0)

    # S = Q K^T, 128 features in 4 steps of 32
    for k_block in cutlass.range(num_k_block, unroll_full=True):
        # shared -> registers, one 32-feature slice of Q and of K
        cute.copy(tiled_copy_s2r_Q, tCsQ_p[None, None, k_block], tCrQ_copy_view[None, None, k_block])
        cute.copy(tiled_copy_s2r_K, tCsK_p[None, None, k_block], tCrK_copy_view[None, None, k_block])

        # 8 tensor core calls, one per tile of diagram 3
        cute.gemm(tiled_mma, tCrC, tCrQ[None, None, k_block], tCrK[None, None, k_block], tCrC)

Everything before this line was setup. This is where bytes actually start to move. k_tile is the strip index, the last axis of gK and gV from block 1. Every turn brings one K strip and one V strip into shared memory, always into stage 0, because this version has only one stage.

The three lines after the copies are three different waits. cp_async_commit_group closes the batch: the K copy, the V copy, and on the first turn the Q copy from block 2. cp_async_wait_group(0) blocks this thread until its own copies land. sync_threads blocks until every thread in the CTA got there. Without the last one, a thread would read K rows that another thread’s copy did not finish yet. Remember, the shared strip is common property and every warp reads all 64 keys.

cp.async, commit, wait, sync

tCrC.fill(0.0) resets S. Every strip gets a fresh score tile. tCrO is not touched, it keeps summing across strips, that is the whole point of the online softmax.

Then the first GEMM. The atom takes 32 features per call and Q has 128, so num_k_block is 4. Each step moves one 32-feature slice of Q and one of K from shared memory into registers, with the ldmatrix copies from block 3, and calls cute.gemm once. One cute.gemm is 8 tensor core instructions, the 8 tiles of diagram 3, all accumulating into tCrC. After 4 steps, tCrC holds this thread’s 32 scores of the (64, 64) S tile. Diagram 4 says which 32.

Watch out: Q does not change, yet the loop reads it from shared memory 473 times. Holding the whole Q fragment in registers instead would keep 16 more registers live for the entire loop. The final kernel keeps the re-read too. I did not measure that trade.

Online softmax, one strip at a time

one query row, one quad, two shuffles

# two query rows per thread: row 0 is query g, row 1 is query g + 8 (diagram 4)
for row in cutlass.range_constexpr(2):
    # this lane's 16 raw scores of that row, (2, 8), scaled
    tCrC_row = tCrC[(None, row), 0, None]
    row_scores = tCrC_row.load() * qk_scale

    # only compiled when the sequence is not a multiple of 64. 30272 = 473 x 64, so not here.
    if cutlass.const_expr(sequence % 64 != 0):
        tCrC_row.store(row_scores)
        for column_group in cutlass.range_constexpr(8):
            for column_pair in cutlass.range_constexpr(2):
                key_row = (
                    k_tile * 64 + (tidx % 4) * 2
                    + column_group * 8 + column_pair
                )
                if key_row >= sequence:
                    tCrC_row[column_pair, column_group] = -cutlass.Float32.inf
        row_scores = tCrC_row.load()

    # max over my 16, then over the 4 lanes of the quad
    local_max = row_scores.reduce(cute.ReductionOp.MAX, -cutlass.Float32.inf, 0)

    neighbor_max = cute.arch.shuffle_sync_bfly(local_max, offset=1)
    pair_max = cute.arch.fmax(local_max, neighbor_max)

    neighbor_max = cute.arch.shuffle_sync_bfly(pair_max, offset=2)
    tile_max = cute.arch.fmax(pair_max, neighbor_max)

    # diagram 6: new running max, and the correction for everything summed so far
    m_new = cute.arch.fmax(m[row], tile_max)
    alpha = cute.math.exp2((m[row] - m_new) * math.log2(math.e), fastmath=True)

    # p = e^(s - m_new), 16 values
    p = cute.math.exp2((row_scores - m_new) * math.log2(math.e), fastmath=True)

    # lane-local sum of the 16 p values: 16 -> 8 -> 4 -> 2 -> 1
    p_partial_sums = cute.make_rmem_tensor((8,), cutlass.Float32)
    for column_group in cutlass.range_constexpr(8):
        p_partial_sums[column_group] = p[0, column_group] + p[1, column_group]

    for level in cutlass.range_constexpr(3):
        for sum_index in cutlass.range_constexpr(4 >> level):
            p_partial_sums[sum_index] = (
                p_partial_sums[2 * sum_index]
                + p_partial_sums[2 * sum_index + 1]
            )
    tile_sum = p_partial_sums[0]

    # l = alpha * l + sum, m = m_new
    l[row] = alpha * l[row] + tile_sum
    m[row] = m_new

    # the scores are done, the same registers now hold p
    tCrC_row.store(p)

    # O = alpha * O, same correction as l
    tCrO_row = tCrO[(None, row), 0, None]
    tCrO_row.store(tCrO_row.load() * alpha)

Okay, we somehow managed to have values in tCrC, the scores. But, there is one thing to discuss. tCrC belongs to one thread and has only 16 scores of that row: 2 columns of each 8-wide tile, 8 tiles, see diagram 4. The row has 30272 values in total. In order to calculate softmax, we have to know the max of the whole row, and the sum. So how do we calculate that? Luckily, we have Online Softmax.

Online Softmax does the calculation on the road. Basically it takes the max of that row and does the calculation, and updates if we find a new max alongside on the road. But, how do we do that?

For one row, when we get a new strip, we have 3 values so far. One, m, the max, second l which is sum, and O, the weighted sum of V. The common thing for those 3 values is that all of them were computed with m. What will happen when we have a new max m value? Only 7 steps:

  1. Find the biggest of my 16 new scores. The other 48 keys of this strip sit in 3 neighbor lanes, two shuffles collect their max. Compare it with m. The bigger one is the new m.
  2. Compute the correction: alpha = e^(old m - new m). If m did not change, alpha is 1.
  3. Multiply l by alpha. Multiply O by alpha. The old work now looks as if we used the new m from the start.
  4. Subtract the new m from each of my 16 new scores and take e to that power. These are p.
  5. Sum the p values and add the sum to l.
  6. Multiply the p values with their V rows and add to O. That part is block 6.
  7. Store the new m.

Watch out, multiply first, then add. The new p values are already computed with the new m. If you add them before the multiply, alpha scales them too, and they are wrong.

P goes through shared memory, then P.V

P through shared memory, then P.V

# p in [0, 1] -> [0, 256] before FP8, pv_scale / 256 undoes it in block 7
tCrP_C.store((tCrC.load() * 256.0).to(mQ.element_type))
# my 32 FP8 p values -> shared, in C layout
cute.autovec_copy(tCrP_C, tCsP_C)
# everyone's p is in sP before anyone reads it in A layout
cute.arch.sync_threads()

# O += P V, 64 keys in 2 steps of 32
for k_block in cutlass.range(num_k_block_PV, unroll_full=True):
    # shared -> registers, one 32-key slice of P and of V
    cute.copy(
        tiled_copy_s2r_P,
        tCsP_p[None, None, k_block],
        tCrP_copy_view[None, None, k_block],
    )
    cute.copy(
        tiled_copy_s2r_V,
        tCsV_p[None, None, k_block],
        tCrV_copy_view[None, None, k_block],
    )
    # 16 tensor core calls, one per tile of the (64, 128) O
    cute.gemm(
        tiled_mma,
        tCrO,
        tCrP[None, None, k_block],
        tCrV[None, None, k_block],
        tCrO,
    )

# nobody may overwrite sK, sV, sP with the next strip while a warp still reads them
cute.arch.sync_threads()

Okay, now we have P, but it’s not complete softmax yet. We still did not do the final division. Dividing at the end gives the exact result. Dividing in the middle is extra work for the same result. We will do it later. Now, we focus on final P @ V, which is the second GEMM operation. This time, our A operand is P, and the B operand is V.

The extra thing we did is the FP8 detour. At the end of block 5, tCrC has 32 values per thread, in the output order of the MMA, and all values are FP32. We need to change it and make it suitable for the A operand. The A operand needs 2 things. First, the values must be FP8, second, the order.

This second GEMM loop is not as fancy as the first one, so take a look at the code block and try to understand.

After the last strip: divide, write

row sum, LSE, O row, global O strip

# who am I inside the warp: 4 lanes of a quad share two query rows
warp_id = tidx // 32
lane_id = tidx % 32
quad_id = lane_id // 4

for row in cutlass.range_constexpr(2):
    # the lane-local l of block 5, summed over the 4 lanes of the quad
    row_sum = l[row]

    neighbor_sum = cute.arch.shuffle_sync_bfly(row_sum, offset=1)
    row_sum += neighbor_sum

    neighbor_sum = cute.arch.shuffle_sync_bfly(row_sum, offset=2)
    row_sum += neighbor_sum

    l[row] = row_sum
    # diagram 6 bottom left
    lse = m[row] + cute.math.log(row_sum)

    # which query row of the head this is: strip, warp, quad, row (diagram 2 + 4)
    query_row = bidx * 64 + warp_id * 16 + quad_id + row * 8
    # only compiled when the sequence is not a multiple of 64
    valid_query = True
    if cutlass.const_expr(sequence % 64 != 0):
        valid_query = query_row < sequence
    # one lane per quad writes LSE
    if valid_query and lane_id % 4 == 0:
        mLSE[query_row, bidz] = lse

    # diagram 6 bottom right: O / l, with the V descale and the /256 folded in
    tCrO_row = tCrO[(None, row), 0, None]
    inverse_row_sum = pv_scale / row_sum
    tCrO_row.store(tCrO_row.load() * inverse_row_sum)

    # each lane writes its own 32 features, no shuffle needed
    if valid_query:
        for feature_group in cutlass.range_constexpr(cute.size(tCrO_row, mode=[1])):
            for feature_pair in cutlass.range_constexpr(2):
                feature = (lane_id % 4) * 2 + feature_group * 8 + feature_pair
                output_value = tCrO_row[feature_pair, feature_group]
                mO[query_row, feature, bidz] = output_value.to(mO.element_type)

Okay, we are finally at the last step. The loop is over. Two things are left: divide O by l, and write O and LSE to global memory.

Before the steps, one rule. Every strange line in this block comes from it.

The rule (diagram 4): the tensor core hands its 16 x 8 output tile to 32 lanes. Lane L holds 4 numbers: rows L / 4 and L / 4 + 8, columns (L % 4) * 2 and (L % 4) * 2 + 1. The 4 lanes with the same L / 4 share the same two rows. Call them a quad. Example: lane 5 holds rows 1 and 9, columns 2 and 3. Lanes 4, 5, 6, 7 are one quad, all on rows 1 and 9.

S is 8 of these tiles side by side (diagram 3), O is 16. So for one row, this lane holds 2 keys per tile, 16 keys per strip, and 32 features of O. The numbers of one row are split across the 4 lanes of a quad, two per tile. That is why each lane’s l is only its own share.

What we hold after 473 strips, per thread:

Five steps:

  1. Add the 4 lanes’ l together, two shuffles. row_sum is now the full sum of the row.
  2. lse = m + log(row_sum). One lane per row writes it to mLSE.
  3. query_row = bidx * 64 + warp_id * 16 + quad_id + row * 8. The global row this thread writes.
  4. O = tCrO * (pv_scale / row_sum). One multiply: the softmax division, the V descale, the /256 from block 6.
  5. Each lane writes its 32 features of the row to mO.

Each step against the rule:

Step 1, why xor 1 and xor 2. The 4 lanes of a quad sit at L % 4 = 0, 1, 2, 3. shuffle_sync_bfly(offset=1) swaps 0 with 1 and 2 with 3, then add. offset=2 swaps 0 with 2 and 1 with 3, then add. After two steps all 4 lanes hold the full sum. offset=4 would reach lane 1 from lane 5, another quad, another row, wrong. Why here and not every strip: m had to be exact every strip, because p = e^(s - m) in all 4 lanes must use the same m. l is a plain sum. The 4 shares were all scaled by the same alpha chain, so adding them once at the end gives the same number as adding them 473 times.

Step 2, LSE. The softmax denominator, stored as a log: log(sum of e^s) = m + log(l). Same number, no overflow. The 4 lanes of the quad hold the same value, so lane_id % 4 == 0 writes and the others skip.

Step 3, query_row. bidx * 64 is the CTA’s strip (diagram 1). warp_id * 16 is the warp’s 16 rows (diagram 2). quad_id = lane_id / 4 is the rule’s row part. row * 8 picks the second row of the rule. Nothing else.

Step 4, the division. pv_scale / row_sum is three factors in one number: 1 / row_sum is the softmax division we postponed in block 5, the V descale undoes V’s FP8 quantization, 1 / 256 undoes the x256 from block 6. Three multiplies on the same 32 values, or one. Same result, one is cheaper.

Step 5, feature. (lane_id % 4) * 2 is the rule’s column part. feature_group * 8 is which of the 16 tiles, they sit 8 features apart. feature_pair is 0 or 1. 16 tiles x 2 = 32 features per lane, and no shuffle: every lane already holds different columns of the row. tCrO_row has shape (2, 16), pair x group, so two loops.

A test, do it by hand before you read on: bidx = 2, warp_id = 0, lane_id = 5, row = 1, feature_group = 3, feature_pair = 0. Which (query_row, feature) does this lane write?

Lane 5 holds rows 1 and 9, columns 2 and 3. query_row = 128 + 0 + 1 + 8 = 137. feature = 2 + 24 + 0 = 26. If you got 137 and 26, the rule is yours, and the epilogue has nothing left to hide.

With the numbers from block 5, for one row:

tCrO = 0.0498 V0 + 0.135 V1 + 1 V2 + 0.368 V3     l = 1.553
O    = tCrO / l = 0.032 V0 + 0.087 V1 + 0.644 V2 + 0.237 V3

The weights add up to 1 now. This row takes 64% of key 2’s V. That is the attention output, and this was the whole kernel.

The Four Horsemen of the Optimization

Okay, so far we did a working FP8 attention kernel, and when I compare it with other production kernels, I came across FlashInfer’s FP8 attention kernel. I thought I had compared with all of them.

RTX 5080, kernel only, same run, msS=30272S=20440
Torch SDPA (FA2, BF16)245.4112.4
cuDNN BF16226.2106.5
this kernel, FP8136.063.1
FlashInfer FP8131.860.4

Okay, we are still behind. I sighed. It was not over. I took a break. Next day, I asked AI about one thing. What does a faster kernel do that mine does not? Between Blake’s suggestions and the ncu counters, I ended up with 4 steps.

1: 2-stage KV Pipeline

This one is a common and classic method. While you compute the data you have, you also load the next group of data at the same time. That’s all, really.

one stage vs two stages

The loop head, as a diff:

 cute.copy(tiled_copy_Q, tQgQ[None, None, None], tQsQ[None, None, None, 0])
+# strip 0 is requested together with Q
+cute.copy(tiled_copy_K, tKgK[None, None, None, 0], tKsK[None, None, None, 0])
+cute.copy(tiled_copy_V, tVgV[None, None, None, 0], tVsV[None, None, None, 0])
+cute.arch.cp_async_commit_group()

 for k_tile in range(k_tile_count):
-    cute.copy(tiled_copy_K, tKgK[None, None, None, k_tile], tKsK[None, None, None, 0])
-    cute.copy(tiled_copy_V, tVgV[None, None, None, k_tile], tVsV[None, None, None, 0])
-    cute.arch.cp_async_commit_group()
+    stage = k_tile % 2
+    next_stage = 1 - stage
+
+    # one wait: this strip landed. The barrier also proves that every thread
+    # finished the previous strip, so the other stage is free.
     cute.arch.cp_async_wait_group(0)
     cute.arch.sync_threads()

+    # request the next strip into the free stage, then compute this one
+    if k_tile + 1 < k_tile_count:
+        cute.copy(tiled_copy_K, tKgK[None, None, None, k_tile + 1], tKsK[None, None, None, next_stage])
+        cute.copy(tiled_copy_V, tVgV[None, None, None, k_tile + 1], tVsV[None, None, None, next_stage])
+    cute.arch.cp_async_commit_group()
+
+    tCsK_p = tCsK_copy_view[None, None, None, stage]
+    tCsV_p = tCsV_copy_view[None, None, None, stage]

     tCrC.fill(0.0)

     # ... QK^T, softmax, P.V, unchanged ...

-    # nobody may overwrite sK, sV, sP with the next strip while a warp still reads them
-    cute.arch.sync_threads()
+    # no trailing barrier: the next strip's leading barrier does this job

The result is promising, but not enough for me.

RTX 5080, kernel only, same run, msS=30272S=20440
before136.063.0
after131.962.1
change-3.0%-1.6%
FlashInfer FP8132.060.0

Shared memory per CTA 28 → 44 KB.

2: P stays in registers

In block 6, P took a trip through shared memory. The reason was the layout. My lane holds keys 2c and 2c+1 of each tile, but the A operand wants 4 keys in a row. Keys 2 and 3 sit in the neighbor lane. So in every strip write P to shared memory, barrier, read it back with ldmatrix. 473 times. No no no.

The question that rules them all, does A slot number 4 have to hold key 4? No. The tensor core computes a sum over slots. Slot i of P meets slot i of V. It does not know what a key is.

So:

  1. Each lane drops its own 32 scores into its own A slots, in a fixed order. No data crosses lanes.
  2. Now slot 4 holds key 8, not key 4. Fine, as long as V’s slot 4 also holds key 8.
  3. So V is stored in that same key order, once, in the prepare step. Not per strip. Once.

P through shared memory vs in registers

For example, lane 0, holds the first 16 keys. From tiles 0 and 1 it holds keys 0, 1, 8, 9. Its first four A slots are positions 0 to 3. Keys 0, 1, 8, 9 go there. V positions 0 to 3 get V rows 0, 1, 8, 9. Lane 1 holds keys 2, 3, 10, 11, they go to positions 4 to 7, and V follows. The full order, inside every 16-key group:

key order inside a 16-key group

Kernel side setup:

 @cute.struct
-class SharedStorageQKVP:
+class SharedStorageQKV:
     q: cute.struct.Align[cute.struct.MemRange[mQ.element_type, cute.cosize(sQ_layout)], 16]
     k: cute.struct.Align[cute.struct.MemRange[mK.element_type, cute.cosize(sK_layout)], 16]
     v: cute.struct.Align[cute.struct.MemRange[mV.element_type, cute.cosize(sV_layout)], 16]
-    probabilities: cute.struct.Align[cute.struct.MemRange[mQ.element_type, cute.cosize(sP_layout)], 16]

-tCsP = thr_mma.partition_A(sP)
-tCsP_C = thr_mma.partition_C(sP[None, None, 0])
-tCrP_C = cute.make_fragment_like(tCrC, mQ.element_type)
-tCrP = tiled_mma.make_fragment_A(tCsP[None, None, None, 0])
+# the PV A fragment, 32 FP8 per lane, in the order mma.sync wants it
+tCrP = cute.make_rmem_tensor(
+    cute.make_layout(((4, 2, 2), 1, 2), stride=((1, 4, 8), 0, 16)),
+    mQ.element_type,
+)
 num_k_block_PV = cute.size(tCrP, mode=[2])

+# two views of the same 32 slots: C order on the left, A order on the right
+pack_shape = ((2, 2), 1, (2, 2, 2))
+tCrC_as_pack = cute.make_tensor(
+    tCrC.iterator,
+    cute.make_layout(pack_shape, stride=((1, 2), 0, (4, 8, 16))),
+)
+tCrP_as_pack = cute.make_tensor(
+    tCrP.iterator,
+    cute.make_layout(pack_shape, stride=((1, 4), 0, (2, 8, 16))),
+)

-tiled_copy_s2r_P = cute.make_tiled_copy_A(atom_copy_s2r_Q, tiled_mma)
-ldmatrix_P = tiled_copy_s2r_P.get_slice(tidx)
-tCsP_copy_view = ldmatrix_P.partition_S(sP)
-tCrP_copy_view = ldmatrix_P.retile(tCrP)
-tCsP_p = tCsP_copy_view[None, None, None, 0]

Kernel side, the loop:

-tCrP_C.store((tCrC.load() * 256.0).to(mQ.element_type))
-cute.autovec_copy(tCrP_C, tCsP_C)
-cute.arch.sync_threads()
+tCrP_as_pack.store((tCrC_as_pack.load() * 256.0).to(mQ.element_type))

 for k_block in cutlass.range(num_k_block_PV, unroll_full=True):
-    cute.copy(
-        tiled_copy_s2r_P,
-        tCsP_p[None, None, k_block],
-        tCrP_copy_view[None, None, k_block],
-    )
     cute.copy(
         tiled_copy_s2r_V,
         tCsV_p[None, None, k_block],
         tCrV_copy_view[None, None, k_block],
     )
     cute.gemm(tiled_mma, tCrO, tCrP[None, None, k_block], tCrV[None, None, k_block], tCrO)

Prepare side, V is written in that key order:

+def key_order_for_positions(positions):
+    group = positions // 16 * 16
+    quad = (positions % 16) // 4
+    element = positions % 4
+    return group + 8 * (element // 2) + 2 * quad + element % 2

-destinations[operand].copy_(padded)
+destinations[operand].copy_(padded[:, :, self.key_order])
RTX 5080, kernel only, same run, msS=30272S=20440
before131.962.1
after127.560.1
change-3.3%-3.1%
FlashInfer FP8132.060.0

Against the first FP8 kernel: -6.3%. Shared memory 44 → 40 KB. Barriers per strip 2 → 1. No shared P, no P ldmatrix.

3: 32 query per-warp

After step 2, ncu said the tensor pipe was at 89.9%. The tensor cores were busy. But every K and V byte that ldmatrix pulls from shared memory feeds only 16 query rows per warp, one m16 tile. The next question becomes, can one ldmatrix feed two? Yes.

  1. The CTA takes 128 query rows instead of 64. Each warp owns 32 rows and two m16 tiles.
  2. Every K and V ldmatrix now feeds two tiles. Shared-memory bytes per MMA drop from 288 to 160.
  3. The price: the score accumulator, the P pack, the O accumulator and the softmax state all double. With a 64-key strip that is 255 registers and a 208-byte spill. Slower.
  4. So the key strip goes from 64 to 32. Score fragment 64 → 32 registers, P pack 16 → 8, V fragment 64 → 32. No spill. Bytes per MMA 192. Same tile shape FlashAttention-2 picks for head_dim 128 on consumer GPUs.

one ldmatrix feeds one tile vs two

One thing changes in the numbers. P is rounded to FP8 against the running max of a 32-key strip, not 64. Output is no longer byte-equal to step 2. Measured 0.3% relative RMS between the two, both 5.4% from cuDNN BF16.

The diff is the tile shape, and every loop that walks it:

-cta_tiler: cutlass.Constexpr = (64, 64, 128),
+cta_tiler: cutlass.Constexpr = (128, 32, 128),

-m = cute.make_rmem_tensor((2,), cutlass.Float32)
-l = cute.make_rmem_tensor((2,), cutlass.Float32)
+# two rows per m-tile, two m-tiles
+m = cute.make_rmem_tensor((4,), cutlass.Float32)
+l = cute.make_rmem_tensor((4,), cutlass.Float32)

-gV = cute.local_tile(mV[None, None, bidz], (128, 64), (0, None))
+gV = cute.local_tile(mV[None, None, bidz], (128, 32), (0, None))

-acc_shape = thr_mma.partition_shape_C((64, 64))
+acc_shape = thr_mma.partition_shape_C((128, 32))
-acc_shape_O = thr_mma.partition_shape_C((64, 128))
+acc_shape_O = thr_mma.partition_shape_C((128, 128))

 tCrP = cute.make_rmem_tensor(
-    cute.make_layout(((4, 2, 2), 1, 2), stride=((1, 4, 8), 0, 16)),
+    # MMA_M = 2, MMA_K = 1: two m-tiles, one 32-key block
+    cute.make_layout(((4, 2, 2), 2, 1), stride=((1, 4, 8), 16, 0)),
     mQ.element_type,
 )

The softmax and the epilogue get one more loop, over the two m-tiles. The body is the same, indexed by state:

-for row in cutlass.range_constexpr(2):
-    tCrC_row = tCrC[(None, row), 0, None]
+for m_tile in cutlass.range_constexpr(2):
+    for row in cutlass.range_constexpr(2):
+        state = m_tile * 2 + row
+        tCrC_row = tCrC[(None, row), m_tile, None]
         ...
-    m_new = cute.arch.fmax(m[row], tile_max)
+        m_new = cute.arch.fmax(m[state], tile_max)
         ...
-    tCrO_row = tCrO[(None, row), 0, None]
+        tCrO_row = tCrO[(None, row), m_tile, None]

-query_row = bidx * 64 + warp_id * 16 + quad_id + row * 8
+query_row = bidx * 128 + m_tile * 64 + warp_id * 16 + quad_id + row * 8
kernel only, same run, ms5080 S=302725080 S=20440PRO 6000 S=30272PRO 6000 S=20440
before127.059.947.523.1
after126.057.945.721.1
change-0.8%-3.5%-3.8%-8.8%

On the RTX PRO 6000 the shared-memory pipe was the wall, on the 5080 it was not.

4: Swizzle shift correction

This one came from a counter, not from an idea. ncu on the kernel after step 3 showed 8.0 shared-memory wavefronts per ldmatrix, 4.03 G bank conflicts. The number should be 4.0 and about 0. Every ldmatrix was paying double.

4 is the target because shared memory serves 128 bytes per cycle, 32 banks of 4 bytes. One ldmatrix.x4 reads 4 matrices of 8 rows x 16 bytes, 128 bytes each. If the 8 rows of a matrix land on 8 different 16-byte bank groups, one matrix is one wavefront, 4 per ldmatrix. The swizzle exists to make that true, it XORs the 16-byte chunk index of a row with the row index, so row 0 uses chunk c, row 1 uses chunk c ^ 1, row 2 chunk c ^ 2, and so on.

chunk 0 of each row after the XOR, shift 4 vs shift 3

The helper that builds the swizzle was written for the BF16 kernel. It takes the XOR shift from the number of elements in a 16-byte chunk, 8 for BF16, 16 for FP8. Shift 3 for BF16 is right, but shift 4 for FP8 is wrong. A 128-byte row starts at address bit 7. With shift 4 the XOR reads bits 8 to 10, that is row / 2, not row. Rows 0 and 1 get the same permutation, rows 2 and 3 too. Two rows per bank group, two wavefronts per matrix, 8 per ldmatrix.

address bits and the XOR source

The shift has to start at bit 7, the 128-byte bank period, for any element width. Swizzle<3,4,3> for the 128-byte Q and K rows, Swizzle<1,4,3> for V’s 32-byte rows.

The fix is three lines in the helper:

-def _make_smem_layout_ab(dtype, copy_bits, smem_tiler):
-    major_size = min(smem_tiler[1], 128 * 8 // dtype.width)
-    swizzle_bits = min(int(math.log2(major_size * dtype.width // copy_bits)), 3)
-    base_bits = int(math.log2(copy_bits // 8))
-    shift_bits = int(math.log2(copy_bits // dtype.width))
+def _make_smem_layout_fp8(dtype, copy_bits, smem_tiler):
+    major_size = smem_tiler[1]
+    row_bytes = major_size * dtype.width // 8
+    chunk_bytes = copy_bits // 8
+    swizzle_bits = min(int(math.log2(row_bytes // chunk_bytes)), 3)
+    base_bits = int(math.log2(chunk_bytes))
+    # the XOR source starts at the 128-byte bank period, not at "elements per chunk"
+    shift_bits = int(math.log2(128 // chunk_bytes))
     swizzle = cute.make_swizzle(swizzle_bits, base_bits, shift_bits)
     atom = cute.make_layout((8, major_size), stride=(major_size, 1))
     layout = cute.tile_to_shape(atom, smem_tiler, (0, 1, 2))
     return layout, swizzle

Nothing else in the kernel changes. Same bytes land in shared memory, in a different order, and every read of them is half the price.

A second effect I did not expect: with the XOR reading bits 8 to 10, the +8-row offset of the next n-tile was not a plain add anymore, so ptxas kept one address register per ldmatrix. With the fix those offsets fold into immediates. The step 3 kernel went from 255 registers plus a 16-byte spill to 244 and no spill.

Counters at S=30272, before and after the fix:

ldmatrixshared wavefrontswavefronts / ldmatrixbank conflicts
step 31.00 G8.05 G8.04.03 G
step 41.21 G4.82 G4.00.001 G
kernel only, same run, ms5080 S=302725080 S=20440PRO 6000 S=30272PRO 6000 S=20440
before126.758.245.721.3
after123.657.043.719.9
change-2.4%-2.1%-4.3%-6.5%
FlashInfer FP8131.860.453.725.0

This is the final kernel. Against the first FP8 kernel: 136.0 → 123.6 ms, 9% in four steps. Against cuDNN BF16: 226.2 → 123.6, 1.83x.

Final Results

Here is the very final result of the kernel and the comparison with the others. Since H3 did not fit on my 16 GB RTX 5080, I rented a RTX PRO 6000 (also SM120) with my limited funds from Verda, and ran some e2e video generations. I can’t describe the joy and excitement I felt when I saw the first results of my kernel. You had to be there.

kernel only, same run, ms5080 S=302725080 S=20440PRO 6000 S=30272PRO 6000 S=20440
FA4 SM120 BF16 (SGLang)251.8115.5
Torch SDPA (FA2, BF16)245.4112.474.434.2
our BF16232.4109.087.542.2
cuDNN BF16226.2106.573.732.3
this kernel, FP8, first136.063.156.326.0
FlashInfer FP8131.860.453.725.0
this kernel, FP8, final123.657.043.719.9
SageAttention*84.939.1

* INT8 QK + FP8 PV, its own quantization included in the time, separate run, 5080 only. FA4: no SGLang tree on the PRO 6000 pod.

FP8 prep (BF16 → FP8 + scales), not in the table above, ms5080 S=302725080 S=20440PRO 6000 S=30272PRO 6000 S=20440
ours, fused Triton4.032.762.401.63
FlashInfer3.732.522.171.47
PRO 6000, 50-step video, same prompt and seedsteady s/itdenoise se2e svs cuDNNvs firstclip vs cuDNN, PSNR / SSIM
cuDNN BF16 (stock)6.87328.0360.2
this kernel, FP8, first6.25301.7330.0-9.0%19.47 / 0.697
this kernel, FP8, final5.61271.4298.4-18.3%-10.2%19.75 / 0.707
cuDNN BF16
this kernel, FP8, first
this kernel, FP8, final

This work is a live PR at SGLang, maybe merged by the time you read this.

Closing

Congrats if you made it this far! I am still trying to understand parts of it myself, so if you did not, don’t worry, the Roman Empire wasn’t built in a day, relax. For the last code block and some optimizations, I also asked AI about how do I explain a thing that I don’t understand yet. This work is the result of a looong month of suffering, sweat and effort. I still may not explain some things but hey, isn’t this part of learning?

I also thank Ezgi, who gave me this whole idea and her support when I was close to bottom, and who lit the spark in me to climb out, at least a little. And thanks to Blake Ledden, who generously walked me through his kernel design and shared his insights, some of which you have just read above.

If you want to ask, share, criticise something or support, feel free to contact me. I wish you all a happy day, and may your GPUs work cool.