# Stable Audio 3 — TensorRT engines Prebuilt TensorRT engines for the Stable Audio 3 stack, organised by CUDA architecture. ``` tensorRT/sm_90/ Hopper — H100, H200 tensorRT/sm_120/ Blackwell — RTX 5090, RTX PRO 6000, RTX PRO 4500 ``` TensorRT engines are **not portable across architectures** — an `sm_90` engine will not load on an `sm_120` GPU. They are also tied to the TensorRT version they were built with (these were built with **TensorRT 10.15.1.29**). If your GPU or TensorRT version is not covered here, build from the ONNX in `onnx/` using the scripts in [`optimized/tensorRT/build/`](https://github.com/Stability-AI/stable-audio-3/tree/main/optimized/tensorRT/build). | directory | contents | |---|---| | `sa3-m/` | medium DiT (the main diffusion transformer) | | `sa3-sm-music/`, `sa3-sm-sfx/` | small domain-specialised DiTs | | `same-l/`, `same-s/` | autoencoder (SAME) encoder + decoder | | `t5gemma/` | text conditioner, plus its tokenizer | --- ## The SAME-L attention kernel: two implementations The SAME-L encoder and decoder use a custom sliding-window-attention plugin (`samel::diff_attn_swa`) for differential attention, which TensorRT has no native op for. That plugin ships **two implementations**, and *which one an engine uses is fixed when the engine is built* — it is baked into the `.trt` file and cannot be changed at load time. **AOT** (ahead-of-time) compiles a block-tiled tensor-core kernel to PTX and embeds it in the engine. No Python runs during inference. **JIT** dispatches a Triton kernel through a Python callback on every enqueue — 12 callbacks per decode, one per attention layer. | | AOT | JIT | |---|---|---| | CUDA-graph capturable | **yes, both architectures** | sm_90 only — **fails on sm_120** | | Python / Triton calls per decode | **0** | 12 | | Triton needed at inference | no | **yes** | | Engine size | +0.1% | baseline | ### Which files use which | published engines | kernel | why | |---|---|---| | `sm_120/same-l/` | **AOT** (block-tiled MMA) | required — the JIT build is not capturable on sm_120 | | `sm_90/same-l/` | **JIT** | works, and AOT measured no faster on Hopper, so these were left alone | So on **sm_90 you need Triton installed to run** the `same-l` engines, and on **sm_120 you do not**. Both implementations stay registered by the runtime, so either kind of engine loads and runs regardless of which GPU you are on. For anything you build yourself, **AOT is the default on every architecture** — it is the only correct choice on sm_120, it is at parity on sm_90, and it drops the runtime Triton dependency. Rebuilding the `sm_90` engines would therefore give you AOT rather than the JIT files published here; that is expected and fine, not a mismatch. ### Why JIT breaks on sm_120 The runtime captures the whole pipeline into one CUDA graph. A JIT engine cannot be captured on sm_120: `enqueueV3` returns `False`, TensorRT logs *"this TRT engine is not stream capturable"*, and **the decode stage is silently omitted from the graph** — no exception, no non-zero exit. What you get is the pre-capture warm-up decode of zero latents: a constant wash of noise, **byte-identical for every prompt and every seed**, with exit code 0. If you ever see that, this is the first thing to check. The diagnostic is to render twice with different seeds and compare the files — identical bytes mean the decode never ran. The same engine captures without complaint on sm_90, which is why the problem did not surface until Blackwell. Re-entering Python inside a captured region was always fragile rather than supported, so do not assume a future architecture will tolerate it either. The decode is inside the captured region on the default path: `sa3_trt.py` captures T5 → DiT loop → decoder → PCM copy as one graph whenever `--cfg 1.0` and no inpaint/init-audio is used, which is the default. `--no-mega-graph` runs eagerly and sidesteps capture entirely, so it is a usable workaround on an affected engine — at the cost of the graph-replay speedup. > **The `sm_120` SAME-L engines were rebuilt AOT on 2026-07-31.** Anything published > before that date was a JIT build and is affected, so if you have an `sm_120` copy cached > from earlier, re-pull it. The `sm_90` files are unchanged — JIT captures correctly on > Hopper, so they were never affected. ### Benchmarks SAME-L decode, milliseconds, median of 7 on an otherwise idle GPU. `L` is the latent sequence length — one latent is 4096 samples, so at 44.1 kHz L=1292 is a two-minute render and L=4096 is about 6m20s. **sm_120 · RTX PRO 4500 Blackwell** | kernel | L=256 | L=1024 | L=1292 | L=2048 | L=4096 | |---|---|---|---|---|---| | JIT *(not capturable here)* | 55.1 | 227.2 | 289.1 | 457.7 | 920.2 | | AOT `ptx` (scalar fallback) | 62.5 | 251.2 | 318.7 | 502.8 | 1008.8 | | **AOT `mma` (default)** | 55.4 | **223.6** | **283.7** | **448.4** | **899.3** | **sm_90 · H200** | kernel | L=256 | L=1024 | L=1292 | L=2048 | L=4096 | |---|---|---|---|---|---| | JIT | 12.4 | 47.0 | 61.1 | 98.5 | 194.4 | | AOT `ptx` (scalar fallback) | 12.7 | 47.3 | 60.7 | 96.6 | 195.5 | | **AOT `mma` (default)** | 12.4 | 48.4 | 62.6 | 100.0 | 196.8 | The default AOT kernel is at parity with JIT on both architectures — within a few percent either way, and faster at long sequence lengths on sm_120 — while remaining capturable. Accuracy against the FP32 decoder at L=1292: AOT `mma` **51.45 dB** PSNR, JIT 51.43, AOT `ptx` 51.11. The default kernel is the most accurate of the three as well as the fastest. Note the ~4.5× gap between the two architectures. That is silicon, not a software regression — a workstation card against a flagship datacenter part. Measured on these two GPUs, the H200 has **5.57×** the copy bandwidth (4084 vs 734 GB/s) and 5.08× the BF16 tensor throughput, so a bandwidth-bound decoder landing at 4.5× is extracting slightly *more* of its silicon on sm_120 than the H200 does of its own. ### How the kernels differ The default AOT kernel (`mma`) is block-tiled: 16 queries per block share a 64-wide K/V tile in shared memory, and both attention products go through `tl.dot`, which lowers to 512 × `mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32` — TF32 tensor cores. It is written in Triton and compiled ahead of time, so Triton generates the fragment layouts rather than us hand-writing them, but the result is a plain PTX blob inside the engine. The fallback AOT kernel (`ptx`) is hand-written scalar FP32: one warp per query, no K/V reuse, zero `mma` instructions. It exists only for targets where the block-tiled kernel will not build, and costs about 10% on sm_120. TF32 in the two attention products matches what the JIT path already did — it is not a new precision loss. The differential subtraction stays in FP32 on FP32 accumulators, so the cancellation-sensitive step is untouched. ### Building your own ```bash SA3_SWA_PLUGIN=aot # default — PTX compiled into the engine SA3_SWA_PLUGIN=jit # Triton through a Python callback per enqueue SA3_SWA_AOT=mma # default AOT kernel — block-tiled, tensor cores SA3_SWA_AOT=ptx # fallback AOT kernel — scalar FP32 ``` Triton is needed to *build* an AOT engine (it is the code generator) but not to run one. Three gotchas, all of which produce wrong output rather than a build failure: - **Kernel parameter order is inputs → runtime scalars → outputs.** TensorRT's AOT launcher passes the `extra` scalars *before* the output pointers. Declaring the output ahead of them makes the kernel dereference a sequence length as an address; ours died with an illegal memory access until the order was fixed. - **`BLOCK_KV` must be ≥ `BLOCK_N + 2*WINDOW`.** Otherwise the K/V tile stops covering every position the block's queries can attend to and attention contributions are *silently dropped*. `BLOCK_N=32` with `BLOCK_KV=64` needs 66 and is therefore wrong, even though it compiles and fits in shared memory. - **Shared memory over 48 KB fails at enqueue, not at build.** `BLOCK_N=64` needs 64 KB; the engine builds cleanly and then reports *"Failed to enqueue status -1"* and returns zeros. `BLOCK_N=16` needs 40 KB, and is faster anyway — with a window of only ±17, a wide K/V tile is mostly masked-out waste. The engine filenames still say `triton_swa` because Triton remains the code generator; only the compilation moved ahead of time. --- ## Other things worth knowing **The medium DiT has three precisions.** `dit_fp16mixed.trt` is the default and the one to use: FP16 attention core with FP32 RMSNorm and RoPE islands, ~4.3× faster than FP32 with no audible quality cost. `dit_fp32.trt` is the reference. `dit_bf16.trt` is kept for comparison — earlier builds of it degraded past about two minutes of audio because the RoPE angle was computed in BF16, where one ULP exceeds 2π at long sequence lengths; the published engine has the RoPE table baked in at FP32 precision. **`dit_fp8.trt` is sm_90 only** and must be built weakly-typed; strong typing breaks TensorRT's attention fusion. This is the *opposite* of the FP16-mixed rule, which requires `STRONGLY_TYPED` or the builder re-casts the FP32 islands.