CuTe DSL fundamentals and primitives [FA3]

Hello all. This post is the companion to the FlashAttention-3 worklog. The worklog walks through eight kernels, K1 to K8, and for each one it gives the idea, what changed, the numbers and the diagrams. It does not re-explain the CuTe DSL and Hopper primitives each kernel is built from. Those live here, one section per primitive, and the worklog links into them by section.

The rule for this post is that every statement about what the hardware or the library does comes from a primary source: the CUDA PTX ISA, the CUTLASS and CuTe DSL source and documentation, NVIDIA’s Hopper material, or the Colfax CUTLASS tutorials. Each section ends with the references it drew on. Where a claim could not be sourced it is marked as unverified rather than asserted. The kernels compile and run against nvidia-cutlass-dsl 4.5.1 on an H100 SXM (confirmed 2026-10-02). Everything quoted from the DSL source was read from the CUTLASS v4.7.0 tree and cross-checked against the v4.5.1 tag; where the two differ it is said in place. Kernel snippets come from the FlashAttention3 repo and are stamped with the date they were read. The repo was restructured on 2026-10-02: the WGMMA helpers moved to kernels/common.py, K4 became k4_1.py and k4_2.py, and K5 to K8 were ported the same day. All eleven kernels compiled and passed the FP32 reference check on the H100 that evening, on every configuration of the worklog’s sweep, so every snippet below is from code that runs; each carries the date it was read.

The plan:

• What CuTe is, and what the Python DSL adds
• The JIT model: cute.compile, compile-time values, const_expr, range_constexpr, static and dynamic layouts
• Layouts: shape, stride, hierarchy, and the handful of operations the kernels use
• Tensors: iterator plus layout, aliasing, identity tensors, register tensors
• Shared memory: structs, the allocator, banks, and swizzled layouts
• Warp reductions
• TMA: the instruction, the atom, partitioning, issuing a copy
• mbarriers: what the object tracks, phases, transaction counts
• Pipelines: PipelineTmaAsync, producer and consumer protocol, pipeline state
• WGMMA: the instruction, operand sources, descriptors, fence, commit and wait, and how cute.gemm dispatches it
• TV layouts: who owns which accumulator element, and the identity-tensor scatter
• The QK to PV register handoff
• Warp specialization and register reallocation
• Ordering two consumer warpgroups with PipelineOrder
• FP8 operand constraints
• Block scales, then wrapping up

Let’s get started.

What CuTe is, and what the DSL adds

CUTLASS is NVIDIA’s open source library of CUDA matrix multiply and related kernels. CuTe is the layout and tensor library inside it: a small algebra for describing how data is arranged in memory and how it is split across threads, with the same vocabulary for global memory, shared memory and registers. In C++ it is a header library of templates, and the CUTLASS 3.x and 4.x GEMM kernels are written on top of it.

The CuTe DSL is the Python front end to the same ideas. It ships as the nvidia-cutlass-dsl package, and its Python source lives in the CUTLASS repository under python/CuTeDSL. Its documentation describes it as a Python-native DSL with JIT compilation for both host and device code, with the GPU kernel compiled through the DSL’s own MLIR-based pipeline down to a CUDA module that is loaded and launched from the Python process. The point of the DSL for this project is that it exposes the Hopper primitives directly: TMA copies, mbarriers, pipelines, warpgroup MMA, and register reallocation all have Python spellings that lower to the matching PTX, and the layout algebra that makes them composable is the same as in C++.

A note on versions, because the API moves. The kernel repo pins nvidia-cutlass-dsl==4.5.1, and that is the version the kernels were compiled and checked with on the H100 (with PyTorch 2.6.0 and CUDA 12.4 wheels, on 2026-10-02). The DSL source quoted in this post was read from the CUTLASS v4.7.0 tree and compared against the v4.5.1 tag: every function named here exists under the same name in both, and the two deprecations that matter are already present in 4.5.1 (warpgroup.OperandMajorMode points at cute.nvgpu.OperandMajorMode, and warpgroup_reg_alloc points at setmaxregister_increase).

References

• CUTLASS repository, nvidia-cutlass-dsl on PyPI
• CuTe DSL code generation overview
• CuTe layouts (C++ docs)

The JIT model and specialization

Two decorators mark DSL code. @cute.jit marks a host function that the DSL compiles and that can call into other DSL functions; @cute.kernel marks a GPU kernel, launched from a jit function with .launch(grid=..., block=..., stream=...). Nothing is compiled when the decorators are applied. Compilation happens either implicitly on the first call, or explicitly with cute.compile(fn, *args), which the documentation calls Zero Compile: it compiles immediately, bypasses the DSL’s own cache, and returns a JIT executor, a callable that holds the compiled host function, the loaded CUDA modules and an argument specification.

The kernels in the worklog use the explicit form, once per sweep configuration (as of 2026-10-02, from kernels/k1.py):

score_compiled = cute.compile(
    score, tensors["q_cute"], tensors["k_cute"], s_cute, scale, stream,
)
...
score_compiled(tensors["q_cute"], tensors["k_cute"], s_cute, scale, stream)

The arguments passed to cute.compile do two different jobs depending on their kind.

• Runtime values become arguments of the compiled function. A cutlass.Float32 scalar, a stream, or a tensor is passed again at every call.
• Compile-time values are folded into the generated code. The documentation calls these Constexpr: a value “passed and computed by python interpreter”. A parameter annotated cutlass.Constexpr, and any attribute of the kernel object that is a plain Python value and is read while the function is traced, fall in this class. The JIT caching document states that arguments with the Constexpr hint are excluded from the argument specification, “since they are evaluated at compile time rather than runtime.” An unannotated Python number passed as an argument is a runtime value: the documentation’s own example compiles add(a, b, print_result: cutlass.Constexpr) with cute.compile(add, 1, 2, True) and then calls the executor with (1, 2).

In the kernels, tile sizes, head dimension, stage counts and the causal flag are attributes like self.d and self.causal, so they are Python ints and bools and get folded. That is the whole specialization story: a kernel compiled for \(d = 64\) is a different binary from one compiled for \(d = 128\), and neither contains a branch on \(d\).

Control flow: what runs in Python and what becomes a loop

The DSL walks the Python AST and decides, statement by statement, whether a loop or branch is evaluated in the Python interpreter at compile time or emitted into the generated code. The control flow document fixes the rules:

const_expr is a check as much as a marker. Its source does one thing: if the expression is a Python value it returns it, and if it is a dynamic (traced) value it raises. So if cutlass.const_expr(self.causal): compiles to one of the two branches and nothing else, and passing a runtime value to it is an error rather than a silent slow path. range_constexpr runs in the interpreter, so a loop like

for k in cutlass.range_constexpr(self.d):
    dot = dot + gQ[i, k].to(cutlass.Float32) * gK[j, k].to(cutlass.Float32)

is emitted as \(d\) straight-line multiply-adds. The KV main loop in every kernel from K2 on uses cutlass.range(n_kv_eff, unroll=1) instead, because the trip count depends on the sequence length and because a fully unrolled loop over hundreds of iterations is not wanted. The two kinds mix freely: the trip count itself is chosen by a compile-time branch and computed at runtime (as of 2026-10-02, from kernels/k3.py):

n_kv = self.N // self.Bc
if cutlass.const_expr(self.causal):
    n_kv_eff = causal_n_kv(q_row_start, self.Br, self.Bc, n_kv)
else:
    n_kv_eff = n_kv

The non-causal binary contains no trace of the first branch, and in the causal binary n_kv_eff is a runtime integer that the loop reads.

Unrolling is a choice, not a free win. A range_constexpr loop becomes straight-line code whose live values all compete for the same per-thread register budget, and the Hopper tuning guide puts the maximum at 255 registers per thread; past that the compiler spills to local memory. K2 ran into exactly this: its two 16-column softmax row scans were first written with range_constexpr, spilled at the cap on the H100, and ran markedly slower than the same scans as runtime loops, so the kernel now uses runtime loops for them on purpose (as of 2026-10-02, from kernels/k2.py; the worklog’s K2 section has the measurements):

for c in cutlass.range(col_group, self.Bc, self.threads_per_row):
    local_max = cute.arch.fmax(local_max, sS[row, c])

The \(d\)-long dot product above stays unrolled because its body holds one accumulator; the row scan did not because its unrolled body held many. Which loops to unroll is a register-budget question, and the kept SASS is where the answer shows up.

The documented limitations of dynamic control flow matter when writing kernel bodies: no break, continue or early return inside an emitted loop or branch, values defined inside a dynamic branch are not visible after it, and a variable cannot change type inside one.

Static and dynamic layouts

The earlier version of this post said layouts are always compile-time objects. That is wrong for the DSL, and the distinction matters for how the kernels are compiled. A tensor converted with from_dlpack(torch_tensor) gets a static layout: its shape and strides are compile-time constants, and the compiled function only accepts tensors of exactly that layout. Calling mark_layout_dynamic(leading_dim=...) on it makes the shape and the non-leading strides dynamic, printed as (?,?):(?,1), so one compiled function serves every shape with the same contiguous mode. The documentation’s rule of thumb is that static layouts let the compiler fold more and unroll loops, and dynamic layouts avoid one compilation per shape.

The worklog’s harness uses the dynamic form for Q, K, V and O (as of 2026-10-02, from utils.py):

def cute_view_ndb(torch_bnd, leading_dim=1, element_type=None):
    """(B, N, d) torch tensor -> CuTe tensor viewing it as (N, d, B)."""
    view = torch_bnd.permute(1, 2, 0)
    t = from_dlpack(view, assumed_align=16)
    if element_type is not None:
        t.element_type = element_type
    t = t.mark_layout_dynamic(leading_dim=leading_dim)
    return t, view

The permute is a view, so the CuTe tensor is (N, d, B) over the same bytes with \(d\) contiguous, and leading_dim=1 names that mode as the one whose stride stays static at 1. So inside a kernel the global tensors have dynamic shapes, while the tile sizes that index into them are compile-time attributes. Both coexist in one kernel: self.N // self.Bc is a Python int, and cute.size(mQ, mode=[0]) would be a runtime value.

Compile options

cute.compile accepts an options string parsed with argparse, or typed option objects. The ones used while debugging the kernels are --keep-ptx, --keep-sass and --keep-cubin, which keep the generated PTX, the SASS disassembly and the cubin as files, and --generate-line-info. Looking at the kept PTX is the fastest way to confirm that a range_constexpr loop really unrolled or that a const_expr branch really disappeared.

References

• CuTe DSL introduction, JIT caching and the JIT executor
• Control flow, const_expr and range_constexpr source
• Static vs dynamic layouts
• JIT compilation options

Layouts

A CuTe layout is a shape and a stride, and together they define a function from a coordinate to an integer offset. The C++ documentation defines it exactly this way and everything in the DSL follows from it. For a flat layout with shape \((s_0, \ldots, s_{n-1})\) and stride \((d_0, \ldots, d_{n-1})\),

In the DSL, cute.make_layout(shape, stride=...) builds one; the stride is keyword-only, and when omitted the layout is compact and column-major, meaning mode 0 has stride 1 (the docstring calls this “left-most”). So make_layout((4, 3), stride=(1, 4)) is a 4 by 3 matrix with rows contiguous, and make_layout((4, 3), stride=(3, 1)) is the same shape with columns contiguous. A layout is only this function. It holds no pointer and touches no memory.

Hierarchical shapes

A mode of a layout can itself be a tuple, and the coordinate for that mode is then a tuple too. The offset is still the sum of coordinate times stride over every leaf. For example ((2, 2), 3):((1, 4), 8) maps the coordinate ((i0, i1), j) to \(i_0 + 4 i_1 + 8 j\). A flat index into a hierarchical mode is unpacked using the shape as place values, innermost first, and the stride is applied to the unpacked digits. This is the mechanism by which one logical axis can walk two different strides, and it is how every hardware layout in this post is written: the GMMA shared memory atoms, the thread and value layouts of the accumulator, and the staged pipeline buffers.

The operations the kernels use

Slicing with None is used constantly. A per-batch 2D view of a (N, d, B) tensor is mQ[(None, None, b)], and a per-stage view of a staged shared memory tensor is sK_full[(None, None, stage)]. The docstring states that None “represents keeping all elements in that mode”, and that slicing preserves the layout structure of the source, so it is a view, not a copy.

flat_divide is how each kernel gets its global tiles. The C++ layout algebra document lists the mode order of every divide flavor; for flat_divide the tile modes come first and the rest modes after. Applied to a (N, d, B) tensor with tiler (Br, d):

gQ = cute.flat_divide(tma_tensor_q, (self.Br, self.d))   # (Br, d, N/Br, 1, B)
gQ_tiles = gQ[(None, None, None, 0, bidx_b)]              # (Br, d, N/Br)

so gQ_tiles[(None, None, m)] is the \(m\)-th row tile of batch bidx_b. The 1 in the fourth mode is \(d / d\), the number of tiles along the head dimension.

References

• CuTe layouts (C++ docs), CuTe layout algebra: divides and their mode orders
• make_layout, slice_, size, select, group_modes, composition, zipped_divide, flat_divide, local_tile docstrings in cute/core.py

Tensors

The DSL defines a tensor as an engine composed with a layout, \(T = E \circ L\), where the engine is a pointer, an array or a counting iterator and the layout maps a logical coordinate to an offset into it. Evaluating tensor[c] computes layout(c) and dereferences the engine at that offset. cute.make_tensor(iterator, layout) builds one; the iterator can be a pointer in any address space, or an integer tuple, in which case the result is a coordinate tensor.

Aliasing

Because the engine and the layout are separate, two tensors can share one engine with different layouts. cute.make_tensor(t.iterator, other_layout) re-views the same bytes or the same registers. No data moves. The kernels use this in three places: to re-view V with its modes reordered before building the PV MMA (cute.make_tensor(mV.iterator, cute.select(mV.layout, [1, 0, 2]))), to view a WGMMA accumulator fragment as a rows-by-columns matrix for the softmax, and inside the register handoff from the QK accumulator to the PV A operand. All three are covered in their own sections.

Identity tensors

cute.make_identity_tensor(shape) builds a tensor that “maps each coordinate to itself”. Its values are coordinates, so after it is partitioned exactly like a data tensor, reading an element tells you which (m, n) position of the original tile that slot corresponds to. The docstring’s example: make_identity_tensor((3, 2)) holds (0,0), (1,0), (2,0), (0,1), (1,1), (2,1) in column-major order. The TV layouts section uses this for the causal mask and the output scatter.

Register tensors

cute.make_rmem_tensor(layout_or_shape, dtype) allocates a tensor in registers; cute.make_rmem_tensor_like(src, dtype) allocates one with the same shape as src and compact column-major strides. The DSL changelog records make_rmem_tensor as the replacement for the older make_fragment name; make_fragment_like still exists as a separate function and is not used in the kernels. The per-row softmax state in K4 is three such tensors of length equal to the number of accumulator rows a thread owns:

s_max      = cute.make_rmem_tensor_like(s_max_layout, cutlass.Float32)
a_sum      = cute.make_rmem_tensor_like(s_max, cutlass.Float32)
s_max_prev = cute.make_rmem_tensor_like(s_max, cutlass.Float32)

A tensor also has .load() and .store(value), which move the whole tensor’s contents as one vector value and convert types on store when the value’s dtype differs from the tensor’s. The register handoff section relies on this pair.

Framework tensors

from_dlpack(torch_tensor, assumed_align=16) wraps a PyTorch tensor without copying, and the 16-byte alignment assumption is what lets the DSL emit vector loads. The resulting tensor’s layout is static unless mark_layout_dynamic is called, as described in the JIT section.

References

• CuTe tensors (C++ docs)
• make_tensor, make_identity_tensor, make_rmem_tensor, make_rmem_tensor_like in cute/tensor.py, CuTe DSL API changelog
• Static vs dynamic layouts and from_dlpack

Shared memory: structs, allocation, banks

Shared memory in a DSL kernel is declared as a cute.struct and materialized once with a SmemAllocator. The struct decorator supports scalar fields, cute.struct.MemRange[dtype, n] for a contiguous region of n elements, nested structs, and cute.struct.Align[field, bytes] to force alignment. Its size and alignment are static (__sizeof__, __alignof__). The allocator’s docstring notes that it computes the kernel’s shared memory usage automatically at launch, so no size has to be passed to .launch, and that it only supports static layouts.

This is the K3 storage, which every later kernel keeps the shape of (as of 2026-10-02, from kernels/k3.py):

@cute.struct
class SharedStorage:
    q_pipeline_array_ptr:  cute.struct.MemRange[cutlass.Int64, self.q_stage  * 2]
    kv_pipeline_array_ptr: cute.struct.MemRange[cutlass.Int64, self.kv_stage * 2]
    sQ: cute.struct.Align[
        cute.struct.MemRange[INPUT_DTYPE, cute.cosize(sQ_layout_staged)],
        self.buffer_align_bytes,
    ]
    sK: cute.struct.Align[
        cute.struct.MemRange[INPUT_DTYPE, cute.cosize(sK_layout_staged)],
        self.buffer_align_bytes,
    ]
    sV: cute.struct.Align[
        cute.struct.MemRange[INPUT_DTYPE, cute.cosize(sV_layout_staged)],
        self.buffer_align_bytes,
    ]

Three things to read off it. The barrier arrays are 64-bit words, two per stage, because a pipeline keeps one full barrier and one empty barrier per stage, and the PTX ISA fixes an mbarrier as a .b64 in shared memory with 8-byte alignment. The tile buffers are sized with cute.cosize of the staged layout, not with a hand computed product, so a swizzled or padded layout gets the right number of elements. And buffer_align_bytes is 1024, for a reason that belongs to the swizzle subsection below.

In the kernel body:

smem = utils.SmemAllocator()
storage = smem.allocate(self.shared_storage)
sQ_full = storage.sQ.get_tensor(sQ_layout_staged)

get_tensor(layout, swizzle=None) turns a MemRange into a tensor over that memory with the given layout; the swizzled form get_tensor(layout.outer, swizzle=layout.inner) is how K4 attaches the hardware swizzle and is explained next. storage.q_pipeline_array_ptr.data_ptr() is the pointer handed to the pipeline as its barrier storage.

Banks

The CUDA Best Practices Guide describes shared memory as divided into equally sized banks that can be accessed simultaneously; on compute capability 5.x and later each bank delivers 32 bits per clock and successive 32-bit words are assigned to successive banks, and when several addresses of one request fall in the same bank the accesses are serialized, with the one exception that identical addresses are broadcast. The programming guide’s compute capability tables give the bank count as 32, so one bank cycle covers 128 contiguous bytes. A plain row-major tile whose row pitch is a multiple of 128 bytes therefore maps the same column of successive rows onto the same bank; the swizzle atoms below are eight rows tall, so a same-column access to the eight rows of one atom would land on one bank. Swizzling exists to break that geometry.

Swizzle

CuTe’s swizzle is a function on offsets, defined in include/cute/swizzle.hpp as Swizzle<BBits, MBase, SShift>: take the BBits bits that sit SShift positions above the low MBase bits, and XOR them into the BBits bits just above MBase. The header draws it as a bit pattern, with MBase the number of least significant bits kept constant, BBits the width of the mask, and SShift the distance between the source bits and the target bits. The function is an involution: applying it twice gives the identity.

The Hopper tensor core and TMA understand exactly four arrangements, and CuTe spells them as Swizzle<B, 4, 3> with \(B \in \{0, 1, 2, 3\}\). The DSL’s make_smem_layout_atom lists them, in units of elements:

The MN-major kinds are the transposes. For BF16, K_SW128 is Swizzle<3,4,3> composed with the layout (8, 64):(64, 1): eight rows of 64 elements, 128 bytes per row, 1024 bytes per atom. The same four atoms appear in the C++ header mma_traits_sm90_gmma.hpp in units of bits, and the PTX ISA’s Table 41 lists the same shapes in 128-bit elements (an 8 by 8 atom of 16-byte chunks for the 128-byte mode in K-major, 8 by 4 for 64-byte, 8 by 2 for 32-byte, 8 by 1 for none). The PTX ISA’s canonical layout table then writes the whole operand tile, for K-major 128-byte swizzle, as Swizzle<3,4,3> o ((8,m),(T,2k)):((8T,SBO),(1,T)) with \(T\) the number of elements per 16 bytes.

What the 4, 3 mean in bytes: the low 4 bits of a byte offset are untouched, so swizzling never splits a 16-byte chunk; the three bits above that, bits 4 to 6, select which 16-byte chunk inside a 128-byte row; and the three bits at 7 to 9, which are the row index inside the 8-row atom, are XORed into them. So row \(r\) of the atom has its 16-byte chunks permuted by XOR with \(r\), and the eight rows of one 16-byte column land in eight different 16-byte positions of their 128-byte rows, which is eight different bank groups rather than one.

Two consequences of the definition shape the kernel code. First, pointer_swizzle.hpp states that the swizzle on a shared memory pointer is applied to the address, which makes the layout “position-dependent”, and notes that the hardware was designed with position-dependent swizzles. That is why the operand buffers are aligned to 1024 bytes: the PTX ISA says the descriptor’s base offset field is zero only when the repeating pattern of the 128-byte mode starts on a 1024-byte boundary, and CUTLASS’s descriptor builder always writes zero and relies on alignment. Second, in the DSL a swizzled layout is a ComposedLayout whose .inner is the swizzle and .outer is the plain layout, and a MemRange is turned into a swizzled tensor by passing them separately:

sQ_full = storage.sQ.get_tensor(sQ_layout_staged.outer, swizzle=sQ_layout_staged.inner)

Building the operand layouts

cutlass.utils.hopper_helpers has make_smem_layout_a and make_smem_layout_b. Each takes a LayoutEnum (row- or column-major, derived from a tensor with LayoutEnum.from_tensor), the MMA tile (M, N, K), the element type and the stage count. The source does four steps: slice the operand’s shape out of the tile (A is (M, K), B is (N, K)), pick the atom kind with get_smem_layout_atom by testing whether the major mode’s size in bits is a multiple of 1024, then 512, then 256 (falling back to the interleaved atom), make the atom, and tile_to_shape it to the staged shape (M or N, K, stages) with an order argument that says which mode is tiled innermost: (0, 1, 2) for K-major, which the WGMMA guide glosses as “K innermost”, and (1, 0, 2) for MN-major. For BF16 with \(d = 64\) or \(128\) the K-major size in bits is a multiple of 1024, so Q and K get K_SW128 atoms. The result is a ComposedLayout with the staging mode outermost, and the per-stage layout handed to the TMA atom is cute.slice_(layout, (None, None, 0)).

References

• cute.struct docstring in cute/core.py, SmemAllocator
• CUDA Best Practices Guide, shared memory and bank conflicts, CUDA C++ Programming Guide, compute capabilities, Hopper tuning guide (228 KB of shared memory per SM)
• include/cute/swizzle.hpp, include/cute/pointer_swizzle.hpp, include/cute/atom/mma_traits_sm90_gmma.hpp
• make_smem_layout_atom in cute/nvgpu/warpgroup/helpers.py, make_smem_layout_a/b and get_smem_layout_atom in utils/hopper_helpers.py, LayoutEnum
• PTX ISA, Asynchronous Warpgroup Level Matrix Instructions: section 9.7.16.5.1.2 (shared memory matrix layout, Table 41, canonical layouts) and 9.7.16.5.1.2.2 (matrix descriptor format, base offset)

Warp reductions

cute.arch.warp_reduction(val, op, threads_in_group=32) reduces a register value across a group of lanes, and warp_reduction_max and warp_reduction_sum are the two partial applications of it the kernels call. The docstring says it works “by shuffle_sync_bfly” and that threads_in_group is the number of threads per reduction group within a warp: 32 reduces the whole warp as one group, 8 splits the warp into four groups of eight. The source is a loop: starting with offset = threads_in_group // 2, each step combines the lane’s value with the value shuffled from lane lane XOR offset and halves the offset, so a group of \(g\) lanes takes \(\log_2 g\) shuffle steps and every lane in the group ends with the full result. The PTX instruction underneath is shfl.sync.bfly, whose semantics compute the source lane as the executing lane XOR the offset.

Two constraints follow from the mechanism. The group must be a power of two that divides 32, and the lanes of one group must be the ones that differ only in the low \(\log_2 g\) bits of the lane id. K2 and K3 arrange threads as row = tidx // 4, col_group = tidx % 4, so the four threads of one row are lanes 4r to 4r + 3 inside a warp, and warp_reduction_max(local_max, threads_in_group=4) reduces each row’s partial maxima in two shuffle steps with no shared memory and no barrier (as of 2026-10-02, from kernels/k3.py):

row_max = cute.arch.warp_reduction_max(local_max, threads_in_group=self.threads_per_row)

K4 cannot choose its grouping. The lanes that share an accumulator row are fixed by the WGMMA accumulator layout, and the TV layouts section shows that they are again four consecutive lanes, which is why the same primitive still applies there.

References

• warp_reduction, warp_reduction_max, warp_reduction_sum in cute/arch/nvvm_wrappers.py
• PTX ISA, shfl.sync, section 9.7.9.6

TMA

The Tensor Memory Accelerator is Hopper’s bulk copy engine. The PTX instruction is cp.async.bulk.tensor, and the ISA describes it as a non-blocking instruction that initiates an asynchronous copy of tensor data between global memory and shared memory. Its operands are the destination shared memory address, a tensor map and a coordinate vector, and an mbarrier. The tensor map is an opaque object created on the host that holds the global tensor’s base pointer, element type, dimensions, strides, the box size to copy and the swizzle mode; the CuTe C++ documentation lists these contents and notes that the hardware walks the tile from them, so no per-element address computation happens in the kernel. The instruction’s completion mechanism for loads is .mbarrier::complete_tx::bytes: when the copy finishes, a complete-tx operation with a count equal to the number of bytes copied is performed on the mbarrier.

TMA runs in the async proxy. The ISA states that cp.async.bulk operations are performed in the asynchronous proxy, that accessing the same memory from different proxies needs a cross-proxy fence (fence.proxy.async), and that the completion of the operation is followed by an implicit fence, so once the mbarrier reports completion the data is visible to ordinary loads. Nothing further is needed on the consumer side after the barrier wait.

The atom

The DSL wraps one tensor map plus the copy instruction as a copy atom. cute.nvgpu.cpasync.make_tiled_tma_atom(op, gmem_tensor, smem_layout, cta_tiler, num_multicast=1) takes the operation (here CopyBulkTensorTileG2SOp(), global to shared in tile mode), the global tensor, the per-stage shared memory layout, and the tile shape, and returns two things: the atom and a TMA tensor. The docstring is specific about the second one: it “maps logical coordinates of the GMEM tensor to coordinates consumed by the TMA unit”, its strides are basis elements rather than integers, and it can be partitioned like any tensor. Every later partitioning for TMA must use this tensor, not the original global tensor. From K3 (as of 2026-10-02):

tma_atom_q, tma_tensor_q = cute.nvgpu.cpasync.make_tiled_tma_atom(
    cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(),
    mQ, sQ_layout, (self.Br, self.d), num_multicast=1,
)

The atom is built from the same shared memory layout, swizzle included, that the WGMMA descriptor is later built from, and the CuTe documentation lists the shared memory swizzling pattern among the contents of the tensor map. That is how TMA writes and tensor core reads agree on the arrangement of bytes: one layout object feeds both.

cute.nvgpu.cpasync.prefetch_descriptor(tma_atom) prefetches the tensor map; the kernels do it from warp 0 once at kernel start.

Partitioning

cute.nvgpu.cpasync.tma_partition(atom, cta_coord, cta_layout, smem_tensor, gmem_tensor) tiles both sides for the atom and returns the shared memory side first and the global side second. The kernels pass 0 and cute.make_layout(1) for the CTA coordinate and layout, meaning a single CTA with no multicast, and group the two tile modes of each tensor into one with cute.group_modes(t, 0, 2), so the copy sees one tile as one unit and the remaining modes enumerate stages on the shared side and tiles on the global side:

tQsQ, tQgQ = cute.nvgpu.cpasync.tma_partition(
    tma_atom_q, 0, cute.make_layout(1),
    cute.group_modes(sQ_full, 0, 2), cute.group_modes(gQ_tiles, 0, 2),
)

The naming convention, inherited from CUTLASS, is t<Op><mem><Operand>: tQsQ is the Q copy’s shared memory side and tQgQ its global side.

Issuing a copy

cute.copy(atom, src, dst, tma_bar_ptr=...) issues one TMA load. The copy docstring describes the general algorithm as operating on tensors with a (V, Rest...) profile and mentions the TMA keyword arguments tma_bar_ptr, mcast_mask and cache_policy. The CopyBulkTensorTileG2SOp docstring adds the rule that matters most: TMA copies must be issued by a single thread, and the DSL adds the election itself, so the copy must not be wrapped in elect_one (the docstring says doing so can deadlock). Barrier initialization and transaction setup are different: those must be done under elect_one, and the pipeline object does that internally.

cute.copy(
    tma_atom_k, tKgK[(None, kv_producer_state.count)],
    tKsK[(None, kv_producer_state.index)],
    tma_bar_ptr=kv_pipeline.producer_get_barrier(kv_producer_state),
)

The global side is sliced at the tile index and the shared side at the stage index (the pipeline state’s count and index, explained in the pipelines section); the mbarrier is the pipeline’s full barrier for that stage. Which thread runs this in the kernels is governed by an if warp_idx == 0: around the producer code, and within that warp the DSL elects one lane.

What TMA buys, in sourced terms: the CuTe documentation says a single instruction copies an entire tile and the hardware no longer computes per-element addresses, the Colfax tutorial says one thread issues the copy and the address and stride work is done “in a single-threaded manner” which is more register efficient, and NVIDIA’s Hopper overview says TMA operations use a descriptor with tensor dimensions and block coordinates instead of per-element addressing. The earlier draft’s claim that TMA reaches memory bandwidth more efficiently than per-thread loads is not stated in any of these sources and is dropped.

References

• PTX ISA, cp.async.bulk.tensor, section 9.7.9.26.5.2, and section 9.7.9.26.2 (async proxy)
• CuTe TMA tensors (C++ docs)
• make_tiled_tma_atom, tma_partition, prefetch_descriptor in cute/nvgpu/cpasync/helpers.py, CopyBulkTensorTileG2SOp in cute/nvgpu/cpasync/copy.py, cute.copy in cute/algorithm.py
• Colfax: Mastering the NVIDIA Tensor Memory Accelerator, NVIDIA Hopper architecture in-depth

mbarriers

An mbarrier is a 64-bit object in shared memory. The PTX ISA lists what it tracks: the current phase, the count of pending arrivals for the current phase, the count of expected arrivals for the next phase, and, from Hopper on, a tx-count, the number of outstanding asynchronous transactions tracked by the current phase. A phase completes when the pending arrival count has reached zero and the tx-count has reached zero; on completion the object atomically moves to the next phase and resets the pending arrival count to the expected count, so it is immediately reusable.

Three operations change the counts.

• Arrive-on decrements the pending arrival count by one (or by a given count). mbarrier.arrive performs it, and with the .expect_tx qualifier it first performs an expect-tx.
• Expect-tx increases the tx-count by an expected number of transaction units. For TMA the unit is bytes.
• Complete-tx decreases the tx-count. It is performed by the asynchronous operation itself when it finishes; for cp.async.bulk.tensor the count is the number of bytes copied.

Waiting is by phase parity. mbarrier.try_wait.parity takes a 0 or 1 and tests whether the phase with that parity has completed; the ISA notes that the test is valid for the current incomplete phase, where it returns false, and the immediately preceding phase, where it returns true, which is why the caller must track the phase bit across the barrier’s lifetime. try_wait may suspend the thread until the phase completes or a time limit passes, so it is wrapped in a retry loop; the DSL’s mbarrier_wait lowers to that loop with a 10 ms hint.

The DSL spellings are cute.arch.mbarrier_init(ptr, arrival_count), mbarrier_arrive_and_expect_tx(ptr, bytes), mbarrier_arrive(ptr) and mbarrier_wait(ptr, phase). The docstrings require the init and the expect-tx to be executed by one thread, under cute.arch.elect_one(), which lowers to PTX elect.sync: a deterministic election of one leader lane from a mask.

One barrier for two copies

Because the tx-count is a total, one barrier can track several TMA loads. K3 initializes the KV pipeline with tx_count = sK_bytes + sV_bytes and issues the K copy and the V copy against the same barrier of the same stage. The expect-tx of the full byte count happens once when the producer acquires the stage, each copy’s completion performs its own complete-tx, and the phase completes only when both have landed and the producer’s single arrival has been counted. Setting the count to the K bytes alone would complete the phase after K arrived, while V was still in flight; that is a correctness error, not a performance one.

References

• PTX ISA, mbarrier: sections 9.7.14.16.3 (contents), 9.7.14.16.5 (phases), 9.7.14.16.6 (expect-tx and complete-tx), 9.7.14.16.8 (phase completion), 9.7.14.16.9 (arrive-on), 9.7.14.16.16 (mbarrier.arrive), 9.7.14.16.19 (test_wait and try_wait); elect.sync, section 9.7.14.15
• cute/arch/mbar.py, elect_one in cute/arch/elect.py
• NVIDIA Hopper architecture in-depth, asynchronous transaction barrier

Pipelines

cutlass.pipeline.PipelineTmaAsync is the DSL’s producer and consumer protocol for a TMA producer and ordinary thread consumers, which its docstring describes as the Hopper mainloop case. It owns two arrays of mbarriers, one full barrier and one empty barrier per stage, and the base class docstring gives the state machine in a table: producer_acquire blocks until the stage’s empty barrier says the consumer released it, producer_commit marks the stage full, consumer_wait blocks until the full barrier says the data landed, and consumer_release marks the stage empty again. The stages form a circular buffer.

The constructor, as used in K3 and K4.2 (as of 2026-10-02, from kernels/k3.py):

kv_pipeline = pipeline.PipelineTmaAsync.create(
    barrier_storage=storage.kv_pipeline_array_ptr.data_ptr(),
    num_stages=self.kv_stage,
    producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
    consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, num_warps),
    tx_count=self.kv_bytes,
    cta_layout_vmnk=cute.make_layout((1, 1, 1, 1)),
    defer_sync=True,
)

Reading the source of create fixes what each field means.

• barrier_storage is the pointer to 2 * num_stages 64-bit words; the full barriers occupy the first num_stages and the empty barriers the next.
• producer_group sets the arrival count of the full barriers. One thread, because one elected thread performs the arrive-and-expect-tx per stage.
• consumer_group sets the arrival count of the empty barriers. The pipeline’s consumer_release arrives only from a signaling thread, and init_empty_barrier_arrive_signal makes that lane 0 of each warp for a single-CTA cluster, so the count must be the number of consumer warps. The kernels pass num_warps with the Thread agent, which produces exactly that count.
• tx_count is the number of bytes one stage expects, used by every arrive on a full barrier.
• cta_layout_vmnk is the cluster layout, (1, 1, 1, 1) for no cluster.
• defer_sync=True skips the built-in fence and block-wide sync after barrier initialization, so the kernel calls pipeline_init_arrive and pipeline_init_wait itself after creating all pipelines.

The four methods, from the PipelineTmaAsync source:

Pipeline state

PipelineState holds a stage index, a phase bit, and a count. make_pipeline_state(PipelineUserType.Producer, stages) returns index 0, count 0, phase 1; the consumer variant returns phase 0. The docstring explains the producer’s flipped bit: producers start with an empty buffer, and waiting for parity 1 on a fresh barrier, whose first incomplete phase is phase 0, returns immediately because the “immediately preceding phase” test succeeds. advance() increments both index and count, and when the index reaches the stage count it wraps to 0 and flips the phase.

So index is the shared memory slot and count is how many stages this side has consumed or produced so far. The kernels exploit the difference: the producer indexes the global tile with count and the shared slot with index:

cute.copy(tma_atom_k, tKgK[(None, kv_producer_state.count)], tKsK[(None, kv_producer_state.index)],
          tma_bar_ptr=kv_pipeline.producer_get_barrier(kv_producer_state))

Prefetch, steady state, drain

K4.2’s main loop has the canonical shape (as of 2026-10-02, from kernels/k4_2.py): before the loop, warp 0 acquires and fills min(kv_stage, n_kv_eff) stages. Each iteration then waits on the consumer side, computes, releases the stage, and finally, if kv_producer_state.count < n_kv_eff, acquires the next stage and issues the next pair of copies. With 5 stages the producer runs up to 4 tiles ahead of the consumer, and the acquire blocks exactly when all 5 slots hold tiles the consumer has not released. There is no explicit drain: the loop ends when the consumer has consumed n_kv_eff stages and the producer has stopped issuing at n_kv_eff. The K3 variants k3_1, k3_2 and k3_3 are the same file instantiated with 1, 2 and 5 KV stages, which is what makes the stage count a one-line experiment in the worklog.

The control flow document also offers cutlass.range(n, prefetch_stages=k) as an experimental attribute that generates the prefetch loop and the main loop automatically. The kernels write the two loops by hand so that the producer and the consumer can later be moved into different warpgroups.

References

• PipelineAsync, PipelineTmaAsync in pipeline/sm90.py, PipelineState, make_pipeline_state, CooperativeGroup, MbarrierArray in pipeline/helpers.py
• Control flow: software pipelining
• Colfax: Mastering the NVIDIA Tensor Memory Accelerator

WGMMA

The Hopper tensor core instruction is wgmma.mma_async. The PTX ISA defines a warpgroup as four contiguous warps whose first warp has a rank that is a multiple of four, and describes the operation as all threads of the warpgroup collectively computing \(D = A B + D\), or \(D = A B\) when the accumulator input is disabled. For BF16 the shape is m64nNk16 with \(N\) any multiple of 8 from 8 to 256 and a 32-bit float accumulator. For FP8 it is m64nNk32. The operation is asynchronous: the thread issues it, the tensor core performs it in the async proxy, and the thread waits for completion explicitly.

Operand sources and major modes

The ISA fixes where operands may live: A may be in registers or in shared memory, B must be in shared memory, and shared memory operands must start on a 16-byte boundary. A shared memory operand is named by a 64-bit matrix descriptor, and the ISA requires that the contents of a descriptor be the same across all warps of the warpgroup. The instruction has two syntax forms:

wgmma.mma_async.sync.aligned.shape.dtype.bf16.bf16  d, a-desc, b-desc, scale-d, imm-scale-a, imm-scale-b, imm-trans-a, imm-trans-b;
wgmma.mma_async.sync.aligned.shape.dtype.bf16.bf16  d, a,      b-desc, scale-d, imm-scale-a, imm-scale-b,              imm-trans-b;

scale-d is a predicate: false gives \(D = AB\), true gives \(D = AB + D\). imm-trans-a and imm-trans-b select MN-major instead of the default K-major for a shared memory operand; the ISA states that the transpose “is only supported for the wgmma.mma_async variants with .f16/.bf16 types on matrices accessed from shared memory using matrix descriptors”, which is why the register-A form has no imm-trans-a and why FP8 has neither. imm-scale-a and imm-scale-b are \(\pm 1\) and negate an operand; there is no scalar scale factor in the instruction.

The DSL enumerations mirror this. OperandSource.SMEM or OperandSource.RMEM applies to A only. OperandMajorMode.K or OperandMajorMode.MN applies to A and B, and LayoutEnum.from_tensor(t).sm90_mma_major_mode() derives it from which mode of a global tensor is contiguous: a row-major (rows, K) operand is K-major. The MmaF16BF16Op docstring lists the constraints: M is 64, N is 8 to 256 in steps of 8, K is 16, and A and B may be K- or MN-major when A is in shared memory, while with A in registers only B may be transposed.

The descriptor

The ISA gives the 64-bit descriptor’s fields:

with \(\text{encode}(x) = (x \,\&\, \text{0x3FFFF}) \gg 4\). The low four bits are dropped, so every address and stride the descriptor can express is a multiple of 16 bytes. For K-major swizzled layouts the ISA says the leading dimension byte offset is “not used, assumed to be 1” and the stride dimension byte offset is the distance from one group of 8 rows to the next. The base offset is zero when the swizzle pattern starts on its natural boundary (1024 bytes for the 128-byte mode), and CUTLASS’s make_gmma_desc always writes zero, checks that the layout is one of the canonical forms, and fills the two strides from the layout’s modes. The swizzle enters the descriptor only as the two-bit mode; the arrangement of bytes it implies is the canonical layout of the previous section.

The DSL never shows the descriptor bits. make_fragment_A and make_fragment_B on a partitioned shared memory tensor produce descriptor tensors, and the WGMMA guide says each element of such a tensor “is an SMEM descriptor, one per (MMA_K, PIPE) pair”. The hardware reads shared memory through them, so no load into registers happens for A or B.

Fence, commit, wait

The ISA defines three warpgroup-collective instructions around mma_async, and their exact scopes are the part most often misstated.

• wgmma.fence orders prior register accesses against a following mma_async. The ISA says only the accumulator registers and the registers holding A fragments need this ordering, that the fence must be issued before the first mma_async and between any register write and an mma_async that reads the same registers, with one exception: accumulator accesses by consecutive mma_async instructions of the same shape are ordered by default. It says nothing about shared memory, and adds that an async proxy fence (fence.proxy.async) must be used to order prior writes to shared memory operands before a mma_async reads them. The kernels issue no such fence for TMA-written tiles, and the reason is a derivation from two ISA statements rather than a sentence in it: cp.async.bulk operations are performed in the same async proxy as mma_async, and the ISA says their completion is followed by an implicit generic-async proxy fence once the mbarrier reports it. A tile the kernel wrote with ordinary stores would need the explicit fence.
• wgmma.commit_group batches all prior uncommitted mma_async instructions of the warp into a group. It does not wait.
• wgmma.wait_group N waits until at most \(N\) of the most recent groups are pending. Reading the accumulator or the A fragment registers of an mma_async without first waiting on a group that includes it is undefined behavior.

All three carry .sync.aligned, meaning every thread of the warpgroup must execute the same instruction, and the MmaF16BF16Op docstring repeats this for the DSL: cute.gemm must be issued warpgroup-uniformly. The DSL spellings are cute.nvgpu.warpgroup.fence(), commit_group() and wait_group(n), each a direct wrapper of the NVVM intrinsic. The DSL’s WGMMA guide carries a bullet saying the fence “orders prior SMEM writes”; that bullet contradicts the ISA and the op’s own docstring, and this post follows the ISA.

The TiledMma and partitioning

cutlass.utils.hopper_helpers.make_trivial_tiled_mma(a_dtype, b_dtype, a_major, b_major, acc_dtype, atom_layout_mnk, tiler_mn, a_source) picks the op from the dtypes (MmaF16BF16Op with K 16 for BF16, MmaF8Op with K 32 for FP8), builds an atom, and tiles it with atom_layout_mnk, the number of atoms along M, N and K, which is (1, 1, 1) in all the kernels, one warpgroup per CTA tile. K4.2 builds two (as of 2026-10-02, from kernels/k4_2.py):

qk_tiled_mma = sm90_utils.make_trivial_tiled_mma(
    INPUT_DTYPE, INPUT_DTYPE, warpgroup.OperandMajorMode.K, warpgroup.OperandMajorMode.K,
    ACC_DTYPE, self.atom_layout_mnk, (self.Br, self.Bc), warpgroup.OperandSource.SMEM)
pv_tiled_mma = sm90_utils.make_trivial_tiled_mma(
    INPUT_DTYPE, INPUT_DTYPE, warpgroup.OperandMajorMode.K, v_layout_enum.sm90_mma_major_mode(),
    ACC_DTYPE, self.atom_layout_mnk, (self.Br, self.d), warpgroup.OperandSource.RMEM)

A TiledMma exposes shape_mnk, the three TV layouts tv_layout_A, tv_layout_B, tv_layout_C, and get_slice(thread_idx), which returns a ThrMma. The thread view has partition_A(tensor_mk), partition_B(tensor_nk) and partition_C(tensor_mn); the WGMMA guide gives the output profiles as (MMA, MMA_M, MMA_K, ...), (MMA, MMA_N, MMA_K, ...) and (MMA, MMA_M, MMA_N, ...), where the first mode is one atom’s worth of values and the others count atoms along each axis, with any extra modes of the input, such as the stage mode, kept at the end. partition_shape_C((M, N)) gives the shape of the accumulator fragment without a tensor, and make_fragment_C(shape) allocates it in registers:

qk_thr_mma = qk_tiled_mma.get_slice(tidx)
tSsQ = qk_thr_mma.partition_A(sQ_full)          # (MMA, MMA_M, MMA_K, q_stage)
tSsK = qk_thr_mma.partition_B(sK_full)          # (MMA, MMA_N, MMA_K, kv_stage)
tSrQ = qk_tiled_mma.make_fragment_A(tSsQ)       # descriptors, same profile
tSrK = qk_tiled_mma.make_fragment_B(tSsK)
qk_acc_shape = qk_thr_mma.partition_shape_C((self.Br, self.Bc))
acc_qk = qk_thr_mma.make_fragment_C(qk_acc_shape)   # registers

For a shared memory operand the per-thread “partition” is not a slice at all. The C++ documentation for the GMMA atom writes its A layout as (T128, V64x16) -> (M64, K16) with shape (128, (64, 16)) and stride (0, (1, 64)): the thread mode has stride 0, so “every thread sees the entire tile of data and the tile is not reordered so that the descriptor can be constructed on it”. partition_A on shared memory therefore returns the same view for all 128 threads. For C the partition is a genuine per-thread slice, covered in the next section.

cute.gemm

cute.gemm(atom, d, a, b, c) computes \(D \leftarrow A B + C\) where \(C\) and \(D\) may alias, and its docstring notes that warpgroup atoms “require manually setting an accumulate boolean field”. The dispatch rules it lists decide how many instructions one call issues: with operands of profile (V, M, K) and (V, N, K) it loops over M, N and K and issues one atom per (m, n, k); with operands sliced down to (V) it issues exactly one. The kernels slice to one K block at a time and set the accumulate flag per issue, which is the pattern the DSL’s own WGMMA guide shows (as of 2026-10-02, from kernels/k4_2.py):

acc_qk = qk_thr_mma.make_fragment_C(qk_acc_shape)
cute.nvgpu.warpgroup.fence()
tSrQ_k = tSrQ[(None, None, None, q_consumer_state.index)]
tSrK_k = tSrK[(None, None, None, kv_consumer_state.index)]
num_k_blocks_qk = cute.size(tSrQ_k, mode=[2])
for k_block_idx in cutlass.range_constexpr(num_k_blocks_qk):
    qk_tiled_mma.set(cute.nvgpu.warpgroup.Field.ACCUMULATE, k_block_idx != 0)
    cute.gemm(qk_tiled_mma, acc_qk,
              tSrQ_k[(None, None, k_block_idx)], tSrK_k[(None, None, k_block_idx)], acc_qk)
cute.nvgpu.warpgroup.commit_group()
cute.nvgpu.warpgroup.wait_group(0)

Field.ACCUMULATE is the DSL’s name for scale-d. Setting it false on the first K block means the fresh accumulator’s contents do not matter, and the ISA’s same-shape exception means no fence is needed between the K blocks; the one fence before the loop covers the accumulator registers that were just allocated. For \(d = 128\) the loop issues \(128 / 16 = 8\) instructions. The PV matmul sets the flag true on every issue because its accumulator persists across KV iterations.

The accumulate-false path also removes the need to zero the accumulator, but it is only safe when the whole tile is overwritten by the first issue, which is the case for the QK scores and not for the running output.

References

• PTX ISA, Asynchronous Warpgroup Level Matrix Instructions: 9.7.16 (overview and warpgroup), 9.7.16.4 (async proxy), 9.7.16.5.1 (register fragments and shared memory layouts), 9.7.16.5.1.2.2 (descriptor format), 9.7.16.5.2 wgmma.mma_async, 9.7.16.7.1 wgmma.fence, 9.7.16.7.2 wgmma.commit_group, 9.7.16.7.3 wgmma.wait_group
• MmaF16BF16Op, MmaF8Op, OperandSource, OperandMajorMode, Field in cute/nvgpu/warpgroup/mma.py, fence, commit_group, wait_group in cute/nvgpu/warpgroup/helpers.py
• MmaAtom, TiledMma, ThrMma, make_tiled_mma in cute/atom.py, make_trivial_tiled_mma in utils/hopper_helpers.py, cute.gemm in cute/algorithm.py
• CuTe DSL WGMMA programming guide, CuTe MMA atoms (C++ docs): the GMMA A and C layouts, make_gmma_desc in include/cute/atom/mma_traits_sm90_gmma.hpp
• Colfax: WGMMA on Hopper

TV layouts and the accumulator

A TV layout is a layout whose two top-level modes are a thread index and a value index, and whose codomain is a position in the tile. The MMA atom documentation builds the GMMA accumulator layout step by step. For one 64x8 core tile, in column-major offsets \(m + 64 n\), each thread of a warp holds values at rows \(\lfloor (t \bmod 32) / 4 \rfloor\) and that plus 8, and at columns \(2(t \bmod 4)\) and that plus 1; the four warps stack down M in steps of 16 rows; and wider N repeats the 8-column pattern. The documented result for 64x128 is

CLayout = ((4, 8, 4), (2, 2, 16)) : ((128, 1, 16), (64, 8, 512))

and for the 64x64 tile the kernels use, the last value digit has extent 8 instead of 16. The DSL’s tiled_mma.tv_layout_C for the BF16 64x64 atom prints as ((4,8,4),(2,2,8)):((128,1,16),(64,8,512)) (printed on an H100 with DSL 4.7.0), the same layout. Reading the strides as rows and columns, with 64 the column stride:

So thread \(t\) owns value \((v_0, v_1, v_2)\) at

Each thread holds \(2 \cdot 2 \cdot 8 = 32\) values, which matches the PTX fragment table for m64nNk16 with a 32-bit accumulator: \(N / 2\) registers per thread, so 32 for \(N = 64\). They cover 2 rows and 16 columns, not a contiguous block: the columns are pairs spaced 8 apart. And the threads sharing a row are the four whose ids differ only in the low two bits, consecutive lanes, which is what makes warp_reduction with threads_in_group=4 the right reduction for a row.

The mn view

K4.2 does its softmax directly on the accumulator fragment, so it needs the fragment viewed as rows by columns. The helpers come from NVIDIA’s own Hopper FMHA example in the CUTLASS tree; the kernels keep them in kernels/common.py as a WgmmaHelpers mixin that every WGMMA kernel class inherits. layout_separate(thr, src, ref) walks the digits of a mode and splits them by whether the reference stride is below the tile’s M extent; with \(M = 64\) the digits with strides 1, 16 and 8 are row digits and those with 128, 64 and 512 are column digits. layout_acc_mn(tiled_mma, acc_layout) applies it to the value mode and appends the atom-count modes, giving a layout whose mode 0 walks the thread’s rows and mode 1 its columns; reduction_target_n(tiled_mma) applies it to the thread mode and returns the column-walking part, (4):(128), whose shape says how many lanes share a row. The kernel aliases the fragment through it:

acc_qk_mn = cute.make_tensor(acc_qk.iterator, self.layout_acc_mn(qk_tiled_mma, acc_qk.layout))
n_rows    = cute.size(acc_qk_mn, mode=[0])     # 2
n_cols_qk = cute.size(acc_qk_mn, mode=[1])     # 16

and the cross-thread reduction loops over the rank of the reduction target, calling warp_reduction_max(s_max[i], threads_in_group=reduction_target_qk.shape[r]).

Identity-tensor scatter and masks

Partitioning an identity tensor the same way as the accumulator yields each thread’s coordinates. cS = cute.make_identity_tensor((Br, Bc)), tScS = qk_thr_mma.partition_C(cS), re-viewed with the same mn layout, gives tScS_mn[i, jj] as the (m, n) of accumulator element (i, jj) inside the tile. The causal mask adds the tile offsets and compares; the output scatter writes gO_tile[tOcO_mn[i, jj][0], tOcO_mn[i, jj][1]]. No thread needs to know the ownership formula; the formula above is what the identity tensor computes.

References

• CuTe MMA atoms (C++ docs): GMMA C layout derivation and A layout
• PTX ISA, section 9.7.16.5.1.1.1, register fragments for m64nNk16 (the per-thread register counts; the placement figures are not reproduced in text, and the formula above is derived from CuTe’s documented layout)
• make_identity_tensor, ThrMma.partition_C
• CUTLASS Hopper FMHA example, layout_separate, layout_acc_mn, reduction_target_n

The QK to PV register handoff

After the softmax, \(P\) sits in the QK accumulator fragment as 32-bit floats in the C layout. The PV matmul wants it as the A operand in registers, as BF16 in the A fragment layout. The DSL’s WGMMA guide describes the sanctioned path for a fused kernel in three steps: the accumulator’s C layout (MMA, MMA_M, MMA_N) is converted to the A layout (MMA, MMA_M, MMA_K) expected by the second TiledMma, the values are type-converted and stored into a register tensor with that layout, and that tensor is passed directly to cute.gemm as A, with no shared memory descriptor involved. It points at the Hopper FMHA example for the full pattern, and the kernels use that example’s two helpers verbatim (as of 2026-10-02, from kernels/common.py; the FP8 branch that follows the store in the same function is covered in the FP8 section):

@staticmethod
def convert_c_layout_to_a_layout(c, a):
    return cute.make_layout(
        (a, c.shape[1], (c.shape[2], cute.size(c, mode=[0]) // cute.size(a))),
        stride=(c.stride[0], c.stride[1],
                (c.stride[2], cute.size(a, mode=[2]) * c.stride[0][2])),
    )

@cute.jit
def make_acc_into_op(self, acc, operand_layout_tv, Element):
    operand = cute.make_rmem_tensor_like(
        self.convert_c_layout_to_a_layout(acc.layout, operand_layout_tv.shape[1]),
        Element,
    )
    operand_as_acc = cute.make_tensor(operand.iterator, acc.layout)
    acc_vec = acc.load()
    operand_as_acc.store(acc_vec.to(Element))
    return operand

The layout arithmetic: a is the value shape of the PV atom’s A TV layout, 8 values per thread per k16 block, matching the PTX fragment table of four .f16x2 registers; the C fragment’s 32 values per thread are regrouped as 8 values times \(32 / 8 = 4\) K blocks, so the result has profile (V_A, MMA_M, MMA_K) with MMA_K = Bc / 16 = 4, and cute.size(tOrP, mode=[2]) is the PV K loop count. The data movement is one load of the accumulator as a vector, a conversion to BF16, and one store through an alias of the new tensor under the old C layout. Each thread touches only its own registers; there is no shuffle.

What this relies on, and how it is established: the regrouping is correct only if the per-thread placement of the C fragment’s first 16 columns coincides with the A fragment’s placement for one k16 block. The PTX ISA gives those placements as figures, not text, so there is no sentence to cite. The same helper is what NVIDIA’s example ships for the BF16 and FP16 paths, and the worklog’s K4 numerics against the PyTorch reference are the test of it. For FP8 the example adds a cross-thread exchange after the store, covered in the FP8 section.

The PV call is then the register-A form of the instruction:

tOrP = self.make_acc_into_op(acc_qk, pv_tiled_mma.tv_layout_A, INPUT_DTYPE)

cute.nvgpu.warpgroup.fence()
tOrV_k = tOrV[(None, None, None, kv_consumer_state.index)]
num_k_blocks_pv = cute.size(tOrP, mode=[2])
for k_block_idx in cutlass.range_constexpr(num_k_blocks_pv):
    pv_tiled_mma.set(cute.nvgpu.warpgroup.Field.ACCUMULATE, True)
    cute.gemm(pv_tiled_mma, acc_pv,
              tOrP[(None, None, k_block_idx)], tOrV_k[(None, None, k_block_idx)], acc_pv)
cute.nvgpu.warpgroup.commit_group()
cute.nvgpu.warpgroup.wait_group(0)

Here the fence is load-bearing, not ceremonial: the ISA requires it between a register write and an mma_async that reads the same registers as A fragments, and tOrP was just written by ordinary stores.

V’s mode order

\(O = P V\) has \(V\) as the B operand with \(N = d\) and \(K\) the KV sequence axis. \(V\) is stored as (N, d, B) with \(d\) contiguous, so for the PV MMA its contiguous axis is \(N\), which makes it MN-major, and the BF16 instruction allows that for a descriptor operand through imm-trans-b. The kernel makes CuTe see it that way by reordering the modes of the view before anything is derived from it:

mV = cute.make_tensor(mV.iterator, cute.select(mV.layout, [1, 0, 2]))   # (d, N, B)
v_layout_enum = utils.LayoutEnum.from_tensor(mV)                         # COL_MAJOR, so MN-major

LayoutEnum.from_tensor looks at which mode has stride 1; with \(d\) now mode 0 it reports column-major, sm90_mma_major_mode() returns MN, make_smem_layout_b picks an MN-major atom, and the TMA tile becomes (d, Bc). The bytes in global memory never move. This re-view is also exactly the thing FP8 cannot do, which is the next section’s subject.

References

• CuTe DSL WGMMA programming guide, “When A comes from registers”
• CUTLASS Hopper FMHA example, convert_c_layout_to_a_layout and make_acc_into_op
• PTX ISA, 9.7.16.5.2 wgmma.mma_async (register-A form, transpose rule) and 9.7.16.7.1 wgmma.fence
• select in cute/core.py, LayoutEnum

Warp specialization

Warp specialization gives different warps different jobs inside one CTA: one warpgroup issues TMA loads and does nothing else, one or two warpgroups run the WGMMA and softmax. The FlashAttention-3 paper describes this producer and consumer split as the first of its three techniques, and NVIDIA’s Hopper FMHA example is structured this way. The primitives it needs are the ones already covered, plus one new instruction.

Roles by warpgroup. A warpgroup index is cute.arch.make_warp_uniform(tidx // THREADS_PER_WARP_GROUP), and the kernel branches on it once at the top. The consumer side must be whole warpgroups because every wgmma instruction is .aligned across the warpgroup. The producer side is a whole warpgroup too, even though only one warp of it issues copies, because the register reallocation below works at warpgroup granularity. Inside the compute branch, thread indices passed to get_slice are the index within the warpgroup, tidx - THREADS_PER_WARP_GROUP.

Pipelines across roles. The same PipelineTmaAsync serves, with consumer_group sized to the number of consumer warps. The empty barrier’s arrival count is per consumer warp, so K5 passes CooperativeGroup(Agent.Thread, NUM_WARPS_PER_WARP_GROUP), 4 arrivals for its one compute warpgroup, and K6 passes NUM_COMPUTE_WGS * NUM_WARPS_PER_WARP_GROUP, 8 arrivals for two compute warpgroups sharing one KV stream, so a stage is released only when both have arrived (verified on the H100 2026-10-02).

Register reallocation. PTX setmaxnreg is a hint to change the number of registers owned by the executing warp. The ISA fixes its contract: a pool of registers is maintained per CTA, .dec releases registers from the warp to the pool and .inc requests them, the count must be between 24 and 256 and a multiple of 8, the same instruction must be executed by all warps of a warpgroup, .inc blocks until enough registers are in the pool, newly obtained registers have undefined contents, and the kernel must have been launched with a maximum register count set at compile time or the instruction may have no effect. The DSL spellings are cute.arch.setmaxregister_decrease(n) and setmaxregister_increase(n); the older warpgroup_reg_dealloc and warpgroup_reg_alloc are deprecated aliases of them. NVIDIA’s FMHA example uses 24 for its load warpgroup and 240 for each of its two compute warpgroups, and K5 uses the same two numbers behind a compile-time switch so the worklog can measure the kernel with and without the reallocation (verified on the H100 2026-10-02, from kernels/k5.py):

NUM_REGS_LOAD = 24
NUM_REGS_MMA = 240
...
if warp_group_idx == LOAD_WG_ID:
    if cutlass.const_expr(REALLOC):
        cute.arch.setmaxregister_decrease(NUM_REGS_LOAD)
    ...
else:
    if cutlass.const_expr(REALLOC):
        cute.arch.setmaxregister_increase(NUM_REGS_MMA)
    wg_tidx = tidx - THREADS_PER_WARP_GROUP      # 0..127 within the warpgroup

Why the producer can live on 24 registers: the TMA path has no per-thread operand data, just addresses, barrier pointers and loop state, and the copies are issued by one elected lane.

Warp-uniform values and election. cute.arch.make_warp_uniform(x) is a compiler hint that x is the same across the warp, used on the warp index before branching on it. cute.arch.elect_one() lowers to elect.sync and picks one lane; the pipeline uses it for barrier setup, and cute.copy applies it for TMA on its own.

Named barriers. When warps of different roles must synchronize without a full __syncthreads, PTX named barriers do it: each CTA has sixteen barriers numbered 0 to 15, bar.sync a, b waits for b threads (a multiple of 32) on barrier a, and bar.arrive a, b signals arrival without waiting, which the ISA presents as the producer and consumer pattern. The DSL’s NamedBarrier wraps a barrier id and a thread count with arrive, wait and arrive_and_wait, and warns that id 0 is what sync_threads uses. cute.arch.barrier_arrive(barrier_id=..., number_of_threads=...) is the raw form; in 4.7.0 it gained an aligned=False option that emits barrier.cta.arrive for the case where only some warps participate, and that keyword does not exist in the 4.5.1 the repo pins.

References

• FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision, Colfax: FlashAttention-3 overview
• PTX ISA, setmaxnreg, section 9.7.20.5; bar, barrier, section 9.7.14.1; elect.sync
• setmaxregister_increase/decrease, barrier_arrive, make_warp_uniform in cute/arch, NamedBarrier in pipeline/helpers.py
• CUTLASS Hopper FMHA example (24 and 240 registers, warpgroup roles)

Ordering two consumers: PipelineOrder

With two compute warpgroups on one KV stream, FlashAttention-3’s “ping-pong” schedule makes them alternate: while one issues its GEMMs, the other runs its softmax, so the tensor cores and the exponential units are busy at the same time. The DSL primitive for the alternation is cutlass.pipeline.PipelineOrder. Its docstring: the pipeline is divided into length groups, each group has depth stages, groups execute in a fixed order, and each group waits for the previous group before proceeding. The constructor is

math_wg_order_barrier = pipeline.PipelineOrder.create(
    barrier_storage=storage.order_barrier_ptr.data_ptr(),
    depth=1, length=NUM_COMPUTE_WGS, group_id=warp_group_idx - 1,
    producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, THREADS_PER_WARP_GROUP),
    defer_sync=True)

(verified on the H100 2026-10-02, from kernels/k6.py; the storage struct reserves NUM_COMPUTE_WGS 64-bit words for it). The source shows how little it is: depth * length mbarriers, each initialized with the producer group’s arrival count (128, every thread of a compute warpgroup). wait() waits on barrier index * length + group_id with the state’s phase; arrive() arrives on barrier index * length + (group_id + 1) % length, the next group’s barrier, and advances the state. The initial phase is group_id == 0, so group 0’s first wait passes immediately and group 1’s first wait blocks until group 0 has arrived. NVIDIA’s example wraps the same call in make_and_init_order_barrier with one stage per math warpgroup.

Where the two calls sit decides what overlaps. K6 brackets each GEMM issue, QK and PV separately: wait() before the fence, arrive() right after commit_group() and before wait_group(0):

math_wg_order_barrier.wait()
cute.nvgpu.warpgroup.fence()
...
cute.nvgpu.warpgroup.commit_group()
math_wg_order_barrier.arrive()
cute.nvgpu.warpgroup.wait_group(0)

So the other consumer may start issuing as soon as this one’s instructions are committed, while this one is still waiting for its results, and then runs its softmax while the other’s instructions execute. The ISA’s model makes this legal: commit does not wait, and only the thread’s own accumulator reads need the wait.

Two Q stages are needed in the same kernel, one per consumer (q_stage = NUM_COMPUTE_WGS), since each compute warpgroup owns a different Q tile while both read the same K and V stages.

References

• PipelineOrder in pipeline/sm90.py, CUTLASS Hopper FMHA example, make_and_init_order_barrier
• FlashAttention-3 paper, pingpong scheduling
• PTX ISA, mbarrier

FP8 operands

The FP8 form of the instruction, from the ISA:

wgmma.mma_async.sync.aligned.shape.dtype.atype.btype  d, a-desc, b-desc, scale-d, imm-scale-a, imm-scale-b;
wgmma.mma_async.sync.aligned.shape.dtype.atype.btype  d, a,      b-desc, scale-d, imm-scale-a, imm-scale-b;

with shapes m64nNk32, atype and btype each .e4m3 or .e5m2 independently, and dtype .f16 or .f32. Three things differ from BF16, all visible in the syntax.

• K is 32 per instruction, so the number of K blocks for the same tile halves, and make_trivial_tiled_mma builds an MmaF8Op with K 32 when both dtypes are FP8.
• There are no transpose arguments. The ISA’s transpose rule limits it to .f16 and .bf16 descriptor operands, and the MmaF8Op docstring states it directly: “Transpose (MN-major) is not supported for A or B. Both operands must be K-major.” For QK nothing changes, since Q and K are K-major already. For PV it means \(V\) must be K-major as the B operand, which is the opposite of how BF16 stores it: the KV sequence axis, the contraction axis of PV, must be the contiguous one. NVIDIA’s example handles this in its data preparation by laying \(V\) out with the sequence axis contiguous when the dtype is FP8. The kernel repo does the same: its harness stores \(V\) as (B, d, N) and views it as (N, d, B) with \(N\) contiguous, and K7 then views that as (d, N, B), passes LayoutEnum.ROW_MAJOR to make_smem_layout_b so the shared memory atom has the \(B_c\) axis contiguous, and names both operands K-major (verified on the H100 2026-10-02, from kernels/k7.py):

# V arrives as (N, d, B) with N contiguous. View it as (d, N, B): the
# PV B operand with N = d rows and K = Bc contiguous columns.
mV = cute.make_tensor(mV.iterator, cute.select(mV.layout, [1, 0, 2]))
...
sV_layout_staged = sm90_utils.make_smem_layout_b(
    utils.LayoutEnum.ROW_MAJOR, pv_mma_tile_mnk, INPUT_DTYPE, self.kv_stage)
...
pv_tiled_mma = sm90_utils.make_trivial_tiled_mma(
    INPUT_DTYPE, INPUT_DTYPE, warpgroup.OperandMajorMode.K, warpgroup.OperandMajorMode.K,
    ACC_DTYPE, self.atom_layout_mnk, (self.Br, self.d), warpgroup.OperandSource.RMEM)

The TMA tile is (d, Bc) as in K4.2, but now the contiguous axis inside it is \(B_c\). The mode re-view trick of the BF16 kernel does not apply, because there is no MN-major option to select; the transpose has to happen in memory.

• The A register fragment packs four elements per register. The fragment table for m64nNk32 gives A as four .b32 registers each holding four .e4m3 or .e5m2 elements. The accumulator stays \(N / 2\) 32-bit registers per thread.

The last point is what makes the register handoff different for FP8. After the C to A reinterpretation and the store with conversion to FP8, NVIDIA’s example runs an extra exchange when the element width is 8: for each pair of 32-bit words it swaps 16-bit halves between lanes of the same 4-lane row group with cute.arch.shuffle_sync_op and reassembles bytes with cute.arch.prmt, with a source comment describing the accumulator order per lane as 00 11 22 33 and the required operand order as 00 00 11 11 22 22 33 33. The placement it corrects for is given in the ISA only as a figure. The kernel repo’s make_acc_into_op carries the same branch, enabled by default behind a compile-time FP8_SHUFFLE switch so the worklog can run both variants, and passes 0x1C1F as the shuffle’s fourth operand: the ISA defines that operand as a packed segment mask in bits 12 to 8 and a clamp in bits 4 to 0, so 0x1C makes the source lane relative to each 4-lane segment and 0x1F is the full clamp, which is what lets one code path serve all 32 lanes of a warp (verified on the H100 2026-10-02, from kernels/common.py).

The exchange is required. The worklog’s correctness test on the H100 ran K7 both ways on sharp attention data against an FP32 reference computed on the dequantized inputs:

The numbers are in the worklog’s K7 section. The reason is the one the comment in NVIDIA’s example gives: per 16 columns of \(P\), the accumulator layout hands a thread of a 4-lane row group the column pairs \(2t\) and \(8 + 2t\), while the FP8 A fragment packs four consecutive \(k\) values into each 32-bit register, so each thread must hold \(k = 4t\) to \(4t + 3\); the register-count side of that is in the fragment table, the placement is in the ISA’s figures, and the test is what confirms it. The earlier project that concluded the opposite had dropped the shuffle’s segment-mask operand, so every lane read from lanes 0 to 3 of the whole warp, and had tested on nearly uniform attention, where the output collapses to the mean of \(V\) whatever layout \(P\) is in. The BF16 kernels pass the same sharp test with the plain reinterpretation, which is the evidence behind the handoff section’s claim.

Everything else in an FP8 kernel is software: the per-tensor or per-block dequantization scales are multiplied into the 32-bit accumulator after the GEMM, since the instruction has no scale operand beyond the \(\pm 1\) negation flags.

References

• PTX ISA, 9.7.16.5.2 wgmma.mma_async (FP8 syntax, transpose rule, element sizes), section 9.7.16.5.1.1.3 (register fragments for m64nNk32)
• MmaF8Op in cute/nvgpu/warpgroup/mma.py, make_trivial_tiled_mma
• CUTLASS Hopper FMHA example: FP8 V permutation and the FP8 branch of make_acc_into_op
• FlashAttention-3 paper, low-precision section, NVIDIA Hopper architecture in-depth, FP8 formats

Block scales

Per-block scaling in the FP8 kernel is plain tensor work on the DSL side; nothing new in the hardware vocabulary is involved. A scale tensor is a small global tensor, one float per (batch, tile), wrapped with from_dlpack(t, assumed_align=16) with a static layout and passed as a kernel argument like any other tensor. K8 is the K7 class compiled with block_scales=True, so the reads sit behind a compile-time branch and vanish from the K7 binary. Inside the kernel they are plain indexing: a per-Q-tile scale once at kernel start, and a per-KV-tile scale once per iteration, both scalars per CTA, and the product with the softmax scale is applied to the 32-bit accumulator in the same loop that already scales the scores (verified on the H100 2026-10-02, from kernels/k7.py):

q_block_scale = cutlass.Float32(1.0)
if cutlass.const_expr(self.block_scales):
    q_block_scale = mQ_scales[bidx_b, my_q_tile]
...
effective_scale = scale
if cutlass.const_expr(self.block_scales):
    k_block_scale = mK_scales[bidx_b, _j]
    effective_scale = scale * q_block_scale * k_block_scale
for i in cutlass.range_constexpr(n_rows):
    for jj in cutlass.range_constexpr(n_cols_qk):
        acc_qk_mn[i, jj] = acc_qk_mn[i, jj] * effective_scale

Because the ISA’s wgmma has no scale-factor operand, there is no way to fold these into the instruction on Hopper; the DSL’s cute.gemm does accept auxiliary scale-factor operands for block-scaled GEMMs, but that form belongs to atoms that carry scale factors, not to MmaF8Op. The worklog’s K8 section covers how the scales are produced offline, including the Hadamard rotation, and what they cost.

References

• cute.gemm docstring, auxiliary operands
• PTX ISA, 9.7.16.5.2 wgmma.mma_async
• FlashAttention-3 paper, block quantization and incoherent processing

Wrapping up

The whole post reduces to a few objects and the contracts the hardware puts on them. A layout is a function from coordinates to offsets, a tensor is a pointer composed with one, and every “partition” and “re-view” in the kernels is a new layout over the same bytes. Shared memory for a tensor core operand is one of four swizzle atoms tiled to the tile, named to the hardware by a 64-bit descriptor whose swizzle field is two bits wide, which is why the buffers are 1024-byte aligned and why TMA and WGMMA can agree on the bytes. A TMA load is one instruction from one thread, completed by bytes counted on an mbarrier, and a pipeline is two arrays of those barriers with a phase bit per side. WGMMA is asynchronous, collective across a warpgroup, and ordered by a fence that covers registers only, a commit that batches, and a wait that is the only thing that blocks. The accumulator’s thread and value layout decides what the softmax code looks like, and the register handoff is a type conversion under a relabeling of that layout. Warp specialization is a branch on the warpgroup index plus a register budget moved from one side to the other, and the ping-pong order is one more array of barriers. FP8 takes the transpose away and halves the K blocks, and everything about scales is ordinary arithmetic on the accumulator.

Where this post stops is where the kernels start: the worklog has the numbers for each of these mechanisms and the per-kernel diagrams, including the FP8 handoff test that settled the one question this post could not answer from the sources alone.

If you find a mistake anywhere in here, please let me know and I’ll fix it.