O Triton trouxe sanidade de volta ao desenvolvimento de kernels de GPU: escreva Python em nível de tile e deixe o compilador cuidar de threads, registradores e pipelining. Mas, conforme as GPUs ficam mais exóticas (clusters de thread blocks, DSMEM) e os modelos ficam mais esparsos (roteamento MoE, atenção esparsa de 1M de tokens), o Triton padrão fica sem espaço expressivo. Ou você escreve CUDA à mão ou luta contra o DSL.
É essa lacuna que o FlagTree — o fork open-source multi-backend do Triton, do stack FlagOS da BAAI — ataca com as TLE (Triton Language Extensions): uma extensão de três camadas que mantém compatibilidade de código-fonte com o Triton enquanto expõe progressivamente controle de hardware. Você não precisa fazer fork dos seus kernels para começar.
As três camadas: escolha sua profundidade
- TLE-Lite — dicas semânticas de alto nível para engenheiros de algoritmos. Pense em fatiamento de sub-tiles, loads assíncronos, pipelines produtor-consumidor, device meshes. Escreva uma vez, rode em qualquer lugar.
- TLE-Struct — controle consciente da arquitetura (GPGPU vs DSA). Você descreve a intenção de memória estruturada; o compilador mapeia um “buffer local” para memória compartilhada em uma GPU ou SRAM scratchpad em uma DSA.
- TLE-Raw — passthrough nativo para especialistas em performance: inline de CUDA/assembly ou intrínsecos do fabricante dentro do programa Triton, para os caminhos que um DSL não consegue expressar.
Todas as três camadas são compiladas para um único kernel via LLVM IR (através de FLIR), então você pode misturar profundidades dentro de um mesmo programa.
Como ficam as primitivas
1. Fatiamento de sub-tiles sem cálculo de endereços
A dor clássica do Triton: extrair um sub-bloco lógico de um tensor grande exige calcular offsets, máscaras e checagens de fronteira manualmente. O TLE-Lite oferece 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])Como o compilador enxerga um acesso regular a tiles, ele pode otimizar layout, vetorização e conflitos de banco — um ganho direto para atenção esparsa, normalização local e estatísticas por bloco.
2. Pipelines produtor-consumidor com tle.pipe
O pipelining em software com tl.range(num_stages=...) é automático, mas opaco. O tle.pipe torna a aresta do fluxo de dados explícita: um escritor faz acquire de um slot, preenche e faz commit; um leitor faz wait, consome e faz release:
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)Barreiras, reutilização de buffers e sincronização passam a ser gerenciadas pelo compilador, mas continuam analisáveis — e é isso que permite ao compilador sobrepor carga e computação em vez de serializá-las.
3. Coordenação em nível de cluster para TopK de sequências longas
O destaque é o seletor TopK usado pela DeepSeek Sparse Attention. Sequências longas com batch=1 têm quase nenhum paralelismo dentro de um único bloco, então o TLE divide uma única linha entre um cluster de blocos:
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 worldAtualizações de histograma, escrita de candidatos e a ordenação final ficam todas on-chip — sem ida e volta pela memória global — enquanto o cluster distribui a varredura da linha única entre os blocos. Kernel completo e funcional: python/tutorials/tle/deepseek_v32/01-topk_selector.py no repositório.
Os números
- Seletor TopK (H800, batch=1, 131K tokens): versão com clusters do TLE em 0,030 ms vs FlashInfer 0,045 ms e TRT-LLM 0,049 ms — até ~2,5× sobre o TRT-LLM; permanece estável em 512K enquanto o conjunto de candidatos do TileLang estoura.
- Radix select: uma reimplementação em TLE do algoritmo do TRT-LLM alcança ~85–97% do TRT-LLM nativo em todos os formatos testados.
- SparseMLA (contexto de 128K): as primitivas de pipeline chegam a ~90% do baseline FlashMLA.
- Autotuning com FlagOSTune: em formatos MM do H20, o espaço de busca colapsa de 620K+ para 4.070 configurações — um speedup de autotune de 120× com perda de performance desprezível. O tuning de kernels de modelos reais os melhorou de 1,21× a 7,35× em chips NVIDIA, Moore Threads e MetaX.
- Passes de compilador: a eliminação de conversões de layout remove 68–79% das ops convert_layout (até 71% de speedup no kernel); o reordenamento de instruções antecipa loads independentes para uma média de 1,19–1,61× (pico de 2×).
Como começar
- Clone o repositório:
git clone https://github.com/flagos-ai/flagtree.gite faça checkout demain(backend NVIDIA, Triton 3.6; Ascend fica emtriton_v3.5.x, Cambricon emtriton_v3.2.x). - Siga o manual do usuário para NVIDIA para compilar e instalar.
- Rode os tutoriais em
python/tutorials/tle/— comece pelo seletor TopK. - Migre incrementalmente: comece com as dicas do TLE-Lite em um kernel existente; use TLE-Struct/Raw apenas onde o profiling mostrar gargalos.
Conclusão
A TLE não é uma linguagem nova — é uma saída para o teto de abstração do Triton. Se você escreve kernels de atenção esparsa, MoE ou multi-chip, vale dedicar um fim de semana: o tutorial do seletor TopK sozinho é a forma mais rápida de ver Triton em nível de cluster feito da maneira certa.
