MegaKittens Internals

00The idea

A normal PyTorch program launches one kernel per operator. Each launch waits for the previous one to drain, and every intermediate makes a round trip through HBM with the GPU partly idle at each boundary. MegaKittens replaces the whole graph with a single persistent kernel that runs an interpreter. The host compiles the graph into a flat list of tile-sized instructions: one 512×256 GEMM output tile, two 128×128 add tiles, a page of RMSNorm rows, and so on. Each SM walks its share of that list in order.

Inside an SM, five specialized warp roles pipeline consecutive instructions. While the consumer warps finish instruction i, the loader is already TMA-loading instruction i+1 into the shared-memory pages that i has released. Dependencies between SMs are tracked per tile with integer counters in global memory. A consumer tile starts as soon as the specific producer tiles it reads have been written, without waiting for the whole previous operator.

01Capture: megakittens.compile

compile() (megakittens/interface.py) is a thin wrapper around TorchDynamo:

torch._dynamo.optimize(
    backend=aot_autograd(fw_compiler=_megakittens_backend, bw_compiler=_megakittens_backend),
    nopython=True,      # any graph break is an error
    dynamic=False,      # every new shape recompiles
    disable=not enable)(fn)

• Dynamo captures the Python function into an FX graph. nopython=True means the entire function must be capturable, because one graph becomes one kernel.
• AOTAutograd turns that into an aten-level FX graph. It functionalizes in-place ops (they reappear as auto_functionalized_v2 nodes), lifts module parameters to placeholders, and flattens list[Tensor] inputs into individual tensors. The same backend is registered for the backward graph, but no backward instruction types exist, so training graphs fail at op resolution.
• The backend (backend.py) prints the FX table if verbose, runs trace() and then schedule(), builds a Dispatcher, and returns make_boxed_func(dispatcher). With dry_run=True it stops after tracing and returns the original graph. The gpt-fast example's --compile megakittens mode does exactly this (dry_run=True, save_dag=True) to dump the Llama DAG without running it.

02Trace: FX graph → DAG

This is the actual aten graph AOTAutograd produced for the test MLP (x = relu(x @ W[i] + b[i]), three layers):

placeholder arg0_1                       # W  (3, 1024, 1024)
placeholder arg1_1                       # x  (1024, 1024)
placeholder arg2_1                       # b  (3, 1024, 1024)
select   = aten.select.int(arg0_1, 0, 0)
mm       = aten.mm.default(arg1_1, select)
select_1 = aten.select.int(arg2_1, 0, 0)
add      = aten.add.Tensor(mm, select_1)
relu     = aten.relu.default(add)
select_2 = aten.select.int(arg0_1, 0, 1)
mm_1     = aten.mm.default(relu, select_2)
…                                         # layers 2 and 3 repeat
output   = output((relu_2,))

tracer.trace() turns it into a node-centric DAG in five steps:

1. Prune. prune_fx_graph runs a DFS back from the output node and drops dead nodes. It also drops every view node (getitem, aten.select.int, aten.slice.Tensor), because views become properties of edges rather than nodes.
2. Classify. placeholder and get_attr become input nodes, each carrying a TensorMeta (dtype, shape of at most 4 dims, device; the tensor must be contiguous). call_function, call_method and call_module nodes go through IType.from_torch. An auto_functionalized_v2 wrapper is first unwrapped back to the mutating op and its real arguments.
3. Resolve the op. Three global maps (function, method name, module type) are filled by each IType subclass at import time. A map entry is either a plain IType instance or a list of resolvers. A resolver returns an IType, or (IType, aten_output_indices) when only some outputs of a multi-output aten op are kept. For example, _fused_rms_norm and cuDNN SDPA keep only output 0.
4. Fold views into ranges. For each input edge, resolve_input_node walks back through getitem/select/slice and records a slice chain. TensorRange.from_slice_chain replays the chain into an absolute, left-padded 4-D half-open range on the producer's tensor. W[1] on a (3,1024,1024) tensor becomes ((0,1),(1,2),(0,1024),(0,1024)) on its padded shape (1,3,1024,1024). Slices with step ≠ 1 are rejected, as are views in front of the output.
5. Validate. DAG.validate checks that each node is exactly one of input/output/compute, that there is exactly one output node, that nodes are in topological order, that edges appear on both endpoints, and that every range fits inside its tensor.

megakittens/tracer.py, megakittens/schema/dag.py, megakittens/schema/tensor.py, megakittens/schema/itype.py

03The IType contract

An instruction type is defined twice: once in Python, which tells the scheduler how to tile and what each tile touches, and once in CUDA, which tells each warp role what to do with one instruction. The scheduler auto-discovers subclasses through IType.__subclasses__(). __init_subclass__ requires a custom op named torch.ops.megakittens.<snake_case_name>, registers it in the dispatch maps, and generates a test_fn.

04Scheduler: DAG → instructions, tensors, barriers

schedule() runs four phases in order. Each phase consumes the previous phase's output. The numbers below are from the MLP above, with M = H = 1024 and 148 SMs (B200).

Phase 1 · count and place

For every compute node in topological order, the scheduler asks the IType for num_instructions and assigns it a contiguous range in the global list. Each range is rounded up to a multiple of cluster_size. This guarantees that instruction pairs (2k, 2k+1) always belong to the same node, which matters because the two CTAs of a cluster run in lockstep (§09).

Phase 2 · assign tensor slots, reuse memory

Every node output gets an index into tensor_metas; there can be at most 256. A freed buffer goes into a pool keyed by its exact TensorMeta. A new output takes a pooled buffer only when its node starts at least #SMs instructions after the last instruction of every consumer of that buffer. In other words the reader must be a full wave behind, so the reuse barrier (Phase 3) is almost always already satisfied, and every old consumer sits earlier than the reuser in every SM's queue. Three kinds of buffer are never pooled: graph inputs, buffers that feed the graph output, and buffers that some consumer mutates in place. An output with an inplace_mapping aliases its input directly and adds a reuse barrier against the input's other readers.

In the MLP this produces 11 slots. T0 = W, T1 = x and T2 = b are the inputs, and T3–T10 are intermediates. The final ReLU writes into T3, the buffer of layer 1's GEMM output. It qualifies because its node starts at instruction 208, and the last reader of T3 (layer 1's add) ends at 48: 208 − 48 = 160 ≥ 148. Layer 3's add, starting at 176, misses the threshold (176 − 48 = 128) and gets a fresh buffer.

Phase 3 · fine-grained barriers

Barriers are global int32 counters. Every barrier has a target, the number of producer instructions that must arrive before a waiter may proceed. They are derived purely from access_regions:

1. Collect each instruction's read and write boxes.
2. For each producer → consumer edge, compute a unit cell: for each dimension, the GCD of every box extent on both sides. Every unit cell is mapped to the producer instructions that write it.
3. For each consumer instruction, the union of producers over its unit cells is its dependency set. Sets are cached per box list.
4. Every distinct (producer node, dependency set) pair becomes one barrier. Its target is the size of the set. Each producer in the set lists the barrier in dst_barriers, and each consumer with that set lists (barrier, target) in src_barriers. Consumers with identical dependencies share one barrier.
5. Reuse barriers are separate: one per reuse event. Every instruction of every old reader arrives, and every instruction of the reusing node waits for the total count before it writes.
6. Each instruction may have at most 16 source and 8 destination barriers.

The GEMM target is 2 because the CTA-pair instructions are duplicates with identical write boxes: both CTAs must finish their half. The totals are 42 + 42 + 40 input barriers plus 1 reuse barrier, which gives the scheduler's reported 125. Layer 2's GEMM row block 0 (instructions 80–87) waits only on the 16 ReLU instructions that produce rows 0–511. It can start while layer 1 is still computing rows 512–1023.

Phase 4 · emit

Each node is validated and assigned an icode, one per unique (itype, src tensor indices, dst tensor indices). Icode 0 is reserved for Noop. Because tensor indices are compile-time template arguments, the three structurally identical MLP layers become nine distinct icodes, and the kernel instantiates nine instruction bodies. The instruction list is then built and padded with Noops at node boundaries to the cluster size.

megakittens/scheduler.py

05The real schedule, SM by SM

Below is the schedule the repository produced for the MLP, laid out as the GPU sees it in the default per-SM mode. Instruction i goes to SM i % 148 in iteration i // 148. Each column pair is one cluster. 240 instructions fill 1.62 waves, and the remaining 56 slots of the second wave are zero-filled, which decodes as icode 0, a Noop. Click any cell to decode it. The instructions it waits on are marked above their cells and the ones that wait on it below.

06The 256-byte instruction

struct instruction_t (csrc/schema.cuh) is exactly 256 bytes, enforced by a static_assert. Python packs it into 64 int32s with _pack_instruction, and the controller warp copies it with two coalesced loads per lane. The grid below shows the 64 words of the instruction selected in §05; select a field to see where it lives.

07Dispatcher, JIT and launch

First call: materialize

1. Validate the inputs against their TensorMetas, then allocate every non-input slot with torch.empty. These buffers belong to the dispatcher, and the returned outputs are those same buffers on every call.
2. Pack the instructions. In per-SM mode they go into an int32 tensor of shape [⌈N/#SMs⌉, #SMs, 64]. In global mode the shape is [N, 64]. The barrier array is zeros(num_barriers).
3. Build a gl descriptor per tensor. The Python gl in jit/pykittens.py mirrors ThunderKittens' kittens::gl byte for byte: raw pointer, a size_t for each runtime dim, then 64-byte-aligned 128-byte CUtensorMaps. For each tensor it collects the union of TMA tile types from every IType that touches it. In the MLP, T3 needs both GEMM's 128×32 epilogue box and the add's 128×128 box. Descriptors are encoded on the host with cuTensorMapEncodeTiled; swizzled tiles use the same 5-D box TK uses.
4. All gls are packed into one C struct (pack_struct applies C alignment) and passed as a single __grid_constant__ kernel parameter, about 3.5 KB for the MLP. The TMA descriptors therefore live in kernel-parameter space, where TMA can read them. The packed bytes are cached per tuple of data_ptrs.

Generated source

This is what Dispatcher._compile_kernel emitted for the MLP, captured by intercepting the NVRTC call (indentation tidied):

NVRTC compiles it with --std=c++20 --use_fast_math -DNDEBUG -lineinfo -DKITTENS_NO_HOST -DKITTENS_BLACKWELL --gpu-architecture=sm_100a, plus include paths for csrc, TK include and prototype, and CUDA/CCCL. The cubin is loaded with cuModuleLoadData, the mangled name of megakittens::kernel<MKConfig, MKGlobals> is resolved, and MAX_DYNAMIC_SHARED_SIZE_BYTES is set to 230,400.

Every call: launch

• barrier_tensor.zero_() runs as a separate small kernel on the current stream.
• cuLaunchKernelEx uses 384 threads per block and 230,400 bytes of dynamic shared memory. The grid is #SMs in per-SM mode, or ⌈N/2⌉·2 in global mode (one CTA per instruction, most of which never launch; see below). Cluster dimension 2 is set both as preferred and required. PROGRAMMATIC_STREAM_SERIALIZATION enables PDL, so the kernel's prologue overlaps the previous kernel's tail, and griddepcontrol.wait runs after semaphore init.
• Global-work-queue mode uses Blackwell Cluster Launch Control. A CTA starts at instruction blockIdx.x. For each next instruction, CTA rank 0 issues clusterlaunchcontrol.try_cancel, which cancels a not-yet-launched cluster and multicasts its first block id to both CTAs. The pair adopts that cluster's instructions (x + cta_rank). When cancellation fails, the grid is exhausted and the CTA stops. This gives dynamic load balancing with the same instruction format.

megakittens/dispatcher.py, megakittens/jit/{pykittens,nvrtc_jit,cuda_utils,c_utils}.py

08Inside one CTA

megakittens::kernel (csrc/megakittens.cuh) sets everything up once and then splits the 12 warps into roles that each loop forever over instructions.

Prologue, in order

1. Static shared memory. This holds instruction_state[2] (each is the 256-byte instruction, a 7-entry pid_order padded to 128 bytes, 32 dynamic semaphores and 128 bytes of scratch), two CLC handles, and the mbarriers clc_arrived[2], instruction_arrived[2], instruction_finished[2], page_finished[7] and tensor_finished.
2. Dynamic shared memory is rounded up to a 1024-byte boundary (needed by the 128-byte swizzle) and viewed as page_t pages[7].
3. TMEM. The tensor_allocator<1, CLUSTER_SIZE> constructor has warp 0 run tcgen05.alloc for all 512 columns and relinquish the allocation permit. Its destructor does a cluster barrier and then tcgen05.dealloc.
4. Semaphores are initialized by distinct threads. instruction_arrived gets count = CLUSTER_SIZE, instruction_finished gets 11 (every warp except the controller), and clc_arrived gets 1. page_finished[p] and tensor_finished get count 1 and are immediately arrived once, so the first instruction finds every page and the TMEM token free.
5. Cluster barrier (barrier.cluster.arrive/wait), or __syncthreads without clusters. Then griddepcontrol.wait for PDL.
6. Role split. Warps 0–7 raise their register limit to 224 and enter consumer_loop. Warpgroup 2 lowers to 56 and dispatches on its warp index: 0 = controller, 1 = loader, 2 = launcher, 3 = storer.
7. On exit, a final cluster barrier, so neither CTA's shared memory or TMEM goes away while its partner still addresses it.

09The instruction pipeline

Two stages (slots) hold decoded instructions, so at most two instructions are in flight per CTA. The controller produces into the stages and the other 11 warps consume from them. Because both CTAs of a cluster arrive on each other's barrier, the pair always works on the same instruction index.

Controller loop (warp 8), per iteration

1. Recycle the slot. From iteration 2 on, wait on instruction_finished[stage] with parity ((iter−2)/2)&1. All 11 workers have then finished the instruction that last used this slot. Its dynamic semaphores are invalidated with mbarrier.inval, one lane each.
2. Fetch. In per-SM mode, if iter ≥ depth (the number of waves), write icode −1, broadcast it and exit. Otherwise copy instructions[iter][blockIdx.x][0..63] with lane l loading words l and l+32. In global mode, the first fetch uses blockIdx.x and later fetches use CLC as described in §07.
3. Rename pages. Lane l computes pid_order[l] = prev.pid_order[prev.lid_release_order(l)], calling the previous instruction's function (§10). The very first instruction gets the identity order.
4. Init semaphores. Dispatch to the new instruction's init_semaphores. Each lane initializes one entry, and the count is remembered for the later invalidation. Then fence.proxy.async.shared::cta makes the inits visible to the TMA and tensor-core async proxy.
5. Publish. warp::sync, then one elected lane runs tensor_commit<CLUSTER_SIZE>(instruction_arrived[stage]). Advance the stage.

Worker loop (every other warp)

for (iter = 0, stage = 0; ; ++iter) {
    wait(instruction_arrived[stage], (iter / 2) & 1);
    icode = instruction_states[stage].instruction.icode;
    if (icode == -1) break;
    dispatch_instruction<WorkerType::ROLE>(icode, g, s);   // JIT switch → Op::ROLE::run
    warp::sync();
    if (warp::elect_leader()) arrive(instruction_finished[stage]);
    stage ^= 1;
}

Roles drift apart freely within the two-slot window. The loader of instruction i+1 can run while consumers are still in i; what bounds that is the page and TMEM tokens (§10, §11) and the global barriers.

csrc/controller.cuh, csrc/workers.cuh, csrc/utils.cuh (dispatch_instruction)

10Pages and page renaming

The seven 32 KB pages are passed between consecutive instructions as tokens. Two rules make this work:

• Every instruction waits on and finishes every physical page exactly once. page_wait(pid) is wait(page_finished[pid], iter&1) and page_finish(pid) is one arrive. Because the count is 1 and each page was pre-arrived at startup, a page's phase advances exactly once per instruction. Pages an instruction does not need are waited on and released at once by a spare lane of the loader warp.
• Instructions address logical pages (lids), not physical ones (pids). Each IType puts the data it needs first in its lowest lids. Through lid_release_order(q), the previous instruction tells the controller which of its lids it releases q-th. The controller then maps new lid q to that page: new.pid_order[q] = old.pid_order[old.release(q)]. The next instruction's first loads therefore land in the pages that free up first.

The simulator runs the real lid_release_order functions. Add instructions to see how physical pages (by colour) move between logical slots. GEMM's order depends on (K/64) mod 4, because that decides which stage's pages its loader frees last.

csrc/schema.cuh (page_wait / page_finish), controller.cuh step 2, each itype's controller::lid_release_order

11The TMEM token and global barriers

TMEM token

tensor_finished works like a page token for the whole of tensor memory. Each instruction arrives on it exactly once (tensor_finish). A launcher that is about to write TMEM first calls tensor_wait with parity iter&1. Types that never touch TMEM (Noop, the elementwise ops, RMSNorm) have the launcher pass the token straight through. GEMM's consumers return it as soon as the accumulator has been copied to registers, before the epilogue stores. The next GEMM's MMAs can therefore overwrite TMEM while the current one is still writing results to HBM. Attention returns it only at the very end.

Global barriers between SMs

• all_input_barrier_wait runs in the loader, before any load from a dependent tensor.
• all_reuse_barrier_wait runs in the consumer, right before the first store (WAR protection for reused or in-place buffers).
• all_barrier_arrive runs after store_async_wait(), when this instruction's writes have completed. It is issued by one elected thread.
• Why this cannot deadlock (per-SM mode): every dependency points to a lower global index, and instruction i sits at iteration i // #SMs on SM i % #SMs. A producer is therefore always earlier in its own SM's queue than any consumer that could be blocking that SM. Reuse barriers are additionally a full wave back (Phase 2).

csrc/utils.cuh (barrier_wait / barrier_arrive), csrc/schema.cuh (tensor_wait / tensor_finish)

12The six instruction types

ElementwiseUnary and ElementwiseBinary

• A tile is st<bf16,128,128>, exactly one 32 KB page. Each of the 8 consumer warps owns 16 rows as an rt_bf<16,128>, loaded with ldmatrix.x4.
• The unary op list is a template parameter pack applied with a fold expression, so ("abs","log","neg") is one pass over registers. NEG, SQRT and RSQRT use bf16_2 intrinsics. The others use TK's register maps.
• Binary is designed for N ops and N+1 inputs (BinaryOps<...>, index layout 4·NUM_INPUTS + 1), but the Python side currently caps MAX_OPS = 1. Inputs must have equal effective shapes; there is no broadcasting. That is why the MLP test uses a bias of shape (L, M, H).
• Index layout (binary): [a.b, a.d, a.row_tile, a.col_tile, b.b, b.d, b.row_tile, b.col_tile, d.b, d.d, d.row_tile, d.col_tile, n_tiles].
• Release order: unused pages first, then B pages, then the A/output pages, which are freed only after their TMA stores complete.

RMSNorm

• Each row is an sv_bf<N>, padded to 128 bytes, and rows are packed into page lid 0. The weight is read directly from global memory (w_gl.raw_ptr, L2-resident) with no TMA.
• Rows are split across 8 warps, ⌈rows/8⌉ each, so with 4 rows per instruction half the warps idle. Each lane reads 4 bf16 (8 bytes) at a time with a stride of 128 elements, and __shfl_xor_sync reduces across the warp.
• EPS is a C++ constant, 1e-6f. The eps captured from the graph is not passed to the kernel.
• Requirements: N must be a multiple of 64, at most 256 or with 2N divisible by 128, and one row must fit in a page.

csrc/itypes/{noop,elementwise_unary,elementwise_binary,rmsnorm}.cuh and the matching megakittens/itypes/*.py

13GEMM in depth

A cluster of two CTAs computes a 512×256 tile of D = A·B (bf16 in, fp32 accumulate, bf16 out). The Python tile is TILE_M = 512. The C++ per-CTA constants are Mb = 256, Nb = 256, Kb = 64, with 2 consumer warpgroups (cid), a 4-stage load ring and an 8-step epilogue.

1. Loader (both CTAs, one lane). Wait on the input barriers. For iterations i = 0 … K/64 + 3: on the first pass through the ring, claim the stage's A page and, for even stages, its B page. After that, wait on inputs_finished[stage] until the MMA has consumed the stage. While i < K/64, issue two A loads (one per cid) and one B load with tma::cluster::load_async(…, mask = 1<<rank, dst_mbar_cta = 0). In the four drain iterations, release each finished stage's pages, except A stage 0, which the epilogue reuses. A second lane releases lid 6 immediately.
2. Launcher (CTA 0 only). Take the TMEM token, then for each K step expect_bytes(inputs_arrived[stage], 96 KB), wait, and issue mm2_AB (first step, overwrite) or mma2_AB (accumulate) for each cid. Each MMA commits to inputs_finished[stage], which has count 2 and is multicast to both CTAs, so both loaders learn the stage is free. Finally tensor_commit<2>(outputs_arrived).
3. Consumers (both CTAs, warpgroup = cid). Wait on outputs_arrived. Eight warpgroup::load_async calls pull the 128×256 accumulator as eight 128×32 chunks into rt_bf<32,32> registers (fp32 → bf16). Then tensor_load_wait, the reuse-barrier wait, and an early tensor_finish. The epilogue loops 8 times: store_async_read_wait<1> frees one of two 8 KB slots, the warpgroup stores registers to the slot, and a TMA store with EVICT_FIRST writes it. Finally wait for all stores, release lid 0, and arrive on the global barriers.
4. Tile order. _swizzled_tile_order visits output tiles in supergroups of 8 column tiles, with serpentine rows, to improve L2 reuse of A and B.

csrc/itypes/gemm.cuh, megakittens/itypes/gemm.py

14Attention in depth

This is FlashAttention-style forward attention on BSHD tensors with head_dim = 128, using a 2-CTA cluster. One instruction pair covers 512 query rows of one (batch, head). Each CTA owns two 128-row Q tiles, one per softmax warpgroup (qid). K and V are streamed in 128-key steps. The softmax scale is fixed at 1/√128.

Roles

• Loader (both CTAs). Load the two Q tiles for this CTA with axis = DEPTH, so the tile's 128 rows run along the sequence dimension of a BSHD tensor without a transpose. Then for each key step, load one K tile (64 keys: this CTA's half of 128) and one V tile (128 keys × this CTA's 64 head-dim columns). Both go into a shared 4-stage ring over lids 2–3 (16 KB per stage), gated by tma::cluster::wait(kv_finished) starting from phase 1.
• Launcher (CTA 0). Wait for Q from both CTAs, then issue the first QKᵀ for both qids. In the steady state, for each qid: wait for P quarter 0, issue PV quarter 0, then quarters 1–3 as their signals arrive (the last one commits kv_finished for V). For qid 0 it also waits for the next K tile, then issues the next QKᵀ and commits scores_arrived. After the last step, the final PV commits tile_arrived.
• Consumers (two softmax warpgroups, both CTAs). Each thread owns one query row and loads the 128 fp32 scores with four tcgen05.ld.32x32b.x32. The row max is a tree of 3-input max.f32.
  • Lazy rescale (as in FA4): if the new max exceeds the old by no more than 8 in log2-scaled units, the old max is kept and O is left alone. Otherwise O is rescaled in TMEM in 16-column chunks (tcgen05.ld → __fmul2_rn → tcgen05.st), and only when some lane in the warp needs it.
  • p = exp2(s·log2e/√d − m·log2e/√d) via __ffma2_rn, converted to bf16 and stored as P in four quarters. After each quarter, warp::tma::cluster::arrive signals CTA 0. The row sum is accumulated with __fadd2_rn.
• Epilogue. Release the Q page, compute rcp.approx of the row sum (guarded against 0/NaN), and claim the O page. Then load O from TMEM, multiply, convert, and st.shared it with TK's 128-byte swizzle written out by hand: addr ^ (((addr & 0x380) >> 7) << 4), into two 64-column panels. Wait on the reuse barrier, TMA store, and finally release the K/V and O pages, return the TMEM token, and arrive on the global barriers.
• Causal. The number of key steps becomes 4·m + 4, enough to reach the diagonal of the cluster's last tile. Steps before a warpgroup's own tile are unmasked. The diagonal step masks col > row using the warpgroup thread index as the row; later steps are fully masked. On the host, instructions are emitted in LPT order (largest m first) with an L2 swizzle that groups heads so their K/V fit in about 50 MB of L2.

csrc/itypes/attention.cuh, megakittens/itypes/attention.py

15The ThunderKittens layer

ThunderKittens (TK, csrc/ThunderKittens/include, header-only) supplies the types and the PTX. MegaKittens uses about a dozen of its constructs. They are compiled under NVRTC with KITTENS_NO_HOST, which swaps the STL for cuda::std and stubs CUtensorMap as a 64-byte-aligned 128-byte blob. Everything TK would do on the host, such as building TMA descriptors, is re-implemented in pykittens.py.

Execution scopes

• kittens::warp = group<1> and kittens::warpgroup = group<4>. MegaKittens adds group<8> as consumer_group. group<N>::laneid() is the thread index within the group (0…32N−1). warpid() is the warp within the group, and groupid() is threadIdx.x / 32N, which gives MegaKittens' cid/qid.
• group<N>::sync(id) emits bar.sync id, 32N. Consumers use ids 1 and 4 across 256 threads, and ids 1 and 2 across 128 threads per warpgroup.
• elect_leader() emits elect.sync, so exactly one lane is true. The loader uses if (elect_leader()) {…} else if (elect_leader_from_active()) {…}. The second branch elects one of the other 31 lanes through activemask, so one lane issues TMA while another releases unused pages, concurrently.
• warpgroup::increase/decrease_registers<n> emits setmaxnreg.inc/dec.sync.aligned.u32 n.

Semaphores (mbarriers)

• init_semaphore(s, n) → mbarrier.init … n. arrive → mbarrier.arrive.release.cta. invalidate_semaphore → mbarrier.inval.
• wait(s, p) spins on mbarrier.try_wait.parity p, which succeeds once the phase with parity p has completed. On a fresh barrier, waiting with p = 1 returns immediately. Producers start at phase 1, as in attention's kv_phase = 1; tokens that start free are instead pre-arrived once.
• tma::expect_bytes(s, n) → mbarrier.arrive.expect_tx. It arrives once and adds n bytes to the transaction count, and TMA completion subtracts them. Every TMA barrier is therefore initialized with count 1.
• Cluster variants map a CTA-local address into the peer CTA with mapa.shared::cluster: tma::cluster::arrive(sem, cta) performs a remote arrive, and tma::cluster::wait uses .acquire.cluster.
• tensor_commit<2> / detail::tcgen05::commit<2> → tcgen05.commit.cta_group::2.mbarrier::arrive::one.shared::cluster.multicast::cluster [sem], 0b11. This arrives on both CTAs' copies once the issuing thread's earlier MMAs finish.

Tiles, layouts, TMA

• st<bf16,R,C> is stored as column panels of 64 elements (128 bytes) × R rows, XOR-swizzled with addr ^ (((addr % 1024) >> 7) << 4). That is why pages are 1024-byte aligned and every in-page offset is a multiple of 1 KB. The swizzle width is picked from C: 128 B for C ∈ {64, 128}, but 64 B for GEMM's st_bf<128,32> epilogue tile. The Python st applies the same rule.
• sv<T,N> is padded to 128 bytes. A TMA on a vector is split into ⌈N/≤256⌉ 4-D boxes.
• gl is raw_ptr, then a size_t per runtime dim (or a 1-byte empty struct for compile-time dims), then a recursive descriptor_dict of CUtensorMaps looked up at compile time by (tile type, axis). A missing pair is a static_assert.
• Coordinates are in tile units: {b, d, r, c} where the axis named by the template's axis is multiplied by the tile's rows. Attention uses axis = DEPTH to tile along the sequence of BSHD.
• Loads emit cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes. The cluster form adds .cta_group::2.multicast::cluster with the CTA's own bit as mask and CTA 0's mbarrier (via mapa). It does not multicast data; it lets CTA 0 count both CTAs' bytes. Stores emit fence.proxy.async + cp.async.bulk.tensor…bulk_group + commit_group. store_async_wait is cp.async.bulk.wait_group 0, and store_async_read_wait<1> waits only until shared memory may be reused.
• group<8>::load(rt_bf<16,128>, st<128,128>) uses ldmatrix.sync.aligned.m8n8.x4. Warps are interleaved over row slabs (warp w of warpgroup g gets rows 32·w + 16·g) to match TMEM lane quarters.

Tensor memory and tcgen05

• tt<T,R,C> is just a 32-bit TMEM address (lane << 16) | column. 16-bit types pack two per column. tensor_allocator::allocate<tt>(col) is pure address arithmetic; no per-instruction allocation happens.
• mm2_AB / mma2_AB / mm2_ABt issue tcgen05.mma.cta_group::2.kind::f16 once per K/16 slice, with shared-memory matrix descriptors (128 B swizzle, SBO = 1024) and an instruction descriptor encoding M, N, dtypes and transposes. mm* clears the accumulator on the first slice, and mma* accumulates. The PV product uses the TS form, where A comes from TMEM.
• warpgroup::load_async(rt, tt) uses tcgen05.ld.sync.aligned.16x256b. Attention bypasses it with raw tcgen05.ld/st.32x32b, where lane i of warp w%4 is TMEM lane 32·(w%4) + i, i.e. one thread per row. tensor_load_wait / tensor_store_wait emit tcgen05.wait::ld/st.

Clusters, CLC, PDL

• everyone::tma::cluster::sync() → barrier.cluster.arrive.release; barrier.cluster.wait.acquire, at kernel start and end.
• clc::schedule → clusterlaunchcontrol.try_cancel.async…multicast::cluster::all.b128, and clc::query → query_cancel.is_canceled + get_first_ctaid.
• pdl::wait() → griddepcontrol.wait. launch_dependents is never called, so dependents start when blocks exit.

TK files: common/util.cuh, ops/group/group.cuh, ops/thread/util/{sync,tma,util}.cuh, ops/thread/memory/tile/tma.cuh, ops/thread/mma/tcgen05.cuh, types/{shared/st,shared/sv,global/gl,tensor/tt,tensor/tensor}.cuh

16Limits and sharp edges

By design, today

• Static shapes only (dynamic=False), whole-graph capture only (nopython=True), and at least one tensor output. No backward pass.
• Contiguous tensors of at most 4 dims, slices with step 1, at most 256 tensors, 16 inputs and 8 outputs per node, 16 indices per instruction, 16 source and 8 destination barriers per instruction.
• Every IType is bf16-only. GEMM needs M % 512, N % 256 and K % 64 to be 0, with full-K ranges. Attention needs head_dim 128, sequence length % 512, BSHD layout, and a fixed scale. Binary ops do not broadcast.
• The dispatcher owns the output buffers, so a second call overwrites the tensors returned by the first.
• Each (op, tensor-index) combination is a separate icode and a separate template instantiation, so code size and compile time grow with graph depth even when the layers are identical.

Observed in the code

• RMSNorm hard-codes EPS = 1e-6f and ignores the traced eps.
• Barrier waits spin without nanosleep (the helper exists but is unused). The storer warp has no work in any type.
• The README points to csrc/itypes/add.cuh, which does not exist; ElementwiseBinary replaced it.
• gemm.cuh's comment says d_st_t is 128×64; the type is 128×32.

Worth verifying (not confirmed either way)

• TMEM handoff across the CTA pair. CTA 0's launcher waits only on CTA 0's tensor_finished before issuing 2-CTA MMAs that also write CTA 1's TMEM. I found no explicit wait on CTA 1's consumers having finished reading the previous accumulator. In practice the next MMA also waits for CTA 1's next loads, but that is timing, not a guarantee.
• No tcgen05.fence::before/after_thread_sync around the mbarrier-based TMEM handoffs, and the ordering of the controller's plain shared-memory writes relative to the tcgen05.commit "broadcast". Both are worth checking against the PTX memory model.
• barrier_arrive is red.relaxed after cp.async.bulk.wait_group 0, with no explicit release fence. The consumer side does use fence.acquire.gpu.

17Source map