Triton Close to the Metal

We selected Triton to analyze close to the metal. Triton sits as perhaps the most compelling way for a developer to code for GPUs, effectively balancing the challenges of enabling hardware-awareness and accuracy while also offering an easy way to develop.

We believe that Gestell’s analyzer-enabled inspection of PTX and SASS offers significant insights into the behavior of Triton and how its code lowers. Lowering patterns and the increasing complexity of NVIDIA hardware require venturing to this often neglected source of evidence. Triton’s own authors have acknowledged this through the introduction of things like Gluon. Instruction-specific behavior, tensor-memory operation, ptxas operations, register allocation and so forth all require understanding how execution pathways influence Triton’s own development.

The main case we analyze is Issue #8328, which displayed a Blackwell FlexAttention performance regression caused not by the source but by a layout-conversion path that appeared in the lowering path. Here, we show that Gestell’s analyzer of PTX/SASS would augment investigation

The secondary case is an examination of Gluon across Blackwell as a positive-control case to demonstrate how Gluon takes advantage of the B200 hardware and verify the compiler’s lowering into PTX and SASS

We attempt to show that the current limitations of analysis presented by a compiler’s lowering pathways itself create challenges for investigation for Triton maintainers. We believe that through statically parsing and bringing forth the actions of the compiler, Triton may be augmented for both developers and users alike

How Triton lowers

Triton compiles in stages, each one fixing detail the stage above left open:

Tile distribution at the top becomes thread, warp, and memory assignment lower down. By the time the code reaches ptxas, the decisions about what actually runs on the machine are made.

As NVIDIA’s hardware has added tensor memory, warp-group MMA, and asynchronous coordination, Triton has had to keep up, enabling high-level kernel writing while still exposing the full instruction set to people who need it. This is why Gluon exists. It is also why reading the compiled output has become increasingly important.

This is a difficult task as some PTX patterns vanish after ptxas, and lowering shifts between versions. SASS is a black box and individual inspection of it is challenging due to the nature of machine code. A finding of true importance names a specific lowering path, shows a before-and-after delta in the compiled output, confirms the consequence in SASS, and points at the place in the stack to look. Which is how we use our analyzer here

Methodology

For each state of the kernel we collected the compiled artifacts, read the emitted PTX and SASS statically, grouped instructions by kernel family, and compared what appeared, disappeared, or changed across states. The Gestell Analyzer does not execute the kernel, run the model, read profiler counters, or benchmark anything. It produces a compiled-output diff: which kernel families showed up or dropped out, and which instruction-level structures moved. The runtime numbers in this post come from running aakhundov’s repro directly; the structural findings come from the diff.

Case Study: Triton FlexAttention Regression

Issue #8328 reports an 18% performance regression on Blackwell for the FlexAttention forward kernel. PR #7565 introduced this regression while generalizing the swizzling algorithm for convert_layout lowering. PR #8353 fixed it by folding a redundant convert_layout op into Blackwell’s tmem_store.

Thomas Raoux started debugging the regression in the issue thread, he suggested the approach “A first debugging step may also be to diff the ptx or llvm generated to see which places changed significantly.” We listened to this advice and did that using Gestell. The diff exposes the nature of the regression and points at a Blackwell matrix-fragment staging path that the regressed build ran through a redundant layout conversion, and that the fix collapses back down

The python source of the Flex Attention forward kernel never changed, it is byte-identical across the baseline, regressed, and the fixed builds, and so is the frontend intent because neither PR changed the kernel source. The relevant evidence lives below source review: the TTGIR conversion site shown in PR #8353, and the PTX/SASS lowering that changed around it. Here a reviewer reading the kernel observes a source-identical kernel, yet sees a performance regression on their profiler. No evidence for this exists to the external viewer. Only in diving into the lowering method of the compiler is it revealed the PTX and SASS have changed from the prior state

Three builds

We compared three commits using the sample flex_attention_fwd script attached in issue #8328. It has the shape B=4, Hq=16, M=8192, Hkv=16, N=8192, D=128.

The baseline used the regular vector shared-memory path. The regressed build replaced that path with Blackwell matrix-fragment staging, extra synchronization, and a shared-memory layout detour around the same compute.

Why only Blackwell

PR #7565, written by Mario Lezcano with the algorithm co-authored by Adam Goucher, generalized the swizzling Triton uses for ldmatrix and stmatrix. The PR is explicit about the cost: “we now require having a dedicated allocator for nvidia, as the required shmem for a convert_layout will now depend on the instructions we can emit.” ldmatrix and stmatrix go back to sm_75. The change touched a path shared by Ampere, Hopper, and Blackwell. Blackwell is where it broke because of a pass that only runs on Blackwell. TMemLoadReducePattern forward-propagates the layout that comes out of a tmem_load. After PR #7565 that layout was #linear, an artifact of the new swizzling. It propagated forward until it hit ttng.tmem_store, which wanted #blocked. The pipeline reconciled the mismatch the only way it could without the fix: it dropped a ttg.convert_layout between the arith.truncf and the store. That convert is what lowered into all the staging traffic. Only through careful analysis of the emitted SASS is this confirmed

pchen7e2 found it by disabling this pass in the issue thread:

The TTGIR diff pchen7e2 posted in the fix PR shows the inserted op:

%acc_316 = arith.truncf %p_308
: tensor<128x64xf32, #blocked> to tensor<128x64xbf16, #blocked>

ttng.tmem_store %acc_316, %acc_35, %true
: tensor<128x64xbf16, #blocked> ->
  !ttg.memdesc<128x64xbf16, #tmem2, #ttng.tensor_memory, mutable>

%acc_315 = arith.truncf %p_307
: tensor<128x64xf32, #linear> to tensor<128x64xbf16, #linear>

%acc_316 = ttg.convert_layout %acc_315
: tensor<128x64xbf16, #linear> -> tensor<128x64xbf16, #blocked3>

ttng.tmem_store %acc_316, %acc_35, %true
: tensor<128x64xbf16, #blocked3> ->
  !ttg.memdesc<128x64xbf16, #tmem2, #ttng.tensor_memory, mutable>

One op, ttg.convert_layout, sitting in the main loop. That is the crux of the regression.

What matrix-fragment staging is, and why the counts move

NVIDIA calls these “warp-level matrix fragments.” A fragment is the slice of a matrix tile that one warp lane holds in registers while a tensor-core MMA runs; no single lane owns the whole tile, which is what lets the warp feed the MMA in parallel. Staging is the use of shared memory as the layout-conversion waypoint: store the tile in one layout, synchronize the warp, load it back in the fragment layout the MMA expects. ldmatrix does that load as a single warp-collective instruction and then stmatrix does the store. In SASS they are LDSM and STSM. They are the purpose-built version of what you would otherwise hand-assemble from ld.shared, st.shared, and a pile of per-lane address math.

The baseline never used this path and PR #7565 addressed that by turning it on.

The Gestell Analyzer Reveals this via a diff between the 3 states

Tensor compute holds at 128 across all three builds and TMEM movement holds at 47. That means the actual tensor-core work is identical in every version and the matmul is unchanged. The things that move are those associated with it: the staging, the barriers that staging needs, the scalar address math that computes the shared-memory offsets.

The regressed count of 296 staging hits and 501 synchronizations is what the path looks like when the conversion is wrong. The fixed count of 40 and 251 is what it looks like when the conversion is gone and the staging only runs where it earns its keep.

The fix

PR #8353 makes two changes.

It does not revert PR #7565. The generalized swizzling stays, which is the point because that work is wanted. What changes is that tmem_store now absorbs a compatible incoming conversion instead of forcing it through shared memory first. The accumulator goes into TMEM in the layout the compute already produced

On our B200 reproduction the fixed kernel runs at 507 TFLOP/s, above the baseline’s 434. Upstream PR #8353 summarizes the standalone repro fixed result as 499 TFLOP/s, and issue #8328’s closing table reports a fixed TritonBench average of 494.838 TFLOP/s. The swizzling improvement survives while the redundant conversion is gone, and the result clears both earlier builds

The PTX agrees. ldmatrix.sync.aligned.m8n8.x1.shared.b16 falls from 70 to 6. The x1 stmatrix store path also falls from 70 to 6, with the regressed build using stmatrix.sync.aligned.x1.m8n8.shared.b16 and the fixed build using stmatrix.sync.aligned.m8n8.x1.shared.b16. bar.sync falls from 176 to 52. In SASS, ULOP3.LUT, the bitwise-lookup op the regressed build leaned on for swizzle address math, falls from 120 to 13, and USHF.R.U32.HI from 114 to 16. Most of that work no longer emits because the redundant conversion is gone.

The fix reverses it without going back to the baseline path:

A Gestell Config to Catch This

The Gestell can fit easily here as either a CI or inspection to help determine these sorts of movements. The regression has a fixed shape in the IR: a convert_layout feeding directly into a tmem_store whose destination encoding the source already satisfies. That is a pattern a review policy can match on

rule: review-convert-layout-before-tmem-store
match: convert_layout -> tmem_store
when: source layout satisfies tmem_store encoding
flag: redundant layout conversion in main loop

Configured this way, the analyzer flags the regressed build and stays quiet on the baseline and the fixed build. The same policy generalizes to any architecture where a layout-anchored store can absorb its incoming conversion; only the destination op changes

Gluon in Context

Gluon presents another useful case study to show how this style of compiled-output analysis can help confirm when lowering patterns are used. Gluon was introduced to give programmers more direct control over the parts of modern GPU execution that Triton often asks the compiler to infer. Particularly for Blackwell, kernels often require and depend upon TMA, TMEM, tcgen05 etc.

Here we analyzed a Gluon attention-forward kernel on a B200. This case serves to show that PTX and SASS review can act as a positive-control artifact to show what intentional Blackwell orchestration looks like and verify it as it is lowered to PTX and SASS

At the PTX level, the Gluon kernel emits:

Here you see the intentional hardware of the orchestration of Gluon at play - tensor compute, TMEM movement, barrier coordination and register-budget control all are coming together at the low level

So we now have a contrast with the FlexAttention case, there, a compiled-output review showed an unintended fragment-staging detour. In the Gluon case, compiled-output review confirms the intended hardware path: the high-level programming model asked for explicit Blackwell orchestration, PTX expressed it, and SASS confirmed it

As a CI, PTX/SASS compiled output can serve Gluon kernels to make explicit the movements that are being requested at higher levels. As a programmer may specify commands to make an optimal kernel, these artifacts can expose those decisions and serve as possible review gates or lens of analyses for the author

Conclusion

Here we’ve analyzed how PTX/SASS compiled-output review may provide a useful tool for helping understand and augmenting the development of the Triton ecosystem. The case study of FlexAttention and how the lowering patterns were made clear showed that PTX/SASS deserves an important status for compiler development. The performance story of the regression lived in the black box of the correlation between PTX and SASS. Specifically, a structural regression where an op appeared in the loop, and the analyzer sees it because it is present in the emitted output and reading what is present is what a static tool does

We do believe that there are limitations to our approach. For example, Triton issue #9433 is a different kind of problem. There, a wgmma.wait_group N-1 fails to guarantee that every in-flight async dot has finished before its input registers get overwritten. The instructions are all present and correct, so there is nothing for a static read to flag. The fault is in when they run relative to each other, which a static tool cannot see

Staying static allows the analyzer to read the emitted artifact instead of running the kernel, which is what makes the diff in this post reproducible from a cached build, on any machine, without a profiler or the hardware in the loop

In the case of our analysis of Gluon, this analysis serves to show a positive-control version of the analyzer. Here, a kernel explicitly enables Blackwell orchestration and PTX/SASS review can verify that

The role that a compiled-output CI plays is as a useful companion to the users and maintainers alike. Source review and profiling serve to expose the intent and ground truth while compiled-output review sits between these things. A useful CI gate doesn’t fail instruction diffs but flags structural changes especially where things like layout conversions enter, PTX intent doesn’t survive to SASS or architecture specific patterns materialize

As Triton continues its progress to help developers maximize their ability to code on GPUs, compiled artifacts become increasingly critical. PTX and SASS now serve as an evidence layer for maintainers and users to understand what their code actually does