cpu-amx: torch-free C++ AMX CPU runtime for Stable Audio 3 medium - #105
Open
Cortexelus wants to merge 7 commits into
Open
Cortexelus wants to merge 7 commits into
Cortexelus wants to merge 7 commits into
Conversation
CPU-native SA3-medium inference on Intel AMX Xeon — the whole pipeline on torch-free C++ engines (T5Gemma -> DiT int8 -> SAME-S/L decoder, + SAME-S/L encoder for audio-to-audio/inpaint). MLX-parity CLI + gradio, full feature parity (CFG, negative prompt, APG, inpainting, audio-to-audio). Compiled .so's + weight blobs (6.5 GB) live on HF (stabilityai/stable-audio-3-optimized/cpu-amx); scripts/weights.py pulls each engine on first use. Includes build tooling (build/), tests (tests/), and docs (README/BUILD/TESTING/LESSONS). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…it_medium_int8, same_*_encoder_bf16); update weights.py to match Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…-precision int8 The int8 SAME-S/L encoders (SmoothQuant+GPTQ, real-audio calib; ear-cleared, round-trip-transparent vs bf16) are now on HF. weights.py pulls them; CppEncoder routes a2a/inpaint init-encode to the int8 encoder when --decoder-precision int8, matching the decoder. Verified end-to-end (a2a int8 → valid audio). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…nt dropdowns
- backends.load_dit(precision, threads): routes the int8 (shipped, 1-thread) or
bf16 (near-lossless fp32 RoPE/RMSNorm islands, ~59/54 dB, runs at --threads) DiT core.
- CLI --dit-precision {int8,bf16} (default int8); full feature parity across
text-to-audio / CFG+neg+APG / audio-to-audio / inpainting, both precisions.
- gradio: "DiT quantization" + "Decoder quantization" dropdowns, threaded into the CLI.
- weights.py: dit_bf16 group (dit_medium_bf16.so + core_bf16 + pin_fp32 + bf16 flash,
reuses the shared AOT kernels tarball) pulled from HF cpu-amx/.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…er pinned to 1) The int8 DiT was pinned to 1 thread (DIT_THREADS=1), so --threads=16 never reached it — a full 380s song took 676s (0.56x RT), ~664s of it DiT sampling at 1902 ms/step. Stress-tested int8 at 16 threads (bit-identical + finite over 48 forwards @ L=172/1356/4160, one length/process): no heap-race, deterministic. Now both precisions honor --threads (default 16) → int8 DiT ~7.6x faster (1902 -> ~250 ms/step). Pass --threads 1 to restore. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Everything resolved from a hardcoded /weka2/... prefix, so the engines only
loaded on the machine they were built on and the published tree leaked a
private directory layout (88 such lines across 32 files).
Paths now come from $SA3_CPUAMX_HOME (default ~/.cache/stable-audio-3/cpu-amx)
in BOTH halves:
* scripts/backends.py derives every DIR_* from the base weights.py uses, so
the loader and the downloader can no longer disagree;
* the seven C++ engines read the same variable via getenv() instead of
baking WBASE / AOT / COREBASE / FLASH_SO_DIR in at link time;
* build.sh takes $ONEDNN_HOME, the weight dumpers take $SA3_REPO,
compile_isa_all.py takes $SA3_TRITON_CPU.
NOTE: engines published before this need a rebuild from build/ to be portable.
Also:
* weights.py places the DiT files in dit_medium_cpu_amx/ with every other
engine, instead of a tritoncpu_sa3/ tree outside the layout.
* Remove the torch fp32 AEEncoder / load_ae_encoder_torch. Nothing called it
(a2a and inpaint use the C++ AMX encoders through load_encoder), and it
pulled in torch + stable_audio_tools + an unpublished loader, contradicting
the torch-free runtime. requirements.txt claimed torch was required for
--init-audio; it is not.
* .gitignore: the repo-root `build/` rule was silently excluding all 21
engine sources — a PR would have shipped the runtime with no way to rebuild
anything. Re-included, with bytecode still ignored.
* Drop tests/, TESTING.md, LESSONS.md: dev material, and every test hardcodes
a local ground-truth path.
Verified after the change: 15/15 python files compile, 5/5 shell scripts parse,
all modules import, 10/10 engine dirs resolve, and ./sa3 renders 8.0 s of audio
in 17.8 s wall (0.45x realtime, 4 steps) to a non-silent 44.1 kHz stereo WAV.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Cortexelus
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds
optimized/cpu-amx— build sources, runtime, CLI and gradio for running Stable Audio 3medium entirely on CPU via AOT C++ AMX engines (Intel AMX / AVX-512) plus numpy. No torch, MLX,
TFLite or stable-audio-tools at runtime.
Why now
The engine binaries have been published at
stabilityai/stable-audio-3-optimized/cpu-amx/(39files) since August; the code that builds and loads them was on this branch but never merged.
Contents
scripts/sa3_cpu_amx.py), gradio, pipeline, ctypes backends, HF weight fetch, self-testbuild/assets/sa3,sa3-gradioscripts/weights.pydownloads the published binaries on first use into$SA3_CPUAMX_HOME(default
~/.cache/stable-audio-3/cpu-amx).Relocatable
Engine paths resolve from
$SA3_CPUAMX_HOMEat load time in both the Python loader and theC++ engines, so nothing absolute is compiled in.
they need a rebuild from
build/to be portable.Verified
End to end on Sapphire Rapids via
./sa3: 8.0 s of audio in 17.8 s wall (0.45× realtime at 4steps), peak RSS 1.77 GB, non-silent 44.1 kHz stereo. 15/15 Python files compile, 5/5 shell
scripts parse, all modules import, 10/10 engine directories resolve.
Not yet included
does not, so output can hard-clip — a smoke render peaked at 32767/32767. Parity with the
other backends is still to do.
Commits
1-5 (Aug): initial release,
<model>_<precision>HF filenames, int8 encoders +--decoder-precision, bf16 DiT tier + gradio quant dropdowns, multi-threaded DiT for bothprecisions.
6: merge current
main.7: make the tree relocatable and strip local paths — 88 lines of a private directory layout
across 32 files, plus a
.gitignorecollision that was silently excluding all 21 engine sourcesfrom the commit, and a dead torch encoder that contradicted the torch-free runtime.
🤖 Generated with Claude Code