github ROCm/aotriton 0.14b
AOTriton 0.14 Beta

7 hours ago

Overview

AOTriton 0.14 is a feature release in the maintenance-feature cadence
introduced in 0.13: it synchronizes Triton with upstream, adds kernel features,
and updates the tuning database.

Headline changes:

FlyDSL SDPA Kernel Ship as the flyc backend

  • Available for gfx950/gfx1201
  • Feature Parity with Triton kernel

Gfx950 Performance Boost from 0.13b

d_head fwd non-causal fwd causal bwd non-causal bwd causal
16 196.78 -> 355.91 (+81%) 115.18 -> 339.91 (+195%) 186.04 -> 231.92 (+25%) 148.19 -> 214.49 (+45%)
32 542.85 -> 698.91 (+29%) 364.04 -> 623.48 (+71%) 346.51 -> 455.55 (+31%) 181.57 -> 449.67 (+148%)
48 265.45 -> 647.24 (+144%) 272.18 -> 571.24 (+110%) 131.94 -> 122.65 (-7%) 93.19 -> 125.85 (+35%)
64 790.92 -> 858.36 (+9%) 548.75 -> 759.05 (+38%) 350.93 -> 586.07 (+67%) 236.17 -> 517.81 (+119%)
80 539.32 -> 610.79 (+13%) 361.53 -> 565.82 (+57%) 119.11 -> 139.41 (+17%) 101.17 -> 170.29 (+68%)
96 691.69 -> 928.44 (+34%) 516.32 -> 897.54 (+74%) 285.32 -> 579.28 (+103%) 243.90 -> 327.71 (+34%)
128 952.89 -> 1176.47 (+23%) 669.70 -> 983.31 (+47%) 294.16 -> 315.98 (+7%) 285.12 -> 314.51 (+10%)
160 688.44 -> 818.54 (+19%) 571.57 -> 691.51 (+21%) 284.52 -> 496.04 (+74%) 279.73 -> 467.17 (+67%)
192 565.49 -> 875.19 (+55%) 587.83 -> 703.51 (+20%) 304.53 -> 545.30 (+79%) 293.61 -> 477.46 (+63%)
224 518.71 -> 872.83 (+68%) 548.30 -> 701.23 (+28%) 188.99 -> 408.06 (+116%) 216.46 -> 330.31 (+53%)
256 596.85 -> 840.32 (+41%) 583.15 -> 685.41 (+18%) 268.15 -> 535.05 (+100%) 260.49 -> 339.81 (+30%)
384 99.37 -> 249.21 (+151%) 90.95 -> 238.81 (+163%) 59.50 -> 51.36 (-14%) 59.01 -> 50.85 (-14%)
512 335.17 -> 467.27 (+39%) 369.86 -> 409.15 (+11%) 161.80 -> 298.72 (+85%) 156.78 -> 312.36 (+99%)

d_head=512 is measured with sq/sk=8192 to match gfx1201's practice.

Gfx1201 Performance Boost from 0.13b

d_head fwd non-causal fwd causal bwd non-causal bwd causal
16 43.05 -> 56.00 (+30%) 26.16 -> 50.87 (+94%) 34.64 -> 40.16 (+16%) 33.88 -> 33.96 (+0%)
32 56.35 -> 68.84 (+22%) 38.10 -> 67.49 (+77%) 43.34 -> 53.46 (+23%) 41.37 -> 49.23 (+19%)
48 58.11 -> 71.50 (+23%) 40.11 -> 71.19 (+77%) 38.89 -> 51.35 (+32%) 39.15 -> 49.51 (+26%)
64 64.93 -> 88.24 (+36%) 52.98 -> 77.99 (+47%) 46.00 -> 59.28 (+29%) 44.14 -> 56.01 (+27%)
80 61.90 -> 91.46 (+48%) 43.89 -> 87.77 (+100%) 40.24 -> 56.36 (+40%) 40.94 -> 53.23 (+30%)
96 64.18 -> 90.39 (+41%) 47.39 -> 86.54 (+83%) 39.36 -> 57.81 (+47%) 40.61 -> 55.09 (+36%)
128 69.22 -> 80.31 (+16%) 48.91 -> 78.57 (+61%) 37.84 -> 57.17 (+51%) 40.51 -> 52.89 (+31%)
160 64.91 -> 85.06 (+31%) 50.31 -> 79.88 (+59%) 25.01 -> 48.84 (+95%) 27.81 -> 47.44 (+71%)
192 61.49 -> 54.60 (-11%) 47.14 -> 46.84 (-1%) 26.64 -> 45.70 (+72%) 27.21 -> 44.60 (+64%)
224 43.21 -> 83.79 (+94%) 39.62 -> 81.09 (+105%) 17.64 -> 39.47 (+124%) 20.58 -> 37.96 (+84%)
256 50.97 -> 84.20 (+65%) 43.85 -> 80.51 (+84%) 20.47 -> 45.43 (+122%) 24.31 -> 29.97 (+23%)
384 13.20 -> 19.93 (+51%) 13.20 -> 19.54 (+48%) 6.44 -> 5.90 (-8%) 6.78 -> 6.10 (-10%)
512 28.46 -> 20.19 (-29%) 27.72 -> 19.72 (-29%) 10.50 -> 26.52 (+153%) 10.75 -> 28.25 (+163%)

d_head=384/512 is measured with sq/sk=8192 due to VRAM limits.

VarlenBits supersedes VarlenType

VarlenBits massively expands the expressibility of specifying varlen inputs:
users can now configure sequence accessing independently for Q and K, along
three choice axes:

  • Input tensor shape is BHSD or 1THD
    • Certain varlen inputs are not stacked (T) but padded to S uniformly
  • Sequence lengths are uniform, given as cumulative sums, or specified
    individually by an array
  • Sequence offsets are uniform, re-used from the cumulative sums above, or
    specified by a second array

Concretely, VarlenBits is defined as (abridged from
include/aotriton/flash.h):

struct VarlenMode {       // one side's three axes, one byte
  uint8_t stacked : 1, length : 2, position : 2, reserved : 3;
};

struct alignas(uint32_t) VarlenBits {
  VarlenMode qmode;         // bits  7:0  -- Q and K decode independently,
  VarlenMode kmode;         // bits 15:8     so mixed addressing is expressible
  uint16_t lse_layout : 2;  // bits 17:16
  uint16_t reserved   : 14; // bits 31:18
};

The extra lse_layout enables both HT and TH memory layouts for the LSE
tensor.

API BREAKING, BUT ABI STABLE. AOTriton has no plan to promise
source-level compatibility, due to the maintenance cost. However, AOTriton
begins delivering ABI compatibility with 0.14b. Applications targeting 0.13b
can upgrade to 0.14b as a drop-in replacement (note that aotriton.images
needs to be replaced in this process as well). The library automatically
translates passed-in 0.13b params structs to 0.14b.

Full Modularization

The source tree is fully modularized: everything flash-specific now lives
under modules/flash/.

pip install . at the root of the source tree installs the Python package
aotriton into the environment, which serves as the family-neutral toolchain
that transforms Triton/FlyDSL kernels into a C++ library with precompiled GPU
kernels.

pytest-gpu-lease: one GPU per pytest-xdist worker

A standalone pytest plugin (#215, #223), shipped in-tree but useful to anyone
running GPU tests under pytest-xdist:

pip install ./python/pytest-gpu-lease

It activates through the pytest11 entry point — no conftest.py wiring.
Each worker leases exactly one GPU through fcntl byte-range locks on a shared
lock file, and tests take it from the gpu_id (int) or gpu_device
('cuda:3') fixture. Set GPU_LEASE_PIN to bypass leasing and pin every
worker to one GPU instead.

The optional watchdog monitors the heartbeat timestamp each pytest-xdist
worker writes into the lock file, and terminates a worker that stops making
progress with a SIGTERM-then-SIGKILL escalation, separated by a grace period.
Monitoring out-of-process sidesteps the limits of in-process stall detection
like pytest-timeout, whose SIGALRM only takes effect at the next bytecode
boundary — which a thread parked in a blocking HIP call never reaches.

export GPU_LEASE_LOCKFILE=/dev/shm/gpu_lease.$$   # both sides must agree
python -m pytest_gpu_lease.watchdog --lockfile "$GPU_LEASE_LOCKFILE" --workers 4 &
pytest -n 4 ...

Major Changes from Release 0.13 Beta

API Changes

  • BREAKING VarlenType is replaced by VarlenBits (#222).
    • Spells configurations the enum could not: mixed per-side addressing,
      seqused_k on a packed KV cache, and TransformerEngine's (T, H)
      logsumexp layout (VarlenLseLayout::TH; default stays HT).
    • Covers torch.nn.attention.varlen.
  • ABI backward compatibility is now formally supported (#222).
  • Backends of the flash module now have names (#230).
    • For use with attn_options.force_backend_index.
    • Declared in the new generated public header <aotriton/flash/backends.h>.
      • OpAttnFwdBackend/OpAttnBwdBackend constants.
    • Also exposed under pyaotriton.v3.flash for the Python testing API.

FlyDSL / flyc Backend (new)

  • gfx1201 and gfx950 SDPA kernels developed in FlyDSL are vendored under
    modules/flash/flyc/, cross-compiled ahead of time by
    python/flyc_compile.py, and packed into aotriton.images (#227, #230).
  • Operator tuning now probes flyc; the shipped operator database selects it
    for gfx950/gfx1201 training and inferencing on non-fp32 inputs (#239).
  • 0.14b must not be built from the released flydsl==0.3.1 wheel — upstream
    LLVM miscompiles register spills and produces wrong kernels rather than
    absent ones. .ci/build_llvm_tarball.sh and .ci/build_flydsl_wheel.sh
    build the pinned LLVM and the compiler wheel from source; the pins live in
    third_party/flydsl-{compiler,kernel,llvm}.txt (#234).

Triton Compiler and Databases

  • Triton bumped to ROCm/triton release/internal/3.8.x (3.8.0, aa3cf1a1).
  • AITER ASM kernels bumped to v0.1.22 (#238), including the backward odo
    argument ABI change — the vendored dispatchers and the pin move together.
  • Kernel tuning database fully re-tuned for gfx950 and gfx1100;
    gfx90a/gfx942/gfx1201 validated against the level-3 UT (#239).
  • Operator tuning database updated for gfx950/gfx1201 (#239).
  • Partial gfx1150/gfx1151 flash databases for hdim=64, dropout=0.0
    (#200) — attn_fwd averages 61–76% faster than the default configuration.

Kernel Bug Fixes

  • attn_fwd: Fix incorrect early exits in persistent dynamic mode (#237,
    #235).
  • bwd_kernel_fuse: Fix addressing dO with Out's strides in the dK/dV
    phase (#237, #236).
  • attn_fwd: Fix the causal early-exit +inf LSE OOB writes, which corrupted
    the next sequence's LSE under a stacked varlen layout (#222).
  • attn_fwd: Fix scaling of bias, which should be cast to fp32 first
    (#228, #232).

Runtime Bug Fixes

  • Thread safety: the following functions no longer rely on the caller holding
    the GIL (#224):
    • getGpuFromStream/isArchExperimentallySupported/isArchTechPreview.
    • getMultiProcessorCount.
    • TritonKernel::invoke/TritonKernel::direct_invoke and
      AiterAsmKernel::launch_kernel
      • Now take the device from the stream (hipStreamGetDevice) instead of the
        calling thread's ambient current device, which can differ from the
        stream's.
    • mha_fwd's AiterAsmKernel cache: thread_local dropped, matching
      mha_bwd. It built a separate copy of every kernel for each new thread.
  • OnDeviceKernel::get_kernel propagates hipError_t; a missing kernel now
    fails with hipErrorInvalidImage (#238).
    • Otherwise a later call would fail due to a null hipFunction_t.

Build and Infrastructure

  • Modularization complete. Flash C++ sources, pybind bindings, tuning
    database, test suite, Triton kernels and tuning logic all live under
    modules/flash/ (#195#199); v3python/ is deleted and the tuning stack is
    the aotriton.tune package under python/tune/ (#212).
  • ROCm installed from pip wheels or a tarball is now supported (available
    since ROCm 7.14): find_package(hip) honours $ENV{ROCM_PATH} instead of
    assuming a system install at /opt/rocm (#207).

Minor Changes from Release 0.13 Beta

  • attn_options::KernelSlot gains the FlyDSL kernels (#230). Tuning builds
    only — the enum lives behind AOTRITON_BUILD_FOR_TUNING.
  • flyc is registered as op_attn_fwd backend 2 and op_attn_bwd backend 3,
    both metros mixing FlyDSL and Triton steps on one stream (#230).
  • FWD_IMPL/BWD_IMPL take a backend name (triton, aiter, flyc,
    triton_split, triton_fuse); V3_API, FORCE_FWD_BACKEND and
    FORCE_BWD_BACKEND are removed (#231).
  • test_forward.py is deprecated — use SKIP_BWD=1 test_backward.py (#231).
  • New level-0 tests: test_prime_hdim and test_memory_layouts (#231),
    test_bottom_right_fully_masked_rows (#237), test_dropout_layout (#226),
    and test_varlen_bits (#222).
  • bwd_preprocess_varlen is fused into bwd_preprocess, changing the kernel
    numbering inside attn_options (#222).
  • fast_philox emitted [r0, r2, r1, r3] because tl.join appends a minor
    axis; the four randoms of one offset now land in generation order (#226).
    The API returns the same mask either way — the win is fewer registers and a
    word order other Philox implementations can match.
  • "Plain Object Notation" (PON) formalizes the existing k=v;k=v wire format
    used by exaid and the flyc build targets (#225).
  • __signature__ gains FLYC_VERSION (#234).
  • Flash kernel tuning configuration generators, left as single-config
    placeholders during the ATI migration, are restored (#216).
  • Widened the gfx1100 attn_fwd tuning space for inference functionals, and
    dropped the stale gfx11 seqlen > 2048 tuning gate (#220). Tuning-time
    generation only; release builds come from database rows.
  • Operator tuning stores a second-best backend (op$best1st/op$best2nd);
    gfx1200 falls back to op$best2nd when the shared gfx1201 row picks flyc,
    since no gfx1200 flyc kernel is shipped (#239).
  • Tuner v3.5: kernel-level and operator-level tuning are unified into one
    module, discriminated by a new tuning_level column; dispatch flash_op is
    now dispatch flash --tuning_mode op (#213). The WebUI follows (#214).
    • CAVEAT: no migration path. Every existing tuning database must be
      re-initialized and re-tuned, and the fleet must be drained and restarted.
      Starting from a fresh work directory is recommended.
  • amd-smi: resolve the device by PCI BDF so HIP_VISIBLE_DEVICES is
    honoured, read vram_total as MB, and fall back between junction and edge
    temperature sensors (#217).
  • Tests can force a CPU reference per test via USE_ADIFFS_TXT=CPUREF or
    AOTRITON_REF_DEVICE_OPTIONS=cpu, working around ROCm 7.14 math_sdp
    defects (#231, #239).
  • AOTRITON_DEBUG_SKIP_TRITON_KERNELS builds only the Triton kernels the flyc
    backend needs (#230).
  • CMake: .nsv added to the AKS2 packing DEPENDS, fixing stale packs when
    the hsaco manifest changes (#239).
  • CI builds CPython >= 3.13 from source in base.Dockerfile; AlmaLinux 8 has
    no such package (#211).

Known Problems

  • gfx1150 and gfx1151 are not fully validated against the updated Triton
    compiler, due to hardware availability.
  • gfx1250 GPU images are not tested. A tech preview release built with an
    internal compiler is planned for better gfx1250 support.
  • The op$best2nd mechanism that protects gfx1200 from gfx1201's flyc
    selection has never been exercised on real gfx1200 hardware.
  • gfx950 flyc kernels natively support head dimensions that are multiples of
    32 only.
  • flyc hsaco build rules have no proper dependency tracking, the same
    limitation Triton hsaco rules have: neither DSL toolchain ships a -MD
    equivalent.
  • On ROCm 7.14.x, Math SDP is unreliable on both gfx950 and gfx1201. It should
    be documented with adiff files, which are not yet written.
  • The AITER forward .co packaging globs every *.co, so affine_kernels.zip
    carries uncallable fp8bf16 kernels (+383.9 KiB on gfx950).
  • third_party/flydsl-llvm.txt must be emptied once the upstream LLVM fix
    lands, in the same change that bumps flydsl-compiler.txt.

What's Changed (Auto-generated)

  • Modularization Step 1: Relocate flash attention specific C++ code to modules/flash/csrc by @xinyazhang in #195
  • Modularization Step 2: Move tuning database to per-family modules/flash/database by @xinyazhang in #196
  • Modularization Step 3: Relocate flash attention API bindings to modules/flash/bindings by @xinyazhang in #197
  • Modularization Step 4: Make python/test self-contained and migrate tritonsrc dev harness to modules/flash/kernel by @xinyazhang in #198
  • Modularization Step 5: Move SDPA test suite to modules/flash/tests by @xinyazhang in #199
  • Add gfx1150 and gfx1151 flash tuning databases for hdim=64, dropout=0.0 by @amd-callumm in #200
  • Port Fixes found in 0.13.50 Tech Preview to Mainline by @xinyazhang in #207
  • ci: build Python >= 3.13 from source in base.Dockerfile by @xinyazhang in #211
  • Modularization Step 6: Tuner v3.5 (Phase 1) by @xinyazhang in #212
  • Add pytest-gpu-lease package for pytest-xdist GPU assignment by @xinyazhang in #215
  • Modularization Step 6: Tuner v3.5 (Phase 2 - Tuning Level Unification) by @xinyazhang in #213
  • Modularization Step 6: Tuner v3.5 (Phase 3 - WebUI) by @xinyazhang in #214
  • Fix amd-smi usage related issues introduced in Tuner V3.5 by @xinyazhang in #217
  • Restore flash kernel tuning config generators in ATI by @nemanjaudovic in #216
  • Widen attn_fwd forward tuning space on gfx1100 by @nemanjaudovic in #220
  • Replace VarlenType with VarlenBits; add ABI version translation by @xinyazhang in #222
  • Fix various race conditions by @xinyazhang in #224
  • Add watchdog to pytest-gpu-lease package by @xinyazhang in #223
  • FlyDSL Integration PR1: Introduce PON and Deserializable ChoiceView by @xinyazhang in #225
  • FlyDSL Integration PR2: Rearrange the Word Order Produced by the Triton fast_philox Kernel by @xinyazhang in #226
  • fixed bias and alibi_block f32 asymmtry between the fwd and bwd by @Chi-Chu319 in #228
  • FlyDSL Integration PR3: Vendor FlyDSL SDPA Kernels and Add the AOT Compiler Driver by @xinyazhang in #227
  • FlyDSL Integration PR4: Prepare ATI Codegen for a Second Language by @xinyazhang in #229
  • FlyDSL Integration PR5: Compile, Pack and Dispatch FlyDSL attention kernels as "Flyc" Backend by @xinyazhang in #230
  • FlyDSL Integration PR6: Test the flyc Backend, and Fix the gfx950 Slab Bound by @xinyazhang in #231
  • FlyDSL Integration PR7: CI Script to Build the FlyDSL/LLVM from Source by @xinyazhang in #234
  • unit test added for bias asymmetry by @Chi-Chu319 in #232
  • Bump AITER to v0.1.22 by @xinyazhang in #238
  • Fix attn_fwd persistent-loop tile skipping and bwd_kernel_fuse dO strides by @xinyazhang in #237
  • 0.14 Final Tuning: Triton 3.8 Databases and the Flyc Operator Backend by @xinyazhang in #239

New Contributors

Full Changelog: 0.13b...0.14b

Don't miss a new aotriton release

NewReleases is sending notifications on new releases.