cpu-amx/README.md

metadata

license: other
tags:
  - stable-audio
  - text-to-audio
  - cpu
  - amx
  - int8
  - bf16

Stable Audio 3 — cpu-amx engines

Torch-free C++ / Intel-AMX inference engines for Stable Audio 3 medium, running the whole pipeline on CPU with no PyTorch, MLX, TFLite, or stable-audio-tools at runtime — just the .sos + numpy.

prompt ─▶ T5Gemma (C++ AMX) ─▶ DiT pingpong (C++ AMX int8 | bf16) ─▶ SAME-S/L decoder (C++ AMX) ─▶ WAV
audio  ─▶ SAME-S/L encoder (C++ AMX) ─▶ latent ─▶ … (audio-to-audio / inpainting)

The runtime code (CLI, gradio, sampler) lives in the GitHub repo under optimized/cpu-amx/; this HF folder holds the compiled engines + their weight blobs, which the repo downloads on first use.

⚠ Target hardware

These .sos are prebuilt for AMX-equipped Intel Xeon — 4th/5th-gen Sapphire Rapids / Emerald Rapids (the CPU must report amx_tile + amx_int8 + amx_bf16, i.e. cpu_isa_avx10_1_512_amx). The headline speed comes entirely from the AMX matrix unit: on a non-AMX CPU these engines are not faster than fp32 — rebuild from source (optimized/cpu-amx/build/) for other ISAs. Linux x86-64, glibc ≥ 2.31.

Model classes & how each was quantized

file(s)	model	precision / method	notes
t5gemma_bf16.*	T5Gemma-b-b-ul2 encoder (text conditioner)	bf16 AMX GEMMs, fp32 islands (RMSNorm/softmax/softcap/RoPE)	standard softmax attention → bf16-safe
dit_medium_int8.* + dit_medium_int8_kernels.tar.gz	medium DiT (24-block rectified-flow)	int8 — fully-fused all-integer, AOT-compiled Triton kernels + oneDNN	naive RTN (weight-PTQ hurts the DiT — see below); default, 1-thread core
dit_medium_bf16.* (reuses the int8 kernels tarball)	medium DiT	bf16 AMX GEMMs (oneDNN, packed weights) + bf16 flash attention, fp32 RoPE/RMSNorm islands	near-lossless (~59/54 dB @ L1292/L4096, cos 0.9997+) — --dit-precision bf16; ~1.24× the int8 latency, runs at --threads
same_{s,l}_decoder_bf16.*	SAME-S (50M) / SAME-L (426M) decoder	bf16 AMX GEMMs, fp32 differential-attention islands	highest-fidelity CPU decoders
same_{s,l}_decoder_int8.*	SAME-S / SAME-L decoder	int8 (w8a8, fused) via SmoothQuant α0.9 → GPTQ ("improved" grid)	½ size; SAME-S +1.2 dB over naive int8 on real music
same_{s,l}_encoder_bf16.*	SAME-S / SAME-L encoder (audio→latent)	bf16 AMX GEMMs, fp32 differential-attention islands	for audio-to-audio / inpainting
same_{s,l}_encoder_int8.*	SAME-S / SAME-L encoder	int8 (w8a8, fused) via SmoothQuant α0.9 → GPTQ (real-audio calibration)	½ size; round-trip transparent vs bf16 (≤0.13 dB)
same_l_encoder_bf16_weights_f32.bin	SAME-L encoder (optional)	fp32 refinement mode	max-fidelity; bf16 is the default

Quantization findings baked into these choices (full write-up in the repo's LESSONS.md):

• The medium DiT's quantized tier stays naive-int8 — GPTQ/SmoothQuant lower its accuracy (its adaLN-modulated qkv + the chaotic 8-step sampler make a calibration-averaged Hessian overfit and generalize worse than RTN). For fidelity-over-speed there's a near-lossless bf16 DiT tier (--dit-precision bf16): bf16 only on the AMX matmuls (GEMMs + flash attention), RoPE/RMSNorm kept in fp32 islands (bf16-RoPE angle breaks long renders). ~59/54 dB vs the int8 tier's ~40/36 dB, at ~1.24× the latency.
• The decoders are activation-limited: SmoothQuant (activation-outlier migration) does the work, GPTQ just lets α go higher. The gain transfers to real music for SAME-S (+1.2 dB), not SAME-L.
• bf16 is used where attention is either standard-softmax (T5Gemma) or where the fp32 differential- attention island absorbs the rounding (decoders/encoders); RoPE stays fp32 (a bf16-RoPE bug clips long renders).

Benchmarks (single AMX Xeon socket, this build)

Per-stage (the encoder/DiT run once per generation; the decoder once):

stage	precision	time	×realtime	vs alternative
T5Gemma encode (256 tok)	bf16	28 ms	~1000×	2.5× faster than TFLite fp16
SAME-L decode (20 s clip)	bf16 / int8	1028 / 618 ms	20× / 32×	int8 1.66× faster than bf16
SAME-S decode (20 s clip)	bf16 / int8	162 / 124 ms	124× / 161×	int8 1.30× faster
SAME-S encode (30 s clip)	bf16	208 ms	143×	(a2a/inpaint)
SAME-L encode (30 s clip)	bf16	1046 ms	28×	fp32 mode 4169 ms / 7×

End-to-end (medium DiT, 8-step): a 3 s clip generates in ~6 s (t2a, CFG, and — with the C++ encoders — audio-to-audio & inpainting alike). All 14 CLI configs (t2a / a2a / inpaint / CFG / negative prompt / APG / step counts / unconditional) pass the release self-test.

Quality (vs the fp32/reference of each stage): T5Gemma cos 0.9997 (61–67 dB); bf16 decoders ~62 dB vs the torch port; int8 decoders ~40–49 dB; encoders cos 0.9995 (SAME-S) / mean-cos 0.99996 (SAME-L bf16, 107–116 dB in fp32 mode). Audio quality is ear-verified, not dB-gated (PSNR is a harsh proxy).

Usage

Use through the GitHub repo (Stability-AI/stable-audio-3 → optimized/cpu-amx/), which pulls these files on first run:

optimized/cpu-amx/sa3 --prompt "warm analog synthwave, 120 bpm" --seconds 10
optimized/cpu-amx/sa3-gradio        # web UI

Files are flat and self-describing: each engine is <name>.so + <name>_weights.bin (+ a _manifest.txt giving the mmap layout the loader reads). The DiT additionally needs dit_medium_int8_core.bin (int8 block weights) and dit_medium_int8_kernels.tar.gz (its AOT Triton kernels).