cpu-amx/README.md

---
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 `.so`s + 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 `.so`s 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:

```bash
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).