fla-triton-to-gluon

name

fla-triton-to-gluon

description

Workflow for porting an existing Triton kernel in `fla/ops/**` to Gluon (`triton.experimental.gluon`) to gain explicit control over tensor layouts, shared memory, async data movement (cp.async / TMA), MMA (WGMMA / tcgen05), and scheduling (persistent kernels, warp specialization). Covers when a port is worth it, an incremental porting sequence that keeps numerical parity at every step, a Triton-to-Gluon API mapping, compile-time / autotune / smem-budget management for heavily unrolled kernels, and a pitfall checklist (proxy fences, mbarrier semantics, layout costs, bitwise-cancellation traps, NaN-poisoned OOB handling). Use when a Triton kernel is register-bound, when `num_stages` pipelining underperforms, or when Hopper/Blackwell features (TMA, TMEM, tcgen05) are needed.

FLA Triton → Gluon Porting Skill

Gluon shares Triton's compiler stack, JIT, and SPMD tile model; host-side launch code is unchanged. The difference: layouts, shared memory, asynchrony, and synchronization are all explicit. Port incrementally: first a literal translation that passes the op's frozen pytest, then upgrade layer by layer driven by profiling, keeping numerical parity after every step.

Related skills:

• fla-optimization-loop — the iteration discipline around this port (frozen test contract, recording, when to stop).
• fla-nvidia-performance — profiling workflow, hardware baselines, MR-ready perf evidence.
• fla-correctness-coverage — test coverage matrix for the op being ported.

When a port is worth it

Worth it:

• Register spills cap the block size (Triton gives you no lever; Gluon's TMA path moves addressing out of registers).
• num_stages pipelining fails to overlap load and compute the way you want.
• The kernel needs TMA features Triton does not expose well (im2col, gather/scatter, multicast).
• Blackwell-specific paths: TMEM accumulators, tcgen05_mma, 2-CTA MMA, CLC dynamic scheduling.
• Load and compute are imbalanced enough to justify warp specialization.

Not worth it:

• The kernel already saturates bandwidth or tensor-core throughput.
• The bottleneck is algorithmic, not scheduling.
• The op must stay portable across vendors — Gluon's nvidia modules are NVIDIA-only (AMD is a separate submodule).

Environment and versions

• Gluon lives under triton.experimental and its API moves between Triton versions. The official tutorials (https://triton-lang.org/main/getting-started/tutorials/gluon/) track the main branch; verify names against the installed Triton with dir() before copying tutorial code. Known examples: in Triton 3.5.1 there is no gluon.aggregate, and the TMA/cp.async load is async_copy_global_to_shared (named async_load on main).
• Since Triton 3.6, a kernel may not read plain module-level Python globals (NameError: ... instantiated as constexpr). Pass such values as constexpr kernel arguments instead — that also keeps host and kernel in sync and makes them visible to autotune key/prune functions.
• Hardware gating: cp.async needs Ampere+; TMA, WGMMA, gl.warp_specialize, CGA clusters need Hopper+; TMEM, tcgen05_*, CLC, and TMA gather/scatter need Blackwell. Follow the hardware baseline rules in fla-nvidia-performance.

from triton.experimental import gluon
from triton.experimental.gluon import language as gl
from triton.experimental.gluon.language.nvidia import ampere, hopper, blackwell
from triton.tools.tensor_descriptor import TensorDescriptor  # TMA, host side

Triton → Gluon mapping

Triton	Gluon	Notes
@triton.jit	@gluon.jit	triton.autotune, triton.cdiv, do_bench are reused as-is
tl.load / tl.store	gl.load / gl.store	every tensor (including pointer tensors) needs an explicit layout
tl.arange	gl.arange(..., layout=gl.SliceLayout(dim, parent))	2D offsets = SliceLayout + expand_dims + broadcast (free)
tl.dot	Hopper: hopper.warpgroup_mma; Blackwell: blackwell.tcgen05_mma	async instructions; explicit wait/commit required
num_stages=N	manual multi-buffering: smem gets a leading [num_buffers, ...] dim	prologue / steady-state / epilogue skeleton below
(compiler-managed smem)	gl.allocate_shared_memory(dtype, shape, layout)	SwizzledSharedLayout / NVMMASharedLayout to avoid bank conflicts
(compiler-managed layout)	gl.BlockedLayout(size_per_thread, threads_per_warp, warps_per_cta, order)	see layout guidance below
tl.trans(b)	b_smem.permute((1, 0))	forwarded to the MMA hardware, zero-copy
tl.static_range	gl.static_range	used for prologue peeling

Porting sequence

1. Freeze the baseline

Keep the Triton kernel and the op's tests/ops/test_<op>.py untouched (they are the frozen contract per fla-optimization-loop). The Gluon kernel is added alongside and must pass the same parity tests (forward and backward, via fla.utils.assert_close) before any optimization.

2. Literal translation (gl.load/gl.store, correctness first)

Add layouts, no async anything. A literal port is a parity scaffold, not a deliverable: it drops Triton's automatic vectorization and num_stages pipelining without adding any manual control, so it usually ties or loses to the Triton kernel. Measured on a bandwidth-bound op (attnres, GB200): the literal port was ≈ Triton; the wins (fwd 1.3–1.45×, reaching 66–72% of HBM peak) all came from the restructuring steps below.

While translating, also restructure what Triton forced you to express dynamically: small runtime dimensions (e.g. a source/tensor count) become constexprs, so runtime pointer-select chains (tl.where(o == i, ptrs_i, p)) turn into static indexing over an unrolled gl.static_range — and gathers over many tensors become contiguous per-tensor block loads that async copy can handle.

Layout starting points:

• 1D: size_per_thread=[1] — each warp issues exactly one 128-byte coalesced access; measured faster than larger per-thread tiles in the tutorials.
• 2D row-major: size_per_thread=[1, N], order=[1, 0]. The layout's contiguous dim must match the tensor's contiguous dim — a mismatch costs an order of magnitude of bandwidth (6.3 → 0.8 TB/s in the tutorial).
• Input and output with opposite contiguity: derive a layout from each tensor's strides, pay one gl.convert_layout in the middle, use square-ish blocks (e.g. 128×128).
• Broadcast waste: a tensor smaller than the layout's block shape still burns the full register budget (redundant copies per thread/warp).

3. Async data movement (the first big jump)

cp.async (Ampere+, small diff): ampere.async_copy.async_copy_global_to_shared → commit_group() → wait_group(N).

TMA (Hopper+, frees registers so blocks can grow): host side TensorDescriptor.from_tensor(t, block_shape); smem must use NVMMASharedLayout; strides 16-byte aligned; loads tracked by an mbarrier (expect(bar, nbytes) → wait(bar, phase)), stores by tma.store_wait(pendings=N). Out-of-bounds masking is automatic.

Pipeline skeleton (same shape for both mechanisms):

smem = allocate([num_buffers, BM, BN]); one mbarrier per buffer
prologue:      issue num_buffers - 1 loads (gl.static_range)
steady state:  issue load i + num_buffers - 1; wait load i; compute; release buffer i
               buffer index = i % num_buffers; mbarrier phase = i // num_buffers & 1
epilogue:      drain with decreasing wait counts

Pick the pipeline depth from the load/compute latency ratio; going deeper past bandwidth saturation buys nothing.

cp.async specifics learned the hard way:

• Same-lane staging: when the cp.async pointer tensor and the smem readback use the same blocked layout, every thread reads back exactly the bytes it copied — smem is pure staging for asynchrony, with no cross-thread exchange. Buffer reuse still gets a gl.thread_barrier() before the refill (WAR safety); it costs ~a barrier, not a pipeline stall.
• Commit groups are one global FIFO counter: wait_group(N) counts every group issued later, so it cannot express "wait for slot l only" once you interleave prefetches for the next loop iteration with consumption of the current one — the wait would also cover the new issues and serialize you again. For per-slot pipelining across iterations, switch to ampere.mbarrier: one barrier per slot, mbarrier.init(bar, count=num_warps * 32), and after each thread's issues async_copy.mbarrier_arrive(bar, increment_count=False) (the noinc form consumes the pre-initialized count; the default self-increments and never completes with a thread-count init). Consumers mbarrier.wait(bar, phase=t & 1) — one fill per iteration flips parity.

4. MMA (if the kernel has a dot)

• WGMMA (Hopper): B must be in smem; accumulator in registers with NVMMADistributedLayout(version=[3, 0]); M ≥ 64 (one warpgroup minimum); results must flow through the return value of warpgroup_mma_wait(deps=...) or ordering is not guaranteed.
• tcgen05 (Blackwell): accumulator must live in TMEM (allocate_tensor_memory + TensorMemoryLayout); TMEM loads/stores need a full warpgroup (each warp sees only 32 of 128 rows); completion via tcgen05_commit + mbarrier; tcgen05_copy moves smem→TMEM without a register round-trip, and same-pipe tcgen05 instructions are implicitly ordered (a copy followed by an MMA needs no wait).
• Both: use_acc=False is the cheapest way to zero-initialize the accumulator.

5. Scheduling layer (profile first, never by default)

Persistent kernels (grid = min(num_sms, num_tiles) + a tile scheduler — add grouped/swizzled tile order or L2 hit rate drops) → gl.warp_specialize (load/MMA/epilogue partitions; a TMA-issue-only partition needs 1 warp and 24 registers; set maxnreg explicitly) → multi-CTA / CLC (Blackwell). Re-autotune after every layer: in the tutorials, the pre-pipelining best config lost >100 TFLOPS after pipelining was added.

Compile time, autotune, and the smem budget

Three interacting constraints that only show up at scale:

• Unroll × configs = compile explosion. gl.static_range(K) fully unrolls; a body unrolled ~30× across two passes, multiplied by ~9 autotune configs, can take tens of minutes per shape. Use fla_cache_autotune(..., prune_configs_by={'early_config_prune': fn}); the prune fn receives all kernel args ({**named_args, **kwargs}), so it can cap the sweep to 1–2 configs when the unroll factor is large.
• Shared memory is a hard cap (228KB/SM on Hopper/Blackwell; budget ~192KB to leave headroom). When a "keep everything resident" design can exceed it for some shapes, put both designs in one kernel behind a gluon.constexpr_function switch (e.g. RESIDENT = (L+1)*BT*BD*ES <= budget): resident path for the common case, streaming double-buffer fallback for the rest. Prune configs whose minimal footprint still exceeds the budget — Triton's autotuner does not reliably skip smem-overflow configs on its own.
• Big smem buys traffic but kills occupancy. A resident design at 192KB runs 1 CTA/SM (12.5% warp occupancy at 8 warps); at that point latency hiding must come from within the CTA — that is exactly what the per-slot mbarrier pipelining above provides. Read the trade-off from NCU: dram__throughput...pct_of_peak low + stalled_long_scoreboard high + warps_active low means serialized loads, not insufficient bandwidth.

Pitfall checklist (check here first when things break)

1. Proxy fences: registers use the generic proxy; TMA/WGMMA/tcgen05 use the async proxy; the two are unordered. Any plain smem access adjacent to a TMA/MMA op on the same buffer needs fence_async_shared() — this holds across warp_specialize partitions and is not waived by mbarrier arrive/wait ordering. Sole exception: after mbarrier.wait on a TMA read barrier, reading that smem needs no fence.
2. mbarrier phase: phase = i // num_buffers & 1. A barrier only tracks the current and previous phase; running more than one phase ahead desynchronizes permanently. Never reuse one mbarrier for both TMA and tcgen05 completion (undefined behavior) — allocate separately or reinitialize.
3. tma.store_wait waits only for the smem read by default, not the global write. If the stored range is read afterwards (e.g. cross-CTA signaling), pass read_only=False.
4. Wrong values without a crash: usually a layout broadcast / conversion misunderstanding. Debug with gl.static_print on layouts and convert_layout(..., assert_trivial=True) to prove a conversion is actually free.
5. Illegal-instruction / driver errors: check TMA alignment first (descriptor strides 16-byte; gather y_offset 16-byte).
6. Slower after "optimizing": register budget blown (warp-specialized total ≈ maxnreg × (num_warps + 4) × 32), a cross-warp convert_layout silently routing through smem, or a persistent schedule tanking L2 hit rate (lts__t_sector_hit_rate in NCU — see fla-nvidia-performance for the profiling workflow).
7. Never rely on two separately-compiled reductions cancelling bitwise. If a gradient is mathematically zero only because sum(a*b) at two program points must agree to the last bit (e.g. softmax bwd over a single source: ds = p*(dp - delta) with dp == delta), Gluon may compile the two reductions differently and leave O(eps) residue that explodes against an exactly-zero reference. Branch on the degenerate constexpr case and emit exact zeros.
8. OOB rows under NaN-poisoned tests: masked cp.async leaves smem uninitialized, and NaN garbage in dead rows leaks through cross-row reductions (0 * NaN = NaN). Clamp indices to a valid row instead of masking the loads, then zero the one tensor (e.g. the incoming gradient) whose zeroing provably kills every masked contribution downstream; keep masks only on stores.

Verification and benchmarking

Run on a GPU worker per fla-nvidia-performance hardware baselines (sm_90+; prefer sm_100/sm_103):

python -m pytest tests/ops/test_<op>.py -q                 # frozen parity gate, fwd + bwd
python benchmarks/ops/run.py --op <op> --base main         # before/after vs the Triton baseline

Record every iteration per the fla-optimization-loop protocol; dense workloads for quick iteration, varlen checked before the MR.

Iteration-speed hygiene (Gluon compiles are expensive):

• Keep one warm worker per optimization loop and a persistent TRITON_CACHE_DIR — a fresh machine per run recompiles every kernel × config from scratch and dominates wall-clock.
• Order each round for fast signal: cheapest bench first, full frozen pytest after; split slow-compiling parameterizations (huge unroll factors) into their own pytest invocation.
• If the backend is selected via a cached dispatch env var (FLA_<OP>_<BACKEND>), benchmark each backend in its own process with the env var set at launch.

References

• Tutorial series (read in order; last six are advanced topics): https://triton-lang.org/main/getting-started/tutorials/gluon/ (intro → layouts → async-copy → tma → wgmma → tcgen05 → persistence → warp-specialization → tma-gather-scatter, tcgen05-copy, tcgen05-mma-scaled, cluster-launch-control, conv-im2col, multicta)
• Tutorial sources live in the Triton repo under python/tutorials/gluon/; complete kernels under python/examples/gluon/ (e.g. 02-convolution.py, a pipelined warp-specialized convolution).
• Ground truth for the installed version: python -c "from triton.experimental.gluon.language.nvidia import hopper; print(dir(hopper.tma))"