sm_120 SAME-L: AOT (graph-capturable) engines + tensorRT/README.md

tensorRT/README.md

@@ -0,0 +1,177 @@
+# 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.

tensorRT/sm_120/same-l/dec_dynamic_triton_swa.trt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ccd5f48922fe46711593e5251fe45ef40a8a5eac52462b6d72ed69595458fd42
-size 1194932516
+oid sha256:e7b030cb0cf0fa075caa8c11c81b5e1bc621ddf4dd6e03e84f6c0acc81450d09
+size 1196159668

tensorRT/sm_120/same-l/enc_dynamic_triton_swa.trt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9a72e57ecee8c0d6db9077474bcd93086a7bd709cfcb1de3ab126bcfdc205c89
-size 1194895540
+oid sha256:d48b4210bff72a3d2f41ad0eb5c2741b2be7f82c6e5a8101927e755f728b7b0c
+size 1196143124