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 toSuniformly
- Certain varlen inputs are not stacked (
- 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-leaseIt 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
VarlenTypeis replaced byVarlenBits(#222).- Spells configurations the enum could not: mixed per-side addressing,
seqused_kon a packed KV cache, and TransformerEngine's(T, H)
logsumexp layout (VarlenLseLayout::TH; default staysHT). - Covers
torch.nn.attention.varlen.
- Spells configurations the enum could not: mixed per-side addressing,
- ABI backward compatibility is now formally supported (#222).
- Backends of the
flashmodule now have names (#230).- For use with
attn_options.force_backend_index. - Declared in the new generated public header
<aotriton/flash/backends.h>.OpAttnFwdBackend/OpAttnBwdBackendconstants.
- Also exposed under
pyaotriton.v3.flashfor the Python testing API.
- For use with
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 intoaotriton.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.1wheel — upstream
LLVM miscompiles register spills and produces wrong kernels rather than
absent ones..ci/build_llvm_tarball.shand.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_fwdaverages 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 addressingdOwithOut's strides in the dK/dV
phase (#237, #236).attn_fwd: Fix the causal early-exit+infLSE OOB writes, which corrupted
the next sequence's LSE under a stacked varlen layout (#222).attn_fwd: Fix scaling ofbias, 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_invokeand
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.
- Now take the device from the stream (
mha_fwd'sAiterAsmKernelcache:thread_localdropped, matching
mha_bwd. It built a separate copy of every kernel for each new thread.
OnDeviceKernel::get_kernelpropagateshipError_t; a missing kernel now
fails withhipErrorInvalidImage(#238).- Otherwise a later call would fail due to a null
hipFunction_t.
- Otherwise a later call would fail due to a null
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
theaotriton.tunepackage underpython/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::KernelSlotgains the FlyDSL kernels (#230). Tuning builds
only — the enum lives behindAOTRITON_BUILD_FOR_TUNING.flycis registered asop_attn_fwdbackend 2 andop_attn_bwdbackend 3,
both metros mixing FlyDSL and Triton steps on one stream (#230).FWD_IMPL/BWD_IMPLtake a backend name (triton,aiter,flyc,
triton_split,triton_fuse);V3_API,FORCE_FWD_BACKENDand
FORCE_BWD_BACKENDare removed (#231).test_forward.pyis deprecated — useSKIP_BWD=1 test_backward.py(#231).- New level-0 tests:
test_prime_hdimandtest_memory_layouts(#231),
test_bottom_right_fully_masked_rows(#237),test_dropout_layout(#226),
andtest_varlen_bits(#222). bwd_preprocess_varlenis fused intobwd_preprocess, changing the kernel
numbering insideattn_options(#222).fast_philoxemitted[r0, r2, r1, r3]becausetl.joinappends 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=vwire format
used byexaidand the flyc build targets (#225). __signature__gainsFLYC_VERSION(#234).- Flash kernel tuning configuration generators, left as single-config
placeholders during the ATI migration, are restored (#216). - Widened the gfx1100
attn_fwdtuning space for inference functionals, and
dropped the stale gfx11seqlen > 2048tuning 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 toop$best2ndwhen the shared gfx1201 row picksflyc,
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 newtuning_levelcolumn;dispatch flash_opis
nowdispatch 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.
- CAVEAT: no migration path. Every existing tuning database must be
amd-smi: resolve the device by PCI BDF soHIP_VISIBLE_DEVICESis
honoured, readvram_totalas MB, and fall back between junction and edge
temperature sensors (#217).- Tests can force a CPU reference per test via
USE_ADIFFS_TXT=CPUREFor
AOTRITON_REF_DEVICE_OPTIONS=cpu, working around ROCm 7.14math_sdp
defects (#231, #239). AOTRITON_DEBUG_SKIP_TRITON_KERNELSbuilds only the Triton kernels the flyc
backend needs (#230).- CMake:
.nsvadded to the AKS2 packingDEPENDS, 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$best2ndmechanism that protects gfx1200 from gfx1201'sflyc
selection has never been exercised on real gfx1200 hardware. - gfx950
flyckernels natively support head dimensions that are multiples of
32 only. flychsaco 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
.copackaging globs every*.co, soaffine_kernels.zip
carries uncallablefp8bf16kernels (+383.9 KiB on gfx950). third_party/flydsl-llvm.txtmust be emptied once the upstream LLVM fix
lands, in the same change that bumpsflydsl-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
- @amd-callumm made their first contribution in #200
- @nemanjaudovic made their first contribution in #216
- @Chi-Chu319 made their first contribution in #228
Full Changelog: 0.13b...0.14b