Technical Deep Dive | How to Use FlagOS’s New Triton-TLE Language to Build a TopK Selector Faster Than FlashInfer

In the landscape of large language model (LLM) inference, particularly as context windows expand from 128K to 1M+ tokens, the TopK Selector has emerged as a critical bottleneck for end-to-end latency.

While 45 microseconds per layer on an H800 might seem negligible, scaling that across a 60-layer model adds 2.7ms of additional latency. In the pursuit of extreme inference speed, every millisecond counts. This article explores how we leveraged Triton-TLE — a language extension from the FlagOS community — to outpace industry standards like FlashInfer and TRT-LLM.

01. Why the TopK Selector Needs Optimization

As context windows grow, the O(N²‌) computational and memory cost of standard Attention becomes prohibitive. Architectures like DeepSeek Sparse Attention (DSA) mitigate this by limiting the number of KV Cache tokens involved in computation.

The DSA workflow follows three steps:

1.Generate logits for the current query against all historical tokens. 2.Select the top-k indices from these logits (the TopK Selector). 3.Perform subsequent attention calculations only on the KV of these k positions.

As sequence length increases, step 2 accounts for a rising percentage of total latency. Unlike torch.topk, the specialized TopK Selector only requires indices, does not require full sorting, and focuses on minimizing latency for small batches (batch=1), where overhead directly impacts end-to-end response time.

02.Traditional GPU Approach: Radix Selection

To perform TopK on a GPU (especially in scenarios where only indices are needed and sorting is not required), Radix Selection is a common solution: GPUs are good at parallel counting and prefix sums, and radix selection does not require full sorting; it can compress the candidate range round by round. The basic process is:

●Mapping: Convert floating-point numbers to comparable unsigned integers. ●Binning: Divide bits into segments from high to low. ●Histogramming: Count frequencies in each bin. ●Thresholding: Identify which bin contains the k-th element and discard the rest. ●Collection: Gather elements exceeding the final threshold.

The pseudocode example is as follows:

# 返回第k大值（k从1开始）
def radix_select_kth_largest(A, k):
    C = A
    rank = k - 1                      # 0-based
    exp = highest_power_of_10(max(A)) # 例如 839 -> 100
    while exp > 0:
        count[0..9] = 0
        for x in C:
            d = (x // exp) % 10
            count[d] += 1
        acc = 0
        chosen = 0
        for d from 9 downto 0:        # 找“第rank个最大”落在哪个桶
            if rank < acc + count[d]:
                chosen = d
                rank = rank - acc
                break
            acc += count[d]
        C = [x in C where ((x // exp) % 10) == chosen]
        exp = exp // 10
    return C[0]
def radix_topk(A, k, need_sorted=true):
    T = radix_select_kth_largest(A, k)
    out = [x in A where x > T]
    need = k - len(out)
    for x in A:
        if x == T and need > 0:
            out.append(x)
            need -= 1
    if need_sorted:
        sort out in descending order
    return out

However, traversing the entire sequence in every round creates massive memory overhead as sequence length grows.

03. Optimizations from TileLang and TRT-LLM

To address the issues of traditional radix selection, TileLang and TRT-LLM have introduced mature optimization strategies: ●TileLang: Reduces memory access by keeping candidate indices on-chip (shared memory) after an initial 8-bit screening. ●TRT-LLM: Uses a four-stage screening process with early stopping (when bucket size ≤ 4096) and vectorized instructions to maximize global memory throughput.

4. The Bottleneck of Standard Triton Triton-TLE (Triton Language Extension) fills these gaps by providing primitives for explicit on-chip buffer management and cluster-level coordination. Key features include: ●tle.gpu.alloc: Explicitly allocate on-chip buffers; ●tle.gpu.local_ptr: Construct pointer views on on-chip buffers to avoid manual address calculation; ●tle.remote: Access on-chip buffers of other blocks; ●tle.device_mesh: Define the organization of block clusters; ●tle.distributed_barrier: Perform scoped synchronization within a cluster. With Triton-TLE, implementations relying on shared memory and clusters can be written naturally without the need to write manual CUDA.

5. Triton-TLE: Extending Triton’s Power When replicating the TRT-LLM selector with Triton-TLE, first place the histogram, thresholds, output indices, counts, and temporary buffers in shared memory, allowing histogram updates, candidate writing, and final sorting to be completed in a closed loop on-chip. The method of allocating shared memory in Triton-TLE:

@triton.jit
def tle_topk_selector_kernel(...):
    ...
    HIST_SIZE: tl.constexpr = 4096
    s_histogram = tle.gpu.alloc(
        [HIST_SIZE],
        dtype=tl.int32,
        layout=None,
        scope=tle.gpu.smem,
        nv_mma_shared_layout=False,
    )

tle.gpu.local_ptr allows the code to naturally express: performing load / store / atomic_add on the histogram in shared memory; writing candidate indices and values into the final bucket buffer; and re-reading these candidates during the final sort stage.

The method of accessing shared memory in Triton-TLE:

@triton.jit
def tle_topk_selector_kernel(...):
    ...
    flush_chunks: tl.constexpr = (TOPK + BLOCK_SIZE - 1) // BLOCK_SIZE
    for flush_chunk in tl.static_range(flush_chunks):
        pos = flush_chunk * BLOCK_SIZE + lane
        mask = pos < TOPK
        out_vals = tl.load(tle.gpu.local_ptr(s_out_indices, (pos, )), mask=mask, other=-1)
        tl.store(out_row + pos * stride_outn, out_vals, mask=mask)

Furthermore, aiming at the problem of insufficient single-block parallelism for batch=1 long sequences, TLE distributed + DSMEM splits one row among multiple blocks for joint processing: use device_mesh to define a cluster, each block is responsible for a portion of tiles, each block first calculates its own local histogram, summarizes the local histogram to Rank 0’s shared memory via remote, then synchronizes using distributed_barrier, and all blocks continue screening according to the new threshold. In this way, the serial nature of single-row scanning is dispersed, and summarization is still completed on-chip.

The gain of this step in histogram statistics based on local_ptr lies in dispersing the workload of a single row: the single-block version is limited by the serial nature of single-row scanning; the cluster version splits single-row scanning across multiple blocks; summarization is still completed on-chip without retreating to global memory. The final performance results are as follows:

●Hardware: Single NVIDIA H800 card ●Sequence length: Covered up to 512K ●Comparison objects: Triton, TRT-LLM prefill, TRT-LLM prefill-1024T, FlashInfer, TileLang, TLE (ours)

Note 1: TRT-LLM uses num_threads=512 by default. We found that using 1024 has higher performance on H800, so we added the TRT-LLM 1024T test.

Note 2: The TileLang algorithm will have candidate set overflow in the seq_len ≥ 262144 test, leading to incorrect results, marked as N/A in the table.

1. Triton-TLE replication of TileLang algorithm performance is close to the native version

Note: TLE-TileLang refers to the TileLang algorithm replicated using Triton-TLE

2. batch=1 Performance (Core Highlight)

The TLE cluster version takes 0.030ms at 131072 length, better than FlashInfer’s 0.045ms and TRT-LLM 1024T’s 0.049ms; at 262144 length, it is 0.038ms, and FlashInfer is 0.048ms. Compared with TRT-LLM, the maximum acceleration is about 2.5 times.

3. batch=64 Performance

The Triton-TLE version is close to TRT-LLM 1024T in performance, with a gap within 10%.

Summary

The secret to an ultra-fast TopK Selector lies in vectorized memory access and rapid candidate reduction. While native Triton struggles with on-chip state management and inter-block collaboration, Triton-TLE provides the necessary toolkit to break these barriers.

For developers looking to squeeze every bit of performance out of heterogeneous AI chips, Triton-TLE offers a path to implement complex CUDA-like optimizations while maintaining the productivity of the Triton ecosystem.

References & Resources

● TileLang’s approach: tilelang/examples/deepseek_v32/topk_selector.py at v0.1.8 · tile-ai/tilelang · GitHub

● TRT-LLM prefill’s approach: TensorRT-LLM/cpp/tensorrt_llm/kernels/indexerTopK.cu at v1.3.0rc10 · NVIDIA/TensorRT-LLM · GitHub

● Using TLE to replicate the TRT-LLM selector: https://github.com/flagos-ai/FlagTree/blob/f9a8d23602a65ec5c1af3b117e1faa46fe6f63b7/python/tutorials/tle/deepseek_v32/01-topk_selector.py#L658

● Using TLE distributed + DSMEM to optimize batch=1: https://github.com/flagos-ai/FlagTree/blob/f9a8d23602a65ec5c1af3b117e1faa46fe6f63b7/python/tutorials/tle/deepseek_v32/01-topk_selector.py#L3055

● Detailed description of Triton-TLE: TLE · flagos-ai/FlagTree Wiki · GitHub

About FlagOS

FlagOS is an open-source, unified system software stack designed for heterogeneous AI chips, initiated by the Beijing Academy of Artificial Intelligence (BAAI) and industry partners.

Official website: https://flagos.io

GitHub Project Address: FlagOS · GitHub

GitCode Project Address:AtomGit | GitCode — 全球开发者的开源社区,开源代码托管平台