On August 1st, the 9th Meet AI Compiler Technical Salon was held in Beijing. This event focused on the latest advancements in AI compilation technology, with several experts from industry and research institutions sharing their insights on programming languages, operator development, compilation optimization, and inference execution, showcasing the collaborative evolution of AI compilers from high-level language expression to hardware execution.

in,Guo Hui and Xiao Hang, researchers at BAAI AI Compiler, presented on "FlagTree: Triton-TLE Language Extension, Tile IR Backend and Compiler Optimization Practices," introducing a series of explorations conducted by the FlagTree team around Triton.

Faced with increasingly complex hardware architectures and more irregular model operators, the team proposed the Triton Language Extension (TLE) to try to establish a progressive channel between high-level DSL and low-level hardware control; at the same time, they carried out compilation optimizations on the FlagTree compiler, focusing on TileIR backend, automatic tuning, data layout, instruction scheduling, and other aspects.

The core goal of this work is to allow developers to delve deeper into hardware details as needed, while preserving Triton's ease of use and community ecosystem, and to cover rapid operator optimization, architecture-aware tuning, and native code-level optimization with a relatively unified development system.

HyperAI has compiled and summarized the shared content without altering its original meaning.

Follow the WeChat official account "HyperAI" and reply with the keyword "" in the background.0801 AI CompilerYou can obtain the authorized speaker's presentation PPT by clicking "...".

Building upon Triton, rebalancing abstraction and performance.

In recent years, AI compilers have faced an increasingly prominent contradiction:Hardware architectures and model operators are becoming increasingly complex, but developers still hope to use higher-level DSLs to achieve performance close to that of expert-written kernels.

If the abstraction is too high, the compiler may not be able to obtain enough information to complete in-depth optimization; if the abstraction is too low, developers will return to the complex development mode of CUDA or vendor proprietary languages.

Triton's success lies in elevating GPU operator development from the thread level to the tile level.Users use the Python DSL to describe the computational relationships between data blocks, while tasks such as thread mapping, register allocation, data layout, pipelining, and synchronization are mainly handled by the compiler.

This approach lowers the development barrier for high-performance operators and has also fostered a large community ecosystem. However, with the continuous development of next-generation GPUs, domain-specific architectures, and domestically produced AI chips, Triton's original abstraction has begun to encounter limitations.

on the one hand,If the compiler backend does not yet support the new hardware's storage structure, communication mechanism, or computing unit, it will be difficult for frontend developers to utilize these capabilities on their own.on the other hand,More and more critical operators require precise control over storage levels, parallel granularity, CTA collaboration, and the overlap of communication and computation, and the original Triton code sometimes struggles to fully express these intentions.

The emergence of new languages and DSLs such as Gluon, TLX, and TileLang reflects the same trend: AI operator development is no longer just about writing a single kernel, but also about expressing data layout, parallel hierarchy, pipeline, communication topology, and hardware characteristics.

TLE does not attempt to replace Triton, but rather to extend its syntax and ecosystem in layers.It consists of three layers: TLE-Lite, TLE-Struct, and TLE-Raw, which correspond to lightweight semantic hints, architecture-aware control, and native code-level optimization, respectively.

TLE-Lite is designed for algorithm engineers and rapid optimization scenarios.Developers don't need to directly concern themselves with the underlying hardware; instead, they can supplement the compiler with more explicit structural information.For example, a tensor might need to be accessed by sub-tiles, a computation might run on a distributed mesh, or a CTA might use a Producer-Consumer Pipeline. Taking sub-tile operations as an example, developers can directly extract a logical sub-block from a larger tensor, perform activation, normalization, or statistical computations, and then write it back, without having to manually calculate offsets, construct masks, and handle boundaries.

Since the compiler can recognize that this is a regular tile access, it can further optimize around data layout, vectorization, bank conflict, and register reuse. This approach is particularly suitable for sparse attention, local normalization, block statistics, and routing operators. In distributed scenarios, TLE uses Device Mesh to describe different levels such as nodes, GPUs, Block Clusters, and Blocks, organizing them into a unified multidimensional topology.

Developers write programs with a Mesh-based approach, and the compiler and runtime then map the logical topology to the actual hardware and communication mechanisms. As a result, Ring communication, Barrier synchronization, and sharded access are no longer just scattered Ranks and communication groups in the code, but become analyzable, structured semantics. Once communication relationships are explicitly expressed, the compiler has the opportunity to perform topology-aware scheduling, communication-computation overlap, Barrier merging, and deadlock checking.

TLE-Lite also abstracts the internal collaboration of the CTA into a Producer-Consumer model through Pipeline primitives. Developers mainly describe who produces the data and who consumes it, while the underlying Barrier, buffer reuse, and synchronization mechanisms are handled by the compiler.This does not mean shielding the underlying capabilities, but rather transforming them into a program structure that can be analyzed, verified, and optimized.

From lightweight semantics to native pass-through, covering different levels of optimization depth.

If TLE-Lite primarily addresses cross-platform semantic representation,TLE-Struct, on the other hand, focuses more on architecture awareness and fine-tuning.

Different GPUs, DSAs, and AI accelerators differ significantly in their storage tiers, execution units, synchronization mechanisms, and on-chip networks. TLE-Struct therefore exposes a hierarchical parallel and storage structure to developers, allowing them to explicitly define data layout, computation mapping, and memory hierarchy, without having to be directly bound to a vendor's proprietary interface.

For example, the same local buffer can be mapped to shared memory on the GPU backend, and to Scratchpad or on-chip SRAM on the DSA backend. The user expresses a structured memory intent, and the compiler is responsible for translating it into an address space and memory access instructions suitable for the target hardware.

Expert Counting in MoE is a typical scenario. This operator essentially counts the number of tokens routed to different Experts, and is easily affected by factors such as shared memory layout, concurrent updates, bank conflicts, and cross-block aggregation.

With TLE-Struct, developers can explicitly organize the counter layout, mapping different Experts or Tokens to different areas of local storage; after the compiler obtains this structural information, it generates appropriate access and synchronization methods.

TLE-Raw is geared towards performance optimization experts, preserving interfaces for vendor-native code.

Some performance paths must use CUDA, assembly, or dedicated intrinsics directly. Forcing them to be repackaged into a higher-level DSL may limit performance or increase migration costs. TLE-Raw allows developers to inline native code within the Triton/TLE architecture or directly access vendor compilation pipelines.

Taking All-Gather GEMM as an example,This operator involves communication, matrix computation, buffering, and synchronization management. Developers can reuse the underlying communication capabilities while continuing to use structured Tensors and Tiles to express computations, and finally, a unified compilation pipeline organizes the different parts into callable operators.

This layered design allows algorithm engineers, operator developers, and performance experts to choose different optimization depths within the same system, without having to start with the lowest level of programming.

Performance tests also show thatThe overall abstraction overhead of TLE is controllable.

In the Radix Select test, TLE replicated the TensorRT-LLM algorithm, achieving performance of approximately 85%–97% across multiple shapes. For teams requiring cross-platform maintenance and rapid iteration, achieving near-expert-level performance with more unified code has significant engineering value.

In a SparseMLA scenario with a 128K context,TLE uses Pipeline primitives to express collaboration between different execution roles, achieving performance of approximately 90%, the baseline for FlashMLA.

The team also tested All-Gather on a single node with eight NVIDIA H100 processors and further explored the integration of GEMM and All-Gather. The focus is not simply on replacing the communication library, but on incorporating communication into operator-level expressions and compilation optimizations, enabling data transmission, local computation, synchronization, and subsequent consumption to form a pipeline.

For inference systems, this kind of communication-computation fusion is often more meaningful than improving the peak performance of a single GEMM in isolation, because what users ultimately perceive is the end-to-end latency.

Flagtree Compilation Optimization Practices

In addition to TLE language extensions, the FlagTree team also undertook two other tasks:First, we integrated with the CUDA Tile IR backend; second, we implemented several Triton compilation optimizations for real-world model workloads.

The core concept of CUDA Tile IR is to allow programs to express tiles, rather than pre-determining the underlying thread mapping.Its input consists of a Tile Program and a three-dimensional Tile Grid. Inside Dialect, computation, data view, and necessary dependencies are expressed through Tile Compute, View, and Token-Ordered Operation (TKO).

Among them, TensorView describes the global pointer, Shape, and Stride; PartitionView adds block mapping capabilities on top of this; Token is used to constrain the dependencies between related TKOs; and Memory Model specifies memory semantics and scope separately.

FlagTree does not completely replace the Triton NVIDIA native backend. Instead, it adds a TileIR path to the existing system and overrides the TileIR View and Token interfaces through TLE primitives. Kernels that meet the requirements can access the TileIR backend; those that are not yet supported will fall back to Native CUDA. Both paths can coexist in the compiler.

Regarding automatic optimization,The team proposed FlagOSTune to resolve the conflict between coverage and cost in real-world models of Triton Autotune.

In actual model inference, the same operator may correspond to a large number of different shapes. The team's statistics show that there are 1994 unique MM Shapes in six models and four types of inference scenarios, which is far greater than the coverage of daily benchmarks. If the candidate configuration is directly expanded, the theoretical search scale will rapidly grow to millions of sets.

FlagOSTune improves performance by expanding the search space and compresses search costs through a combination of model prediction and limited real-world testing. The system first uses XGBoost to rank candidate configurations, then sends only a small number of the most promising configurations to real compilation and GPU testing, before continuing the search using a genetic algorithm.

The performance of multiple operators has been improved on various computing power systems from NVIDIA, Moore Threads, and Muxi, resulting in speedups of 1.21 to 7.35x.In multiple shape experiments with the NVIDIA H20 MM operator, the search space configurations were compressed from over 620,000 to 4,070, resulting in a 120-fold speedup in tuning efficiency with almost no performance loss.

In terms of data layout, the team focused on optimizing the performance overhead of data layout transformation.

In Triton, `convert_layout` typically involves cross-thread data rearrangement, requiring writing to shared memory, synchronization, and then reading back, which is not a zero-cost operation. For small operators or memory-intensive operators, even a small number of layout transformations can become a performance bottleneck.

FlagTree enhances the Remove Layout Conversions mechanism by using techniques such as cost models, backpropagation, and local joint solutions to reduce unnecessary data layout conversions.In tests with more than 100 operators, the net conversion elimination rate was approximately 68%–79%, with a performance improvement of up to 71%.

Another optimization is instruction reordering. Loop unrolling simply copies the loop body and does not automatically fire multiple loads ahead of time. With compiler reordering enabled, independent loads can be executed ahead of time, interleaving subsequent calculations, thus utilizing the hardware pipeline to hide memory access latency.In the three typical operators, this optimization brings an average speedup of about 1.19 to 1.61 times, with a peak speedup of 2 times.

The team also optimized the performance of the Fused Marlin MoE operator for real shapes in DeepSeek-V4-Flash, reducing data loading time through fragment merging. Combined with algorithm-level optimizations,Compared to vLLM CUDA, FlagGems achieved speedups on all 53 tested shapes, with an average improvement of approximately 1.208 times when weighted by shape frequency.

In the DeepSeek-V4-Flash end-to-end test of the NVIDIA H20,When TP=4, the latency of the first token decreases by 20.221 TP3T, and the total throughput increases by 12.631 TP3T; when TP=8, the latency of the first token decreases by 12.871 TP3T, and the total throughput increases by 7.651 TP3T.

Future work prospects

The next stage,FlagTree will focus on advancing the NUMA programming model and the MegaKernel Compiler.

For multi-node, multi-GPU setups, and GPU-internal clusters and local SRAM, data access exhibits significant differences in proximity. TLE aims to explicitly describe how Tensors are segmented along the mesh and how they transition between different distribution methods, enabling the compiler to determine which data should be proximityed for computation and which communications can be merged or overlapped with computation.

The MegaKernel Compiler attempts to extend the optimization scope from a single kernel to a model execution process. The compiler identifies computations that can be co-optimized, breaks them down into specific tasks, and then uniformly decides how to allocate, parallelize, and pipeline these tasks, ultimately generating one or a few MegaKernels.

TLE acts as a bridge between front-end semantics and hardware architecture. Tile, Mesh, Pipeline, Memory Layout, and native capability interfaces allow the compiler to see not just a series of independent operators, but a task graph that can be scheduled, placed, synchronized, and merged.

From TLE-Lite's lightweight semantic hints to TLE-Struct's architecture-aware control, and TLE-Raw's native capability pass-through, combined with various compilation optimization techniques, FlagTree is attempting to establish a more flexible layered system between Triton's ease of use, cross-platform migration, and ultimate performance, and to enable new hardware capabilities to be transformed into the actual performance of models and operators more quickly.