TLX - Triton Low-level Language Extensions

Introduction

TLX (Triton Low-level Language Extensions) is a low-level, warp-aware, hardware-near extension of the Triton DSL. It offers intrinsics and warp-specialized operations for fine-grained GPU control, hardware-oriented primitives for advanced kernel development, and explicit constructs for GPU memory, computation, and asynchronous control flow. TLX is designed for expert users pushing Triton closer to the metal.

Primarily targeting NVIDIA GPUs (for now), TLX extends Triton to support:

Hardware-specific intrinsics (e.g., wgmma, async_copy, barrier)
Shared and local memory allocation
Instruction-level scheduling and control
Cross-warpgroup synchronization

While this approach places more responsibility on the user, it reduces the compiler's role as a performance bottleneck. Although it may introduce divergence across hardware platforms, it empowers users to perform deeper, architecture-specific optimizations without relying solely on compiler heuristics.

Nightly builds (fbtriton)

Nightly .dev wheels are published to a self-managed index (not PyPI):

pip install --pre fbtriton \
  --index-url https://facebookexperimental.github.io/triton/nightly/simple/

Each nightly is built from the newest main commit whose GPU/CI checks are all green. triton.__version__ reports 3.8.0.dev<YYYYMMDD>+fb.git<hash>. Nightlies are retained for ~30 days. Formal releases remain on PyPI (pip install fbtriton).

Gluon support

Gluon (python/triton/experimental/gluon/) is upstream-synced and not a first-class DSL here (TLX is the focus) — do Gluon feature/bug work upstream, not in this fork. But since fbtriton is a secondary Triton, we run fundamental Gluon CI so we don't silently break it.

CI: the b200-gluon-test / mi350-gluon-test jobs run pytest python/test/gluon/ (every test_*.py). It is green: ~200 passed, ~1600 skipped, 0 failed. The real signal is the compile-only, target-agnostic frontend suite (test_frontend.py, one run covers NVIDIA + AMD codegen via mock GPUTarget); the version-skewed cases below are skipped via python/test/gluon/conftest.py rather than failing the job.

Fork-side fixes (Gluon frontend itself unmodified): core.py reduce reduction_ordering compat, semantic.py dot() allow_tf32 default, a test_core.py collection fix (a bad cherry-pick left an IndentationError), and regenerated test_frontend.py goldens (upstream-synced — overwritten on next sync).

Skipped, needs upstream Gluon re-sync — TODO(gluon-ci): the GPU-execution suites (test_core, test_lowerings, test_consan, test_fpsan, one kernel in test_layout_format_view) were synced from much newer upstream (e.g. test_fpsan.py via bundle #1956) than the pinned Gluon frontend (_semantic.py ~2026-06-29). They require Gluon behavior the frontend doesn't have yet — raw pointer gl.load/gl.store inferring a distributed layout — so they fail wholesale with expected ... distributed_type but got block_type. The fix is to sync the Gluon frontend forward (upstream cherry-pick / re-sync), not a local patch. Also skipped: ~9 frontend per-target golden tests (single inline golden can't match every parametrized target) and the create_lds_barrier_wait pybind mismatch. conftest.py lists these; remove/trim it after the re-sync.

The DSL Extension

Hardware availability tags. Each op below is tagged with the targets it runs on: **[Hopper+]** = NVIDIA Hopper and newer; **[MI300+]** = AMD MI300 (CDNA3) and MI350 (CDNA4); **[MI350]** = AMD MI350 (CDNA4) only. A trailing ? (e.g. **[MI300+?]**) marks AMD availability that has not been confirmed yet and needs verification. Note: async copies (async_load and its commit/wait groups) and the AMD buffer ops require MI350 — they are not available on MI300. barrier_arrive on AMD requires arrive_count == 1.

Local buffer operations

buffers = tlx.local_alloc(shape, dtype, NUM_BUFFERS) [Hopper+, MI300+]

Allocate NUM_BUFFERS buffers in local memory per thread block, each of the specified size. The memory layout is inferred from its consumers.
buffers = tlx.local_alloc(shape, dtype, NUM_BUFFERS, tlx.storage_kind.tmem) [Blackwell]

Allocate NUM_BUFFERS of buffers in the tensor memory per thread block, each with size size. The memory layout is inferred from its consumers.
buffers = tlx.local_alloc(shape, dtype, NUM_BUFFERS, reuse=other_buffers) [Hopper+, MI300+]

Alias this allocation to an existing buffered_tensor so multiple logical buffers reuse the same underlying local storage (SMEM or TMEM) without reallocation.
buffer = tlx.local_view(buffers, buffer_idx) or buffer = buffers[buffer_idx] [Hopper+, MI300+]

Return a subview of the buffer indexed by buffer_idx from buffers. Both the explicit local_view() call and the indexing syntax [] are supported.
distributed_tensor = tlx.local_load(buffer, token=None, layout=None, relaxed=False, rematerialize_coordinates=False, rematerialize_coordinates_group=None) [Hopper+, MI300+]

Loads the buffer from local or tensor memory. layout pins a requested register layout. On AMD, the rematerialization options start fresh address coordinate live ranges either per load or for a named group of nearby loads; the two options are mutually exclusive.
tlx.local_store(buffer, distributed_tensor) [Hopper+, MI300+]

Store a distributed tensor into a buffer in local memory or tensor memory.
distributed_tensor = tlx.local_gather(src, indices, axis, optional_token) [Hopper+, MI300+?]

Gather elements from shared memory along a specified axis using an indices tensor. The output shape matches the indices shape, and elements are gathered from src at positions specified by indices along the given axis.
tlx.local_scatter(dst, src, indices, axis, optional_token) [Hopper+, MI300+?]

Scatter elements to shared memory along a specified axis using an indices tensor. Elements from src are written to dst at positions specified by indices along the given axis.
buffer = tlx.local_trans(buffer, dims) [Hopper+, MI300+]

Permutes the dimensions of a tensor.
buffer = tlx.local_slice(buffer, offsets=[m, n], shapes=[M, N]) [Hopper+, MI300+]

Slice a M x N tensor at a m x n offset.

Buffer Reuse

TLX provides you the ability to reuse the same allocated buffer across multiple disjoint steps in your kernel. This is useful to allow additional pipelining when you may not have enough isolated SMEM or TMEM.

tlx.storage_alias_spec(storage=storage_kind) [Hopper+, MI300+]

Defines a buffer that you will want to share across multiple aliases. The storage can be either SMEM or TMEM. To use this in an allocation you should provide the spec in the reuse argument for local_alloc. Here is the example from the FA kernel.

# Create the storage alias spec for all shared buffers. Cannot be directly
# indexed.
qk_storage_alias = tlx.storage_alias_spec(storage=tlx.storage_kind.tmem)

# Allocate all buffers referencing the same spec
qk_tiles = tlx.local_alloc(
    (BLOCK_M_SPLIT, BLOCK_N), qk_dtype, NUM_MMA_GROUPS,
    tlx.storage_kind.tmem, reuse=qk_storage_alias,
)
p_tiles = tlx.local_alloc(
    (BLOCK_M_SPLIT, BLOCK_N // NUM_MMA_SLICES), tlx.dtype_of(desc_v),
    NUM_MMA_GROUPS * NUM_MMA_SLICES, tlx.storage_kind.tmem,
    reuse=qk_storage_alias,
)
alpha_tiles = tlx.local_alloc(
    (BLOCK_M_SPLIT, 1), tl.float32, NUM_MMA_GROUPS * NUM_BUFFERS_QK,
    tlx.storage_kind.tmem, reuse=qk_storage_alias,
)
l_tiles = tlx.local_alloc(
    (BLOCK_M_SPLIT, 1), tl.float32, NUM_MMA_GROUPS * NUM_BUFFERS_QK,
    tlx.storage_kind.tmem, reuse=qk_storage_alias,
)
m_tiles = tlx.local_alloc(
    (BLOCK_M_SPLIT, 1), tl.float32, NUM_MMA_GROUPS * NUM_BUFFERS_QK,
    tlx.storage_kind.tmem, reuse=qk_storage_alias,
)

tlx.reuse_group(*tensors, group_type=REUSE_TYPE, group_size=SUBTILE_SIZE) [Hopper+, MI300+]

A reuse group expresses how you intend to access the shared buffer. There are two types: Shared or Distinct. A shared buffer wants to occupy the same memory and each index should not be accessed at the same time. A distinct buffer will be accessible at the same index at the same time. The compiler will isolate buffer locations and potentially expand the buffer allocation to enforce this guarantee, which is helpful with buffers of unequal sizes.

The group_size is used to enable subtiling a buffer. This ensures that for every 1 index of a buffer that SUBTILE_SIZE indices of this other buffer/group can be accessed. Reuse groups can be nested to allow expressing more complex relationships. Currently a reuse group is not applied unless you assign it to a buffer with spec.set_buffer_overlap.

Here is the example implementation for Flash Attention. In this kernel as the comment suggests, QK is shared with P, l, m, and alpha, and P is potentially subtiling.

# Define the buffer overlap strategy:
#   QK : |                                                   BLK_M/2 * BLOCK_N * fp32                         |
#   P:   |  BLK_M/(2*SLICES) * fp16| BLK_M/(2*SLICES) * fp16|...
# Alpha:                                                        |BLK_M/2*1*fp32|
#   l  :                                                                        |BLK_M/2*1*fp32|
#   m  :                                                                                       |BLK_M/2*1*fp32|
qk_storage_alias.set_buffer_overlap(
    tlx.reuse_group(
        qk_tiles,
        tlx.reuse_group(
            tlx.reuse_group(p_tiles, group_size=NUM_MMA_SLICES),
            alpha_tiles, l_tiles, m_tiles,
            group_type=tlx.reuse_group_type.distinct,
        ),
        group_type=tlx.reuse_group_type.shared,
    )
)

Compiler Pipeline Inspection Steps To introspect the pipeline add_stages, before running your kernels, simply set the add_stages_inspection_hook like so:

def inspect_stages(_self, stages, options, language, capability):
    # inspect or modify add_stages here
triton.knobs.runtime.add_stages_inspection_hook = inspect_stages

Examples of how to use this for out of tree plugin passes is here

Binary wheels are available for CPython 3.10-3.14.

Remote buffer operations

buffer = tlx.remote_view(buffer, remote_cta_rank) [Hopper+]

Return a remote view of the buffer living in another CTA in the same cluster with ID remote_cta_rank. NOTE: for now we only support barrier as buffer, not general SMEM.
tlx.remote_shmem_store(dst, src, remote_cta_rank) [Hopper+]

Store a distributed tensor into a buffer in the remote shared memory of a cluster (synchronous).

Parameters:
- dst: The destination buffer in local shared memory (will be internally mapped to the remote CTA)
- src: The source distributed tensor to store
- remote_cta_rank: The rank (unique ID) of the remote CTA within the cluster
Example:
```
# Allocate shared memory buffer
buffer = tlx.local_alloc((BLOCK_M, BLOCK_N), tl.float16, 1)

# Store to remote CTA's shared memory (synchronous)
tlx.remote_shmem_store(buffer[0], src_tensor, remote_cta_rank=1)
```

Async memory access

tlx.async_descriptor_load(desc, buffer, offsets, barrier, pred=None, cache_modifier="", eviction_policy="", multicast_targets=[]) [Hopper+]

Load a chunk of data from global memory into a local memory buffer using TMA. The global address, strides, and buffer size are defined by the tensor descriptor. A barrier object is provided and signaled upon completion of the operation.

Parameters:
- desc: Tensor descriptor for the source
- buffer: Destination buffer in shared memory
- offsets: List of offsets for each dimension
- barrier: mbarrier to signal upon completion
- pred: Optional predicate to guard the load
- cache_modifier: Cache modifier hint (e.g., "", "evict_first")
- eviction_policy: L2 cache eviction policy ("", "evict_first", "evict_last")
- multicast_targets: Optional list of multicast targets for cluster-wide loads
tlx.async_descriptor_prefetch_tensor(memdesc, [offsets], pred, eviction_policy) [Hopper+]

Hint hardware to load a chunk of data from global memory into a L2 cache to prepare for upcoming async_descriptor_load operations.
tlx.async_descriptor_store(desc, source, offsets, eviction_policy="", store_reduce="") [Hopper+]

Store a chunk of data from shared memory into global memory using TMA. The global address, strides, and buffer size are defined by the tensor descriptor.

Supports optional atomic reduction (store_reduce) and L2 cache eviction hints (eviction_policy). Both regular stores and atomic reduce stores support cache eviction policies.

Parameters:
- desc: Tensor descriptor for the destination
- source: Source buffer in shared memory
- offsets: List of offsets for each dimension
- eviction_policy: L2 cache eviction policy ("", "evict_first", "evict_last")
- store_reduce: Atomic reduction kind ("", "add", "min", "max", "and", "or", "xor")
Example:
```
# Regular TMA store with L2 evict_first hint
tlx.async_descriptor_store(desc_c, c_buf[0], [offs_m, offs_n], eviction_policy="evict_first")

# TMA atomic reduce-add with L2 evict_first hint
tlx.async_descriptor_store(desc_c, c_buf[0], [offs_m, offs_n],
                           eviction_policy="evict_first", store_reduce="add")
```
tlx.async_remote_shmem_store(dst, src, remote_cta_rank, barrier) [Hopper+]

Store a distributed tensor into a buffer in the remote shared memory of a cluster asynchronously. Signals the provided mbarrier when the store completes.

Parameters:
- dst: The destination buffer in local shared memory (will be internally mapped to the remote CTA)
- src: The source distributed tensor to store
- remote_cta_rank: The rank (unique ID) of the remote CTA within the cluster
- barrier: mbarrier to signal when the store completes
Example:
```
# Allocate shared memory buffer and barrier
buffer = tlx.local_alloc((BLOCK_M, BLOCK_N), tl.float16, 1)
barrier = tlx.alloc_barriers(num_barriers=1, arrive_count=1)

# Store to remote CTA's shared memory
tlx.async_remote_shmem_store(buffer[0], src_tensor, remote_cta_rank=1, barrier=barrier[0])
```
tlx.remote_shmem_copy(dst, src, remote_cta_rank) [Hopper+]

Store a local shared memory buffer into a buffer in the remote shared memory of a cluster asynchronously.

Parameters:
- dst: The destination buffer in local shared memory (will be internally mapped to the remote CTA)
- src: The source distributed tensor to store
- remote_cta_rank: The rank (unique ID) of the remote CTA within the cluster
- barrier: mbarrier to signal when the store completes (will be internally mapped to the remote CTA)
Example:
```
# Allocate shared memory buffer
buffer0 = tlx.local_alloc((BLOCK_M, BLOCK_N), tl.float16, 1)
buffer1 = tlx.local_alloc((BLOCK_M, BLOCK_N), tl.float16, 1)
barrier = tlx.alloc_barriers(num_barriers=1, arrive_count=1)

# Copy to remote CTA's shared memory
tlx.remote_shmem_store(buffer0[0], buffer1[0], remote_cta_rank=1, barrier=barrier[0])
```
desc_ptrs = tlx.allocate_tensor_descriptor(num) [Hopper+]

Allocates global memory for tensor descriptor storage with built-in parameters (nbytes=128, alignment=128 per descriptor). Returns a tensor_descriptor_ptr with 128-byte stride semantics that supports indexing.

Parameters:
- num: Number of tensor descriptors to allocate (must be a constexpr)
Returns:
- A tensor_descriptor_ptr where indexing (e.g., desc_ptrs[0], desc_ptrs[1]) advances by 128 bytes per index
Example:
```
# Allocate storage for 4 tensor descriptors
desc_ptrs = tlx.allocate_tensor_descriptor(num=4)

# Access individual descriptors using indexing
desc_ptr_0 = desc_ptrs[0]  # First descriptor
desc_ptr_1 = desc_ptrs[1]  # Second descriptor (128 bytes offset)
```

tlx.make_tensor_descriptor(desc_ptr, base, shape, strides, block_shape, padding_option) [Hopper+]

Create a TMA (Tensor Memory Accelerator) descriptor for efficient asynchronous data movement on Hopper and Blackwell GPUs.

Parameters:

desc_ptr (optional): Tensor descriptor pointer from allocate_tensor_descriptor(). Pass None for automatic allocation.
base: Base pointer to the tensor in global memory
shape: List of tensor dimensions (dynamic, runtime values)
strides: List of tensor strides (dynamic, runtime values)
block_shape: Shape of the block to be loaded/stored (compile-time constants)
padding_option: Padding option for out-of-bounds accesses (default: "zero")

Example:

# Create a 2D tensor descriptor with automatic scratch allocation
desc = tlx.make_tensor_descriptor(
    desc_ptr=None,  # Compiler allocates scratch memory automatically
    base=tensor_ptr,
    shape=[M, N],
    strides=[N, tl.constexpr(1)],
    block_shape=[64, 64],
)

# Or with explicit descriptor allocation for advanced use cases (e.g., pipelining)
desc_ptrs = tlx.allocate_tensor_descriptor(num=2)

# Create descriptor at index 0
tlx.make_tensor_descriptor(
    desc_ptr=desc_ptrs[0],
    base=tensor_ptr,
    shape=[M, N],
    strides=[N, tl.constexpr(1)],
    block_shape=[64, 64],
)

# Reinterpret the descriptor for TMA operations
desc = tlx.reinterpret_tensor_descriptor(
    desc_ptr=desc_ptrs[0],
    block_shape=[64, 64],
    dtype=tl.float16,
)

# Use with async TMA operations
tlx.async_descriptor_load(desc, buffer, offsets=[m_offset, n_offset], barrier=mbar)

desc = tlx.reinterpret_tensor_descriptor(desc_ptr, block_shape, dtype) [Hopper+, MI300+]

Reinterpret a tensor descriptor pointer as a TMA-backed tensor descriptor object.

Parameters:

desc_ptr: A tensor_descriptor_ptr pointing to the TMA descriptor (from allocate_tensor_descriptor)
block_shape: Shape of the block to be loaded/stored (compile-time constants)
dtype: Data type of the tensor elements

Example:

# Allocate and create descriptor
desc_ptrs = tlx.allocate_tensor_descriptor(num=2)
tlx.make_tensor_descriptor(desc_ptr=desc_ptrs[0], base=a_ptr, shape=[M, K], strides=[K, 1], block_shape=[128, 64])

# Reinterpret for use with TMA
a_desc = tlx.reinterpret_tensor_descriptor(desc_ptr=desc_ptrs[0], block_shape=[128, 64], dtype=tl.float16)
tlx.async_descriptor_load(a_desc, buffer, offsets=[offs_m, offs_k], barrier=mbar)

tlx.async_load(tensor_ptr, buffer, optional_mask, optional_other, cache_modifier, eviction_policy, is_volatile) [Hopper+, MI350]

Load a chunk of data from global memory into a local memory buffer asynchronously.

The operation returns a token object which can be used to track the completion of the operation.

MI350X caveat: When mask is provided, callers should also provide other—typically 0.0 for numerical kernels. Elements for which mask=True are copied from global memory into the destination local buffer. For elements where mask=False, other is written only if it is provided. Otherwise, the corresponding LDS locations retain unspecified, potentially stale contents from an earlier use of the buffer.

For example, in FlashAttention, when the sequence length is not a multiple of BLOCK_N, the final V tile contains masked-out rows. If other is omitted, those rows may contain unspecified values, including bit patterns representing NaN or infinity. These values can propagate through the subsequent matrix multiplication and produce incorrect output. Use other=0.0 for such padded tiles.
tlx.async_load_commit_group(tokens) [Hopper+, MI350]

Commits all prior initiated but uncommitted async_load ops an async group. Optionally, each token represents a tracked async load operation.
tlx.async_load_wait_group(pendings, tokens) [Hopper+, MI350]

Wait for completion of prior asynchronous copy operations. The pendings argument indicates the number of in-flight operations not completed. Optionally, each token represents a tracked async commit group operation.

Async tensor core operations

acc = tlx.async_dot(a[i], b[i], acc) [Hopper+]
acc = tlx.async_dot(a_reg, b[i], acc) [Hopper]
acc[i] = tlx.async_dot(a[i], b[i], acc[i], barrier) [Blackwell]

acc[i] = tlx.async_dot_scaled(a[i], b[i], acc[i], a_scale[i], a_format, b_scale[i], b_format, use_acc, two_ctas, mBarriers) [Blackwell]

Parameters:

a[i]: A tile in shared memory (FP8 format)
b[i]: B tile in shared memory (FP8 format)
acc[i]: Accumulator tile in tensor memory (TMEM)
a_scale[i]: Per-block scaling factors for A (E8M0 format in SMEM)
a_format: FP8 format string for A: "e4m3", "e5m2", or "e2m1"
b_scale[i]: Per-block scaling factors for B (E8M0 format in SMEM)
b_format: FP8 format string for B: "e4m3", "e5m2", or "e2m1"
use_acc: If True, compute D = A@B + D; if False, compute D = A@B
two_ctas: If True, enables 2-CTA collective MMA (generates tcgen05.mma.cta_group::2)
mBarriers: Optional list of mbarriers for MMA completion signaling

2-CTA Scaled MMA: When two_ctas=True, the scaled MMA operates across two CTAs in a cluster. Key considerations:

B data is split: Each CTA loads half of B (BLOCK_N // 2)
B scale is NOT split: Both CTAs need the full B scale for correct MMA computation
CTA synchronization: Use "Arrive Remote, Wait Local" pattern before MMA
MMA predication: Compiler auto-generates predicate so only CTA 0 issues the MMA

Example: 2-CTA Scaled MMA

# B data split across CTAs, but B scale is full
desc_b = tl.make_tensor_descriptor(b_ptr, ..., block_shape=[BLOCK_K, BLOCK_N // 2])
desc_b_scale = tl.make_tensor_descriptor(b_scale_ptr, ..., block_shape=[BLOCK_N // 128, ...])  # Full scale

# Load B with CTA offset, B scale without offset
tlx.async_descriptor_load(desc_b, b_tile[0], [0, cluster_cta_rank * BLOCK_N // 2], bar_b)
tlx.async_descriptor_load(desc_b_scale, b_scale_tile[0], [0, 0, 0, 0], bar_b_scale)  # Full B scale

# CTA sync: "Arrive Remote, Wait Local"
tlx.barrier_arrive(cta_bars[0], 1, remote_cta_rank=0)
tlx.barrier_wait(cta_bars[0], phase=0, pred=pred_cta0)

# 2-CTA scaled MMA with mBarriers for completion tracking
tlx.async_dot_scaled(
    a_tile[0], b_tile[0], c_tile[0],
    a_scale_tile[0], "e4m3",
    b_scale_tile[0], "e4m3",
    use_acc=False,
    two_ctas=True,
    mBarriers=[mma_done_bar],
)
tlx.barrier_wait(mma_done_bar, tl.constexpr(0))

Alternative: Using tcgen05_commit for MMA completion

# Issue MMA without mBarriers
tlx.async_dot_scaled(..., two_ctas=True)

# Use tcgen05_commit to track all prior MMA ops
tlx.tcgen05_commit(mma_done_bar, two_ctas=True)
tlx.barrier_wait(mma_done_bar, tl.constexpr(0))

TMEM-backed MX Scales:

For scaled MMA operations on Blackwell GPUs, scales can be stored in Tensor Memory (TMEM) for efficient access. TLX provides automatic layout resolution for TMEM scale buffers.

Allocating TMEM Scale Buffers:

When allocating TMEM buffers for uint8/int8 types (used for MX scales), TLX uses a placeholder layout (DummyTMEMLayoutAttr) that gets automatically resolved to TensorMemoryScalesEncodingAttr during compilation when the buffer is used with async_dot_scaled.

# Allocate TMEM buffers for scales (layout is automatically resolved)
a_scale_tmem = tlx.local_alloc((128, 8), tl.uint8, num=1, storage=tlx.storage_kind.tmem)
b_scale_tmem = tlx.local_alloc((256, 4), tl.uint8, num=1, storage=tlx.storage_kind.tmem)

Copying Scales from SMEM to TMEM:

Use tlx.tmem_copy [Blackwell] to efficiently transfer scale data from shared memory to tensor memory:

# Copy scales from SMEM to TMEM (asynchronous, uses tcgen05.cp instruction)
tlx.tmem_copy(a_scale_smem, a_scale_tmem)
tlx.tmem_copy(b_scale_smem, b_scale_tmem)

Using TMEM Scales with Scaled MMA:

# TMEM scales are automatically detected and used with the correct layout
tlx.async_dot_scaled(
    a_smem, b_smem, acc_tmem,
    A_scale=a_scale_tmem, A_format="e4m3",
    B_scale=b_scale_tmem, B_format="e4m3",
    use_acc=True,
    mBarriers=[mma_bar],
)

Complete Example: TMEM-backed Scaled GEMM:

@triton.jit
def scaled_gemm_kernel(...):
    # Allocate TMEM for accumulator and scales
    acc = tlx.local_alloc((BLOCK_M, BLOCK_N), tl.float32, num=1, storage=tlx.storage_kind.tmem)
    a_scale_tmem = tlx.local_alloc((BLOCK_M // 128, BLOCK_K // 32), tl.uint8, num=1, storage=tlx.storage_kind.tmem)
    b_scale_tmem = tlx.local_alloc((BLOCK_N // 128, BLOCK_K // 32), tl.uint8, num=1, storage=tlx.storage_kind.tmem)

    # Load scales from global memory to SMEM
    tlx.async_descriptor_load(a_scale_desc, a_scale_smem, [...], barrier=bar)
    tlx.async_descriptor_load(b_scale_desc, b_scale_smem, [...], barrier=bar)
    tlx.barrier_wait(bar, phase)

    # Copy scales from SMEM to TMEM
    tlx.tmem_copy(a_scale_smem[0], a_scale_tmem[0])
    tlx.tmem_copy(b_scale_smem[0], b_scale_tmem[0])

    # Perform scaled MMA with TMEM scales
    tlx.async_dot_scaled(
        a_smem[0], b_smem[0], acc[0],
        A_scale=a_scale_tmem[0], A_format="e4m3",
        B_scale=b_scale_tmem[0], B_format="e4m3",
        use_acc=False,
    )

Note: Multibuffering is automatically cancelled for scale buffers since TMEM scales don't support multibuffering. 3D allocations (1×M×K) are automatically flattened to 2D (M×K).

acc = tlx.async_dot_wait(pendings, acc) [Hopper+]

Wait for completion of prior asynchronous dot operations. The pendings argument indicates the number of in-flight operations not completed.

Example:
```
acc = tlx.async_dot(a_smem, b_smem)
acc = tlx.async_dot_wait(tl.constexpr(0), acc)
tl.store(C_ptrs, acc)
```

Barrier operations

barriers = tlx.alloc_barrier(num_barriers, arrive_count=1) [Hopper+]

Allocates buffer in shared memory and initialize mbarriers with arrive_counts.

Input:
- num_barriers: The number of barriers to allocate.
- arrive_counts: The number of threads that need to arrive at the barrier before it can be released.
tlx.barrier_wait(bar, phase) [Hopper+]

Wait until the mbarrier phase completes
tlx.barrier_arrive(bar, arrive_count=1) [Hopper+]

Perform the arrive operation on an mbarrier
tlx.named_barrier_wait(bar_id, num_threads) [Hopper+]

Wait until num_threads threads have reached the specified named mbarrier phase.
tlx.named_barrier_arrive(bar_id, num_threads) [Hopper+]

Signal arrival at a named mbarrier with the given thread count.
tlx.barrier_expect_bytes(bar, bytes) [Hopper+]

Signal a barrier of an expected number of bytes to be copied.
tlx.barrier_arrive(bar, arrive_count=1, remote_cta_rank=None) [Hopper+]

Perform the arrive operation on an mbarrier. If remote_cta_rank is provided, signals the barrier in the specified remote CTA's shared memory (useful for multi-CTA synchronization).
tlx.amd_sched_barrier(mask=0) [AMD]

Prevents selected AMD machine-instruction classes from crossing a source boundary. It is a scheduling marker, not a workgroup barrier or memory fence.

Memory Fences

tlx.fence(scope) [Hopper+] issues a memory fence. The scope argument is required:

Scope	PTX	Description
`"gpu"`	`fence.acq_rel.gpu`	Device-scope fence. Orders prior global/shared memory writes to be visible to all GPU threads.
`"sys"`	`fence.acq_rel.sys`	System-scope fence. Like `"gpu"` but also visible to the host CPU.
`"async_shared"`	`fence.proxy.async.shared::cta`	Proxy fence for async shared memory. Required between `local_store` and a subsequent TMA store (`async_descriptor_store`) to the same shared memory.

Example:

tlx.local_store(smem_buf, data)
tlx.fence("async_shared")
tlx.async_descriptor_store(desc, smem_buf, offsets)

tlx.fence_mbarrier_init_cluster(scope) [Hopper+] issues a memory fence to make mbarrier init visible to cluster.

Example:

bars = tlx.alloc_barriers(num_barriers=1, arrive_count=1)
tlx.fence_mbarrier_init_cluster()
tlx.cluster_barrier()

# now bars is ready for cross CTA use
tlx.barrier_arrive(bar=bars[0], remote_cta_rank=1)

Cluster Launch Control (CLC)

CLC (Cluster Launch Control) is a Blackwell-specific feature [Blackwell] that enables dynamic persistent kernel execution with efficient work stealing across thread blocks. It allows CTAs to dynamically acquire tile IDs from a hardware-managed work queue, enabling load balancing without explicit inter-CTA communication.

CLC API

context = tlx.clc_create_context(num_consumers=num_consumers) [Blackwell]

Create a CLC pipeline context with the specified number of stages and expected consumer count.

Parameters:
- num_consumers: Number of consumers that will signal completion per tile (typically 3 async tasks × num_CTAs)
tlx.clc_producer(context, p_producer=phase, multi_ctas=False) [Blackwell]

Issue a CLC try_cancel request to acquire a new tile ID.

Parameters:
- context: CLC pipeline context from clc_create_context
- phase: Current barrier phase (0 or 1, alternates each iteration)
- multi_ctas: Set to True for 2-CTA mode (cluster of 2 CTAs). When enabled, pred_cta0 is computed internally from cluster_cta_rank().
tile_id = tlx.clc_consumer(context, p_consumer=phase, multi_ctas=False, k=0, return_3d=False) [Blackwell]

Decode the tile ID from a CLC response and signal completion.

Parameters:
- context: CLC pipeline context from clc_create_context
- phase: Current barrier phase
- multi_ctas: Set to True for 2-CTA mode. When enabled, pred_cta0 is computed internally.
- return_3d: Set to True to return (ctaIdX, ctaIdY, ctaIdZ) tuple instead of scalar tile_id.
Returns: The tile ID (already offset by cluster_cta_rank() for unique tile assignments), or -1 if no work available. With return_3d=True, returns (ctaIdX, ctaIdY, ctaIdZ) tuple.

How CLC Works

CLC uses hardware-assisted work stealing via the PTX instruction:

clusterlaunchcontrol.try_cancel.async.shared::cta.mbarrier::complete_tx::bytes.multicast::cluster::all.b128

The .multicast::cluster::all qualifier means the response is asynchronously written to all CTAs in the cluster. This enables efficient multi-CTA execution where all CTAs in a cluster receive the same base tile ID.

CLC Synchronization Flow

┌─────────────────────────────────────────────────────────────────┐
│                    CLC Producer (clc_producer)                  │
├─────────────────────────────────────────────────────────────────┤
│  1. WAIT:   barrier_wait(bar_empty)      ← Wait for consumers   │
│  2. EXPECT: barrier_expect_bytes(bar_full, 16)                  │
│  3. ISSUE:  clc_issue(response, bar_full) ← Hardware request    │
└─────────────────────────────────────────────────────────────────┘
                              ↓
                    [Hardware processes CLC]
                    [Multicasts response to all CTAs]
                              ↓
┌─────────────────────────────────────────────────────────────────┐
│                    CLC Consumer (clc_consumer)                  │
├─────────────────────────────────────────────────────────────────┤
│  1. WAIT:   barrier_wait(bar_full)       ← Wait for response    │
│  2. QUERY:  tile_id = clc_query(response) ← Extract tile ID     │
│  3. SIGNAL: barrier_arrive(bar_empty)    ← Release producer     │
└─────────────────────────────────────────────────────────────────┘

Multi-CTA Mode (2-CTA Clusters)

In multi-CTA mode (multi_ctas=True), multiple CTAs in a cluster work together on adjacent tiles. The key constraint is: you can arrive at a remote mbarrier, but you cannot wait on a remote mbarrier (per NVIDIA specification).

Key Principle: "Arrive Remote, Wait Local"

Operation	Local mbarrier	Remote mbarrier
`barrier_wait`	✅ Allowed	❌ Undefined behavior
`barrier_arrive`	✅ Allowed	✅ Allowed (via `remote_cta_rank`)

Example: Multi-CTA GEMM with CLC

@triton.jit
def matmul_kernel(..., PAIR_CTA: tl.constexpr):
    # Create CLC context: 6 consumers for 2-CTA mode (3 tasks × 2 CTAs)
    clc_context = tlx.clc_create_context(num_consumers= 6 if PAIR_CTA else 3)

    with tlx.async_tasks():
        with tlx.async_task("default"):  # Epilogue consumer
            clc_phase_producer = 1
            clc_phase_consumer = 0
            tile_id = start_pid

            while tile_id != -1:
                # Producer: acquire next tile
                tlx.clc_producer(clc_context, p_producer=clc_phase_producer, multi_ctas=PAIR_CTA)
                clc_phase_producer ^= 1

                # ... process tile ...

                # Consumer: get tile ID and signal completion
                tile_id = tlx.clc_consumer(clc_context, p_consumer=clc_phase_consumer, multi_ctas=PAIR_CTA)
                clc_phase_consumer ^= 1
        with tlx.async_task(num_warps=1, num_regs=24):  # MMA consumer
            clc_phase_consumer = 0
            tile_id = start_pid

            while tile_id != -1:
                # ... process tile ...

                # Consumer: get tile ID and signal completion
                tile_id = tlx.clc_consumer(clc_context, p_consumer=clc_phase_consumer, multi_ctas=PAIR_CTA)
                clc_phase_consumer ^= 1
        with tlx.async_task(num_warps=1, num_regs=24):  # producer, TMA load
            clc_phase_consumer = 0
            tile_id = start_pid

            while tile_id != -1:
                # ... process tile ...

                # Consumer: get tile ID and signal completion
                tile_id = tlx.clc_consumer(clc_context, p_consumer=clc_phase_consumer, multi_ctas=PAIR_CTA)
                clc_phase_consumer ^= 1

Examples: how mbarriers are communicated in warp specialization

    phase = 0
    with tlx.async_tasks():
        with tlx.async_task("default"):

            tlx.barrier_wait(bar=b1, phase=phase ^ 1)

            # Placeholder block to do something

            tlx.barrier_arrive(bar=b0)  # Release

        with tlx.async_task(num_warps=4):

            tlx.barrier_wait(bar=b0, phase=phase)  # Wait

            # Some arith ops TODO. add WS
            offsets = block_start + tl.arange(0, BLOCK_SIZE)
            mask = offsets < n_elements
            x = tl.load(x_ptr + offsets, mask=mask)
            z = x * x
            tl.store(z_ptr + offsets, z, mask=mask)

            tlx.barrier_arrive(bar=b0)  # Wait

Warp Specialization operations

tlx.async_tasks and tlx.async_task [Hopper+]

    with tlx.async_tasks()
        with tlx.async_task("default")
            ...
        with tlx.async_task(num_warps=4)
            ...

tlx.async_tasks opens a multi-tasking region where independent asynchronous tasks can be declared. Each task executes in parallel using a dedicated subset of warps within the thread block.

tlx.async_task("default") defines the default task, also known as the trunk. It uses the available warps not explicitly reserved by other tasks.

tlx.async_task(num_warps=4) defines a warp-specialized asynchronous task that explicitly reserves 4 warps in addition to those used by the trunk task.

async_tasks Parameters

Parameter	Description
`exclusive`	Assert this is the only one `tlx.async_tasks` in the kernel for more efficient PTX. Default to False.
`no_ending_cluster_sync`	This suppresses compiler generated cluster sync at end of Warp Spec. Should only be used if user guarantees all cross CTA SMEM/TMEM access are done by end of WS default task. Default to False.
`mbarrier_try_wait_suspend_ns`	On Blackwell, use the four-operand `mbarrier.try_wait.parity` form with this suspend hint for waits in the kernel. `None` is unspecified, `0` explicitly disables the hint, and positive values enable it. If multiple `async_tasks` regions specify a value, the minimum explicit value is used module-wide. Default to None.

async_task Parameters

Parameter	Description
`"default"`	First positional argument to mark this as the default/trunk task
`num_warps`	Number of warps to reserve for this task
`num_regs`	Number of registers per thread (optional, for register allocation tuning). It is supported by both default and non-default tasks and must be divisible by 8.
`replicate`	Number of replicas for this task (default: 1). Creates multiple copies of the task region
`warp_group_start_id`	Starting warp ID for this task (optional). Allows explicit control over warp assignment

Default Task Register Budget

Register budgets are specified per thread. When the default task does not set num_regs, register allocation keeps the original donation model: non-default warp groups receive their requested budgets, and the default task receives the remaining registers.

Setting num_regs on the default task selects a fixed budget instead:

with tlx.async_tasks():
    with tlx.async_task("default", num_regs=80):
        ...
    with tlx.async_task(num_warps=4, num_regs=24):
        ...

In this example, the default and non-default tasks receive 80 and 24 registers per thread, respectively. The default task does not absorb unused registers. Non-default warp groups without an explicit budget evenly share the remaining register pool. If every task has a fixed budget, any surplus is left unused. The compiler may raise a request to the hardware or instrumentation safety minimum when required.

Explicit Warp Assignment with warp_group_start_id

By default, the compiler automatically assigns warp IDs to each task. However, you can use warp_group_start_id to explicitly specify which warps each task should use. This is useful for:

Fine-grained control over warp-to-task mapping
Ensuring specific hardware resource allocation
Advanced optimization scenarios

Example:

with tlx.async_tasks():
    with tlx.async_task("default"):  # Uses warps 0-3 (from num_warps=4 kernel param)
        # Producer task
        ...
    with tlx.async_task(num_warps=2, warp_group_start_id=4, replicate=2):
        # Two replicas, each using 2 warps
        # Replica 0: warps 4-5
        # Replica 1: warps 6-7
        ...
    with tlx.async_task(num_warps=1, warp_group_start_id=8):
        # Consumer task using warp 8
        ...

Validation Rules:

Warp ranges must not overlap between tasks
Non-default tasks must not overlap with the default region (warps 0 to kernel's num_warps)
When using warp_group_start_id, it must be specified for ALL non-default tasks or NONE

CUDA Thread Block Clustering

TLX supports CUDA Thread Block Clustering (available on SM90+ Hopper/Blackwell GPUs) through the ctas_per_cga parameter. This provides explicit control over cluster dimensions for multi-CTA cooperative kernels.

Usage

Pass ctas_per_cga as a tuple when launching a kernel:

kernel[(grid_x, grid_y)](
    ...,
    ctas_per_cga=(2, 1, 1),  # 2x1x1 cluster of CTAs
    **kwargs
)

Using ctas_per_cga with Autotune

You can specify ctas_per_cga in triton.Config for autotuning:

@triton.autotune(
    configs=[
        triton.Config(
            {"BLOCK_M": 128, "BLOCK_N": 128},
            num_warps=4,
            ctas_per_cga=(2, 1, 1),  # 2x1x1 cluster
        ),
        triton.Config(
            {"BLOCK_M": 64, "BLOCK_N": 64},
            num_warps=4,
            ctas_per_cga=(1, 1, 1),  # No clustering
        ),
    ],
    key=["M", "N", "K"],
)
@triton.jit
def matmul_kernel(...):
    ...

TLX vs Triton Semantics

TLX uses CUDA-native cluster semantics which differs from Triton's approach:

Aspect	Triton's way (`num_ctas`)	TLX way (`ctas_per_cga`)
Grid interpretation	Grid × cluster_dims = total CTAs	Grid = total CTAs
Cluster definition	Multiplicative	Regrouping
`num_ctas` value	`product(cluster_dims)`	Always 1
`launch_cluster`	Can be False (enabled by `num_ctas != 1`)	Always True

Other operations

tlx.cluster_cta_rank() [Hopper+]

Returns the rank (unique ID) of the current CTA within the cluster.
tlx.thread_id(axis) [Hopper+]

Returns the id of the current thread instance along the given axis.
tlx.dtype_of(v) [Hopper+]

Returns the dtype of a tensor or tensor descriptor.

tlx.size_of(dtype) [Hopper+]

Returns the size in bytes of a given Triton dtype. This is useful for dynamically computing memory sizes based on dtype, especially in barrier synchronization code.

Example:

# Instead of hardcoding size values
tlx.barrier_expect_bytes(barrier, 2 * BLOCK_M * BLOCK_K)  # Assumes float16

# Use size_of for dtype-aware computation
tlx.barrier_expect_bytes(barrier,
                       tlx.size_of(tlx.dtype_of(desc)) * BLOCK_M * BLOCK_K)

tlx.clock64() [Hopper+]

Returns the current 64-bit hardware clock value. E.g,

    start = tlx.clock64()
    # ... kernel code ...
    end = tlx.clock64()
    elapsed = end - start  # Number of clock cycles elapsed

tlx.stoch_round(src, dst_dtype, rand_bits) [Blackwell]

Performs hardware-accelerated stochastic rounding for FP32→FP8/BF16/F16 conversions on Blackwell GPUs (compute capability ≥ 100). Uses PTX cvt.rs.satfinite instructions for probabilistic rounding.

Why Use Stochastic Rounding:
- Reduces bias in low-precision training/inference by randomly rounding up or down
- Improves numerical accuracy compared to deterministic rounding (e.g., round-to-nearest-even)
- Particularly beneficial when accumulating many small updates in FP8/FP16
Performance Characteristics:
- Hardware-accelerated: Uses native Blackwell instructions (cvt.rs.satfinite)
- Minimal overhead: Similar throughput to deterministic rounding
- Memory bandwidth: Requires additional random bits (uint32 per element)
Parameters:
- src: Source FP32 tensor
- dst_dtype: Destination dtype (FP8 E5M2, FP8 E4M3FN, BF16, or FP16)
- rand_bits: Random bits (uint32 tensor) for entropy, same shape as src
  - Important: Use n_rounds=7 with tl.randint4x() for sufficient entropy
  - Fewer rounds may result in biased rounding behavior
  - Different seeds produce different rounding decisions for better statistical properties
Example:
```
    # Generate random bits for entropy
    # n_rounds=7 provides sufficient randomness for unbiased stochastic rounding
    offsets = tl.arange(0, BLOCK_SIZE // 4)
    r0, r1, r2, r3 = tl.randint4x(seed, offsets, n_rounds=7)
    rbits = tl.join(tl.join(r0, r1), tl.join(r2, r3)).reshape(x.shape)

    # Apply stochastic rounding
    y = tlx.stoch_round(x, tlx.dtype_of(y_ptr), rbits)
```
tlx.vote_ballot_sync(mask, pred) [Hopper+]

Collects a predicate from each thread in the warp and returns a 32-bit mask where each bit represents the predicate value from the corresponding lane. Only threads specified by mask participate in the vote.
```
    ballot_result = tlx.vote_ballot_sync(0xFFFFFFFF, pred)
```
tlx.prefetch(pointer, level="L2", mask=None, tensormap=False) [Hopper+] issues a non-blocking prefetch hint for pointer-based scattered/gather loads. This complements tlx.async_descriptor_prefetch_tensor (which works on TMA tensor descriptors) by supporting raw pointer tensors. Additionally, if tensormap is specified to True, the API instead does a prefetch of tensor map object (TMA descriptor) and ignores other parameters other than pointer.

Level PTX Description
"L1" prefetch.global.L1 Prefetch into L1 and L2 cache
"L2" prefetch.global.L2 Prefetch into L2 cache only (default)

Example:
```
offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements
tlx.prefetch(input_ptr + offsets, level="L2", mask=mask)
x = tl.load(input_ptr + offsets, mask=mask)

...
# desc_in can be host side descriptor or device side like this:
desc_in = tl.make_tensor_descriptor(
          input_ptr,
          shape=[M, N],
          strides=[N, 1],
          block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_N],
      )
tlx.prefetch(desc_in, tensormap=True)
```
tlx.dump_layout(x) [Hopper+, MI300+]

Compile-time diagnostic that prints the resolved layout of a value to the compiler log. x may be a register tensor or a shared/tensor-memory buffer (memdesc). It emits no device code and returns nothing — this is a static, host-side diagnostic, distinct from the runtime tl.device_print / tl.print. The op is rendered at the end of the TTGIR pipeline, so the printed layout reflects all compiler optimizations, then it is erased.

The layout is printed in CuTe (CUTLASS) Shape:Stride notation (_N marks a static integer):
- Register tensors → a thread-value (TV) layout ((thread...),(value...)):((thread...),(value...)), where the thread group comes from the hardware lane/warp/block dims and the value group from the per-thread registers (stride _0 denotes a broadcast).
- Shared/tensor-memory buffers → a single strided layout, e.g. _64:_1.
- Swizzled shared buffers → Swizzle<B,M,S> o (base):(stride).
In all cases the layout maps a coordinate to the logical tensor's row-major element index (its codomain): for a register tensor a (thread, value) coordinate → the logical element index it holds, and for a buffer an offset → the buffer element index. The strides are offsets in that logical index space, not physical byte/bank addresses.

Layouts that are not representable as a CuTe layout fall back to the raw linear-layout string.

Example:
```
x = tl.load(x_ptr + offs)          # register tensor
tlx.dump_layout(x)                  # -> // cute: ((_32,_2,_2),_1):((_1,_32,_0),_0)

buf = tlx.local_alloc((BLOCK,), tl.float32, 1)
v = tlx.local_view(buf, 0)
tlx.dump_layout(v)                  # -> // cute: _64:_1
```
x = tlx.require_layout(x, layout, pin=True, late_address_compute=False) [Hopper+, MI300+]

Require a register tensor x to use layout, expressed as a tlx.layout(...) (Shape:Stride). With the default pin=True, the #linear encoding is wrapped as #tlx.no_verify_layout(#tlx.user_layout(...)). The inner #tlx.user_layout carries PinnedEncodingTrait, so downstream passes (tritongpu-coalesce, remove-layout-conversions, AMD optimize-epilogue) treat it as fixed and never rewrite it; the outer #tlx.no_verify_layout defers operand-layout verification until the pin is peeled by TLXResolvePlaceholderLayouts in make_ttgir. Example: pin an FP16 epilogue tl.store to a coalesced layout so OptimizeEpilogue keeps the wide buffer_store_dwordx4 instead of narrowing it to the MMA-accumulator store, without staging the value through LDS.

Pass pin=False for an optimizer-flexible requirement. This emits the requested encoding without the #tlx.user_layout hard anchor, allowing later layout passes to propagate the requirement or materialize a layout conversion. The default remains pin=True for existing callers.

On AMD, late_address_compute=True asks shared-memory-backed layout conversions to compute their addresses at this use. The backend does this by rematerializing inexpensive lane/warp coordinates, shortening their live ranges across register-heavy regions.

Pair with tlx.assert_same_layout(x, layout) (below) to statically verify the pin survived to the final TTGIR.
tlx.assert_same_layout(lhs, rhs) [Hopper+, MI300+]

Compile-time assertion that two layouts are equivalent after layout propagation and all other TTGIR layout optimizations have completed. Like tlx.dump_layout, it emits no device code and is consumed at the end of the TTGIR pipeline.

rhs supports two forms:
- Value/value: rhs is another register tensor or shared/tensor-memory buffer. The frontend emits tlx.assert_same_layout, whose two operands retain their independently resolved final types.
- Value/layout: rhs is a constant tlx.layout_encoding. The frontend lowers the constant to an encoding attribute and emits tlx.assert_same_layout_expected. At assertion time, the pass combines that encoding with lhs's shape, element type, and (for buffers) memory properties to construct an expected tensor or memdesc type.
These are separate internal operations only because an SSA value is an MLIR operand while a constant layout is an MLIR attribute. They share the same comparison path and the same public Python API.

Before comparison, both final types are converted with ttg::toLinearLayout. The assertion compares the resulting LinearLayouts, not the original encoding attributes. Consequently, structurally different encodings pass if they describe the same logical mapping. A mismatch reports both normalized LinearLayouts and fails compilation.

Example:
```
x = tlx.local_load(x_buf, layout=REGISTER_LAYOUT)
y = tlx.local_load(y_buf, layout=REGISTER_LAYOUT)

tlx.assert_same_layout(x, y)                # value/value
tlx.assert_same_layout(x, REGISTER_LAYOUT)  # value/layout
```

Level	PTX	Description
`"L1"`	`prefetch.global.L1`	Prefetch into L1 and L2 cache
`"L2"`	`prefetch.global.L2`	Prefetch into L2 cache only (default)

Buffer Operations (AMD)

[MI350] — available on AMD MI350 (CDNA4) only; not available on MI300.

Buffer operations access global memory via a scalar base pointer and a tensor of i32 element offsets, rather than a tensor of pointers. This maps directly to AMD's hardware buffer instructions, which use a resource descriptor and byte offsets, enabling the hardware to do out-of-bounds checking and cache optimization.

`tlx.buffer_load`

Load a tensor of values from global memory.

result = tlx.buffer_load(
    ptr, offsets, mask=None, other=None, cache=None, contiguity=1
)

Argument	Type	Description
`ptr`	scalar pointer	Base address in global memory.
`offsets`	i32 tensor	Per-element byte offsets from `ptr`.
`mask`	bool tensor, optional	When `mask[i]` is `False`, the element is not loaded.
`other`	tensor or scalar, optional	Value used for masked-out elements (where `mask[i]` is `False`).
`cache`	str, optional	Cache modifier (e.g. `".ca"`, `".cg"`).
`contiguity`	constexpr int	Trusted positive power-of-two vectorization width; it must divide each thread's element count.

Returns: A tensor with the same shape as offsets and element type matching the pointee type of ptr.

Lowers to amdg.buffer_load, which is eventually lowered to rocdl.raw.ptr.buffer.load.

`tlx.buffer_store`

Store a tensor of values to global memory.

tlx.buffer_store(stored_value, ptr, offsets, mask=None, cache=None)

Argument	Type	Description
`stored_value`	tensor	Values to write.
`ptr`	scalar pointer	Base address in global memory.
`offsets`	i32 tensor	Per-element byte offsets from `ptr`.
`mask`	bool tensor, optional	When `mask[i]` is `False`, the element is not written.
`cache`	str, optional	Cache modifier.

Returns: Nothing.

Lowers to amdg.buffer_store, which is eventually lowered to rocdl.raw.ptr.buffer.store.

`tlx.buffer_atomic_add`

Atomically add a tensor through an AMD scalar resource descriptor.

previous = tlx.buffer_atomic_add(
    ptr, offsets, value, mask=None, sem=None, scope=None, contiguity=1
)

contiguity is the same trusted per-thread adjacency width accepted by buffer_load; values greater than one also preserve the selected layout. The operation returns the values observed before the additions. FP16 and BF16 use packed two-element instructions, so contiguity must be at least two and any mask must be uniform for each adjacent pair. Compilation fails if mask analysis reduces the legal vector width below two.

`tlx.buffer_load_to_local`

Async load from global memory directly into shared (local) memory, bypassing registers. This is useful for producer warps that prefetch data into shared memory for other warps to consume.

token = tlx.buffer_load_to_local(dest, ptr, offsets, mask=None, other=None, cache_modifier="")

Argument	Type	Description
`dest`	`tlx.buffered_tensor`	Destination slice in shared memory.
`ptr`	scalar pointer	Base address in global memory.
`offsets`	i32 tensor	Per-element byte offsets from `ptr`.
`mask`	bool tensor, optional	When `mask[i]` is `False`, the element is not loaded.
`other`	tensor or scalar, optional	Value used for masked-out elements.
`cache_modifier`	str, optional	Cache modifier string (default `""`).

Returns: A tlx.async_token that can be used with tlx.async_load_wait_group() to synchronize on the completion of the transfer.

Lowers to amdg.buffer_load_to_local, which is eventually lowered to rocdl.raw.ptr.buffer.load.async.lds — a single hardware instruction that moves data from global memory to LDS without going through VGPRs.

Requirements. The direct-to-LDS copy has the following hardware constraints:

Vector width. Each thread's load must reach a supported direct-to-LDS width (32 or 128 bits). If it can only be vectorized to a smaller width, it cannot be lowered.
Provable pointer/offset alignment. The compiler must be able to prove the alignment that vector width needs.
Mask alignment. If mask is given it must be aligned to the vector width: each group of (vector width) consecutive mask values must be identical. The copy transfers each lane's whole vector in one transaction, so a mask whose True/False boundary cannot be proven vector-aligned (e.g. offs < K for a runtime K) forces per-element vectorization and cannot lower.

Example

import triton.language.extra.tlx as tlx

@triton.jit
def kernel(src_ptr, dst_ptr, BLOCK_SIZE: tl.constexpr):
    offsets = tl.arange(0, BLOCK_SIZE).to(tl.int32)
    mask = offsets < BLOCK_SIZE

    # Load from global memory using buffer semantics
    data = tlx.buffer_load(src_ptr, offsets, mask=mask, other=0.0)

    # Store to global memory using buffer semantics
    tlx.buffer_store(data, dst_ptr, offsets, mask=mask)

For the async global-to-shared variant, see the warp-pipeline GEMM example (third_party/amd/python/examples/gluon/f16_gemm_warp_pipeline_gfx1250.py).

Wave Uniformity (AMD)

`tlx.assume_uniform`

Assert that a scalar holds the same value in every lane of the wave.

value = tlx.assume_uniform(value)

Argument	Type	Description
`value`	scalar pointer, or 16/32/64-bit int or float	Value asserted to be wave-uniform. Narrower types are not supported.

Returns: value, unchanged.

The main use is buffer operations, which keep their base pointer in the scalar (SGPR) resource descriptor, so it has to be wave-uniform. When the backend cannot prove that it is — most commonly because the pointer was loaded from memory — it falls back to a per-lane waterfall loop around every access. tlx.assume_uniform tells the backend to take uniformity as given:

base = tl.load(ptr_array + gid).to(tl.pointer_type(tl.float16))
base = tlx.assume_uniform(base)
data = tlx.buffer_load(base, offsets)

Lowers to amdg.assume_uniform, which is eventually lowered to llvm.amdgcn.readfirstlane — that makes the result uniform by construction as far as LLVM's uniformity analysis is concerned. On non-AMD backends it is a no-op that returns its argument. Nothing verifies the assertion: if the value is not actually uniform, every lane silently gets lane 0's value.

Explicit MFMA Scheduling (AMD)

[MI350] — the source-scheduled operations are restricted to CDNA4 (gfx950) native BF16 MFMA layouts.

tl.dot(a, b, acc) preserves a TLX-pinned accumulator layout, so whole dots need no AMD-specific wrapper.
tlx.extract_slice(source, shape, offsets) selects an aligned register fragment without cross-thread movement.
tlx.rematerialized_range(start, end, anchor, placement=None) recreates inexpensive distributed coordinates near a use instead of carrying them through a long software pipeline.
tlx.amd_register_resident(value, register_class="agpr", registers_per_group=1) keeps a distributed value in allocator-visible native register tuples.
tlx.amd_scheduled_mfma(...) exposes independent native MFMA accumulator chains in deterministic N-major, M-minor, K-reduction source order.
tlx.amd_mfma_commit(value, preserve) applies the CDNA4 MFMA result hazard boundary while threading a live dot-operand dependency.
tlx.amd_sched_barrier(mask=0) prevents selected AMD machine-instruction classes from crossing a source boundary. It is a scheduling marker, not a workgroup barrier or memory fence.

These primitives describe fragments, lifetimes, and ordering without assigning physical registers. Their verifiers reject unsupported targets, layouts, element types, and native fragment widths before lowering.

Unlike tl.dot, amd_scheduled_mfma carries an explicit accumulator_role. transient selects the latency-aware intrinsic path for phase-local work; persistent selects register-constrained lowering for a chain carried across phases. Neither role changes the numerical matrix operation.

AMD TDM Descriptor Loads

tlx.update_tensor_descriptor(desc, add_offsets=None, set_bounds=None, pred=None, clamp_bounds=False) produces a positioned descriptor SSA value. add_offsets advances the tile position without changing bounds; set_bounds rewrites absolute bounds; and pred replaces the inherited predicate. Use clamp_bounds=True with add_offsets to derive the remaining OOB extent of an advanced tile.

tlx.async_amd_descriptor_load(desc, result, offsets=None, pred=None, clamp_bounds=True) issues an AMD TDM descriptor load from global memory to a TLX local buffer. If offsets is omitted, desc is used as already positioned and its predicate is preserved. The analogous store accepts offsets=None for the same reason. Both operations are available on TDM-capable AMD targets (gfx1250+) and should be synchronized with tlx.async_amd_descriptor_wait.

tlx.async_amd_descriptor_load_group(descs, results, offsets, warp_masks, preds=None) groups multiple AMD TDM descriptor loads behind one static hardware TDM instruction. Each list entry is one arm:

Argument	Description
`descs[i]`	Tensor descriptor for arm `i`.
`results[i]`	Local buffer or local view receiving arm `i`.
`offsets[i]`	Offset list for arm `i`; all arms must have the same rank.
`warp_masks[i]`	Bitmask selecting the waves that use arm `i`.
`preds[i]`	Optional predicate for arm `i`; defaults to true.

The warp masks must be non-empty, disjoint, axis-aligned, and cover all waves in the CTA exactly once. The grouped operation currently requires one CTA, the same rank and element bitwidth for every arm, one shared cache modifier, and shared layouts supported by AMD TDM lowering. This is useful for kernels where different wave groups load different inputs, such as A/B GEMM tiles or A/B plus scale tiles in MXFP GEMM, while keeping the assembly to one TDM instruction per load group.

Example:

a_tok = tlx.async_amd_descriptor_load_group(
    [a_desc, b_desc],
    [tlx.local_view(a_buf, slot), tlx.local_view(b_buf, slot)],
    [[off_m, k * BLOCK_K], [k * BLOCK_K, off_n]],
    [0b0011, 0b1100],
)
tlx.async_amd_descriptor_wait(0, [a_tok])

Warp Pipeline (AMD)

[MI350] — AMD MI350 (CDNA4); not available on MI300.

tlx.warp_pipeline_stage(label, *, priority=None) is a context manager that marks explicit pipeline stage boundaries inside a loop. The compiler partitions the loop body at these boundaries and inserts conditional barriers so that one warp group executes one stage ahead of the other, overlapping memory latency with compute.

This is an explicit partitioning marker, not an automatic optimization. Correctness depends on the user's buffering and synchronization structure. In particular:

Use multi-buffered shared memory (typically triple buffering with NUM_BUFFERS=3) to prevent data races between warp groups accessing the same buffer.
Use explicit tlx.async_load_wait_group() to ensure data is ready before consumption.
Handle prologue (prefetch) and epilogue (drain) around the main loop.

See the gfx1250 warp-pipeline GEMM example (third_party/amd/python/examples/gluon/f16_gemm_warp_pipeline_gfx1250.py) for the full pattern.

Parameter	Type	Description
`label`	`str`	Stage name for diagnostics (e.g. `"load"`, `"compute"`)
`priority`	`int` (0-3), optional	Hardware scheduling hint, maps to `s_setprio`. Higher = more urgent.

Auto software pipelining is automatically disabled on loops that contain warp pipeline stages.

Example (simplified — see gfx1250 example for production pattern):

import triton.language.extra.tlx as tlx

@triton.jit
def gemm_kernel(..., BLOCK_K: tl.constexpr, NUM_BUFFERS: tl.constexpr):
    buf_A = tlx.local_alloc((BLOCK_M, BLOCK_K), tl.float16, NUM_BUFFERS)
    buf_B = tlx.local_alloc((BLOCK_K, BLOCK_N), tl.float16, NUM_BUFFERS)

    # Prologue: prefetch NUM_BUFFERS-1 tiles into shared memory
    for i in tl.range(0, NUM_BUFFERS - 1, loop_unroll_factor=NUM_BUFFERS - 1):
        tlx.async_load(a_ptrs, tlx.local_view(buf_A, i), mask=...)
        tlx.async_load(b_ptrs, tlx.local_view(buf_B, i), mask=...)
        tlx.async_load_commit_group()
        a_ptrs += BLOCK_K * stride_ak; b_ptrs += BLOCK_K * stride_bk
    tlx.async_load_wait_group(NUM_BUFFERS - 2)

    # Main loop with warp pipelining
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in tl.range(NUM_BUFFERS - 1, K_ITERS):
        consumer = (k - (NUM_BUFFERS - 1)) % NUM_BUFFERS
        producer = k % NUM_BUFFERS
        with tlx.warp_pipeline_stage("lds_load", priority=1):
            a_tile = tlx.local_load(tlx.local_view(buf_A, consumer))
            b_tile = tlx.local_load(tlx.local_view(buf_B, consumer))
        tlx.async_load_wait_group(0)
        with tlx.warp_pipeline_stage("compute_and_load", priority=0):
            tlx.async_load(a_ptrs, tlx.local_view(buf_A, producer), mask=...)
            tlx.async_load_commit_group()
            acc = tl.dot(a_tile, b_tile, acc)

    # Epilogue: drain remaining buffers
    ...

Scaled Dot (AMD)

tlx.dot_scaled(lhs, lhs_scale, lhs_format, rhs, rhs_scale, rhs_format, acc=None, *, fast_math=False, lhs_k_pack=True, rhs_k_pack=True, out_dtype=tl.float32, tiles_per_warp=None) is a thin wrapper around tl.dot_scaled. Without tiles_per_warp it is exactly equivalent to tl.dot_scaled — pass it only when you need the AMD-specific WMMA scheduling hint described below.

`tiles_per_warp` — what it controls

Concept	Controlled by	What it means
Warp distribution	`warpsPerCTA` (chosen automatically by `AccelerateAMDMatmul::planWarps`)	How the total result tile is split across warps along M/N.
Per-warp tiling	`tiles_per_warp` (this hint)	How many `instrShape`-sized WMMA tiles each warp covers contiguously before the layout repeats.

So tiles_per_warp=[2, 2] does not mean "distribute 4 tiles across 4 warps." It means each warp emits a 2×2 block of WMMA instruction tiles, holding the corresponding 2×2 accumulator registers. Concretely, for a tt.dot_scaled lowered to gfx1250 WMMA (instrShape = [16, 16, K]), 4 warps, warpsPerCTA = [2, 2]:

`tiles_per_warp`	Per-warp coverage (M × N)	Per-CTA coverage before repeat (M × N)
`[1, 1]` (default)	`16 × 16`	`32 × 32`
`[2, 2]`	`32 × 32`	`64 × 64`

For a 256 × 256 result, [1, 1] repeats the layout 8 × 8 times, [2, 2] repeats it 4 × 4. Larger tiles_per_warp gives each warp more contiguous accumulator state (better register reuse for preshuffled MXFP scales, fewer warp-level reductions), at the cost of more registers per warp.

Together, instrShape, warpsPerCTA, and tiles_per_warp define the M/N extent of one CTA-level WMMA layout period: period[d] = instrShape[d] * warpsPerCTA[d] * tiles_per_warp[d]. If the result tile is larger than this period, the period repeats. The K entry of instrShape is the per-instruction reduction depth and is handled separately from this M/N tiling.

tiles_per_warp is validated by AccelerateAMDMatmul: it must have one entry per result-tile dim, each entry must be positive, and instrShape[d] * warpsPerCTA[d] * tiles_per_warp[d] must fit in the result tile shape.

Example

import triton.language.extra.tlx as tlx

acc = tlx.dot_scaled(
    a, a_scale, "e5m2",
    b, b_scale, "e5m2",
    acc,
    tiles_per_warp=[2, 2],   # pack 2x2 WMMA tiles per warp for preshuffled MXFP
)

Mechanism (for IR-level users)

The wrapper attaches amdg.wmma_tiles_per_warp = array<i32: m, n> on the resulting tt.dot_scaled op. ScaledBlockedToScaledWMMAF8F6F4 reads the attribute and substitutes m, n for the default 1, 1 when building the WMMA encoding. Setting the attribute directly on a tt.dot[_scaled] op in MLIR has the same effect; the wrapper just spares Python kernels from hand-poking attributes.

Currently consumed only by the scaled-WMMA pattern (gfx1250). Regular tt.dot WMMA and the MFMA patterns do not read it.

Explicit WMMA layout pinning

tlx.require_amd_wmma_layout(x, version=3, transposed=True, warp_bases=..., reg_bases=..., instr_shape=(16, 16, 128)) pins a tensor to an explicit AMD WMMA register/warp layout. This is useful for tuned gfx1250 epilogues that must retain the accumulator ownership chosen by tiles_per_warp across otherwise layout-neutral tensor operations. The bases contain one linear-layout basis vector per register or warp bit and must match the tensor rank.

The helper lowers to a pinned tlx.require_layout; omit it when automatic layout propagation is sufficient.

Kernels Implemented with TLX

Build and install TLX from source

git clone https://github.com/facebookexperimental/triton.git
cd triton

pip install -r python/requirements.txt # build-time dependencies
pip install -e .

Run the tutorials after the build finishes, e.g,

python third_party/tlx/tutorials/hopper_fa_ws_pipelined_pingpong_test.py

To run Blackwell GEMM tutorial kernels, you can use the following command:

Change 2: One correctness test script

[TLX_VERSION=<kernel_name>] pytest third_party/tlx/tutorials/testing/test_correctness.py

By default only one autotune config will be used by correctness test.

All kernels — Hopper, Blackwell, and AMD — share this one file; each test is arch-gated with @pytest.mark.skipif, so on any given GPU only the relevant cases run and the rest skip. To run just the AMD/IKBO cases:

pytest third_party/tlx/tutorials/testing/test_correctness.py -k "amd or ikbo"

(on gfx950 the gfx1250-only GEMM cases skip automatically).

Change 3: One performance test script per op × arch (Hopper, Blackwell, AMD)

third_party/tlx/denoise.sh third_party/tlx/tutorials/testing/test_hopper_gemm_perf.py [--version {ws|pipelined}]

third_party/tlx/denoise.sh third_party/tlx/tutorials/testing/test_hopper_fa_perf.py [--version {ws|ws_pipelined|ws_pipelined_pingpong|ws_pipelined_pingpong_persistent}]

third_party/tlx/denoise.sh third_party/tlx/tutorials/testing/test_blackwell_gemm_perf.py [--version {ws|pipelined|clc|2cta}]

third_party/tlx/denoise.sh third_party/tlx/tutorials/testing/test_blackwell_fa_perf.py [--version {ws|ws_persistent|ws_pipelined|ws_pipelined_persistent|clc}]

denoise.sh wraps AMD runs too (it applies NUMA pinning and runs the benchmark; the GPU clock/power lock is NVIDIA-only and is simply skipped on AMD). gfx950 / CDNA4:

third_party/tlx/denoise.sh python third_party/tlx/tutorials/testing/test_amd_gemm_perf.py [--version {warp_pipeline|pipelined}]

third_party/tlx/denoise.sh python third_party/tlx/tutorials/testing/test_amd_fa_perf.py [--version {simple|prefetch|persistent|cluster}]

Without --version, AMD FA perf runs simple, prefetch, and persistent; select cluster explicitly because it is D=128-only.

third_party/tlx/denoise.sh python third_party/tlx/tutorials/testing/test_amd_addmm_glu_perf.py [--version {tlx_baseline|tlx_simple_async|tlx_optimized_async|tlx_optimized|tlx_persistent}]

third_party/tlx/denoise.sh python third_party/tlx/tutorials/testing/test_amd_ikbo_fa_perf.py (IKBO Flash Attention)

third_party/tlx/denoise.sh python third_party/tlx/tutorials/testing/test_amd_ikbo_lce_perf.py (IKBO LCE — distinct op, not attention)

gfx1250:

third_party/tlx/denoise.sh python third_party/tlx/tutorials/testing/test_amd_mxfp_gemm_perf.py [--transpose-b]

TLX-AMD CI

AMD tutorial kernels are exercised by .github/workflows/mi350.yml on a gfx950 (MI350 / CDNA4) runner, mirroring the H100 job in .github/workflows/h100.yml:

mi350-tlx-test — TLX unit tests (python/test/unit/language/test_tlx_*.py) plus the tutorial correctness suite (third_party/tlx/tutorials/testing/test_correctness.py). AMD and IKBO cases run; Hopper/Blackwell and gfx1250 cases auto-skip via the arch gates.
mi350-meta-triton-test — TritonBench performance coverage (the AMD perf scripts above are for local runs; perf-regression tracking lives in TritonBench).

Both run on push, PR, and the nightly schedule; nightly failures are filed as issues via report-nightly-failure.yml.

facebookexperimental/triton

TLX - Triton Low-level Language Extensions

Introduction

Nightly builds (fbtriton)

Gluon support

The DSL Extension

Local buffer operations

Buffer Reuse

Remote buffer operations

Async memory access

Async tensor core operations

Barrier operations

Memory Fences

Cluster Launch Control (CLC)

CLC API

How CLC Works

CLC Synchronization Flow

Multi-CTA Mode (2-CTA Clusters)

Key Principle: "Arrive Remote, Wait Local"

Example: Multi-CTA GEMM with CLC

Warp Specialization operations

async_tasks Parameters

async_task Parameters

Default Task Register Budget

Explicit Warp Assignment with warp_group_start_id

CUDA Thread Block Clustering

Usage

Using ctas_per_cga with Autotune

TLX vs Triton Semantics

Other operations

Buffer Operations (AMD)

tlx.buffer_load

tlx.buffer_store

tlx.buffer_atomic_add

tlx.buffer_load_to_local

Example

Wave Uniformity (AMD)

tlx.assume_uniform

Explicit MFMA Scheduling (AMD)

AMD TDM Descriptor Loads

Warp Pipeline (AMD)

Scaled Dot (AMD)

tiles_per_warp — what it controls

Example

Mechanism (for IR-level users)

Explicit WMMA layout pinning

Kernels Implemented with TLX

GEMM kernels

Attention kernels

AMD kernels (gfx950 / CDNA4)

AMD kernels (gfx1250)

Build and install TLX from source

Change 2: One correctness test script

Change 3: One performance test script per op × arch (Hopper, Blackwell, AMD)

TLX-AMD CI

More reading materials

facebookexperimental triton

`tlx.buffer_load`

`tlx.buffer_store`

`tlx.buffer_atomic_add`

`tlx.buffer_load_to_local`

`tlx.assume_uniform`

`tiles_per_warp` — what it controls