Fla Triton To Gluon

Triton to Gluon

Use Gluon when profiling points to register pressure, layout conversion, or load/compute overlap that needs explicit control. A kernel already saturating its useful bandwidth or compute limit may not benefit. Keep the portable default path when adding hardware-specific implementations.

Follow fla-optimization-loop for the baseline and correctness gate, and fla-nvidia-performance for NVIDIA measurements.

Check the target environment

Gluon is experimental. Check the installed Triton API before copying upstream examples; their names and signatures can differ from the installed release.

Gate architecture-specific code appropriately: Ampere supports cp.async, Hopper adds TMA/WGMMA, and Blackwell adds TMEM/tcgen05. Code under Gluon's nvidia namespace requires NVIDIA hardware. Keep compile-time choices in explicit constexpr arguments so the kernel, launcher, and autotune pruning use the same values.

from triton.experimental import gluon
from triton.experimental.gluon import language as gl

Port incrementally

1. Record the baseline commit and retain the validated tests, reference, and tolerances. Keep the Triton fallback available.
2. Translate loads, stores, and arithmetic with explicit layouts first. Match each tensor's contiguous dimension and preserve masking. Pass forward, state, and supported backward comparisons before adding asynchrony.
3. Add async transfers only where load latency or addressing pressure limits the kernel. Define buffer ownership, completion, and reuse for the prologue, steady state, and drain.
4. Add architecture-specific MMA or scheduling when profiling supports it. Recheck register/shared-memory use and retune after each change.
5. Run the full relevant correctness gate and same-hardware benchmarks, including supported variable-length and boundary cases. Measure production entry points as well as individual kernels.

A literal translation establishes parity; measure it before assuming a performance gain. Preserve the validated numerical algorithm and precision while changing execution layout or scheduling.

Layout and resource constraints

• Pointer tensors and values need compatible layouts. A cross-warp gl.convert_layout may use shared memory; use assert_trivial=True when a conversion is intended to be free.
• Broadcasts and oversized layout tiles can duplicate values and increase register use. Choose layouts from actual tensor strides rather than copying a tutorial's block sizes.
• Account for all live shared-memory buffers and the target device's allocation limit. Prune infeasible autotune configurations; provide a streaming path when a resident design cannot fit supported shapes.
• Large static unrolls multiplied by many autotune configurations can dominate compilation. Prune from the actual shape and live-memory budget, and reuse the worker and Triton cache during iteration.
• Persistent schedules can trade launch overhead for worse L2 locality. Warp specialization also changes the total register budget; measure occupancy and stalls after changing either.

Async correctness

Transfers and buffer reuse

cp.async groups are ordered as a queue. wait_group(N) bounds all outstanding groups, so it cannot identify a particular buffer after unrelated prefetches are interleaved. For per-buffer scheduling, use separate mbarriers and match their arrival counts to the participating threads. In the installed API, check whether mbarrier_arrive increments the count; a preinitialized thread count requires the non-incrementing form.

Wait for a load before consuming its buffer, and protect reuse until all readers finish. Even same-lane shared-memory staging needs protection against overwriting data still being read. Track each mbarrier's phase as its buffer is reused; do not run more than one phase ahead or share a completion barrier between TMA and tcgen05 without reinitializing it.

TMA descriptors must satisfy the target's alignment and stride requirements. TMA stores may distinguish completion of the shared-memory read from completion of the global write. Check the installed wait API and require global completion before another operation reads the stored range.