Debugging Warp-Specialized Kernels

GEMM Steps 7-9 in Scaling GEMM with Warp Specialization and Clusters overlap TMA load, tcgen05 MMA, and TMEM/SMEM writeback. The same debugging method applies to Flash Attention handoffs: identify the roles, identify the storage each role owns, then verify the generated CUDA against that model.

Do not start by rewriting the kernel. First make sure the run is valid, then inspect the generated CUDA. After environment and compile-time issues are ruled out, runtime failures in these kernels usually reduce to a broken handoff: an uninitialized barrier, the wrong arrival count, a collective hidden inside a role guard, a stale barrier phase, or storage reused before the producer has made its writes visible.

Before Debugging the Kernel

Rule out the runtime context first:

python -c "import tvm, tvm.tirx; print(tvm.__file__, tvm.__version__)"
python -c "import torch; print(torch.cuda.get_device_name(), torch.cuda.get_device_capability())"

These kernels target Blackwell (sm_100a). If Python imports a stale TVM checkout, or the GPU is not Blackwell-class, fix that before changing the kernel. Then run the kernel’s smallest correctness check, such as run_correctness(), before looking at performance.

Debugging Workflow

1. Reproduce the failure at the smallest shape that still fails. If the failure is an illegal memory access, restart Python before the next run.
2. If compilation fails, check the installed API, target, dispatch=, and buffer scopes before reading the runtime synchronization code.
3. Save inspect_source("cuda") output. Search it for role guards, mbarrier_init, tcgen05, cp.async.bulk.tensor, and cta_sync() before reading the Python again.
4. Write the roles / storage / handoff / lifetime table for the kernel path that failed.
5. Check the generated CUDA against that table: barrier inits before role branches, expected TMA producer, MMA issuer(s), writeback group(s), and no CTA-wide collective inside a warpgroup-only branch.
6. Classify the run as a deadlock, crash, wrong result, or correct-but-slow run, then use the matching section below.
7. Change one handoff at a time: init count, arrive/wait phase, role guard, fence, TMA store drain, TMEM alloc/dealloc, or tile-scheduler advance.
8. Re-run correctness before measuring performance.

What Transfers

For any asynchronous kernel, make a small worksheet before changing code:

Item	What to write down
Roles	The exact threads, warps, warpgroups, or CTAs that issue each async operation.
Storage	The live location of each tile at each step: GMEM, SMEM, TMEM, or registers.
Handoff	The producer, the consumer, the signal object, the arrival count, the phase, and the fence or drain that makes the data visible.
Lifetime	The earliest point where each storage slot can be reused, read back, or freed.

Then verify the generated CUDA against the worksheet:

• Role guards match the roles table.
• Barrier inits appear before guarded role branches.
• Collective operations are not accidentally narrowed by lane, warp, or warpgroup guards.
• Arrive/wait phases match the handoff table.
• TMA store drains, TMEM dealloc, and SMEM reuse happen only after the lifetime table says they are legal.

Use the same worksheet for TMA->MMA->writeback GEMM pipelines and for the score/softmax/value/correction handoffs in Flash Attention.

If Compilation Fails

Fix compile-time failures before debugging runtime synchronization:

Symptom	Likely area	First check
Unknown TIRx API or attribute error	Installed wheel does not match the tutorial code	Print tvm.__file__ and tvm.__version__; compare the API name with [[GPU/modern-gpu-programming-for-mlsys/reference/TIRx Language Reference
Unsupported dispatch=	The selected target or primitive does not support that path	Check the dispatch argument and target capability; tcgen05 paths in this tutorial require Blackwell.
Buffer scope mismatch	A buffer is being used through the wrong hardware path	Check the storage row of the worksheet: TMEM must be accessed through tcgen05, and TMA operands must use compatible GMEM/SMEM layouts.
Compile succeeds but generated CUDA lacks the expected path	Dispatch did not lower the way you expected	Inspect the generated CUDA for tcgen05 and cp.async.bulk.tensor before changing the algorithm.

Inspecting Generated Code

For any compiled kernel, save the CUDA so you can search and diff it:

from pathlib import Path
 
cuda_source = ex.mod.imports[0].inspect_source("cuda")
Path("artifacts").mkdir(exist_ok=True)
Path("artifacts/my_kernel.cu").write_text(cuda_source, encoding="utf-8")
print(cuda_source)

The generated code maps TIRx constructs to CUDA like this:

TIRx	Generated CUDA
wg_id == 0	(warp_id_in_cta >> 2) == 0
wg_id == 1	(warp_id_in_cta >> 2) == 1
warp_id == 0	(warp_id_in_cta & 3) == 0
warp_id == 3	(warp_id_in_cta & 3) == 3
lane_id == 0	(((int)threadIdx.x) % 32) == 0
.init() internal guard	((int)threadIdx.x) < 1 (CTA thread 0 only)
elect_sync()	tvm_builtin_elect_one_sync_op()

Scan for these strings before reading the full kernel:

Generated CUDA	Check
if (threadIdx.x < 1)	Single CTA-thread guard, often barrier initialization
mbarrier_init	Barrier initialization exists and appears before role branches
tcgen05	The Tensor Core path was generated
cp.async.bulk.tensor	The copy lowered to TMA
cta_sync();	CTA-wide barrier; it must not sit inside a wg_id branch

Step 7 Reference Skeleton

A correctly compiled Step 7 kernel has this top-level shape. The guards below are written with role names for readability; in generated CUDA, search for the corresponding expressions from the table above.

// (1) Barrier inits: top level, CTA thread 0 only
if (threadIdx.x < 1) {
  mbarrier_init(tma2mma[0..1], 1);
  mbarrier_init(mma2tma[0..1], 1);
  mbarrier_init(mma2ld, 1);
  mbarrier_init(ld2mma, 128);   // arrived by all 128 WG0 threads
}
 
// (2) TMEM alloc: WG0 warp 0, all lanes of the issuing warp
if (wg_id == 0 && warp_id == 0) tcgen05_alloc(..., 512);
 
// (3) Fences + cta_sync, then phase init: producer=1, consumer=0
 
// (4) Warp-specialized loop
if (wg_id == 1 && warp_id == 3 && elect_sync) { /* TMA  */ while(valid){ ... next_tile(); } }
if (wg_id == 1 && warp_id == 0 && elect_sync) { /* MMA  */ while(valid){ ... next_tile(); } }
if (wg_id == 0)                                { /* WB   */ while(valid){ ... next_tile(); } }
 
// (5) Cleanup: issuing warp, no lane guard
cta_sync();
if (warp_id == 0) { tcgen05_relinquish_alloc_permit(); tcgen05_dealloc(..., 512); }

Check these before changing the algorithm:

• Barrier inits sit at top level, not inside a wg_id guard.
• tcgen05_alloc and tcgen05_dealloc have a warp guard but no lane guard; all lanes of the issuing warp participate.
• TMA and MMA loops both iterate K_TILES times.
• Phase init is producer=1, consumer=0.

Symptom Map

Start from the symptom, but treat it as a clue rather than a final diagnosis:

Clue	Likely area	First check
Kernel hangs, then the runtime reports an unspecified launch failure	Deadlock	Barrier init placement, arrival count, cta_sync() placement, and next_tile() participation
Illegal memory access, XID, or later unrelated CUDA calls also fail	Crash / poisoned context	Restart Python, then check pointer ranges, storage lifetime, and collective participation
Wrong rows appear in 128-row or tile-sized stripes	Sync race or tile-index mismatch	Producer/consumer phases, scheduler advance, and which warpgroup owns each row stripe
NaN or obviously invalid values	Descriptor, operand setup, or uninitialized accumulation	SMEM/TMEM descriptor setup, swizzle/layout, and accumulator initialization
Finite but patterned wrong values	Stale or partially visible data	Missing fence, missing TMA store drain, or storage reused before the lifetime table allows it
Correct output but no expected speedup	Dispatch or resource issue	Generated CUDA path, pipeline depth, occupancy, and register spill

When to Restart Python

A CUDA error does not always clean up after itself. After an illegal memory access, XID, or “CUDA context poisoned” error, later unrelated calls such as torch.randn may keep failing. Restart the Python process before testing the next fix, otherwise you may be debugging the previous crash instead of the current code.

Deadlocks

Check these in order:

• Arrival count does not match init count. Common case: MBarrier.init(128) but arrive is guarded by if warp_id == 0: if lane_id == 0:, so only 1 thread arrives and the wait never returns.

  Barrier	init(count)	Who arrives	Arrivals
  TMABar (tma->mma)	1	TMA engine via arrive(stage, bytes)	1
  TCGen05Bar (mma->tma, mma->ld)	1	MMA warp via tcgen05.commit	1
  MBarrier (ld->mma)	128	All WG0 threads via arrive	128

• Barrier init nested inside a wg_id guard. .init() lowers to if threadIdx.x < 1:, meaning CTA thread 0. CTA thread 0 lives in WG0, so if wg_id == 1: prevents every thread from running the init. Inits must be at top level; grep mbarrier_init in inspect_source() to verify.

• cta_sync() inside a warpgroup branch. cta_sync is __syncthreads(), which requires all CTA threads. Inside if wg_id == 0:, WG1 never reaches it. Use T.cuda.warpgroup_sync(10) for a single-warpgroup barrier.

• tile_scheduler.next_tile() skipped by some consumer-warpgroup threads. The scheduler tracks per-thread state; threads that skip it can loop forever.

• TMA and MMA disagree on K-tile count. If MMA does K_TILES - 1 instead of K_TILES, barrier phases drift and the second outer tile deadlocks.

• PipelineState initial phase wrong. Producer starts at phase=1 so the first wait passes; consumer starts at phase=0 so the first wait blocks. If both start from the same phase, the first handoff can deadlock immediately.

Crashes and Context Poisoning

Common causes:

• pool.alloc after pool.commit(). Barrier wrappers call alloc internally. Correct order: tmem_addr -> barrier wrappers -> move_base_to(1024) -> Asmem / Bsmem / Dsmem -> commit().
• tcgen05.alloc or tcgen05.dealloc with a lane guard. The issuing warp must participate with all lanes. if lane_id == 0: runs one thread, which is undefined behavior.
• Missing cta_sync() before tcgen05.dealloc. TMEM is freed while writeback is still reading.
• Out-of-range GMEM or SMEM access. Shrink to one tile, check the scheduler’s m_idx / n_idx, and check that the current shape is a multiple of the kernel’s tile or cluster tile.

Wrong Results

Classify wrong output by pattern before guessing. Whole row stripes often point to a producer/consumer phase, tile-index, or role-ownership mismatch. NaN output often points to descriptor setup, operand setup, or uninitialized accumulation. Finite but patterned wrong values often mean the consumer read an old tile, a partially written tile, or data whose store had not drained yet.

• tcgen05.commit outside elect_sync. All 32 threads create commit groups; the 31 empty groups signal the mbarrier immediately. TMA can overwrite SMEM before MMA reads it.
• Missing fence.proxy_async("shared::cta") before TMA store. The TMA engine may not see SMEM writes from threads.
• Missing cp_async.bulk.commit_group() plus wait_group(0) after TMA store. The next tile can reuse Dsmem before the store drains.
• Persistent kernel fails intermittently at small sizes such as 1024x1024. Larger sizes can mask the race with longer K-loops. Re-check phase reset between tiles and the TMA-store commit/wait.
• fence.after_thread_sync() is usually not the fix. The MMA-completion mbarrier already carries release-acquire semantics. Steps 8 and 9 add it conservatively on the writeback edge, after mma2ld.wait and before the first tcgen05.ld; do not add it routinely on the TMA-to-MMA edge.

Correct but Slow

If the output is correct but performance is far below expectation, use the same inspection loop:

Clue	Likely area	First check
Generated CUDA has no cp.async.bulk.tensor	Copy did not lower to TMA	Check dispatch="tma", target capability, and operand layout
Generated CUDA has no tcgen05 path	MMA did not lower to Blackwell Tensor Core instructions	Check dispatch="tcgen05", target capability, and operand layouts
TMA and MMA do not overlap	Pipeline too shallow or phases serialize producer/consumer	Inspect the order of wait/arrive/advance in the generated CUDA
Good small-shape correctness but poor large-shape speed	Register spill, occupancy, or staging-buffer pressure	Check the compiler resource report; reduce tile size, chunk writeback, or lower pipeline depth

Filing a Good Issue

If the failure survives the checks above, reduce it before filing an issue on the Apache TVM GitHub repository. Include:

• the tvm.__file__ / tvm.__version__ output and GPU capability;
• the smallest shape that reproduces the failure;
• whether the failure is compile-time, deadlock, crash, wrong result, or correct-but-slow;
• the minimal kernel or notebook cell plus its correctness check;
• the saved inspect_source("cuda") output, or the smallest excerpt that shows the suspicious guard, barrier, or dispatch path.

---

Previous · Book index · Next