Triton-TLE: 3-Layer Triton Extension for Faster GPU Kernels

AI Guide

Triton made GPU kernel development sane again: write tile-level Python, let the compiler handle threads, registers, and pipelining. But as GPUs get weirder (thread block clusters, DSMEM) and models get sparser (MoE routing, 1M-token sparse attention), stock Triton runs out of expressive room. You either hand-write CUDA or fight the DSL.

That's the gap FlagTree — the open-source multi-backend Triton fork from BAAI's FlagOS stack — attacks with TLE (Triton Language Extensions): a three-layer extension that stays source-compatible with Triton while progressively exposing hardware control. No fork of your kernels needed to get started.

The three layers: pick your depth

  • TLE-Lite — high-level semantic hints for algorithm engineers. Think sub-tile slicing, async loads, producer-consumer pipelines, device meshes. Write once, run anywhere.
  • TLE-Struct — architecture-aware control (GPGPU vs DSA). You describe structured memory intent; the compiler maps a "local buffer" to shared memory on a GPU or scratchpad SRAM on a DSA.
  • TLE-Raw — native passthrough for performance experts: inline CUDA/assembly or vendor intrinsics inside the Triton program, for the paths a DSL can't express.

All three lower into one kernel through LLVM IR (via FLIR), so you can mix depths within a single program.

What the primitives look like

1. Sub-tile slicing without address math

The classic Triton pain: pulling a logical sub-block out of a big tensor means hand-computing offsets, masks, and boundary checks. TLE-Lite gives you extract_tile / insert_tile:

# x is [4, 4]; split into 2x2 sub-tiles, take tile [0, 0]
z = x.extract_tile(index=[0, 0], shape=[2, 2])

# y is [2, 2]; write it back into x at sub-tile [0, 0]
z = x.insert_tile(y, index=[0, 0])

Because the compiler sees a regular tile access, it can optimize layout, vectorization, and bank conflicts — a direct win for sparse attention, local normalization, and per-block statistics.

2. Producer-consumer pipelines via tle.pipe

Software pipelining with tl.range(num_stages=...) is automatic but opaque. tle.pipe makes the dataflow edge explicit: a writer acquires a slot, fills it, commits; a reader waits, consumes, releases:

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, reader = pipe.writer(), pipe.reader()
offs = tl.arange(0, BLOCK)

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

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

Barriers, buffer reuse, and sync become compiler-managed but still analyzable — which is what lets the compiler overlap load and compute instead of serializing them.

3. Cluster-level coordination for long-sequence TopK

The showcase is the TopK selector used by DeepSeek Sparse Attention. Batch=1 long sequences have almost no parallelism inside one block, so TLE splits a single row across a cluster of blocks:

topology = {
    "node": [("node_x", 2), ("node_y", 2)],
    "device": 4,
    "block_cluster": [("cluster_x", 2), ("cluster_y", 2)],
    "block": 4,
}
mesh = tle.device_mesh(topology=topology)

# each block keeps a local histogram in shared memory
s_histogram = tle.gpu.alloc([4096], dtype=tl.int32, scope=tle.gpu.smem)

# summarize local histograms on rank 0 via remote, then scope-sync
tle.remote(...)               # read another block's on-chip memory
tle.distributed_barrier(mesh) # sync only the cluster, not the world

Histogram updates, candidate writes, and the final sort all stay on-chip — no round trip through global memory — while the cluster spreads the single-row scan across blocks. Full working kernel: python/tutorials/tle/deepseek_v32/01-topk_selector.py in the repo.

The numbers

  • TopK selector (H800, batch=1, 131K tokens): TLE cluster version 0.030 ms vs FlashInfer 0.045 ms and TRT-LLM 0.049 ms — up to ~2.5× over TRT-LLM; it stays stable at 512K while TileLang's candidate set overflows.
  • Radix select: a TLE reimplementation of the TRT-LLM algorithm lands at ~85–97% of native TRT-LLM across shapes.
  • SparseMLA (128K context): pipeline primitives reach ~90% of the FlashMLA baseline.
  • FlagOSTune autotuning: on H20 MM shapes the search space collapses from 620K+ to 4,070 configs — a 120× autotune speedup with negligible perf loss. Tuning real-model kernels improved them 1.21–7.35× across NVIDIA, Moore Threads, and MetaX chips.
  • Compiler passes: layout-conversion elimination removes 68–79% of convert_layout ops (up to 71% kernel speedup); instruction reordering hoists independent loads for 1.19–1.61× average (2× peak).

Getting started

  1. Clone the repo: git clone https://github.com/flagos-ai/flagtree.git and check out main (NVIDIA backend, Triton 3.6; Ascend lives on triton_v3.5.x, Cambricon on triton_v3.2.x).
  2. Follow the NVIDIA user manual to build and install.
  3. Run the tutorials under python/tutorials/tle/ — start with the TopK selector.
  4. Migrate incrementally: start with TLE-Lite hints on an existing kernel; reach for TLE-Struct/Raw only where profiling shows a bottleneck.

Bottom line

TLE isn't a new language — it's an off-ramp from Triton's abstraction ceiling. If you're writing sparse-attention, MoE, or multi-chip kernels, it's worth a weekend: the TopK selector tutorial alone is the fastest way to see cluster-level Triton done right.

Resources

Scroll to Top