Tlx API Reference

TLX API Quick Reference

Warp Specialization

Warp specialization skeleton

with tlx.async_tasks():
    with tlx.async_task([0]):       # Producer
        # TMA loads
    with tlx.async_task([1, 2]):    # Consumers
        # MMA compute

Memory Barriers

mbarrier (shared-memory allocated)

Ordinary barrier count rules:

• alloc_barriers(..., arrive_count=N) initializes the number of logical arrivals required to complete a phase. Transaction bytes registered with barrier_expect_bytes are an additional completion condition; do not infer the software arrival topology from the transaction byte count.
• For TMA full/data-ready barriers where one task registers the transaction, use arrive_count=1 unless the surrounding protocol explicitly requires additional software arrivals.
• For local software barriers shared by replicated tlx.async_task consumers, the ordinary count is typically the number of consumer replicas that arrive once per phase.
• For cross-CTA software arrivals using remote_cta_rank, derive the count from the exact CTA protocol. Do not assume the local-task rule applies.

Warp-barrier semantics and safe conversion

alloc_warp_barrier changes the physical arrival protocol, not merely the spelling of the allocation. With an ordinary barrier, TLX selects a leader to perform the arrival and may synchronize the participating threads first. With a warp barrier, every participating thread performs its own arrival, avoiding that leader-path synchronization but issuing more mbarrier arrival operations.

The initialized count is:

expected arrivals per phase = num_warps * 32 * num_arrivals

Here num_warps is the number of warps in each participating task replica, and num_arrivals is the number of such all-thread arrival events targeting the same barrier in one phase. For example, two replicated four-warp consumers that each release one shared buffer slot use num_warps=4, num_arrivals=2, for 256 arrivals. A consumer-owned slot released by one four-warp replica uses num_warps=4, num_arrivals=1, for 128 arrivals. Keep the corresponding barrier_arrive calls at their original final-use points and normally use the default unit arrival count.

Before converting an ordinary barrier, classify its role and arrival source:

Safe-conversion checklist:

1. Identify the protected buffer and confirm the barrier is signaled by explicit software barrier_arrive calls rather than TMA, MMA, multicast, or remote completion.
2. Enumerate every task replica, warp, lane, predicate, and arrival call that targets the barrier during one phase.
3. Verify num_warps * 32 * num_arrivals equals the exact number of unit arrivals.
4. Prove every counted lane reaches the arrival after its final access to the buffer, including prologue, tail, and persistent-loop iterations.
5. Preserve buffer indices, phase calculations, wait sites, and arrival sites; change only the allocator for the first A/B experiment.
6. Keep full/TMA barriers ordinary even when the paired empty/reuse barriers are converted.
7. Run correctness and benchmark both forms. Warp barriers trade additional per-thread arrivals for removal of leader-path synchronization, so conversion is an optimization candidate, not a universal rule.

Named barriers (hardware-allocated, indices 0–15)

num_threads is the total number of threads required to flip the barrier phase: num_waiting_threads + num_arriving_threads. Wait and Arrive calls for the same barrier phase must use the same value. The count must be a multiple of 32 (warp size); it is typically num_warp_groups * warps_per_group * 32.

Used for PingPong scheduling to prevent tensor core contention between consumer warp groups.

Memory Operations

SMEM / TMEM allocation

Storage kinds: tlx.storage_kind.smem, tlx.storage_kind.tmem (Blackwell), tlx.storage_kind.smemCluster

TMA (Tensor Memory Accelerator)

Matrix Multiply (MMA)

Minimum tile sizes for async_dot: M ≥ 64, K ≥ 16, N ≥ 32

Pair-CTA MMA (two_ctas=True): M must be 128 per CTA.

Multi-CTA (Cluster) Kernels

ctas_per_cga=(N,1,1) in triton.Config sets the cluster size. The grid specifies total CTAs; hardware divides by ctas_per_cga to get the number of clusters. E.g., grid=(2,1,1) with ctas_per_cga=(2,1,1) = 1 cluster of 2 CTAs.

2-CTA tile scheduling in attention backward

Each CTA in a cluster gets its own program_id and its own tile_id from CLC. Two CTAs in a cluster naturally get consecutive tiles (pid 0, pid 1). No special tile scheduling is needed for 2-CTA — start_n = pid works as-is. Grid size and n_tile_num do NOT change between 1-CTA and 2-CTA.

Think of 2-CTA as two independent 1-CTAs that handle their own K/V tiles and share Q/dO via multicast. For L2 efficiency, they process consecutive N-blocks.

2-CTA MMA semantics (two_ctas=True)

• A operand (TMEM): per-CTA, each CTA has different data
• B operand (SMEM): split across CTAs and combined by hardware via multicast
• Output (TMEM): split across CTAs along the M dimension, written to both CTAs
• Leader MMA writes to both leader TMEM and peer TMEM

2-CTA barrier patterns

• TMA loads with two_ctas=True: only leader calls barrier_expect_bytes (guarded by if is_leader:). Use arrive_count=1.
• Software arrives (barrier_arrive with remote_cta_rank=0): both CTAs arrive on leader's barrier. Use arrive_count=NUM_CTAS.
• MMA mBarriers with two_ctas=True: hardware signals when input reads complete. The TMEM output write may still be in-flight.

input_precision options: tf32, tf32x3, ieee

CLC (Cluster Launch Control) — Blackwell only

For 2-CTA mode: set multi_ctas=True (uses "arrive remote, wait local" pattern).

Utility

Common patterns

Producer-consumer with mbarrier (pipelined GEMM)

# Full/data-ready barriers track TMA transaction completion and remain ordinary.
bars_full = tlx.alloc_barriers(num_stages, arrive_count=1)

# Ordinary software-release form: one logical arrival per consumer replica.
bars_empty = tlx.alloc_barriers(
    num_stages,
    arrive_count=num_consumers,
)

# Optional optimized form when each consumer replica has four warps and every lane
# is statically guaranteed to execute one unit arrival after its final buffer read.
bars_empty_warp = tlx.alloc_warp_barrier(
    num_barriers=num_stages,
    num_warps=4,
    num_arrivals=num_consumers,
)

# Producer: wait for reuse, then start TMA load tracked by the ordinary full barrier.
tlx.barrier_wait(bar_empty, empty_phase)
tlx.barrier_expect_bytes(bar_full, nbytes)
tlx.async_descriptor_load(desc, indices, barrier=bar_full)

# Consumer: wait for TMA completion, consume the buffer, then release it.
tlx.barrier_wait(bar_full, full_phase)
acc = tlx.async_dot(A, B, acc)
acc = tlx.async_dot_wait(0, acc)  # Prove the final buffer read completed.
tlx.barrier_arrive(bar_empty)

When evaluating the optimized form, substitute bars_empty_warp for bars_empty at both the producer wait and consumer arrival sites. Do not convert bars_full.

PingPong with named barriers

# Consumer 0 waits for Consumer 1, then issues MMA
tlx.named_barrier_wait(9, 256)   # 256 = 2 warp groups * 4 warps * 32 threads
qk = tlx.async_dot(q, k)
tlx.named_barrier_arrive(10, 256)

# Consumer 1 waits for Consumer 0's MMA to finish
tlx.named_barrier_arrive(9, 256)
tlx.named_barrier_wait(10, 256)
qk = tlx.async_dot(q, k)

Deep-dive docs

• API reference: third_party/tlx/README.md
• Barriers: third_party/tlx/doc/tlx_barriers.md
• Placeholder layouts: third_party/tlx/doc/PlaceholderLayouts.md
• Storage alias design: third_party/tlx/doc/storage_alias_spec_design.md