sm_120 SAME-L: AOT (graph-capturable) engines + tensorRT/README.md
Replaces the two sm_120 SAME-L engines with AOT builds, and fills in the
empty tensorRT/README.md.
Why the engines change
The sm_120 SAME-L engines published here are JIT builds, and a JIT-built SWA plugin engine
is not stream-capturable on sm_120. Verified against the file currently in this repo:
[1] capture with NO prior warmup enqueue -> False
[3] capture AFTER warmup enqueue -> False
"this TRT engine is not stream capturable" / "The CUDA Graph is empty"
sa3_trt.py captures the decoder inside the mega-graph (T5 -> DiT loop -> decoder -> PCM),
which is the default path. So the decode is silently dropped from the graph, enqueueV3
returns False, nothing raises, and the decoder returns the pre-capture warm-up decode of
zero latents: a constant wash of noise, byte-identical for every prompt and seed, exit code
0.
The replacements embed an ahead-of-time block-tiled tensor-core kernel, so no Python runs
during inference and the engine captures.
Validation
| capture | replay wrote output | cos(replay, eager) | |
|---|---|---|---|
| encoder | True | True | 1.0000000 |
| decoder | True | True | 1.0000000 |
End-to-end on an RTX PRO 4500, 30 s render through the mega-graph path: 95-97x realtime, and
two seeds produce different audio (the old engines gave byte-identical output). Decode is
also slightly faster than the JIT build at long sequence lengths - 283.7 vs 289.1 ms at
L=1292, 899.3 vs 920.2 at L=4096 - and marginally more accurate against the FP32 decoder
(51.45 vs 51.43 dB PSNR).
sm_90 is deliberately untouched
JIT captures correctly on Hopper and the AOT kernel measured no faster there, so those files
are left as they are. The published set is therefore mixed, which the new README states
outright, along with the consequence: Triton is needed at inference on sm_90 but not on
sm_120. Both implementations stay registered by the runtime, so either kind of engine runs
on either GPU. Confirmed that the published sm_90 engine gives byte-identical output under
the old and new runtime code.
README
tensorRT/README.md was a 0-byte placeholder. It now covers the directory layout, the two
kernel implementations and which files use which, the failure mode above with its
two-seed diagnostic, benchmarks for both architectures, the build-time switches
(SA3_SWA_PLUGIN, SA3_SWA_AOT), and three gotchas that produce wrong output rather than
a build failure.
Code side: https://github.com/Stability-AI/stable-audio-3/pull/85