Skip to content

tle.md

Latest commit

 

History

History
1709 lines (1307 loc) · 79.1 KB

tle.md

File metadata and controls

1709 lines (1307 loc) · 79.1 KB
Raw

TLE Architecture Design

1. Introduction

Triton is an operator programming language in the form of a Python DSL. It follows a block-based programming model that abstracts away hardware details such as memory hierarchy, layout, pipelining, and synchronization, while achieving strong operator performance through compiler optimization. These advantages have attracted a large developer community and ecosystem.

In recent years, however, Triton has faced growth challenges:

  • Adaptation to DSA platforms and new GPU architectures has progressed slowly.
  • Compared with emerging languages like TileLang, Triton lacks abstractions for fine-grained control of memory hierarchy and parallel granularity, which can lead to weaker performance in some cases.

To address these issues, we propose TLE (Triton Language Extentions), which extends Triton across three levels to meet urgent needs from users with different skill profiles.

2. Observations and Proposed Solutions

We analyzed mainstream DSLs in the industry (Triton, TileLang, and cuTile) and summarized a target language design.

2.1 Pythonic

All three are Python-syntax-based DSLs, indicating that developers prefer Python-like syntax for kernel development, even if only a subset of Python is available.

2.2 Tile Programming

All three support block-level programming. In essence, current block programming mainly performs tiling on global memory. cuTile goes further by supporting multi-level tiling, making it possible to design a unified language across multiple memory hierarchy architectures.

Triton, however, does not explicitly model tile/slice concepts, so users can only tile at the global memory level, limiting further language evolution.

TileLang is similar to Triton in that it does not provide explicit tiling primitives. In addition, except for copy and GEMM, it lacks higher-level tensor ops, which makes GPU programming less convenient. Without automatic vectorization, utilizing SIMD hardware well often requires adding many SIMD-specific ops.

2.3 Memory Hierarchy Abstraction

To address the memory wall, modern hardware uses multi-level memory hierarchies.

  • Triton/cuTile expose only two levels: global memory and local tensor.
  • TileLang directly exposes native hardware memory hierarchy without abstraction.

Problems:

  • Exposing too few levels pushes tiling and buffer promotion work to the compiler.
  • Directly exposing native hierarchy significantly hurts portability.

Preferred direction:

  • Developers perform tiling, but do not explicitly select memory levels.
  • Compiler performs buffer promotion.
  • Developers may provide hints; tile sizes are treated as hyperparameters.

This keeps portability while leaving room for further optimization.

2.4 Parallelism Abstraction

  • Triton/cuTile expose only block-level parallelism, and intra-block parallelism is fully compiler-controlled.
  • TileLang lets developers explicitly control intra-block parallelism (Parallel and Vectorize), improving expressiveness but reducing portability and reuse across hardware.

2.5 Distributed Abstraction

None of these languages directly covers cross-block or cross-node communication, which limits compute-communication fusion (ongoing external work includes Triton Distributed and TileScale).

2.6 Ideal Language Design

  • Level 1: Numpy/PyTorch-like algorithm-level programming. Users focus on algorithm logic only; compiler handles hardware mapping and communication.
  • Level 2: cuTile-like tile-level programming plus distributed descriptions. Users explicitly provide tiling and sharding, while compiler handles memory hierarchy, parallelism, and communication, with optional hardware/scenario hints.
  • Level 3: Hardware-specific extensions (memory hierarchy, thread binding, vectorize, etc.). This level is confined to specific regions with explicit interaction contracts with Level 2. Compiler performs only essential optimizations.

Detailed principles:

  • Tile semantics to avoid manual address arithmetic.
  • Do not require tensor shapes to be powers of two.

Open question: what other strong design ideas should be added?

3. Architecture Design

3.1 Architecture Overview

TLE sits in the middle layer of the AI software stack:

  • Upstream: serves AI frameworks through graph compilers and operator libraries.
  • Downstream: integrates with various hardware runtimes.

Content not available outside Feishu document yet.

TLE is split into three layers:

  • TLE-Lite: lightweight extension over Triton. Features are backend-compatible, and only small changes to existing Triton kernels are needed to gain significant speedups. Targets algorithm engineers and fast optimization workflows.
  • TLE-Struct: architecture-clustered abstractions (e.g., GPGPU, DSA) for deeper performance tuning. Requires moderate hardware knowledge.
  • TLE-Raw: direct hardware control, including vendor-native programming languages for maximum performance. Targets expert performance engineers.

Lowering paths:

  • TLE-Lite and TLE-Struct lower to LLVM IR via FLIR.
  • TLE-Raw lowers to LLVM IR via language-specific pipelines (e.g., vendor private compilers).
  • All parts are finally linked into a complete kernel loaded/executed by runtime.

3.1.1 Optional GPU CommonIR lowering

The default NVIDIA build keeps the native TLE GPU lowering path. An opt-in build can instead preserve the structured buffer operations as CommonIR TileIR before converting them to native TTGIR:

tle.gpu.* -> tile.* / !tile.buf -> CommonIRToTTGIR -> native TTGIR

This path makes the frontend buffer semantics explicit without changing the final NVIDIA backend contract. The CommonIRToTTGIR pass must eliminate all tile.*, !tile.buf, and temporary buffer-to-memdesc bridges before the remaining TTGIR pipeline runs.

The path is selected at build time and is disabled by default. It currently supports only the default NVIDIA backend. Place a compatible checkout of the official FLIR main branch at third_party/flir, then build FlagTree with:

git clone https://github.com/flagos-ai/flir.git third_party/flir
FLAGTREE_COMMON_IR=1 python -m pip install -e . --no-build-isolation

The environment variable FLAGTREE_COMMON_IR sets the internal CMake variable FLAGTREE_COMMON_IR_ENABLED. When enabled, the C++ build defines __FLAGTREE_COMMON_IR__, exposes one capability query to Python, and registers the CommonIR dialect and conversion pass only for this build. Do not set FLAGTREE_BACKEND at the same time. Rebuild FlagTree when switching between the native and CommonIR paths; this is not a per-kernel runtime option.

Supported GPU buffer forms are:

TLE/frontend form CommonIR form Native TTGIR result
tle.gpu.alloc (SMEM, no alias) tile.alloc / !tile.buf<..., #shared> ttg.local_alloc / !ttg.memdesc
tle.gpu.copy (full-buffer pointer copy) tile.copy synchronous or asynchronous TTGIR copy operations
buf.load() tile.to_tensor ttg.local_load
buf.store(value) tile.store_tensor ttg.local_store
buffered-tensor slot/view tile.subview memdesc subview
tle.gpu.local_ptr on a local buffer temporary !tile.buf to !ttg.memdesc bridge existing TLE local-pointer operations
tle.gpu.wgmma shared operands, including transpose buffer-to-memdesc bridge existing descriptor views and WGMMA lowering

buf is a tle.gpu.buffered_tensor. Its load() / store(value) methods use the same spelling in native and CommonIR builds. Native builds lower them through tle.gpu.local_ptr and tl.load / tl.store; CommonIR builds preserve the full-buffer TileIR operations until conversion. The former experimental tle.gpu.to_tensor(buf) / tle.gpu.store_tensor(value, buf) entry points are replaced by these methods, not retained as aliases.

TMA descriptor copies retain the existing TMA operation after the buffer-to-memdesc bridge. Global-to-shared TMA copies also preserve a user-provided completion barrier and its expected byte count; barrier validation follows the native path.

The CommonIR build currently rejects the following forms with an explicit frontend diagnostic instead of silently bypassing CommonIR:

  • aliased tle.gpu.alloc buffers;
  • tle.gpu.copy with a mask, or a completion barrier outside global-to-shared TMA copy;
  • remote-buffer tle.gpu.local_ptr;
  • normal pointer copies with offsets (offset the pointer operands instead).

After building, the focused checks are:

python -m pytest -q test/CommonIR/test_gpu_semantics.py
python -m pytest -q test/CommonIR/test_gpu_wgmma_bridge.py python/test/tle/integration/test_tle_tma_copy.py
python -m pytest -q python/test/tle/unit/test_tle_whitelist.py
python -m pytest -q python/test/tle/unit/test_tle_gpu_buffer_access.py
lit -sv --filter='gpu-tileir' build/cmake.*/test

The lit tests check both the intermediate TileIR contract and the absence of CommonIR operations or unresolved bridge casts in final TTGIR.

3.2 TLE-Lite

  • Design philosophy: write once, run anywhere.
  • Core idea: use high-level semantic hints (instead of hard constraints) to guide compiler heuristics. Keep backward compatibility and achieve cross-platform speedups with minimal code changes.

3.2.1 Memory Management

3.2.1.1 tle.load

Extension of tl.load with async hint support:

x = tle.load(..., is_async=True)

3.2.2 Tensor Slicing

3.2.2.1 tle.extract_tile

Split input tensor into a sub-tile grid using a child-tile shape and extract tile at specified coordinates.

  • GPU: supports extraction from registers and shared memory.
# x is [4, 4]
# z is [2, 2]
# Split x into shape=[2, 2] sub-tiles and return tile at [0, 0]
z = x.extract_tile(index=[0, 0], shape=[2, 2])
3.2.2.2 tle.insert_tile

Split input tensor into a sub-tile grid using child-tile shape and update tile at specified coordinates.

  • GPU: supports updates in registers and shared memory.
# x is [4, 4], y is [2, 2], z is [4, 4]
# Split x into shape=[2, 2] sub-tiles, update tile [0, 0] with y,
# and return full updated [4, 4] tensor
z = x.insert_tile(y, index=[0, 0])

3.2.3 Pipeline

3.2.3.1 tle.pipe

tle.pipe is TLE's explicit dataflow edge: a writer fills one stage of a group of shared-memory buffers and commits it; readers wait until that stage is ready, consume it, then release it. It combines "which buffer stage contains the data" and "when producer/consumer visibility is established" into one typed pipe descriptor, and is intended for CTA-local load/compute overlap, single-producer/single-consumer (SPSC), and single-producer/multi-consumer (SPMC) pipelines.

Automatic software pipelining can still be triggered with tl.range(..., num_stages=...):

for yoff in tl.range(0, ynumel, YBLOCK, num_stages=2):
    Q = tl.load(...)
    K = tl.load(...)
    KT = tl.trans(K)
    V = tl.dot(Q, KT)

Use an explicit pipe when the producer/consumer split should be visible in the program:

stage_buf = tle.gpu.alloc([2, BLOCK], dtype=tl.float32, scope=tle.gpu.smem)
pipe = tle.pipe(capacity=2, scope="cta", name="x_pipe", x=stage_buf)
writer = pipe.writer()
reader = pipe.reader()
offs = tl.arange(0, BLOCK)

# producer partition
for k in tl.range(0, n_tiles):
    slot = writer.acquire(k)
    tl.store(tle.gpu.local_ptr(slot.x), tl.load(x_ptr + k * BLOCK + offs))
    writer.commit(k)

# consumer partition
for k in tl.range(0, n_tiles):
    wait = reader.wait(k)
    x = tl.load(tle.gpu.local_ptr(wait.slot.x))
    acc += x
    reader.release(k)

A mixed pipe is the same tle.pipe abstraction used for a payload with multiple fields whose bytes are produced by different mechanisms. For example, one field may be filled by a TMA copy, another by cp.async, and another by a CUDA-core tl.store through tle.gpu.local_ptr. Users still create one logical pipe for the chunk lifecycle and call one commit after producing the fields. The compiler infers each field's transport from the producer-side IR; transport is not a user-facing pipe attribute.

The motivation is to keep synchronization aligned with the logical dataflow edge instead of the implementation transport. Without mixed pipe support, a kernel often has to split one logical chunk into q_pipe, v_pipe, scale_pipe, and similar transport-specific pipes, which duplicates ready/free barriers and makes warp-specialized code diverge from the intended algorithm. Mixed pipe lets these fields share the same slot, phase, reader set, and ready/free protocol while still preserving field-level provenance and compile-time hazard checks.

3.2.4 Distributed

Triton distributed API has four core parts: device mesh definition, sharding specification, resharding (collective communication), and remote access (point-to-point communication).

3.2.4.1 Device Mesh

tle.device_mesh defines physical device topology and serves as the context foundation for distributed operations.

class device_mesh:
    def __init__(self, topology: dict):
        """
        Initialize DeviceMesh.

        Args:
            topology (dict): Hardware hierarchy description.
                             Keys are hierarchy names; values are int (1D)
                             or tuple lists (multi-dimensional).
        """
        self._physical_ids = ... # Internal flattened physical IDs (0..N-1)
        self._shape = ...        # Current logical shape, e.g. (2, 2, 4, 2, 2, 4)
        self._dim_names = ...    # Current dimension names

    @property
    def shape(self):
        """Return logical mesh shape."""
        return self._shape

    @property
    def ndim(self):
        """Return number of dimensions."""
        return len(self._shape)

    def flatten(self):
        """Flatten mesh to 1D, typically for ring communication."""
        return self.reshape(prod(self._shape))

    def __getitem__(self, key):
        """
        Supports slicing and returns a sub-mesh.
        Supports standard slice and integer indexing.
        """
        return sub_mesh

    def __repr__(self):
        return f"DeviceMesh(shape={self._shape}, names={self._dim_names})"


# Define complex hardware hierarchy
topology = {
    # Cross-node hierarchy (2x2 = 4 nodes)
    "node": [("node_x", 2), ("node_y", 2)],
    # In-node GPUs (4 devices)
    "device": 4,
    # In-GPU cluster (2x2)
    "block_cluster": [("cluster_x", 2), ("cluster_y", 2)],
    # In-cluster blocks (4 blocks)
    "block": 4,
}

# mesh.shape -> (2, 2, 4, 2, 2, 4)
# total size = 256
mesh = tle.device_mesh(topology=topology)
3.2.4.2 Sharding Specification

tle.sharding declares tensor distribution state on the device mesh:

  • splits: how each tensor axis is partitioned on mesh axes.
  • partials: whether tensor is partial-sum state.
  • Unspecified mesh axes are treated as broadcast.

Symbols:

  • tle.S(axis): split.
  • tle.B: broadcast/replicate.
  • tle.P(axis): partial; requires reduce on specified axis.
def sharding(tensor, splits, partials):
    """
    Annotation only: marks tensor state, emits no direct code,
    but guides compiler checks and optimizations.
    """
    return tensor

# Split axis0 on cluster, axis1 on device, and partial on block axis
x_shard = tle.sharding(
    mesh,
    split=[["cluster_x", "cluster_y"], "device"],
    partial=["block"],
)

# Define a sharded tensor
x = tle.make_sharded_tensor(x_ptr, sharding=x_shard, shape=[4, 4])
3.2.4.3 Synchronization

In complex distributed kernels (e.g., ring all-reduce or row/column-independent pipelines), only “same-row” or “same-column” blocks often need synchronization rather than the whole cluster. Global synchronization introduces unnecessary waiting.

tle.distributed_barrier is a collective synchronization primitive and returns None. With the default barrier_kind="sync", all participants in the selected synchronization group rendezvous at this point before continuing.

All participants in one barrier instance must execute matching barrier calls in the same order. If a branch or loop makes some participants skip a call, the remaining participants can wait forever and the kernel can deadlock. If the communication only needs a one-way producer-to-consumer notification rather than a collective rendezvous, use tle.signal and tle.signal_wait instead.

def distributed_barrier(
    mesh=None,                    # optional device_mesh; omitted -> full-cluster barrier;
                                  # required with space; a sliced cluster mesh selects sub-mesh first
    device_dptr=None,             # optional normally; required with space;
                                  # DistributedRtContext returned by create_dist_tensor
    space=None,                   # optional FlagCX team selector; None -> mesh/local barrier path;
                                  # "device" | "inter" | "world" (aliases described below)
    group_kind="block",           # optional; effective only on the cross-device/cross-node communicator path; "thread" | "warp" | "block"
    barrier_kind="sync",          # optional; effective only on the cross-device/cross-node communicator path; "arrive" | "wait" | "sync"
    order="acqrel",               # optional; effective only on the cross-device/cross-node communicator path;
                                  # "relaxed" | "acquire" | "release" | "acqrel"
    index=0,                      # optional; effective only on the cross-device/cross-node communicator path; non-negative barrier channel
    context_id=0,                 # optional; effective only on the cross-device/cross-node communicator path; compile-time int32 context index
    memory_scope="system",        # optional; effective only on the cross-device/cross-node communicator path;
                                  # "system" | "device" | "block" | "thread"
):
    ...

Every public argument has a default, so none is unconditionally required. The requirements depend on the selected path:

  • tle.distributed_barrier() is valid and emits the default full-cluster barrier.
  • tle.distributed_barrier(mesh) uses mesh to select/infer a cluster, cluster sub-mesh, or cooperative-grid barrier; no device_dptr is needed.
  • An explicit FlagCX communicator barrier requires all three of space, mesh, and device_dptr. The remaining arguments are optional and use the defaults shown above.

mesh and space select the synchronization mode in the following order:

  • A static slice of a cluster mesh selects a cluster sub-mesh barrier. For example, full_mesh[0, :] synchronizes only the selected cluster members. The current implementation supports non-scalar meshes derived by slicing a launch mesh; every member of that sub-mesh must call the barrier, and CTAs outside it must not call the same sub-mesh barrier.
  • Setting space selects a FlagCX communicator barrier. It requires both mesh and device_dptr. "device" (also "intra" and "intra_node") selects the intra-node team and requires a device topology axis. "inter" (also "inter_node") selects the inter-node team, while "world" selects the world team; both require a node topology axis.
  • A mesh with only block axes selects a cooperative-grid barrier. It synchronizes the entire launched grid, not an arbitrary subset, and enables a cooperative-grid launch.
  • In all other cases, a full cluster mesh (or omitting mesh) selects a full-cluster barrier. When a mesh is passed, its cluster dimensions are inferred for the launch.

A sliced cluster mesh has priority over space. In all other cases, an explicit space has priority over the grid and full-cluster modes.

Example: pass a cluster sub-mesh and a distributed context to the same barrier call:

# On the host.
full_mesh = tle.device_mesh({
    "block_cluster": [("cluster_x", 2), ("cluster_y", 2)],
})
row_mesh = full_mesh[0, :]
device_dptr = tle.create_dist_tensor(buffer)

# In a JIT kernel. `device_dptr` may also be used by FlagCX operations elsewhere.
tle.distributed_barrier(row_mesh, device_dptr=device_dptr, space = "device")

Because row_mesh is a sliced cluster mesh, this call follows the 'submesh' dispatch rule and emits a cluster sub-mesh barrier. device_dptr is accepted but is not consumed by this barrier path; adding space=... does not turn this call into a FlagCX communicator barrier.

For the cross-device/cross-node communicator path (implemented with FlagCX):

  • mesh is required and validates the selected communicator topology. The "device" team requires a device launch axis; the "inter" and "world" teams require a node launch axis.

  • device_dptr is the distributed runtime context returned by tle.create_dist_tensor(...).

  • space selects the participants: "device"/"intra"/"intra_node" for the intra-node team, "inter"/"inter_node" for the inter-node team, or "world" for the world team.

  • group_kind is the collective execution scope: "thread", "warp", or "block" (default).

  • barrier_kind is "arrive", "wait", or "sync" (default). "arrive" only reports arrival; "wait" waits for the matching arrivals; "sync" performs both.

  • order controls the memory ordering of accesses around the barrier and is effective only on the cross-device/cross-node communicator path. Its accepted values are:

    • "relaxed": performs the barrier operation without adding acquire/release ordering constraints to memory accesses around it.
    • "acquire": prevents memory accesses after the barrier from being reordered before it; it is typically used by a waiter so that it can observe data published before a matching release.
    • "release": prevents memory accesses before the barrier from being reordered after it; it is typically used by an arriver to publish completed writes before notifying other ranks.
    • "acqrel": provides both acquire and release semantics. This is the default and is typically paired with a full "sync" barrier.

    Common pairings are barrier_kind="arrive" with order="release", barrier_kind="wait" with order="acquire", and barrier_kind="sync" with order="acqrel". These are semantic recommendations, not parameter-combination restrictions.

  • memory_scope specifies the visibility scope of the memory-ordering guarantee and is effective only on the cross-device/cross-node communicator path. It does not select the participating ranks; space selects them. The API semantics of its accepted values are:

    • "system": system scope, visible to all threads in the system. This is the default and the usual choice for a distributed barrier.
    • "device": device scope, visible to all threads on the current device.
    • "block": thread-block (CTA) scope, visible only to threads in the current block.
    • "thread": thread scope, visible only to the current thread.

    The bundled FlagCX unified-barrier API currently accepts this argument, but its flagcxDevBarrierArrive, flagcxDevBarrierWait, and flagcxDevBarrierSync implementations do not yet use it to change barrier behavior. Consequently, the four values currently have no behavioral difference for distributed_barrier; the definitions above describe the API semantics.

  • index selects a non-negative barrier channel. context_id selects a pre-created FlagCX context (notably the network context for "inter" and "world") and must be a non-negative compile-time int32. Both values must match across all participants.

Example: synchronize the FlagCX world team represented by a node/device mesh:

# On the host: ctx = tle.create_dist_tensor(buffer)
mesh = tle.device_mesh({"node": 2, "device": 8})

# In a JIT kernel, `ctx` and `mesh` are passed as compile-time arguments.
tle.distributed_barrier(
    mesh,
    device_dptr=ctx,
    space="world",
    group_kind="block",
    barrier_kind="sync",
    index=0,
    context_id=0,
    order="acqrel",
    memory_scope="system",
)

Cluster/sub-mesh and cooperative-grid modes currently use the NVIDIA lowering. The FlagCX mode additionally requires a TLE build and runtime with FlagCX enabled.

3.2.4.4 Remote Access

tle.remote obtains a handle for data shards located on other devices. This maps to point-to-point communication or direct memory access (RDMA/NVLink load). The current implementation distinguishes three target domains via space:

def remote(
    tensor,                    # required; cluster path: shared-memory pointer or buffered_tensor;
                               # device/node paths: DistributedRtContext from create_dist_tensor
    shard_id,                  # required; compile-time int or runtime int32 scalar;
                               # cluster path: target block id; device path: intra-node peer rank;
                               # node path: world rank; a compile-time mesh coordinate requires scope
    scope=None,                # device_mesh; required for a mesh-coordinate shard_id;
                               # otherwise optional; also sets cluster/device launch dimensions
    space="cluster",           # optional string; "cluster" | "device" | "node";
                               # defaults to "cluster"
    dtype=None,                # tl.dtype; optional on cluster path, inferred when omitted;
                               # required on device/node paths
    offset=None,               # Python int or scalar integer tensor; required on device path;
                               # unsupported on cluster/node paths (omit)
    coopkind=None,             # optional on node path: "thread" | "warp" | "block" or GroupKind;
                               # default GroupKind.BLOCK; unsupported on cluster/device paths (omit)
    context_id=0,              # optional on node path: compile-time int in [0, INT32_MAX], default 0;
                               # unsupported on cluster/device paths (omit)
):
    """
    For cluster, a pointer input returns a cluster remote pointer and a buffered_tensor
    returns a remote-marked buffer. Device/node paths return a global remote pointer.
    Pointer results participate in tl.load / tl.store; materialize a cluster buffer
    pointer first with tle.gpu.local_ptr.
    """
Cluster Path (intra-cluster communication, DSMEM)

The cluster path targets shared-memory access across blocks within one thread block cluster (CTA cluster, DSMEM on Hopper and later); it is the default when space is omitted. Two input forms are supported:

  • A shared-memory pointer (scalar or tensor): returns a remote pointer in the cluster address space directly, usable in tl.load / tl.store; a block-tensor input keeps its shape.
  • A tle buffered_tensor: returns a remote-marked buffer; materialize remote pointer views with tle.gpu.local_ptr(...).

shard_id is the id of the target block inside the cluster; with scope, coordinates are linearized through the mesh, and launch cluster dimensions are inferred from that mesh (requiring num_ctas=1, one program per block). For pointers already in cluster-shared space, shard_id=0 returns the local access as-is. The cluster / device paths do not support the node-only arguments coopkind / context_id.

Example (read a SMEM tile from a neighbor block):

# Here smem is a buffered_tensor, for example one returned by tle.gpu.alloc.
remote_smem = tle.remote(smem, shard_id=(node_rank, next_device), scope=mesh)
remote_ptr = tle.gpu.local_ptr(remote_smem, (rows, cols))
vals = tl.load(remote_ptr)
Device Path (intra-node cross-GPU communication, NVLink P2P)

The device path targets direct memory access between GPUs within one node through a FlagCX registered window (symmetric memory, NVLink P2P). Like the node path, tensor is the DistributedRtContext returned by create_dist_tensor, but the returned remote pointer lives in the ordinary global address space, and the offset argument is required — it fixes the element offset inside the remote window (Python int or scalar integer tensor, normalized to i64 internally).

Key difference from the node path: the device-path remote pointer is an ordinary global pointer, so subsequent tl.load / tl.store support arbitrary tiled / multi-dimensional access with no "contiguous transfer only" restriction; the offset base is fixed once by offset, and further pointer arithmetic behaves like normal global-pointer math.

Example (scatter a local input shard to a peer GPU in the same node):

remote_base = tle.remote(scatter_ctx, space="device",
                         dtype=input_ptr.dtype.element_ty,
                         shard_id=target_rank,
                         offset=SCATTER_NODE_SLICE_OFFSET_ELEMS + source_slot_offset_elems)
scatter_ptrs = (remote_base +
                (scatter_row + row_offs[:, None]) * N + input_col + col_offs[None, :])
tl.store(scatter_ptrs, values, mask=scatter_row_mask & col_mask)
Node Path (inter-node communication, FlagCX/RDMA)

The node path targets cross-node point-to-point transfers, using FlagCX registered memory (symmetric window) for the data plane.

context_id has the same user-level meaning as tle.distributed_barrier(..., context_id=...): it selects a pre-created FlagCX network context. It is neither a peer rank nor a transfer/barrier sequence number. The value is emitted as an i32 operation attribute rather than a runtime operand. The compiler validates that it is a compile-time integer in [0, INT32_MAX]; the selected context must also exist in the runtime communicator, and corresponding peers must use compatible context assignments. The default is context 0.

The current version only supports contiguous data transfer. Caveats:

  • On the host, tle.create_dist_tensor(comm_buf) must register the communication buffer into a FlagCX symmetric memory window first; the returned DistributedRtContext is passed to the kernel as an argument. Device-side scatter and inter-node P2P may share one registered window. The local-side buffer must be the same buffer registered by that context and must be passed directly to the kernel as a global-pointer argument.
  • dtype is required. offset is not accepted — add offsets to the returned pointer instead.
  • Transfer shapes are restricted: a scalar copy moves one element; both sides of a tensor copy must reuse the same contiguous range value tl.arange(0, N), with either no mask or the same shared prefix mask offsets < valid_n (1 <= valid_n <= N).
  • A single transfer is subject to the compiler's tensor size limit: the number of elements N transferred per load/store must be a power of two and no greater than 33554432 (2^25); exceeding it fails at compile time. To transfer more data, split it into chunks and call load/store multiple times in a loop, one chunk per iteration:
for start in range(0, nelems, CHUNK_N):           # transfer one chunk per iteration
    chunk_n = tl.minimum(nelems - start, CHUNK_N) # the last chunk may be shorter than CHUNK_N
    offsets = tl.arange(0, CHUNK_N)               # CHUNK_N: power of two, <= 33554432
    mask = offsets < chunk_n                      # must be written as arange result < scalar
    vals = tl.load(comm_buf + src_offset + start + offsets, mask=mask)
    tl.store(remote_dst + dst_offset + start + offsets, vals, mask=mask)

When looping, note: the mask must be written as offsets < scalar — where offsets is the bare tl.arange result, not an expression like start + offsets < nelems, otherwise compilation fails. The scalar may be computed at runtime (e.g. the tl.minimum above handles the last, possibly shorter chunk).

  • The load result must be consumed directly and exclusively by its paired store. Both operations must be in the same basic block, with the load before the store and no intervening memory access, atomic operation, barrier, or other communication operation; exactly one side must use a node remote pointer.
  • Sparse access, multi-dimensional tiled load/store, strided access, non-zero-start ranges, and mismatched ranges on the two sides are rejected at compile time; dynamic violations trigger a device assertion.
  • coopkind supports thread / warp / block (default block: the whole CTA convergently issues one transfer). context_id follows the network-context rules above.
  • Since device code has no rank() / num_ranks(), topology is passed into the kernel as constexpr by the host. Pass shard_id as a host-precomputed world rank int (recommended), or as a mesh coordinate tuple with scope (resolved to a world rank via physical_ids). On the node path, scope does not affect launch configuration.

Usage: first obtain a remote pointer to the peer node with tle.remote, then complete the transfer with a load/store pair — the transfer direction is determined by which side of the pair is remote:

  • PUT (local → remote): load from the local buffer first, then store through the remote pointer; the data is written from this node to the remote node.
  • GET (remote → local): the usage is exactly symmetric, with the direction reversed — load through the remote pointer, then store into the local buffer; the data is read back from the remote node.

PUT example (local load + remote store):

remote_dst = tle.remote(ctx, space="node", dtype=DTYPE,
                        shard_id=remote_rank, coopkind=tle.GroupKind.BLOCK,
                        context_id=0)
offsets = tl.arange(0, BLOCK_SIZE)
mask = offsets < nelems
vals = tl.load(comm_buf + src_offset + offsets, mask=mask)     # local load
tl.store(remote_dst + dst_offset + offsets, vals, mask=mask)   # remote store = PUT: data written to the remote node

GET example (remote load + local store) — same usage as PUT, with the load/store sides swapped:

remote_src = tle.remote(ctx, space="node", dtype=DTYPE,
                        shard_id=remote_rank, coopkind=tle.GroupKind.BLOCK,
                        context_id=0)
offsets = tl.arange(0, BLOCK_SIZE)
mask = offsets < nelems
vals = tl.load(remote_src + src_offset + offsets, mask=mask)   # remote load = GET: data read from the remote node
tl.store(comm_buf + dst_offset + offsets, vals, mask=mask)     # local store

Common mistakes (do not write it this way):

# Wrong: multi-dimensional / tiled access
vals = tl.load(remote_dst + rows[:, None] * N + cols[None, :])  # rejected at compile time
# Why: the node path lowers the load/store pair into a single RDMA put/get.
# The interface only supports contiguous data (a 1D linear range) and cannot
# express the scattered addresses of a 2D tile. To move 2D data, lay it out
# contiguously in a local buffer first and transfer that.

# Wrong: passing an offset to remote
remote_dst = tle.remote(ctx, space="node", dtype=DTYPE, shard_id=r, offset=1024)  # error
# Why: on the node path tle.remote only resolves the target peer and returns a
# pointer to the base of the remote window; the source/destination offsets are
# expressed as src_offset / dst_offset in the following load/store.
# Correct: add the offset to the returned pointer
tl.store(remote_dst + dst_offset + offsets, vals, mask=mask)

# Wrong: mismatched ranges on the two sides, or a mask that is not a shared prefix
vals = tl.load(local + tl.arange(0, N), mask=mask)             # arange(0, N)
tl.store(remote + tl.arange(0, M), vals, mask=mask2)           # a different range -> rejected
# Why: the transfer length and both offsets are jointly inferred from the
# pointer expressions of the load/store; the arange range is the transfer size
# (the mask gives the valid length). If the two sides disagree, how much to
# transfer and from where to where are undefined, so the compiler rejects it.
3.2.4.5 Resharding

tle.reshard is the entrypoint for collectives. Compiler compares source and target specs and inserts communication primitives automatically.

def reshard(tensor, spec):
    """
    Action: transform tensor to a new distribution state.

    Typical transitions:
    1. [ ] -> [S]: Scatter
    2. [S] -> [ ]: Gather
    3. [P] -> [ ]: Reduce
    4. [B] -> [S]: Local slice (no communication)
    5. [S] -> [B]: All-gather
    6. [P] -> [B]: All-reduce
    7. [B] -> [P]: Error
    """
3.2.4.6 Distributed GEMM

NVIDIA Hopper (H100) and newer architectures introduce Thread Block Cluster, allowing groups of CTAs to cooperate via DSMEM for high-bandwidth, low-latency exchange.

tle.distributed_dot is designed to use this feature so developers can write cross-block matrix multiplication without manually handling DSMEM barriers and data movement.

def distributed_dot(a, b, c=None):
    """
    Execute distributed matrix multiplication within current
    Thread Block Cluster scope.

    Behavior depends on sharding specs of input tensors `a` and `b`
    over the cluster mesh.

    Args:
        a (Tensor): left operand with cluster-level sharding annotation.
        b (Tensor): right operand with cluster-level sharding annotation.
        c (Tensor, optional): accumulator.

    Returns:
        Tensor: result tensor with distribution inferred from inputs.
    """

Open question: what additional distributed primitives are needed?

3.2.4.7 tle.signal

tle.signal atomically updates a synchronization slot on a remote peer. It only sends a signal; it neither transfers data nor waits for completion on the receiving peer.

def signal(device_dptr, peer, slot_id, value=None, op="inc",
          space="intra_node", group_kind="block", context_id=0,
          scope="system"):
    """
    Atomically update a remote peer's synchronization slot.

    :param device_dptr: distributed communicator handle (required)
    :param peer: target peer rank (int32 scalar) (required)
    :param slot_id: signal slot index (uint32 scalar) (required)
    :param value: optional uint64 scalar; required when op="add", must be omitted otherwise
    :param op: optional; "inc" to increment by one, "add" to add value; default "inc"
    :param space: optional; "intra_node", "inter_node", or "world"; default "intra_node"
    :param group_kind: optional; "thread", "warp", or "block"; default "block"
    :param context_id: optional; compile-time int selecting a pre-allocated network context; default 0
    :param scope: optional; visibility scope of the operation ("system" or "device"); default "system"
    """
    pass

op="inc" increments the selected signal slot by one. op="add" adds value to the selected signal slot; value is required in that case and must be omitted otherwise. space selects the communication scope (intra_node, inter_node, or world), and peer is a rank within that scope. context_id selects a pre-allocated network context.

scope controls the visibility of the signal operation to the threads on the node: "system" means visible to all threads on all devices, and "device" means only visible to the threads on the current device. "device" may only be meaningful when working on a single node, and most of the time "system" is the correct choice.

For group_kind="block" (the default), every thread in the CTA must execute this operation convergently; the group collectively emits one remote update.

3.2.4.8 tle.signal_wait

tle.signal_wait waits until a local synchronization slot reaches its target value.

def signal_wait(device_dptr, slot_id, wait_kind, target=None,
               group_kind="block", context_id=0, order="acquire"):
    """
    Wait until a local synchronization slot reaches its target.

    :param device_dptr: distributed communicator handle (required)
    :param slot_id: signal slot index (int32 scalar) (required)
    :param wait_kind: "signal", "counter", or "shadow" (required)
    :param target: optional; required for "signal"/"counter", must be omitted for "shadow"
    :param group_kind: optional; "thread", "warp", or "block"; default "block"
    :param context_id: optional; compile-time int selecting a pre-allocated network context; default 0
    :param order: optional; memory ordering constraint ("relaxed" or "acquire"); default "acquire"
    """
    pass

wait_kind selects the waiting mode: "signal" waits until the slot value reaches target; "counter" waits until a counter at the slot reaches target; "shadow" reads the target from the runtime's locally maintained shadow buffer, so target must be omitted. slot_id is interpreted in the same signal slot namespace as tle.signal. group_kind and context_id have the same semantics as in tle.signal.

order constrains the memory order of the wait operation. As tle.signal_wait is a read operation, only "relaxed" and "acquire" are allowed, and most of the time the default value "acquire" is the correct choice.

3.2.5 API Reference and Practical Examples

3.2.5.1 tle.load
  • Signature: tle.load(ptr, mask=None, other=None, is_async=False)
  • Use case: Keep tl.load semantics while adding async scheduling hints.
  • Practical guidance:
    • Use is_async=True for global-memory reads that are later reused in compute-heavy regions.
    • Keep mask and other explicit on boundary tiles to avoid undefined values.

Example: guarded async load for tail tiles

offs = base + tl.arange(0, BLOCK)
mask = offs < n_elements
x = tle.load(x_ptr + offs, mask=mask, other=0.0, is_async=True)

Example: async load + compute overlap pattern

for k in tl.range(0, K, BK, num_stages=2):
    a = tle.load(a_ptr + k * stride_a, is_async=True)
    b = tle.load(b_ptr + k * stride_b, is_async=True)
    acc = tl.dot(a, b, acc)
3.2.5.2 tle.extract_tile and tle.insert_tile
  • extract_tile: read a sub-tile view from a larger tile tensor.
  • insert_tile: write a processed sub-tile back to a larger tile tensor.
  • Typical use: local transforms (activation, quant/dequant, normalization) on sub-regions without manual pointer arithmetic.

Example: tilewise post-processing in registers

# x: [4, 4]
sub = x.extract_tile(index=[1, 0], shape=[2, 2])  # rows [2:4], cols [0:2]
sub = tl.maximum(sub, 0.0)  # ReLU on the sub-tile
x = x.insert_tile(sub, index=[1, 0])
3.2.5.3 tle.pipe
  • Signature: tle.pipe(*, capacity, scope="cta", name=None, readers=None, one_shot=False, **fields)
  • Purpose: create a typed pipe for explicit CTA-local producer/consumer dataflow, ring-buffer stage reuse, and synchronization edges.
  • Parameters:
    • capacity: compile-time positive integer, the number of pipe stages. The leading dimension of every payload field must equal capacity.
    • scope: the current MVP supports "cta".
    • name: optional pipe name for IR/diagnostics; if provided, it must be a string.
    • readers: optional reader-name list. Omit it for the default SPSC reader; pass values such as ("left", "right") for SPMC.
    • one_shot: whether the pipe is a single ready/full edge, useful for one-time broadcast data. one_shot=True does not support close.
    • **fields: one or more payload buffers. Each field must be a shared-memory buffered tensor returned by tle.gpu.alloc(..., scope=tle.gpu.smem) or a static subslice of one, with rank >= 2.
  • Endpoint API:
    • pipe.writer() -> pipe_writer
    • pipe.reader(name=None, fields=None) -> pipe_reader
    • writer.acquire(iter) -> pipe_slot
    • writer.commit(iter) -> None
    • writer.close(iter) -> None
    • reader.wait(iter) -> pipe_wait_result
    • reader.release(iter) -> None
  • Semantics:
    • iter maps to stage = iter % capacity, with a phase bit distinguishing ring-buffer reuse rounds.
    • writer.acquire returns the slot for the current stage; fields are exposed by name on the slot, for example slot.kv.
    • reader.wait returns {slot, is_closed}. Normal reads use wait.slot; close-aware flows inspect is_closed.
    • reader(..., fields=("kv_r",)) subscribes to only a subset of fields, reducing unnecessary dependencies.
    • Current lowering maps CTA-scoped SMEM pipes to GPU NVWS token/mbarrier synchronization.
    • The NVIDIA lowering recognizes a restricted multi-writer form when each writer contributes pure-TMA copies to a disjoint field set and all writers follow the same acquire/commit cadence. The union must cover every field. Static non-overlapping field subslices may share one backing SMEM allocation; overlapping fields or aliases that cannot be proven disjoint are rejected.

Example 1: SPSC double-buffered load/compute

smem = tle.gpu.alloc([2, BLOCK], dtype=tl.float32, scope=tle.gpu.smem)
pipe = tle.pipe(capacity=2, scope="cta", name="tile_pipe", tile=smem)
writer = pipe.writer()
reader = pipe.reader()
offs = tl.arange(0, BLOCK)

# producer partition
for i in tl.range(0, n_tiles):
    slot = writer.acquire(i)
    ptr = tle.gpu.local_ptr(slot.tile)
    tl.store(ptr, tl.load(gmem_ptr + i * BLOCK + offs))
    writer.commit(i)

# consumer partition
for i in tl.range(0, n_tiles):
    ready = reader.wait(i)
    tile = tl.load(tle.gpu.local_ptr(ready.slot.tile))
    acc += tile
    reader.release(i)

Example 2: SPMC readers and partial-field subscription

kv = tle.gpu.alloc([PIPE_CAPACITY, BK, D], dtype=tl.float16, scope=tle.gpu.smem)
meta = tle.gpu.alloc(
    [PIPE_CAPACITY, BK],
    dtype=tl.int32,
    scope=tle.gpu.smem,
    nv_mma_shared_layout=False,
)
pipe = tle.pipe(
    capacity=PIPE_CAPACITY,
    scope="cta",
    name="kv_pipe",
    readers=("qk", "value"),
    kv=kv,
    meta=meta,
)

writer = pipe.writer()
qk_reader = pipe.reader("qk")
value_reader = pipe.reader("value", fields=("kv",))

slot = writer.acquire(k)
tl.store(tle.gpu.local_ptr(slot.kv), kv_tile)
tl.store(tle.gpu.local_ptr(slot.meta), valid_mask.to(tl.int32))
writer.commit(k)

qk_slot = qk_reader.wait(k).slot
qk = tl.dot(q, tl.trans(tl.load(tle.gpu.local_ptr(qk_slot.kv))))
qk_reader.release(k)

value_slot = value_reader.wait(k).slot
acc = tl.dot(prob, tl.load(tle.gpu.local_ptr(value_slot.kv)), acc)
value_reader.release(k)

Example 3: mixed pipe with inferred TMA and local-store transports

q = tle.gpu.alloc([PIPE_CAPACITY, BM, BK], dtype=tl.float16, scope=tle.gpu.smem)
scale = tle.gpu.alloc(
    [PIPE_CAPACITY, BM],
    dtype=tl.float32,
    scope=tle.gpu.smem,
    nv_mma_shared_layout=False,
)

pipe = tle.pipe(
    capacity=PIPE_CAPACITY,
    scope="cta",
    name="mixed_inputs",
    readers=("mma", "epilogue"),
    q=q,
    scale=scale,
)

writer = pipe.writer()
mma_reader = pipe.reader("mma", fields=("q",))
epilogue_reader = pipe.reader("epilogue", fields=("q", "scale"))

slot = writer.acquire(k)
tle.gpu.copy(q_desc, slot.q, [BM, BK], [q_block, k_block])
tl.store(tle.gpu.local_ptr(slot.scale, (tl.arange(0, BM),)), scale_values)
writer.commit(k)

q_ready = mma_reader.wait(k).slot
q_tile = tl.load(tle.gpu.local_ptr(q_ready.q))
mma_reader.release(k)

epilogue_ready = epilogue_reader.wait(k).slot
scale_tile = tl.load(tle.gpu.local_ptr(epilogue_ready.scale, (tl.arange(0, BM),)))
epilogue_reader.release(k)

The example above is one logical pipe, not one TMA pipe plus one local-store pipe. slot.q is classified as a TMA-produced field because it is the destination of tle.gpu.copy. slot.scale is classified as a local-store field because it is written through tle.gpu.local_ptr and tl.store. Both fields become visible to subscribed readers after the same writer.commit(k) and the readers' wait(k).

Usage rules for mixed pipes:

  • Split pipes by logical lifecycle and reader protocol, not by transport.
  • Allocate one shared-memory field per payload component; all fields use the same leading capacity dimension.
  • Produce every field between the matching writer.acquire(iter) and writer.commit(iter).
  • Use the natural operation for each field: descriptor/TMA copy, cp.async-style copy, or tle.gpu.local_ptr plus tl.store.
  • Let the compiler infer the transport. Do not encode transport in pipe names, field names, or user attributes.
  • Use pipe.reader(name, fields=(...)) when a reader only consumes a subset of fields; this narrows the reader view without creating another token.
  • Keep pipe-field provenance visible. Opaque shared-memory pointer escapes, untracked shared stores, or unprovable overlapping writes are rejected instead of lowered with a silent fallback.

The current NVIDIA lowering maps accepted mixed pipes to NVWS/mbarrier synchronization. It supports provable TMA/local-store and TMA/cp.async mixed payloads at pipe-field root granularity and fails fast when the payload window, field ownership, participant count, or source-order safety cannot be proven.

3.2.5.4 tle.device_mesh + tle.sharding + tle.reshard
  • Recommended workflow:
    1. Define topology with tle.device_mesh.
    2. Mark tensor layout with tle.sharding.
    3. Transform layout with tle.reshard.
    4. Keep compute kernels operating on logical tensor views.

Example: split-by-device input, then all-gather before compute

mesh = tle.device_mesh({"node": 2, "device": 4})
x_spec = tle.sharding(mesh, split=["device"], partial=[])
x = tle.make_sharded_tensor(x_ptr, sharding=x_spec, shape=[M, K])

# [S] -> [B] on device axis (all-gather)
x_full = tle.reshard(x, spec=tle.sharding(mesh, split=[], partial=[]))
3.2.5.5 tle.shard_id
  • Signature: tle.shard_id(mesh, axis)
  • Returns current program's coordinate on a mesh axis.
  • axis can be a mesh-axis name (e.g. "node", "device", "cluster_x") or an axis index.
  • Typical use: build peer shard IDs for ring exchange, staged all-reduce, and cluster-cooperative kernels.

Example: query current program coordinates on node/device axes

mesh = tle.device_mesh({"node": 2, "device": 4})
node_rank = tle.shard_id(mesh, "node")      # 0..1
device_rank = tle.shard_id(mesh, "device")  # 0..3
3.2.5.6 tle.remote + tle.distributed_barrier
  • tle.remote reads/writes explicit remote shards.
  • tle.distributed_barrier synchronizes only the mesh/sub-mesh you pass in.

Example: remote read from neighbor shard (ring-like exchange)

node_rank = tle.shard_id(mesh, "node")
device_rank = tle.shard_id(mesh, "device")
next_device = (device_rank + 1) % mesh.shape[1]
remote_x = tle.remote(x, shard_id=(node_rank, next_device), scope=mesh)
tle.distributed_barrier(mesh)
neighbor_vals = tl.load(remote_x)
3.2.5.7 tle.signal + tle.signal_wait
  • tle.signal: atomically update a remote peer's synchronization slot without transferring data.
  • tle.signal_wait: block until a local slot reaches its target value.
  • Typical use: lightweight cross-device synchronization in pipelined producer/consumer kernels where the full cost of a collective barrier is unnecessary.
  • wait_kind="signal" waits for the signal in the slot to reach the target value. wait_kind="shadow" reads the target from the runtime's locally maintained shadow buffer, so no explicit target is needed.

Example: cross-device one-shot synchronization

@triton.jit
def signal_kernel(device_dptr, peer: tl.constexpr):
    # ... do work ...
    tle.signal(
        device_dptr, peer, slot_id=0,
        op="inc", space="inter_node",
    )


@triton.jit
def wait_kernel(device_dptr):
    tle.signal_wait(
        device_dptr, slot_id=0,
        wait_kind="signal", target=1,
    )
    # ... safe to proceed ...

3.3 TLE-Struct

  • Design philosophy: architecture-aware, fine-grained tuning.
  • Core idea: classify backends by hardware-topology families (e.g., GPGPU, DSA), expose common hierarchical parallel/storage structures, and let developers explicitly define structured compute/data mappings (e.g., warp-group control, pipeline scheduling). This decouples algorithm logic from hardware physical implementation at the abstraction level.

3.3.1 GPU

3.3.1.1 Memory Management
3.3.1.1.1 tle.gpu.memory_space

Specify tensor memory_space:

x = ...
x = tle.gpu.memory_space(x, "shared_memory")
3.3.1.1.2 tle.gpu.alloc

Allocate memory:

a_smem = tle.gpu.alloc(
    [XBLOCK, YBLOCK],
    dtype=tl.float32,
    layout=None,
    scope=tle.gpu.storage_kind.smem,
)
3.3.1.1.3 tle.gpu.local_ptr

Get memory pointers:

# pointers for a_smem[0, :]: [(0, 0), (0, 1), ..., (0, YBLOCK-1)]
a_smem_ptrs = tle.gpu.local_ptr(
    a_smem,
    indices=(tl.broadcast(0, [YBLOCK]), tl.arange(0, YBLOCK)),
)
  • Signature: tle.gpu.local_ptr(buffer, indices=None) -> tl.tensor | tl.ptr
  • Purpose: Build arbitrary-shaped pointer views over shared memory buffers for tl.load/tl.store/tl.atomic*.
  • Parameters:
    • buffer: buffered tensor returned by tle.gpu.alloc (SMEM/TMEM).
    • indices: optional tuple of integer tensors. Tuple length must equal rank(buffer), and all tensors must have identical shapes. If omitted/None, backend treats it as full indices.
  • Semantics:
    • If indices is provided: output pointer tensor shape equals common shape of index tensors.
    • For each logical output index (i0, i1, ...), pointer value corresponds to buffer[indices0(i0,...), indices1(i0,...), ...].
    • If indices=None: build full-view pointers over buffer shape (rank>0 returns pointer tensor with shape(buffer), rank=0 returns scalar pointer).
    • Returned pointers live in shared-memory address space (LLVM addrspace=3). Indices must be integers (i32/i64, etc.; lowered to i32).
    • Linearization is row-major (last dimension fastest); shared-memory layout/encoding follows buffer memdesc.

Example 1: 1D slice

smem = tle.gpu.alloc([BLOCK], dtype=tl.float32, scope=tle.gpu.smem)
# Slice [offset, offset + SLICE)
idx = offset + tl.arange(0, SLICE)
slice_ptr = tle.gpu.local_ptr(smem, (idx,))
vals = tl.load(slice_ptr)

Example 2: K-dimension tiling (matrix slice)

smem_a = tle.gpu.alloc([BM, BK], dtype=tl.float16, scope=tle.gpu.smem)
# Slice (BM, KW), where KW is K-dimension slice
rows = tl.broadcast_to(tl.arange(0, BM)[:, None], (BM, KW))
cols = tl.broadcast_to(tl.arange(0, KW)[None, :] + k_start, (BM, KW))
a_slice = tle.gpu.local_ptr(smem_a, (rows, cols))
a_vals = tl.load(a_slice)

Example 3: arbitrary gather view

smem = tle.gpu.alloc([H, W], dtype=tl.float32, scope=tle.gpu.smem)
# Take an offset column per row
rows = tl.broadcast_to(tl.arange(0, H)[:, None], (H, SLICE))
cols = tl.broadcast_to(1 + tl.arange(0, SLICE)[None, :], (H, SLICE))
gather_ptr = tle.gpu.local_ptr(smem, (rows, cols))
out = tl.load(gather_ptr)

Supported downstream ops:

  • tl.load
  • tl.store
  • tl.atomic_add/and/cas/max/min/or/xchg/xor

Practical notes:

  • Atomic ops require element dtype/backend support; use integer/float types supported by target hardware.
  • For local-pointer load-after-store hazards, TLE backend pass TleInsertLocalPointerBarriers inserts barriers automatically; add manual barriers only for custom synchronization patterns outside pass coverage.

Example 4: load/store/atomic on the same local_ptr

smem_i32 = tle.gpu.alloc([BLOCK], dtype=tl.int32, scope=tle.gpu.smem)
ptr = tle.gpu.local_ptr(smem_i32, (tl.arange(0, BLOCK),))

tl.store(ptr, tl.zeros([BLOCK], dtype=tl.int32))
tl.atomic_add(ptr, 1)
vals = tl.load(ptr)
3.3.1.1.4 tle.gpu.local_ptr (for remote)
  • Signature: tle.gpu.local_ptr(remote_buffer, indices=None) -> tl.tensor | tl.ptr
  • Purpose: materialize pointer views for remote shared/local buffers returned by tle.remote(...).
  • Inputs:
    • remote_buffer: result of tle.remote(buffer, shard_id, scope), where buffer is typically allocated by tle.gpu.alloc.
    • indices: same rules as local mode (None for full view, or tuple of integer tensors with identical shapes).
  • Semantics:
    • Pointer shape/linearization rules are identical to local tle.gpu.local_ptr.
    • Address resolution targets the remote shard selected by shard_id.
    • Use tle.distributed_barrier(...) when cross-shard producer/consumer ordering is required.

Example: read remote SMEM tile from neighbor shard

smem = tle.gpu.alloc([BM, BK], dtype=tl.float16, scope=tle.gpu.storage_kind.smem)
remote_smem = tle.remote(smem, shard_id=(node_rank, next_device), scope=mesh)

rows = tl.broadcast_to(tl.arange(0, BM)[:, None], (BM, BK))
cols = tl.broadcast_to(tl.arange(0, BK)[None, :], (BM, BK))
remote_ptr = tle.gpu.local_ptr(remote_smem, (rows, cols))

vals = tl.load(remote_ptr)
3.3.1.1.5 tle.gpu.copy

Memory copy:

tle.gpu.copy(a_ptrs + ystride_a * yoffs[None, :], a_smem, [XBLOCK, YBLOCK])
  • Signature: tle.gpu.copy(src, dst, shape, offsets=None, barrier=None, mask=None, cache_policy=None)
  • cache_policy: optional backend-specific compile-time policy object; the common API does not prescribe its fields.
  • MThreads accepts tle.gpu.mthreads.TMECachePolicy(inner="normal", outer="normal"). Each field accepts "none", "once", "normal", or "persist". This immutable object has value-based equality and hashing for JIT specialization.
  • Explicit policies are supported only by the MThreads backend for tensor-descriptor → shared-memory copies. Other backends, CommonIR, pointer copies, and shared-to-global copies reject a non-None policy.
  • Omitting the argument or passing Python None preserves the existing behavior (normal/normal for MUSA TME). Python None is different from the string "none".
  • The two hints apply to the same copy operation. A and B copies can use different policy objects; hints do not change barrier completion or synchronization.

Cache persistence is performance-sensitive. Start from the default and measure each non-default pair on the target kernel, shape, dtype, and layout. Streaming access alone does not justify inner="none": it can severely degrade GEMM performance. A benefit from outer="none" does not imply a benefit from inner="none". The names do not guarantee cache bypass, a fixed reuse count, or locked residency, and the physical L1/L2 mapping is not specified by this API.

MUSA example using the default persistence explicitly:

completion = tle.gpu.alloc_barrier(expect_bytes=BM * BK * 2)  # FP16/BF16
tle.gpu.copy(a_desc, a_smem, [BM, BK], [m_offset, k_offset],
             barrier=completion,
             cache_policy=tle.gpu.mthreads.TMECachePolicy(inner="normal", outer="normal"))
tle.gpu.barrier_wait(completion, phaseIdx=0)

A TMECachePolicy can be constructed on the host or inside a JIT kernel and passed as a tl.constexpr argument to a warp_specialize producer worker. Bare tuples and strings are not MThreads cache-policy objects.

3.3.1.2 Execution Orchestration
3.3.1.2.1 tle.gpu.warp_specialize

tle.gpu.warp_specialize creates an explicit warp-specialized region inside one CTA, placing different JIT functions into separate warp partitions. Typical uses include splitting TMA/cp.async producers, WGMMA consumers, epilogues, and reductions, with shared-memory data passed through tle.pipe or other explicit synchronization primitives.

  • Signature: tle.gpu.warp_specialize(functions_and_args, worker_num_warps, worker_num_regs)
  • Parameters:
    • functions_and_args: [(fn0, args0), (fn1, args1), ...]. Entry 0 is emitted into the default partition; later entries are emitted into worker partitions.
    • worker_num_warps: warp counts for worker partitions. Length must equal len(functions_and_args) - 1.
    • worker_num_regs: requested register counts for worker partitions. Length must equal len(functions_and_args) - 1.
  • Semantics:
    • Each args value must be a tuple. Plain Python int/float/bool/str/tuple/tl.dtype values and objects exposing a tl.constexpr_type through .type are passed as constexpr arguments.
    • The default partition may return values; the return value of tle.gpu.warp_specialize(...) comes from the default partition. Worker partitions perform side effects and end with warp return.
    • Worker callees receive the corresponding "ttg.num-warps" attribute, and the region records requestedRegisters.
    • Captured worker arguments are deduplicated in IR, so multiple workers can share the same pipe endpoint or buffer handle.
    • warp_specialize does not itself provide data-visibility ordering. Producer/consumer ordering should be expressed with tle.pipe commit/wait/release, barriers, or other synchronization primitives.

Example: one producer partition loads shared memory, one consumer worker computes.

@triton.jit
def producer(writer, x_ptr, n_tiles: tl.constexpr, BLOCK: tl.constexpr):
    offs = tl.arange(0, BLOCK)
    for i in tl.range(0, n_tiles):
        slot = writer.acquire(i)
        vals = tl.load(x_ptr + i * BLOCK + offs)
        tl.store(tle.gpu.local_ptr(slot.tile), vals)
        writer.commit(i)


@triton.jit
def consumer(reader, out_ptr, n_tiles: tl.constexpr, BLOCK: tl.constexpr):
    offs = tl.arange(0, BLOCK)
    acc = tl.zeros([BLOCK], dtype=tl.float32)
    for i in tl.range(0, n_tiles):
        ready = reader.wait(i)
        tile = tl.load(tle.gpu.local_ptr(ready.slot.tile))
        acc += tile
        reader.release(i)
    tl.store(out_ptr + offs, acc)


@triton.jit
def kernel(x_ptr, out_ptr, n_tiles: tl.constexpr, BLOCK: tl.constexpr):
    smem = tle.gpu.alloc([2, BLOCK], dtype=tl.float32, scope=tle.gpu.smem)
    pipe = tle.pipe(capacity=2, scope="cta", name="x_pipe", tile=smem)

    tle.gpu.warp_specialize(
        [
            (producer, (pipe.writer(), x_ptr, n_tiles, BLOCK)),
            (consumer, (pipe.reader(), out_ptr, n_tiles, BLOCK)),
        ],
        [4],      # consumer worker uses 4 warps
        [168],    # consumer worker requested registers
    )

Example: multiple workers with an SPMC pipe.

tile = tle.gpu.alloc([2, BM, BK], dtype=tl.float16, scope=tle.gpu.smem)
pipe = tle.pipe(
    capacity=2,
    scope="cta",
    name="spmc_tile",
    readers=("qk", "value"),
    tile=tile,
)

tle.gpu.warp_specialize(
    [
        (load_tile_producer, (pipe.writer(), a_desc, b_desc)),
        (qk_consumer, (pipe.reader("qk"), acc_qk)),
        (value_consumer, (pipe.reader("value", fields=("tile",)), acc_v)),
    ],
    [4, 4],
    [240, 168],
)

3.3.2 DSA

This section is rewritten from triton_v3.2.x (python/triton/experimental/tle/language/dsa and its README). DSA APIs are split into:

  • Generic DSA APIs under tle.dsa.*
  • Backend-specific address spaces under tle.dsa.ascend.*
3.3.2.1 Memory and Data Movement
3.3.2.1.1 tle.dsa.alloc
  • Signature: tle.dsa.alloc(shape, dtype, mem_addr_space)
  • Purpose: allocate DSA local buffers in a target memory space.

Ascend memory spaces exposed in source:

  • tle.dsa.ascend.UB
  • tle.dsa.ascend.L1
  • tle.dsa.ascend.L0A
  • tle.dsa.ascend.L0B
  • tle.dsa.ascend.L0C
a_ub = tle.dsa.alloc([XBLOCK, YBLOCK], dtype=tl.float32, mem_addr_space=tle.dsa.ascend.UB)
b_l1 = tle.dsa.alloc([XBLOCK, YBLOCK], dtype=tl.float32, mem_addr_space=tle.dsa.ascend.L1)
3.3.2.1.2 tle.dsa.copy
  • Signature: tle.dsa.copy(src, dst, shape, inter_no_alias=False)
  • Purpose: explicit movement between GMEM pointers and DSA local buffers (both directions).
tle.dsa.copy(x_ptrs, a_ub, [tail_m, tail_n])          # GMEM -> local buffer
tle.dsa.copy(a_ub, out_ptrs, [tail_m, tail_n])        # local buffer -> GMEM
3.3.2.1.3 tle.dsa.local_ptr
  • Signature: tle.dsa.local_ptr(buffer, indices=None) -> tl.tensor | tl.ptr
  • Purpose: build pointer views over DSA local buffers (for example UB/L1) for explicit local-memory access patterns.
  • Parameters:
    • buffer: DSA buffered tensor, typically from tle.dsa.alloc.
    • indices: optional tuple of integer tensors. If omitted/None, backend treats it as full indices.
  • Semantics:
    • Shape and indexing behavior follow tle.gpu.local_ptr (same pointer-view model).
    • Intended for DSA-local data access paths that require explicit pointer materialization.

Example:

a_ub = tle.dsa.alloc([BM, BK], dtype=tl.float16, mem_addr_space=tle.dsa.ascend.UB)
rows = tl.broadcast_to(tl.arange(0, BM)[:, None], (BM, BK))
cols = tl.broadcast_to(tl.arange(0, BK)[None, :], (BM, BK))
a_ptr = tle.dsa.local_ptr(a_ub, (rows, cols))
a_val = tl.load(a_ptr)
3.3.2.1.4 tle.dsa.local_ptr (for remote)
  • Signature: tle.dsa.local_ptr(remote_buffer, indices=None) -> tl.tensor | tl.ptr
  • Purpose: materialize pointer views over remote DSA local buffers obtained from tle.remote(...).
  • Inputs:
    • remote_buffer: result of tle.remote(dsa_buffer, shard_id, scope).
    • indices: same rules as local DSA mode.
  • Semantics:
    • Same pointer-view semantics as local DSA mode.
    • Pointer dereference is routed to the remote shard selected by shard_id.
    • Pair with tle.distributed_barrier when cross-shard ordering is required.

Example:

a_ub = tle.dsa.alloc([BM, BK], dtype=tl.float16, mem_addr_space=tle.dsa.ascend.UB)
remote_a_ub = tle.remote(a_ub, shard_id=peer_rank, scope=mesh)

rows = tl.broadcast_to(tl.arange(0, BM)[:, None], (BM, BK))
cols = tl.broadcast_to(tl.arange(0, BK)[None, :], (BM, BK))
remote_ptr = tle.dsa.local_ptr(remote_a_ub, (rows, cols))
remote_val = tl.load(remote_ptr)
3.3.2.1.5 tle.dsa.to_tensor / tle.dsa.to_buffer
  • tle.dsa.to_tensor(buffer, writable=True): convert a DSA buffer to a tensor view for tensor expressions.
  • tle.dsa.to_buffer(tensor, space): convert a tensor value back to a buffer in a target DSA address space.
c_val = tle.dsa.to_tensor(c_ub, writable=True)
result = c_val * 0.5
d_ub = tle.dsa.to_buffer(result, tle.dsa.ascend.UB)
tle.dsa.copy(d_ub, out_ptrs, [tail_m, tail_n])
3.3.2.2 Elementwise Compute Ops (buffer-based)

Builtins provided by source:

  • tle.dsa.add

  • tle.dsa.sub

  • tle.dsa.mul

  • tle.dsa.div

  • tle.dsa.max

  • tle.dsa.min

  • Common signature: tle.dsa.<op>(lhs, rhs, out)

  • Compute model: elementwise binary op over DSA local buffers.

  • Shape rules:

    • lhs, rhs, out must have the same rank and shape.
    • No implicit broadcast is assumed in this API layer.

  • Dtype rules:

    • Three operands should use the same dtype in practice.
    • Integer dtypes are typical for index/count paths; float dtypes are typical for activation/math paths.

  • Memory-space rules:

    • Buffers should be allocated in compatible DSA local spaces (for example UB/L1 combinations allowed by backend).
    • Keep hot operands/results in local space to avoid extra GMEM traffic.

Per-op semantics:

  • tle.dsa.add(lhs, rhs, out): out = lhs + rhs
  • tle.dsa.sub(lhs, rhs, out): out = lhs - rhs
  • tle.dsa.mul(lhs, rhs, out): out = lhs * rhs
  • tle.dsa.div(lhs, rhs, out): out = lhs / rhs (backend-dependent precision/rounding)
  • tle.dsa.max(lhs, rhs, out): out = max(lhs, rhs)
  • tle.dsa.min(lhs, rhs, out): out = min(lhs, rhs)

In-place usage:

  • You can reuse the same output buffer across steps, for example tle.dsa.mul(tmp, b, tmp).
  • Avoid aliasing inputs/outputs unless backend semantics explicitly allow it.

Example 1: arithmetic chain ((a - b) * b) / scale

a_ub = tle.dsa.alloc([BM, BK], dtype=tl.float16, mem_addr_space=tle.dsa.ascend.UB)
b_ub = tle.dsa.alloc([BM, BK], dtype=tl.float16, mem_addr_space=tle.dsa.ascend.UB)
scale_ub = tle.dsa.alloc([BM, BK], dtype=tl.float16, mem_addr_space=tle.dsa.ascend.UB)
tmp_ub = tle.dsa.alloc([BM, BK], dtype=tl.float16, mem_addr_space=tle.dsa.ascend.UB)
out_ub = tle.dsa.alloc([BM, BK], dtype=tl.float16, mem_addr_space=tle.dsa.ascend.UB)

tle.dsa.copy(a_ptrs, a_ub, [BM, BK])
tle.dsa.copy(b_ptrs, b_ub, [BM, BK])
tle.dsa.copy(scale_ptrs, scale_ub, [BM, BK])

tle.dsa.sub(a_ub, b_ub, tmp_ub)      # tmp = a - b
tle.dsa.mul(tmp_ub, b_ub, tmp_ub)    # tmp = tmp * b
tle.dsa.div(tmp_ub, scale_ub, out_ub)  # out = tmp / scale

tle.dsa.copy(out_ub, out_ptrs, [BM, BK])

Example 2: clamp by max + min

x_ub = tle.dsa.alloc([BM, BK], dtype=tl.float16, mem_addr_space=tle.dsa.ascend.UB)
floor_ub = tle.dsa.alloc([BM, BK], dtype=tl.float16, mem_addr_space=tle.dsa.ascend.UB)
ceil_ub = tle.dsa.alloc([BM, BK], dtype=tl.float16, mem_addr_space=tle.dsa.ascend.UB)
tmp_ub = tle.dsa.alloc([BM, BK], dtype=tl.float16, mem_addr_space=tle.dsa.ascend.UB)
y_ub = tle.dsa.alloc([BM, BK], dtype=tl.float16, mem_addr_space=tle.dsa.ascend.UB)

tle.dsa.copy(x_ptrs, x_ub, [BM, BK])
tle.dsa.copy(floor_ptrs, floor_ub, [BM, BK])
tle.dsa.copy(ceil_ptrs, ceil_ub, [BM, BK])

tle.dsa.max(x_ub, floor_ub, tmp_ub)  # tmp = max(x, floor)
tle.dsa.min(tmp_ub, ceil_ub, y_ub)   # y = min(tmp, ceil)

tle.dsa.copy(y_ub, y_ptrs, [BM, BK])
3.3.2.3 Loop and Hint APIs

Source includes:

  • tle.dsa.pipeline(...)
  • tle.dsa.parallel(...)
  • tle.dsa.hint(...) (used as with tle.dsa.hint(...) compile-time hints)
with tle.dsa.hint(inter_no_alias=True):
    tle.dsa.copy(x_ptr + offs, a_ub, [tail_size], inter_no_alias=True)
3.3.2.4 Slice/View Utilities

Source includes:

  • tle.dsa.extract_slice
  • tle.dsa.insert_slice
  • tle.dsa.extract_element
  • tle.dsa.subview
sub = tle.dsa.extract_slice(full, offsets=(0, k0), sizes=(BM, BK), strides=(1, 1))
full = tle.dsa.insert_slice(full, sub, offsets=(0, k0), sizes=(BM, BK), strides=(1, 1))
elem = tle.dsa.extract_element(sub, indice=(i, j))

3.3.3 Struct API Cookbook

3.3.3.1 Shared-memory staging (alloc + copy + local_ptr)

Use this pattern when data is reused across multiple math operations.

# 1) Allocate SMEM tile
a_smem = tle.gpu.alloc([BM, BK], dtype=tl.float16, scope=tle.gpu.storage_kind.smem)

# 2) Copy GMEM -> SMEM
tle.gpu.copy(a_ptrs, a_smem, [BM, BK])

# 3) Build local pointer view and load
rows = tl.broadcast_to(tl.arange(0, BM)[:, None], (BM, BK))
cols = tl.broadcast_to(tl.arange(0, BK)[None, :], (BM, BK))
a_ptr_local = tle.gpu.local_ptr(a_smem, (rows, cols))
a_tile = tl.load(a_ptr_local)
3.3.3.2 Shared-memory atomics with local_ptr

Useful for histogram, bucketization, and radix-select style counting.

bins = 256
counts = tle.gpu.alloc([bins], dtype=tl.int32, scope=tle.gpu.storage_kind.smem)
idx = tl.arange(0, BLOCK) % bins
count_ptr = tle.gpu.local_ptr(counts, (idx,))
tl.atomic_add(count_ptr, 1)
3.3.3.3 DSA local-buffer flow (dsa.alloc + dsa.copy + dsa.to_tensor/to_buffer)

Use this for DSA backends that expose dedicated local buffer spaces.

a_ub = tle.dsa.alloc([BM, BK], dtype=tl.float16, mem_addr_space=tle.dsa.ascend.UB)
b_ub = tle.dsa.alloc([BM, BK], dtype=tl.float16, mem_addr_space=tle.dsa.ascend.UB)
c_ub = tle.dsa.alloc([BM, BK], dtype=tl.float16, mem_addr_space=tle.dsa.ascend.UB)

tle.dsa.copy(a_ptrs, a_ub, [BM, BK])
tle.dsa.copy(b_ptrs, b_ub, [BM, BK])
tle.dsa.add(a_ub, b_ub, c_ub)

c_val = tle.dsa.to_tensor(c_ub, writable=True)
out_ub = tle.dsa.to_buffer(c_val, tle.dsa.ascend.UB)
tle.dsa.copy(out_ub, out_ptrs, [BM, BK])

3.4 TLE-Raw

  • Design philosophy: native passthrough and maximal control.
  • Core idea: break DSL abstraction boundaries and support inlined vendor-native code. Target instructions are generated through vendor-private pipelines, bypassing general compiler middle layers and giving experts strong control over instruction scheduling, register allocation, and low-level synchronization primitives.

Content not available outside Feishu document yet.

Open question: should Raw integration be limited to Python DSL only?

3.4.1 Language Extensions

3.4.1.1 MLIR
from typing import Annotated
from mlir import ir
from mlir.dialects import arith, nvvm, tensor
import triton.language as tl
from triton.experimental.flagtree.edsl import dialect
import triton.experimental.flagtree.language as fl

# 1. Dialect declaration
@tle.raw.language(name="mlir")
# 2. Hardware constraints
@tle.hardware_constraint(threads_dim=1, sync_scope="block")
# 3. Function implementation
def vector_add_tile(
    x: Annotated[ir.RankedTensorType, "tensor<1024xf32>"],
    y: Annotated[ir.RankedTensorType, "tensor<1024xf32>"],
    output: Annotated[ir.RankedTensorType, "tensor<1024xf32>"]
):
    tidx = nvvm.ThreadIdXOp(ir.IntegerType.get_signless(32)).res
    bidx = nvvm.BlockIdXOp(ir.IntegerType.get_signless(32)).res
    bdimx = nvvm.BlockDimXOp(ir.IntegerType.get_signless(32)).res
    idx = arith.addi(arith.muli(bidx, bdimx), tidx)
    idx = arith.index_cast(ir.IndexType.get(), idx)
    xval = tensor.extract(x, [idx])
    yval = tensor.extract(y, [idx])
    result = arith.addf(xval, yval)
    tensor.insert(result, output, [idx])

@tle.jit
def add_kernel(
    x_ptr, y_ptr, output_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    output = tl.zeros_like(x)

    # 4. Function call
    tle.call(
        vector_add_tile,
        args=[x, y, output],
        hardware={
            "threads": (BLOCK_SIZE,),
        },
        layout={
            x: {"space": "shared", "order": [0]},
            y: {"space": "shared", "order": [0]},
            output: {"space": "shared", "order": [0]},
        }
    )
    tl.store(output_ptr + offsets, output, mask=mask)

4. Examples and Evaluation

4.1 SparseMLA

Optimization and tests have been conducted for SparseMLA in DSA on H800.

4.1.1 DeepSeek V3.2 SparseMLA Prefill

The cases match the FlashMLA V3.2 sparse prefill performance fixture, with attn_sink omitted because the local Triton, TLE, and TileLang kernels do not implement it: B=1, S=4096, H=128, HKV=1, DQK=576, DV=512, topk=2048.

Latency in milliseconds:

SKV Triton TLE TLE-Pipe-Pipelined TLE-FlashMLA-Prefill TileLang TileLang-Pipelined TileLang-Seesaw FlashMLA
8192 9.896 6.927 4.832 4.273 62.409 5.432 5.160 3.850
32768 11.210 7.624 5.321 4.834 75.160 6.428 5.577 4.117
65536 11.655 8.378 5.731 5.305 84.432 6.865 5.786 4.348
98304 11.835 8.658 5.972 5.599 86.561 7.139 5.873 4.447
131072 11.923 8.863 6.122 5.887 87.448 7.143 5.916 4.534

Speedup summary:

SKV TLE-Pipe over Triton TLE-FlashMLA over Triton TLE-FlashMLA over TLE-Pipe TLE-FlashMLA over FlashMLA TLE-FlashMLA over TileLang-Seesaw
8192 2.05x 2.32x 1.13x 0.90x 1.21x
32768 2.11x 2.32x 1.10x 0.85x 1.15x
65536 2.03x 2.20x 1.08x 0.82x 1.09x
98304 1.98x 2.11x 1.07x 0.79x 1.05x
131072 1.95x 2.03x 1.04x 0.77x 1.00x

4.2 MoeAlignBlockSize

With shared-memory extensions in tle-struct, it is possible to implement vllm/sglang-style moe_align_block_size and improve performance.

4.2.1 RTX 5060 Ti

num_tokens triton triton_atomic tle_atomic_fused [ours] tle_cluster_fused [ours] sglang_cuda Speedup (sglang_cuda / min(tle_atomic_fused, tle_cluster_fused))
256 0.0348 0.0302 0.0323 0.0097 0.0138 1.42x
512 0.0369 0.0301 0.0240 0.0117 0.0138 1.18x
1024 0.0369 0.0313 0.0179 0.0117 0.0139 1.19x
2048 0.0368 0.0313 0.0158 0.0131 0.0138 1.05x
4096 0.0369 0.0301 0.0138 0.0143 0.0148 1.07x
8192 0.0369 0.0313 0.0138 0.0164 0.0179 1.30x
16384 0.0369 0.0301 0.0158 0.0205 0.0240 1.52x
32768 0.0389 0.0322 0.0179 0.0301 0.0312 1.74x
65536 0.0430 0.0374 0.0225 0.0486 0.0507 2.25x
163840 0.0609 0.0512 0.0384 0.1036 0.1001 2.61x

4.2.2 H800

num_tokens triton triton_atomic tle_atomic_fused [ours] tle_cluster_fused [ours] sglang_cuda Speedup (sglang_cuda / min(tle_atomic_fused, tle_cluster_fused))
256 0.0260 0.0408 0.0445 0.0133 0.0160 1.20x
512 0.0262 0.0399 0.0315 0.0140 0.0162 1.16x
1024 0.0274 0.0401 0.0239 0.0158 0.0163 1.03x
2048 0.0509 0.0422 0.0226 0.0169 0.0173 1.02x
4096 0.0265 0.0412 0.0200 0.0177 0.0187 1.06x
8192 0.0476 0.0416 0.0192 0.0211 0.0230 1.20x
16384 0.0548 0.0441 0.0219 0.0256 0.0286 1.31x
32768 0.0443 0.0441 0.0221 0.0358 0.0401 1.81x
65536 0.0361 0.0481 0.0273 0.0561 0.0645 2.36x
163840 0.0509 0.0626 0.0451 0.1177 0.1323 2.93x

4.2.3 H800 Real Data (build/gems/moe_topk_ids.pt)

  • Runtime config: num_tokens=163840, num_experts=512, block_size=16, source=real.
num_tokens num_experts block_size triton triton_atomic tle_atomic_fused [ours] tle_cluster_fused [ours] sglang_cuda Speedup (sglang_cuda / min(tle_atomic_fused, tle_cluster_fused))
163840 512 16 0.0471 0.0535 0.0387 0.0750 0.1467 3.79x

4.2.4 RTX 5060 Ti Real Data (build/gems/moe_topk_ids.pt, Local Measurement)

  • Runtime config: num_tokens=163840, num_experts=512, block_size=16, source=real.
  • Runtime command: conda run -n flagtree python python/tutorials/tle/02-moe_align_block_size.py --skip_correctness --real_data build/gems/moe_topk_ids.pt --num_experts 512 --block_size 16
num_tokens num_experts block_size triton triton_atomic tle_atomic_fused [ours] tle_cluster_fused [ours] sglang_cuda Speedup (sglang_cuda / min(tle_atomic_fused, tle_cluster_fused))
163840 512 16 0.0507 0.0395 0.0261 0.0532 0.1060 4.06x

4.3 TopK

With shared-memory extensions in tle-struct, radix-select-based TopK can improve performance in MoE scenarios with large N and small K.

4.3.1 RTX 5060 Ti (tle-topk-radix-vs-torch)

M N K Triton-RadixSelect Torch-TopK Speedup (Torch / Triton-RadixSelect)
64 128 8 0.008192 0.010240 1.25x
64 1024 32 0.008192 0.020480 2.50x
64 8192 128 0.026624 0.059392 2.23x
128 32768 256 0.124928 0.192512 1.54x

4.3.2 H800 (tle-topk-radix-vs-torch)

M N K Triton-RadixSelect Torch-TopK Speedup (Torch / Triton-RadixSelect)
64 128 8 0.008384 0.017536 2.09x
64 1024 32 0.010688 0.024304 2.27x
64 8192 128 0.029952 0.057184 1.91x
128 32768 256 0.092256 0.117856 1.28x

4.4 TopK Selector

TopK selector performance is evaluated with python/tutorials/tle/deepseek_v32/01-topk_selector.py (plot_name=tle-radix-topk-selector).

4.4.1 RTX 5060 Ti (Local Measurement)

  • Runtime: local benchmark (GeForce RTX 5060 Ti), --skip_correctness --warmup 10 --rep 80.
batch seq_len topk Torch-TopK Triton-Radix TileLang TLE-Radix Speedup (Torch-TopK / TLE-Radix)
64 4096 128 0.038912 0.039456 0.020480 0.015808 2.46x
64 8192 256 0.088624 0.053248 0.028672 0.023936 3.70x
64 32768 1024 0.158272 0.131616 0.073728 0.062912 2.52x
64 32768 2048 0.163264 0.133120 0.075776 0.065536 2.49x

4.4.2 H800 (tle-radix-topk-selector)

batch seq_len topk Torch-TopK Triton-Radix TileLang TLE-Radix Speedup (Torch-TopK / TLE-Radix)
64 4096 128 0.045728 0.054256 0.017200 0.017472 2.62x
64 8192 256 0.097344 0.072512 0.020960 0.020928 4.65x
64 32768 1024 0.125008 0.176768 0.043088 0.041856 2.99x
64 32768 2048 0.125072 0.179264 0.044256 0.041984 2.98x