test(neuron): shrink the ACE-Step VAE tiling test to 115 MB

The vendored tiled-decode test built its random decoder with the production 1920x hop and peaked at 4.6 GB in a debug build, which got the whole neuron test binary SIGKILLed on a 4 GB CI runner. Five 2x blocks keep the decoder's structure and the multi-chunk tiling it checks, at 115 MB and 1.2 s. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_017ohK2pULMVJctL5B1Gixdp
ci: line-table debuginfo for the Test job
2026-10-17 16:42:28 +03:00 · 2026-10-17 16:18:54 +03:00 · 2026-10-17 15:56:41 +03:00 · 2026-10-17 15:22:51 +03:00 · 2026-10-17 15:19:07 +03:00 · 2026-10-17 12:05:16 +00:00
322 changed files with 351778 additions and 274 deletions
@@ -525,7 +525,14 @@ jobs:
            cuda_home: /usr/local/cuda-13.0
            build_jobs: 8
            nvcc_threads: 4
-            cargo_features: "cuda cudnn"
+            # flash-attn on ampere (#337): Laya's encoder runs every
+            # attention through flash_attn_varlen_windowed, matching the
+            # PyTorch reference on a 3060 (36 vs 39 ms at 512 tokens) where
+            # the eager path is 1.4-1.7x slower. This also moves the
+            # flavour's other flash-capable paths onto flash -- qwen3_5
+            # full attention and Z-Image -- which validate-neuron.sh and
+            # validate-image.sh gate. NEURON_FLASH_ATTN=0 rolls back.
+            cargo_features: "cuda cudnn flash-attn"
          - flavour: ada
            compute_cap: "89"
            runner: cuda-13.0
@@ -540,8 +547,8 @@ jobs:
            build_jobs: 8
            nvcc_threads: 4
            # flash-attn on blackwell first (#95): beast carries the
-            # agentic prefill pain; ada/ampere follow once the win is
-            # measured. NEURON_FLASH_ATTN=0 is the runtime rollback.
+            # agentic prefill pain. NEURON_FLASH_ATTN=0 is the runtime
+            # rollback.
            cargo_features: "cuda cudnn flash-attn"
    runs-on: ${{ matrix.runner }}
    env:
@@ -242,10 +242,20 @@ jobs:

  test:
    name: Test
-    timeout-minutes: 25
+    # The compile alone took 19 of the old 25 minutes on main (2026-10-17,
+    # 75% Rust cache hits), so any growth in a workspace crate tipped a
+    # healthy run into a timeout. 45 leaves room without hiding a hang.
+    timeout-minutes: 45
    runs-on: rust
    needs: changes
    if: needs.changes.outputs.rust == 'true'
+    # Line tables only: full debuginfo for the whole graph made the neuron
+    # test binary's compile and link the memory peak of the job, and on the
+    # 15 GB rust runners a Test sharing the host with a Clippy was killed
+    # mid-compile with no log (2026-10-17). Backtraces keep file:line.
+    env:
+      CARGO_PROFILE_DEV_DEBUG: line-tables-only
+      CARGO_PROFILE_TEST_DEBUG: line-tables-only
    steps:
      - uses: actions/checkout@v4
      # See script/ci-cargo-escalate.sh for the escalation rationale.
@@ -244,10 +244,14 @@ jobs:
            flavour: blackwell
            load_timeout: 900
            max_prompt_tokens: 131072
-          - host: benjy.hanzalova.internal
-            flavour: ada
-            load_timeout: 300
-            max_prompt_tokens: 16384
+          # benjy is out of the matrix while its boot NVMe is replaced
+          # (dead 2026-09-29): every deploy failed at SSH init against a
+          # host that is not there. Restore these four lines when it is
+          # back on the mesh.
+          # - host: benjy.hanzalova.internal
+          #   flavour: ada
+          #   load_timeout: 300
+          #   max_prompt_tokens: 16384
          - host: quadbrat.hanzalova.internal
            flavour: ampere
            load_timeout: 300
@@ -28,3 +28,6 @@ doc/plan/*
 .claude/
 helexa-upstream.toml
 helexa-angels.toml
+
+# Local scratch space for playground example checks (never committed).
+.pg-scratch/
@@ -1143,3 +1143,102 @@ constructors on the MoE and MLP.
  budget matches a naive oracle — are closed.
 - **MTP (#313), vision (#314), quantisation and rollout (#315)** are
  untouched.
+
+## 2026-09-29 addendum: System-One decisions (Laya, epic #333)
+
+A third modality beside text and images: **typed decisions**.
+`convaiinnovations/laya` is a 421M non-autoregressive decision model, a
+ModernBERT encoder plus an option-scoring head. It answers `choice` /
+`score` / `noul` questions over a *state* in one encoder forward, with
+calibrated probabilities, on the TypeSafe Jev wire protocol:
+`POST /v1/systemone`. It never generates text. Resident on quadbrat's
+3060 since #351, replacing the Qwen3-1.7B pre-warm.
+
+**The reference is the `laya` pip SDK** (`common.py`, `agent.py`,
+`router.py`, `serve.py`), **not** the `rl_common.py` in the model repo,
+which has diverged. Fixtures recorded from SDK 0.3.21 at model revision
+`55cf4c4e` live in `crates/neuron/src/harness/testdata/laya/`:
+- `reference.json`: 28 cases, f32 CPU.
+- `reference_cuda_bf16.json`: the SDK's own GPU bf16 run, the envelope
+  our bf16 is judged against.
+- `lang_corpus.json`: 924 detector cases.
+
+Recorders are `script/laya-reference.py` and `script/laya-lang-reference.py`.
+
+**Where it lives:**
+- `cortex-core/src/decisions.rs`: wire types, and `Json`, an
+  order-preserving, correctly rounded, `json.loads`-compatible value.
+  Request data must not go through `serde_json::Value`: the workspace
+  has no `preserve_order`, and without `float_roundtrip` serde_json
+  misreads floats (`…445` → `…443`). Key order changes token ids and
+  what the language detector reads.
+- `neuron/src/harness/decision/`: question validation, sequence
+  assembly under every budget rule, calibration in numpy's float32
+  order, answer shaping, and `route.rs` (checkpoint choice, ported from
+  upstream `Router._route`, reason strings included).
+- `neuron/src/harness/lang_detect.rs`: the script/language detector.
+  Python-semantics Unicode classes: `[^\W\d_]` never matches a combining
+  mark.
+- `neuron/src/harness/arch/laya/`: the engine. The ModernBERT backbone
+  is **vendored**, because candle's only runs in f32.
+- `decision_model.rs`, plus `CandleHarness::systemone` in `candle.rs`.
+
+**One repo, three checkpoints:** English at the root, `multilingual/`
+and `typed-decisions/`, all resident. Precedence: a `model` pin (a
+checkpoint name, an alias, or a published id such as
+`convaiinnovations/laya-multilingual`), else the detected script and
+language, else English. The English checkpoint collapses outside
+English (Khmer: 0.000 accuracy at 0.952 confidence upstream), so
+language routing is correctness, not polish.
+
+**Surface semantics:**
+- On this endpoint only, a missing or unknown `model` resolves to
+  **`helexa/one`** at every hop, because Jev clients send their own
+  service's name (`jev-1`).
+- The body is never rewritten, since `model` may be a checkpoint pin;
+  the routed model travels in `x-helexa-decision-model`.
+- Metering follows Jev: **input tokens only**. That is the sum of
+  encoded sequence lengths, and the state is encoded once per question.
+  Reserved as questions × `max_len`, settled to `usage.input_tokens`,
+  priced at $0.042/M.
+- Errors carry the #63 envelope plus FastAPI's `detail`.
+
+**GPU performance (3060, #337).** Laya needs the ampere flavour's
+`flash-attn` (enabled in #350; `NEURON_FLASH_ATTN=0` rolls back). What
+the profile found:
+- candle's `LayerNorm` only fuses when a bias exists, and ModernBERT's
+  norms have none: six kernels per norm, about 25% of GPU time. Fixed
+  with an explicit zero bias.
+- Mask add and scale over the full score tensor. Elided when nothing is
+  masked; scale folded into q.
+- Eager attention replaced by `flash_attn_varlen_windowed` over an
+  unpadded pack (sliding layers ±64).
+
+Measured, bf16: 36 ms for one question over a 512-token state (reference
+39), 32 seq/s batched. Live: 10.6 ms median at the neuron over the mesh.
+bf16 drift is autocast-style: residual, norms and attention scores stay
+in f32. Beyond that, per-case drift is chaotic; judge it by
+mean-vs-reference-envelope, not single cases.
+
+**Validation:**
+- `script/validate-decision.py <url>` replays the corpus against any
+  hop: checkpoint and usage exact, answers except on reference
+  near-ties, probabilities within a bf16 tolerance.
+- The public edge is `https://helexa.ai/v1/systemone`. It is rate-limited
+  (10 r/min per IP), so replay against the router (gallumbits:8088).
+  Probing the public names has network constraints; see
+  `doc/infra/edge-ingress.md` in the private helexa/security repo.
+
+**Sharp edges:**
+- **rustingface#6:** a neuron cold load of anything rf.internal hasn't
+  cached fails on a 503 to hf-hub's bounded-range size probe. A Laya
+  family member then loads *without* that checkpoint, logged but
+  silent. Warm with full GETs, then reload.
+- helexa-bench chat scenarios must gate on `serves_chat()`, which
+  excludes both `image` and `decision` models.
+
+**Open:** #353 (revision pinning, per-language temperatures, fetch
+retries, a bench `decision:` scenario, a dashboard panel,
+micro-batching); #342 (serving fine-tuned checkpoints); #352 (Z-Image on
+the ampere flash build, unrelated to Laya but enabled by the same
+flavour change).
@@ -40,7 +40,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "cabdc9d845d08ec7ed2d0c9de1ae4a1b198301407d55855261572761be90ec9f"
 dependencies = [
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -261,7 +261,7 @@ checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -272,13 +272,13 @@ checksum = "8b75356056920673b02621b35afd0f7dda9306d03c79a30f5c56c44cf256e3de"

 [[package]]
 name = "async-trait"
-version = "0.1.89"
+version = "0.1.92"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb"
+checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 3.0.9",
 ]

 [[package]]
@@ -394,7 +394,7 @@ checksum = "7aa268c23bfbbd2c4363b9cd302a4f504fb2a9dfe7e3451d66f35dd392e20aca"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -508,7 +508,7 @@ checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -697,7 +697,7 @@ dependencies = [
 "heck",
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -853,6 +853,7 @@ dependencies = [
 "reqwest",
 "serde",
 "serde_json",
+ "tempfile",
 "thiserror 2.0.18",
 "tokio",
 "tokio-stream",
@@ -861,6 +862,7 @@ dependencies = [
 "tracing",
 "url",
 "urlencoding",
+ "uuid",
 ]

 [[package]]
@@ -966,7 +968,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "10a2a99df6e410a8ff4245aa2006499ea662245f967cc7c0a38c83ef8eb44dbf"
 dependencies = [
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -1039,7 +1041,7 @@ dependencies = [
 "proc-macro2",
 "quote",
 "strsim",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -1052,7 +1054,7 @@ dependencies = [
 "proc-macro2",
 "quote",
 "strsim",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -1063,7 +1065,7 @@ checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead"
 dependencies = [
 "darling_core 0.20.11",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -1074,7 +1076,7 @@ checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d"
 dependencies = [
 "darling_core 0.23.0",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -1125,7 +1127,7 @@ dependencies = [
 "darling 0.20.11",
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -1135,7 +1137,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c"
 dependencies = [
 "derive_builder_core",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -1157,7 +1159,7 @@ dependencies = [
 "proc-macro2",
 "quote",
 "rustc_version",
- "syn",
+ "syn 2.0.117",
 "unicode-xid",
 ]

@@ -1202,7 +1204,7 @@ checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -1337,7 +1339,7 @@ dependencies = [
 "heck",
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -1667,7 +1669,7 @@ checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -3096,7 +3098,7 @@ checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -3166,6 +3168,9 @@ dependencies = [
 "toml",
 "tracing",
 "tracing-subscriber",
+ "unicode-normalization",
+ "unicode-properties",
+ "zip",
 ]

 [[package]]
@@ -3372,7 +3377,7 @@ checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -3471,7 +3476,7 @@ dependencies = [
 "proc-macro2",
 "proc-macro2-diagnostics",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -3540,7 +3545,7 @@ dependencies = [
 "phf_shared",
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -3569,7 +3574,7 @@ checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -3692,7 +3697,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b"
 dependencies = [
 "proc-macro2",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -3712,7 +3717,7 @@ checksum = "af066a9c399a26e020ada66a034357a868728e72cd426f3adcd35f80d88d88c8"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 "version_check",
 "yansi",
 ]
@@ -4064,7 +4069,7 @@ checksum = "b7186006dcb21920990093f30e3dea63b7d6e977bf1256be20c3563a5db070da"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -4189,7 +4194,7 @@ dependencies = [
 "proc-macro2",
 "quote",
 "serde_json",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -4392,7 +4397,7 @@ dependencies = [
 "proc-macro2",
 "quote",
 "serde_derive_internals",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -4482,7 +4487,7 @@ checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -4493,7 +4498,7 @@ checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -4579,7 +4584,7 @@ dependencies = [
 "darling 0.23.0",
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -4812,7 +4817,7 @@ dependencies = [
 "quote",
 "sqlx-core",
 "sqlx-macros-core",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -4835,7 +4840,7 @@ dependencies = [
 "sqlx-mysql",
 "sqlx-postgres",
 "sqlx-sqlite",
- "syn",
+ "syn 2.0.117",
 "tokio",
 "url",
 ]
@@ -5020,7 +5025,7 @@ dependencies = [
 "heck",
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -5040,6 +5045,17 @@ dependencies = [
 "unicode-ident",
 ]

+[[package]]
+name = "syn"
+version = "3.0.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
 [[package]]
 name = "sync_wrapper"
 version = "1.0.2"
@@ -5057,7 +5073,7 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -5143,7 +5159,7 @@ checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -5154,7 +5170,7 @@ checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -5280,7 +5296,7 @@ checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -5438,7 +5454,7 @@ checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -5798,7 +5814,7 @@ dependencies = [
 "bumpalo",
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 "wasm-bindgen-shared",
 ]

@@ -5988,7 +6004,7 @@ checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -5999,7 +6015,7 @@ checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -6313,7 +6329,7 @@ dependencies = [
 "heck",
 "indexmap 2.14.0",
 "prettyplease",
- "syn",
+ "syn 2.0.117",
 "wasm-metadata",
 "wit-bindgen-core",
 "wit-component",
@@ -6329,7 +6345,7 @@ dependencies = [
 "prettyplease",
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 "wit-bindgen-core",
 "wit-bindgen-rust",
 ]
@@ -6423,7 +6439,7 @@ checksum = "2380878cad4ac9aac1e2435f3eb4020e8374b5f13c296cb75b4620ff8e229154"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 "synstructure",
 ]

@@ -6435,7 +6451,7 @@ checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 "synstructure",
 ]

@@ -6456,7 +6472,7 @@ checksum = "70e3cd084b1788766f53af483dd21f93881ff30d7320490ec3ef7526d203bad4"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -6476,7 +6492,7 @@ checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 "synstructure",
 ]

@@ -6516,7 +6532,7 @@ checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555"
 dependencies = [
 "proc-macro2",
 "quote",
- "syn",
+ "syn 2.0.117",
 ]

 [[package]]
@@ -52,21 +52,10 @@ curl -s 'http://localhost:26559/api/v1/targets' | jq '.data.activeTargets[]|sele

 `grafana-helexa-fleet.json` is a raw dashboard model with a `datasource`
 template variable, so it binds to whichever Prometheus data source you pick
-at import. Import via the API (creds in `/etc/grafana/grafana.env`):
-
-```sh
-# on golgafrinchans; GF_SECURITY_ADMIN_USER/PASSWORD live in the root-only
-# env file, Grafana is published on :28767
-set -a; source <(sudo cat /etc/grafana/grafana.env); set +a
-jq -n --slurpfile d grafana-helexa-fleet.json \
-  '{dashboard: $d[0], overwrite: true, folderUid: null}' \
-| curl -s -u "$GF_SECURITY_ADMIN_USER:$GF_SECURITY_ADMIN_PASSWORD" \
-    -H 'Content-Type: application/json' \
-    -d @- http://localhost:28767/api/dashboards/db | jq '{status,uid,version}'
-```
-
-The dashboard lands at uid `helexa-fleet`. Re-importing with `overwrite:true`
-updates it in place.
+at import. Import it through the Grafana dashboards API with admin
+credentials (`overwrite: true` updates it in place; it lands at uid
+`helexa-fleet`). The exact command and where the admin credentials live are
+in `doc/infra/monitoring-access.md` in the private helexa/security repo.

 ## What the dashboard shows

@@ -7,9 +7,10 @@
 # ─────────────────────────────────────────────────────────────────────
 # neuron.toml for quadbrat.hanzalova.internal
 #
-# 1x RTX 3060 (12 GB) — small / quantised tier. Pre-warms Qwen3-1.7B
-# (bf16, ~4 GB), leaving ~7 GB for KV cache so long contexts on a small
-# model still have plenty of room.
+# 1x RTX 3060 (12 GB). Resident: Laya, the System-One decision model
+# (#333, ~2.3 GB bf16 for its three checkpoints). It replaced the
+# Qwen3-1.7B pre-warm (operator decision 2026-09-29, #341); the 1.7B is
+# still cold-loadable beside it for helexa/small.
 #
 # Synced to /etc/neuron/neuron.toml by script/infra-setup.sh.

@@ -21,10 +22,21 @@ name = "candle"
 [harness.candle]

 [[default_models]]
-model_id = "Qwen/Qwen3-1.7B"
+model_id = "convaiinnovations/laya"
 harness = "candle"
 devices = [0]

+# A decision request is one forward of tens of milliseconds (36 ms for
+# one question over a 512-token state, 162 ms for five), so a one-slot
+# queue of 8 rejects ordinary bursts. Admit
+# several at once -- the device worker still runs forwards one at a
+# time -- and let a deep, short-deadline queue absorb the rest. Applies
+# to every model on this host, including a cold-loaded helexa/small.
+[harness.candle.admission]
+max_in_flight = 4
+max_queue_depth = 64
+max_wait_secs = 10
+
 # Text-to-image (#204): the 12 GB tier serves the q8_0-quantized DiT
 # (8.7 GB resident incl. VAE). Ceiling is 768² — naive f32 attention
 # transients OOM 1024² beside the resident model; lifts to 1024² if
@@ -34,8 +34,8 @@ server {

 server {
    # TCP 443 on the edge is owned by the stream SNI router
-    # (streams-enabled/sni-router.conf, architecture reverse-proxies.md
-    # §5), which preads the SNI and hands every local name to this http
+    # (streams-enabled/sni-router.conf; see doc/infra/edge-ingress.md in
+    # the private helexa/security repo), which preads the SNI and hands every local name to this http
    # tier over PROXY protocol. Listening on :443 directly here would
    # never be reached — the router would answer first with whichever
    # vhost happened to be default, which is exactly the symptom seen
@@ -1,6 +1,7 @@
-# Internal bench UI vhost — https://bench.internal, reachable from inside
-# the WireGuard mesh (the public bench.helexa.ai dead-ends at the OPNsense
-# LAN interface, which only port-forwards :443 from the WAN). Same SPA +
+# Internal bench UI vhost — https://bench.internal, for clients inside
+# the WireGuard mesh, where the public bench.helexa.ai doesn't resolve
+# to a usable path (see doc/infra/edge-ingress.md in the private
+# helexa/security repo). Same SPA +
 # /api→bob proxy as bench.helexa.ai, but with an internal-CA cert
 # (smallstep "lair", renewed by step@bench.timer). Mirrors the
 # *.internal vhost convention on oolon.kosherinata.internal.
@@ -4,7 +4,7 @@
 # Installed on BOTH edge proxies (oolon.kosherinata.internal and
 # hanzalova.internal); Cloudflare DNS load balancing routes the public
 # name to the two site WAN IPs. Mesh clients use helexa.internal instead
-# (public names don't hairpin — see architecture/reverse-proxies.md §2).
+# (see doc/infra/edge-ingress.md in the private helexa/security repo).
 #
 # Same-origin backends (#167): /v1 (+ /health) reverse-proxies to
 # helexa-router and /api to helexa-upstream, both on gallumbits — so the
@@ -1,7 +1,7 @@
-# Internal helexa website vhost — https://helexa.internal, reachable from
-# inside the WireGuard mesh (the public helexa.ai dead-ends at the
-# OPNsense LAN interface, which only port-forwards :443 from the WAN —
-# see architecture/reverse-proxies.md §2). Same SPA webroot as the
+# Internal helexa website vhost — https://helexa.internal, for clients
+# inside the WireGuard mesh, where the public name doesn't resolve to a
+# usable path (see doc/infra/edge-ingress.md in the private
+# helexa/security repo). Same SPA webroot as the
 # public vhost, but with an internal-CA cert (smallstep "lair", renewed
 # by step@helexa.timer). Installed on hanzalova.internal only; the
 # split-horizon helexa.internal record points at hanzalova's mesh IP on
@@ -113,3 +113,19 @@ enabled = false
 # timeout_secs = 5
 # How often to flush served-usage counters to upstream for reconciliation (#58).
 # served_usage_report_interval_secs = 60
+
+# -- Async audio generation jobs (#382) -----------------------------------
+# Music renders take seconds to minutes, so /v1/audio/generations is a job
+# API: submit, poll, fetch the audio. Job records and rendered audio live
+# under `dir` and survive a restart (the packaged unit's StateDirectory
+# creates /var/lib/cortex). All values below are the defaults.
+[audio_jobs]
+# dir = "/var/lib/cortex/audio-jobs"
+# Finished jobs and their audio are deleted after this many hours.
+# retention_hours = 72
+# Renders running at once across the fleet; the rest queue in cortex.
+# max_concurrent = 4
+# Dispatch attempts per job; a busy (429/503) model does not use one up.
+# max_attempts = 2
+# How long a job may wait on a busy model before failing, in seconds.
+# queue_timeout_secs = 3600
@@ -0,0 +1,353 @@
+//! Music / audio generation surface (#382).
+//!
+//! OpenAI has no music endpoint, so this is a helexa extension. It follows
+//! the images surface where that makes sense (#197) but is **asynchronous**:
+//! a song takes seconds to minutes to render, so the public API is a job API.
+//!
+//! - `POST   /v1/audio/generations`               submit → `202` + [`AudioJob`]
+//! - `GET    /v1/audio/generations`               list the caller's jobs
+//! - `GET    /v1/audio/generations/{id}`          poll → [`AudioJob`]
+//! - `GET    /v1/audio/generations/{id}/content`  the rendered audio bytes
+//! - `DELETE /v1/audio/generations/{id}`          cancel
+//!
+//! cortex owns the job record. A neuron renders exactly one generation per
+//! dispatch, synchronously, on `POST /v1/audio/generations`: the response
+//! body is the audio (`audio/wav`) and the result metadata travels in the
+//! [`RESULT_HEADER`] header as JSON ([`NeuronAudioResult`]). Keeping the
+//! bytes out of a JSON envelope matters here — ten minutes of 48 kHz stereo
+//! is over 100 MB of PCM.
+
+use serde::{Deserialize, Serialize};
+
+/// Header carrying [`NeuronAudioResult`] JSON on a neuron's audio response.
+pub const RESULT_HEADER: &str = "x-helexa-audio-result";
+
+/// Capability string a model advertises when it serves this surface.
+pub const CAPABILITY: &str = "audio";
+
+/// `POST /v1/audio/generations` request body. The same body is accepted by
+/// cortex (which turns it into a job) and by neuron (which renders it).
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct AudioGenerationRequest {
+    pub model: String,
+    /// Style / caption: what the music should sound like.
+    pub prompt: String,
+    /// Lyrics for models that sing. Absent, empty, or `[Instrumental]`
+    /// means an instrumental.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub lyrics: Option<String>,
+    /// Requested length in seconds. Defaults to the model's advertised
+    /// default; validated against its advertised range.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub duration: Option<f64>,
+    /// Fixed RNG seed for reproducible renders. The seed actually used is
+    /// always reported back, so a random render can be reproduced too.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub seed: Option<u64>,
+    /// Container for the delivered audio. Only `"wav"` today.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub response_format: Option<String>,
+
+    // ── model-specific extensions ───────────────────────────────
+    /// Diffusion steps (diffusion models only).
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub steps: Option<u32>,
+    /// Guidance / CFG scale.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub guidance_scale: Option<f64>,
+    /// Language of the lyrics, as a code such as `"en"` or `"es"`. Models
+    /// that use it default to `"unknown"`; a song planner may infer it.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub language: Option<String>,
+    /// Song planner: `"off"`, `"auto"` (model default) or `"on"`, for
+    /// models that have one (ACE-Step's LM).
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub planner: Option<String>,
+    /// Sampling temperature (autoregressive models only).
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub temperature: Option<f64>,
+    /// Top-k sampling (autoregressive models only).
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub top_k: Option<u32>,
+}
+
+/// Response formats the surface delivers.
+pub const SUPPORTED_FORMATS: &[&str] = &["wav"];
+
+impl AudioGenerationRequest {
+    /// True when the request asks for an instrumental (no singing).
+    pub fn is_instrumental(&self) -> bool {
+        match self.lyrics.as_deref().map(str::trim) {
+            None | Some("") => true,
+            Some(l) => l.eq_ignore_ascii_case("[instrumental]"),
+        }
+    }
+
+    /// Transport-level validation every hop can do without knowing the
+    /// model: required fields, sane numbers, a known format. Model-specific
+    /// limits ([`AudioLimits`]) are checked separately.
+    pub fn validate_shape(&self) -> Result<(), String> {
+        if self.model.trim().is_empty() {
+            return Err("'model' must be non-empty".into());
+        }
+        if self.prompt.trim().is_empty() {
+            return Err("'prompt' must be non-empty".into());
+        }
+        if let Some(d) = self.duration
+            && !(d.is_finite() && d > 0.0)
+        {
+            return Err(format!("'duration' must be a positive number, got {d}"));
+        }
+        if let Some(f) = self.response_format.as_deref()
+            && !SUPPORTED_FORMATS.contains(&f)
+        {
+            return Err(format!(
+                "response_format '{f}' is not supported; use one of {SUPPORTED_FORMATS:?}"
+            ));
+        }
+        if let Some(p) = self.planner.as_deref()
+            && !matches!(p, "off" | "auto" | "on")
+        {
+            return Err(format!("planner must be 'off', 'auto' or 'on', got '{p}'"));
+        }
+        Ok(())
+    }
+
+    /// Validate against a model's advertised limits.
+    pub fn validate_limits(&self, limits: &AudioLimits) -> Result<(), String> {
+        let d = self.duration.unwrap_or(limits.default_duration);
+        if d < limits.min_duration || d > limits.max_duration {
+            return Err(format!(
+                "duration {d}s is outside this model's range {}..={}s",
+                limits.min_duration, limits.max_duration
+            ));
+        }
+        if !limits.vocals && !self.is_instrumental() {
+            return Err("this model does not sing; omit 'lyrics' or use [Instrumental]".into());
+        }
+        if let (Some(s), Some(max)) = (self.steps, limits.max_steps)
+            && (s == 0 || s > max)
+        {
+            return Err(format!("steps must be in 1..={max}, got {s}"));
+        }
+        Ok(())
+    }
+}
+
+/// What a music model can do, advertised in `/v1/models` metadata so
+/// clients (and the `/music` page) can build their forms from it.
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct AudioLimits {
+    pub min_duration: f64,
+    pub max_duration: f64,
+    pub default_duration: f64,
+    /// Whether the model sings (accepts lyrics).
+    pub vocals: bool,
+    pub sample_rate: u32,
+    pub channels: u16,
+    /// Upper bound for `steps`, when the model takes one.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub max_steps: Option<u32>,
+}
+
+/// Lifecycle of a job.
+#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
+#[serde(rename_all = "snake_case")]
+pub enum JobStatus {
+    Queued,
+    Running,
+    Succeeded,
+    Failed,
+    Cancelled,
+}
+
+impl JobStatus {
+    pub fn is_terminal(self) -> bool {
+        matches!(self, Self::Succeeded | Self::Failed | Self::Cancelled)
+    }
+}
+
+/// The public job object returned by submit, poll, list and cancel.
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct AudioJob {
+    pub id: String,
+    /// Always `"audio.generation"`.
+    pub object: String,
+    pub status: JobStatus,
+    pub model: String,
+    /// The submitted style / caption, echoed so a job list is readable on
+    /// its own.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub prompt: Option<String>,
+    /// The submitted lyrics, if any.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub lyrics: Option<String>,
+    /// Unix seconds.
+    pub created_at: u64,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub started_at: Option<u64>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub completed_at: Option<u64>,
+    /// Requested duration, echoed so a client can show progress context.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub duration: Option<f64>,
+    /// Present when `status == succeeded`.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub result: Option<AudioResult>,
+    /// Present when `status == succeeded`.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub usage: Option<AudioUsage>,
+    /// Present when `status == failed`: the #63 envelope's inner `error`.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub error: Option<serde_json::Value>,
+}
+
+pub const JOB_OBJECT: &str = "audio.generation";
+
+/// The rendered audio's properties.
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct AudioResult {
+    /// Actual length in seconds.
+    pub duration: f64,
+    pub sample_rate: u32,
+    pub channels: u16,
+    /// The seed the render used (random when the request gave none).
+    pub seed: u64,
+    /// Container of the stored audio (`"wav"`).
+    pub format: String,
+    /// Size of the stored audio in bytes.
+    pub bytes: u64,
+}
+
+/// Metering block. `helexa_audio_units` is audio-seconds generated (#386).
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct AudioUsage {
+    pub helexa_audio_units: f64,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub helexa_timing: Option<AudioTiming>,
+}
+
+/// Phase timing for one render, measured on the neuron.
+#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq)]
+pub struct AudioTiming {
+    /// Text / lyric encoding.
+    pub encode_ms: u64,
+    /// Song planner (0 when it did not run).
+    pub plan_ms: u64,
+    /// Denoising, or autoregressive token generation.
+    pub generate_ms: u64,
+    /// Latent / codec decode to waveform.
+    pub decode_ms: u64,
+}
+
+/// What a neuron reports, in [`RESULT_HEADER`], beside the audio bytes.
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct NeuronAudioResult {
+    pub duration: f64,
+    pub sample_rate: u32,
+    pub channels: u16,
+    pub seed: u64,
+    pub usage: AudioUsage,
+}
+
+/// `GET /v1/audio/generations` response.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AudioJobList {
+    /// Always `"list"`.
+    pub object: String,
+    pub data: Vec<AudioJob>,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn req() -> AudioGenerationRequest {
+        serde_json::from_str(r#"{"model": "m", "prompt": "lofi"}"#).unwrap()
+    }
+
+    fn limits() -> AudioLimits {
+        AudioLimits {
+            min_duration: 10.0,
+            max_duration: 600.0,
+            default_duration: 30.0,
+            vocals: true,
+            sample_rate: 48_000,
+            channels: 2,
+            max_steps: Some(20),
+        }
+    }
+
+    #[test]
+    fn minimal_request_deserializes_and_validates() {
+        let r = req();
+        assert!(r.validate_shape().is_ok());
+        assert!(r.validate_limits(&limits()).is_ok());
+        assert!(r.is_instrumental());
+    }
+
+    #[test]
+    fn instrumental_detection() {
+        let mut r = req();
+        r.lyrics = Some("  [Instrumental] ".into());
+        assert!(r.is_instrumental());
+        r.lyrics = Some("[verse]\nla la".into());
+        assert!(!r.is_instrumental());
+    }
+
+    #[test]
+    fn shape_rejects_bad_values() {
+        let mut r = req();
+        r.prompt = " ".into();
+        assert!(r.validate_shape().is_err());
+        let mut r = req();
+        r.duration = Some(f64::NAN);
+        assert!(r.validate_shape().is_err());
+        let mut r = req();
+        r.duration = Some(-1.0);
+        assert!(r.validate_shape().is_err());
+        let mut r = req();
+        r.response_format = Some("ogg".into());
+        assert!(r.validate_shape().is_err());
+        let mut r = req();
+        r.planner = Some("maybe".into());
+        assert!(r.validate_shape().is_err());
+    }
+
+    #[test]
+    fn limits_reject_out_of_range() {
+        let l = limits();
+        let mut r = req();
+        r.duration = Some(5.0);
+        assert!(r.validate_limits(&l).is_err());
+        r.duration = Some(601.0);
+        assert!(r.validate_limits(&l).is_err());
+        r.duration = Some(600.0);
+        assert!(r.validate_limits(&l).is_ok());
+        r.steps = Some(21);
+        assert!(r.validate_limits(&l).is_err());
+        r.steps = Some(0);
+        assert!(r.validate_limits(&l).is_err());
+    }
+
+    #[test]
+    fn lyrics_refused_by_non_singing_model() {
+        let mut l = limits();
+        l.vocals = false;
+        let mut r = req();
+        r.lyrics = Some("sing this".into());
+        assert!(r.validate_limits(&l).is_err());
+        r.lyrics = Some("[Instrumental]".into());
+        assert!(r.validate_limits(&l).is_ok());
+    }
+
+    #[test]
+    fn job_status_serializes_snake_case() {
+        assert_eq!(
+            serde_json::to_string(&JobStatus::Succeeded).unwrap(),
+            "\"succeeded\""
+        );
+        assert!(JobStatus::Cancelled.is_terminal());
+        assert!(!JobStatus::Running.is_terminal());
+    }
+}
@@ -89,6 +89,11 @@ pub struct ModelProfile {
    /// are unioned with this set in the gateway's `/v1/models` response.
    #[serde(default)]
    pub capabilities: Vec<String>,
+    /// What an audio model can do (#382): duration range, vocals, sample
+    /// rate. Advertised in `/v1/models` and checked by cortex before a job
+    /// is accepted. `[models.audio]` in TOML.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub audio: Option<crate::audio::AudioLimits>,
 }

 fn default_min_devices() -> u32 {
@@ -107,6 +112,27 @@ pub const DEFAULT_RESIDENCY_PRIORITY: u32 = 100;
 /// grant on its own.
 pub const PINNED_RESIDENCY_PRIORITY: u32 = 1000;

+/// The alias a `/v1/systemone` decision request routes to when its
+/// `model` is absent or names nothing this fleet serves.
+///
+/// Decision clients written against TypeSafe's Jev API send a `model`
+/// naming *their* service (e.g. `jev-1`), and Laya's own server treats
+/// any unrecognised value as "let the router choose". Every hop of the
+/// serving chain applies the same rule, scoped to the decision endpoint
+/// only — a chat request for an unknown model is still a 404. The
+/// operator points the alias at a decision model in `[aliases]`.
+pub const DECISION_DEFAULT_ALIAS: &str = "helexa/one";
+
+/// Header carrying the concrete decision model a request was routed to.
+///
+/// A decision request's `model` field can name a checkpoint *inside* a
+/// served model family (`multilingual`, `typed-decisions`) or a foreign
+/// service, so the gateway forwards the body untouched and states the
+/// routing outcome here instead of rewriting `model`. The serving node
+/// uses it to pick the model and still reads `model` for a checkpoint
+/// pin.
+pub const DECISION_MODEL_HEADER: &str = "x-helexa-decision-model";
+
 /// The full model catalogue.
 #[derive(Debug, Clone, Serialize, Deserialize, Default)]
 pub struct ModelCatalogue {
@@ -262,6 +288,7 @@ mod tests {
            limit: None,
            cost: None,
            capabilities: vec![],
+            audio: None,
        }
    }

@@ -28,6 +28,63 @@ pub struct GatewayConfig {
    /// — a single operator runs purely local.
    #[serde(default)]
    pub upstream: UpstreamClientConfig,
+    /// Async music/audio generation jobs (#382).
+    #[serde(default)]
+    pub audio_jobs: AudioJobsConfig,
+}
+
+/// `[audio_jobs]` — where cortex keeps async audio generation jobs (#382).
+///
+/// A job record and its rendered audio are written under `dir`, so finished
+/// results and the job list survive a cortex restart. If `dir` cannot be
+/// created the store runs in memory only and says so at startup.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AudioJobsConfig {
+    #[serde(default = "default_audio_jobs_dir")]
+    pub dir: String,
+    /// How long a finished job and its audio are kept, in hours.
+    #[serde(default = "default_audio_retention_hours")]
+    pub retention_hours: u64,
+    /// How many renders cortex runs at once across the fleet. Each one
+    /// occupies a neuron device for its whole duration; queued jobs wait
+    /// in cortex rather than piling onto neuron admission queues.
+    #[serde(default = "default_audio_max_concurrent")]
+    pub max_concurrent: usize,
+    /// Dispatch attempts per job before it fails. A busy neuron (429/503)
+    /// does not consume an attempt; an unreachable or failing one does.
+    #[serde(default = "default_audio_max_attempts")]
+    pub max_attempts: u32,
+    /// How long a job may wait for a busy model before it fails, in seconds.
+    #[serde(default = "default_audio_queue_timeout")]
+    pub queue_timeout_secs: u64,
+}
+
+impl Default for AudioJobsConfig {
+    fn default() -> Self {
+        Self {
+            dir: default_audio_jobs_dir(),
+            retention_hours: default_audio_retention_hours(),
+            max_concurrent: default_audio_max_concurrent(),
+            max_attempts: default_audio_max_attempts(),
+            queue_timeout_secs: default_audio_queue_timeout(),
+        }
+    }
+}
+
+fn default_audio_jobs_dir() -> String {
+    "/var/lib/cortex/audio-jobs".into()
+}
+fn default_audio_retention_hours() -> u64 {
+    72
+}
+fn default_audio_max_concurrent() -> usize {
+    4
+}
+fn default_audio_max_attempts() -> u32 {
+    2
+}
+fn default_audio_queue_timeout() -> u64 {
+    3600
 }

 /// `[upstream]` — the helexa-upstream authority client (#57). Locally
@@ -200,6 +257,7 @@ impl Default for GatewayConfig {
            models_config: default_models_path(),
            entitlements: EntitlementsConfig::default(),
            upstream: UpstreamClientConfig::default(),
+            audio_jobs: AudioJobsConfig::default(),
        }
    }
 }
@@ -1,7 +1,9 @@
 pub mod anthropic;
+pub mod audio;
 pub mod build_info;
 pub mod catalogue;
 pub mod config;
+pub mod decisions;
 pub mod discovery;
 pub mod entitlements;
 pub mod error_envelope;
@@ -186,6 +186,14 @@ pub struct CortexModelEntry {
    /// the absent (not priced) vs `0.0` (intentionally free) distinction.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub cost: Option<ModelCost>,
+    /// What an audio model can do (#382), from its catalogue profile:
+    /// duration range, vocals, sample rate. Clients build forms from it.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub audio: Option<crate::audio::AudioLimits>,
+    /// Ledger tokens charged per second of generated audio (#386), for
+    /// audio models with a configured rate.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub audio_tokens_per_second: Option<u64>,
    /// `true` when any neuron reports this model supports tool calls.
    #[serde(default)]
    pub tool_call: bool,
@@ -27,6 +27,8 @@ eventsource-stream.workspace = true
 bytes = "1"
 urlencoding = "2"
 url = "2"
+uuid = { version = "1", features = ["v4"] }

 [dev-dependencies]
 tokio = { workspace = true, features = ["test-util"] }
+tempfile = "3"
@@ -0,0 +1,886 @@
+//! Async music / audio generation jobs (#382).
+//!
+//! A render takes seconds to minutes, so the public surface is a job API
+//! (see [`cortex_core::audio`] for the routes). cortex owns the job: it
+//! records it, queues it, dispatches it to a neuron serving the model, keeps
+//! the rendered audio, and meters it. A neuron renders one generation per
+//! dispatch, synchronously, and knows nothing about jobs.
+//!
+//! Records and audio are written under `[audio_jobs] dir`, so the job list
+//! and finished results survive a cortex restart. Jobs that were queued or
+//! running when cortex stopped are re-queued at startup. A job whose neuron
+//! fails is retried on whatever node the router picks next, up to
+//! `max_attempts`; a neuron that is merely busy (429/503) is waited on, up
+//! to `queue_timeout_secs`, without spending an attempt.
+//!
+//! Ownership: a job belongs to the account that submitted it, and only that
+//! account can see, fetch or cancel it. Anonymous jobs (no principal) are
+//! reachable by id alone — the id is an unguessable capability — and never
+//! appear in a listing.
+
+use crate::metering::ReservationGuard;
+use crate::state::CortexState;
+use axum::Router;
+use axum::body::{Body, Bytes};
+use axum::extract::{Path, State};
+use axum::http::{HeaderMap, HeaderValue, StatusCode, header};
+use axum::response::{IntoResponse, Json, Response};
+use axum::routing::get;
+use cortex_core::audio::{
+    AudioGenerationRequest, AudioJob, AudioJobList, AudioResult, JOB_OBJECT, JobStatus,
+    NeuronAudioResult, RESULT_HEADER,
+};
+use cortex_core::config::AudioJobsConfig;
+use cortex_core::entitlements::{HEADER_ACCOUNT_ID, HEADER_KEY_ID, Principal};
+use cortex_core::error_envelope::OpenAiError;
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
+use std::collections::HashMap;
+use std::path::PathBuf;
+use std::sync::{Arc, Mutex};
+use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
+use tokio::sync::{RwLock, Semaphore};
+use tokio::task::JoinHandle;
+
+/// Path of the job collection; `{id}` and `{id}/content` hang off it.
+pub const JOBS_PATH: &str = "/v1/audio/generations";
+
+/// Wait between dispatches to a busy model when the neuron sends no
+/// `Retry-After`.
+const DEFAULT_BUSY_WAIT: Duration = Duration::from_secs(5);
+
+/// How often expired jobs are swept.
+const SWEEP_INTERVAL: Duration = Duration::from_secs(600);
+
+/// A job as cortex stores it: the public object plus what is needed to run
+/// (or re-run) it.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct JobRecord {
+    job: AudioJob,
+    request: AudioGenerationRequest,
+    /// The submitting principal; `None` for an anonymous submission.
+    owner: Option<Principal>,
+    /// Ledger tokens per audio-second this job was reserved at (#386), so
+    /// it settles at the price it was accepted at even if the rate changes
+    /// while it runs. `None` for anonymous (unmetered) jobs.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    tokens_per_second: Option<u64>,
+}
+
+/// The job store and dispatcher, held in [`CortexState`].
+pub struct AudioJobs {
+    config: AudioJobsConfig,
+    /// See [`AudioJobs::dir`].
+    dir: std::sync::OnceLock<Option<PathBuf>>,
+    jobs: RwLock<HashMap<String, JobRecord>>,
+    /// Worker tasks for jobs that are queued or running, for cancellation.
+    workers: Mutex<HashMap<String, JoinHandle<()>>>,
+    /// Bounds concurrent renders across the fleet.
+    slots: Arc<Semaphore>,
+    /// Rendered audio, when there is no job directory to write it to.
+    memory_audio: Mutex<HashMap<String, Bytes>>,
+    /// Client for neuron dispatches. A neuron sends nothing until the render
+    /// is finished, so the read timeout has to cover a whole render.
+    http: reqwest::Client,
+}
+
+impl AudioJobs {
+    pub fn new(config: &AudioJobsConfig) -> Self {
+        let http = reqwest::Client::builder()
+            .read_timeout(Duration::from_secs(config.queue_timeout_secs.max(1800)))
+            .build()
+            .expect("failed to build audio HTTP client");
+        Self {
+            config: config.clone(),
+            dir: std::sync::OnceLock::new(),
+            jobs: RwLock::new(HashMap::new()),
+            workers: Mutex::new(HashMap::new()),
+            slots: Arc::new(Semaphore::new(config.max_concurrent.max(1))),
+            memory_audio: Mutex::new(HashMap::new()),
+            http,
+        }
+    }
+
+    /// The job directory, created on first use. `None` (memory only) when it
+    /// cannot be created. Lazy so that building a gateway does not touch the
+    /// filesystem until audio jobs are actually used.
+    fn dir(&self) -> Option<&PathBuf> {
+        self.dir
+            .get_or_init(|| {
+                let dir = PathBuf::from(&self.config.dir);
+                match std::fs::create_dir_all(&dir) {
+                    Ok(()) => Some(dir),
+                    Err(e) => {
+                        tracing::warn!(
+                            dir = %self.config.dir,
+                            error = %e,
+                            "audio job directory unavailable; jobs will not survive a restart"
+                        );
+                        None
+                    }
+                }
+            })
+            .as_ref()
+    }
+
+    fn record_path(&self, id: &str) -> Option<PathBuf> {
+        self.dir().map(|d| d.join(format!("{id}.json")))
+    }
+
+    fn audio_path(&self, id: &str) -> Option<PathBuf> {
+        self.dir().map(|d| d.join(format!("{id}.wav")))
+    }
+
+    async fn persist(&self, record: &JobRecord) {
+        let Some(path) = self.record_path(&record.job.id) else {
+            return;
+        };
+        let bytes = match serde_json::to_vec_pretty(record) {
+            Ok(b) => b,
+            Err(e) => {
+                tracing::warn!(job = %record.job.id, error = %e, "audio job serialise failed");
+                return;
+            }
+        };
+        // Write-then-rename so a crash never leaves a torn record.
+        let tmp = path.with_extension("json.tmp");
+        if let Err(e) = tokio::fs::write(&tmp, &bytes).await {
+            tracing::warn!(job = %record.job.id, error = %e, "audio job write failed");
+            return;
+        }
+        if let Err(e) = tokio::fs::rename(&tmp, &path).await {
+            tracing::warn!(job = %record.job.id, error = %e, "audio job rename failed");
+        }
+    }
+
+    async fn update<F: FnOnce(&mut JobRecord)>(&self, id: &str, f: F) -> Option<JobRecord> {
+        let snapshot = {
+            let mut jobs = self.jobs.write().await;
+            let record = jobs.get_mut(id)?;
+            f(record);
+            record.clone()
+        };
+        self.persist(&snapshot).await;
+        Some(snapshot)
+    }
+
+    async fn get(&self, id: &str) -> Option<JobRecord> {
+        self.jobs.read().await.get(id).cloned()
+    }
+
+    fn forget_worker(&self, id: &str) {
+        if let Ok(mut w) = self.workers.lock() {
+            w.remove(id);
+        }
+    }
+}
+
+/// Load persisted jobs, re-queue unfinished ones, and start the retention
+/// sweep. Called once at startup.
+pub async fn start(fleet: Arc<CortexState>) {
+    let Some(dir) = fleet.audio_jobs.dir().cloned() else {
+        spawn_sweeper(fleet);
+        return;
+    };
+    let mut resumed = Vec::new();
+    if let Ok(mut entries) = tokio::fs::read_dir(&dir).await {
+        while let Ok(Some(entry)) = entries.next_entry().await {
+            let path = entry.path();
+            if path.extension().and_then(|e| e.to_str()) != Some("json") {
+                continue;
+            }
+            let record: JobRecord = match tokio::fs::read(&path)
+                .await
+                .ok()
+                .and_then(|b| serde_json::from_slice(&b).ok())
+            {
+                Some(r) => r,
+                None => {
+                    tracing::warn!(path = %path.display(), "skipping unreadable audio job record");
+                    continue;
+                }
+            };
+            if !record.job.status.is_terminal() {
+                resumed.push(record.job.id.clone());
+            }
+            fleet
+                .audio_jobs
+                .jobs
+                .write()
+                .await
+                .insert(record.job.id.clone(), record);
+        }
+    }
+    for id in resumed {
+        let Some(record) = fleet
+            .audio_jobs
+            .update(&id, |r| {
+                r.job.status = JobStatus::Queued;
+                r.job.started_at = None;
+            })
+            .await
+        else {
+            continue;
+        };
+        // The reservation died with the old process: take a fresh one, and
+        // fail the job if the budget no longer allows it.
+        let guard = match reserve(&fleet, record.owner.as_ref(), &record.request).await {
+            Ok((g, rate)) => {
+                fleet
+                    .audio_jobs
+                    .update(&id, |r| r.tokens_per_second = rate)
+                    .await;
+                g
+            }
+            Err(env) => {
+                fail(&fleet, &id, envelope_error(&env)).await;
+                continue;
+            }
+        };
+        tracing::info!(job = %id, "re-queued audio job after restart");
+        spawn_worker(Arc::clone(&fleet), id, guard);
+    }
+    spawn_sweeper(fleet);
+}
+
+fn spawn_sweeper(fleet: Arc<CortexState>) {
+    tokio::spawn(async move {
+        loop {
+            tokio::time::sleep(SWEEP_INTERVAL).await;
+            sweep_expired(&fleet).await;
+        }
+    });
+}
+
+/// Delete finished jobs older than the retention window, with their audio.
+pub async fn sweep_expired(fleet: &CortexState) {
+    let jobs = &fleet.audio_jobs;
+    let cutoff = now_secs().saturating_sub(jobs.config.retention_hours * 3600);
+    let expired: Vec<String> = jobs
+        .jobs
+        .read()
+        .await
+        .values()
+        .filter(|r| r.job.status.is_terminal())
+        .filter(|r| r.job.completed_at.unwrap_or(r.job.created_at) < cutoff)
+        .map(|r| r.job.id.clone())
+        .collect();
+    for id in expired {
+        jobs.jobs.write().await.remove(&id);
+        if let Ok(mut m) = jobs.memory_audio.lock() {
+            m.remove(&id);
+        }
+        for path in [jobs.record_path(&id), jobs.audio_path(&id)]
+            .into_iter()
+            .flatten()
+        {
+            let _ = tokio::fs::remove_file(path).await;
+        }
+        tracing::debug!(job = %id, "expired audio job removed");
+    }
+}
+
+/// Routes for the job API.
+pub fn routes() -> Router<Arc<CortexState>> {
+    Router::new()
+        .route(JOBS_PATH, get(list_jobs).post(submit))
+        .route(
+            "/v1/audio/generations/{id}",
+            get(get_job).delete(cancel_job),
+        )
+        .route("/v1/audio/generations/{id}/content", get(get_content))
+}
+
+/// `POST /v1/audio/generations` — validate, reserve budget, queue.
+async fn submit(
+    State(fleet): State<Arc<CortexState>>,
+    headers: HeaderMap,
+    body: Bytes,
+) -> Response {
+    let request: AudioGenerationRequest = match serde_json::from_slice(&body) {
+        Ok(r) => r,
+        Err(e) => {
+            return invalid_params(format!("request body is not a valid audio request: {e}"));
+        }
+    };
+    if let Err(msg) = request.validate_shape() {
+        return invalid_params(msg);
+    }
+
+    // The model must be one this fleet knows. Resolving it fully would
+    // cold-load it at submit time, so only check that it exists here; the
+    // worker resolves (and loads) when the job actually runs.
+    let model = fleet.catalogue.resolve_alias(&request.model).to_string();
+    let known = fleet.catalogue.get(&model).is_some()
+        || fleet
+            .nodes
+            .read()
+            .await
+            .values()
+            .any(|n| n.models.contains_key(&model));
+    if !known {
+        return crate::error::envelope_response(OpenAiError::new(
+            404,
+            "invalid_request_error",
+            "model_not_found",
+            format!(
+                "model '{}' not found on any node and not in catalogue",
+                request.model
+            ),
+        ));
+    }
+    if let Some(profile) = fleet.catalogue.get(&model)
+        && let Some(limits) = &profile.audio
+        && let Err(msg) = request.validate_limits(limits)
+    {
+        return invalid_params(msg);
+    }
+
+    let owner = crate::metering::principal_from_headers(&headers);
+    let mut request = request;
+    request.model = model.clone();
+    let (guard, tokens_per_second) = match reserve(&fleet, owner.as_ref(), &request).await {
+        Ok(r) => r,
+        Err(env) => return crate::error::envelope_response(env),
+    };
+
+    let id = format!("audiogen_{}", uuid::Uuid::new_v4().simple());
+    let job = AudioJob {
+        id: id.clone(),
+        object: JOB_OBJECT.into(),
+        status: JobStatus::Queued,
+        model,
+        prompt: Some(request.prompt.clone()),
+        lyrics: request.lyrics.clone(),
+        created_at: now_secs(),
+        started_at: None,
+        completed_at: None,
+        duration: request.duration,
+        result: None,
+        usage: None,
+        error: None,
+    };
+    let record = JobRecord {
+        job: job.clone(),
+        request,
+        owner,
+        tokens_per_second,
+    };
+    fleet
+        .audio_jobs
+        .jobs
+        .write()
+        .await
+        .insert(id.clone(), record.clone());
+    fleet.audio_jobs.persist(&record).await;
+    metrics::counter!("cortex_audio_jobs_submitted_total", "model" => job.model.clone())
+        .increment(1);
+    tracing::info!(job = %id, model = %job.model, "audio job queued");
+
+    spawn_worker(Arc::clone(&fleet), id, guard);
+    (StatusCode::ACCEPTED, Json(job)).into_response()
+}
+
+/// `GET /v1/audio/generations` — the caller's jobs, newest first.
+async fn list_jobs(State(fleet): State<Arc<CortexState>>, headers: HeaderMap) -> Response {
+    let caller = crate::metering::principal_from_headers(&headers);
+    let mut data: Vec<AudioJob> = match &caller {
+        None => Vec::new(),
+        Some(p) => fleet
+            .audio_jobs
+            .jobs
+            .read()
+            .await
+            .values()
+            .filter(|r| {
+                r.owner
+                    .as_ref()
+                    .is_some_and(|o| o.account_id == p.account_id)
+            })
+            .map(|r| r.job.clone())
+            .collect(),
+    };
+    data.sort_by(|a, b| b.created_at.cmp(&a.created_at).then(b.id.cmp(&a.id)));
+    Json(AudioJobList {
+        object: "list".into(),
+        data,
+    })
+    .into_response()
+}
+
+/// `GET /v1/audio/generations/{id}`.
+async fn get_job(
+    State(fleet): State<Arc<CortexState>>,
+    headers: HeaderMap,
+    Path(id): Path<String>,
+) -> Response {
+    match visible_record(&fleet, &headers, &id).await {
+        Some(r) => Json(r.job).into_response(),
+        None => job_not_found(&id),
+    }
+}
+
+/// `GET /v1/audio/generations/{id}/content` — the rendered audio.
+async fn get_content(
+    State(fleet): State<Arc<CortexState>>,
+    headers: HeaderMap,
+    Path(id): Path<String>,
+) -> Response {
+    let Some(record) = visible_record(&fleet, &headers, &id).await else {
+        return job_not_found(&id);
+    };
+    if record.job.status != JobStatus::Succeeded {
+        return crate::error::envelope_response(OpenAiError::new(
+            409,
+            "invalid_request_error",
+            "job_not_complete",
+            format!(
+                "audio job '{id}' is {}; content is available once it has succeeded",
+                status_word(record.job.status)
+            ),
+        ));
+    }
+    let bytes = match fleet.audio_jobs.audio_path(&id) {
+        Some(path) => tokio::fs::read(&path).await.ok(),
+        None => fleet.audio_jobs.memory_audio(&id),
+    };
+    let Some(bytes) = bytes else {
+        return crate::error::envelope_response(OpenAiError::new(
+            410,
+            "invalid_request_error",
+            "content_expired",
+            format!("the audio for job '{id}' is no longer stored"),
+        ));
+    };
+    let mut resp = Response::new(Body::from(bytes));
+    resp.headers_mut()
+        .insert(header::CONTENT_TYPE, HeaderValue::from_static("audio/wav"));
+    if let Ok(v) = HeaderValue::from_str(&format!("attachment; filename=\"{id}.wav\"")) {
+        resp.headers_mut().insert(header::CONTENT_DISPOSITION, v);
+    }
+    resp
+}
+
+/// `DELETE /v1/audio/generations/{id}` — cancel a queued or running job.
+/// Cancelling a finished job is a no-op that returns it unchanged.
+async fn cancel_job(
+    State(fleet): State<Arc<CortexState>>,
+    headers: HeaderMap,
+    Path(id): Path<String>,
+) -> Response {
+    let Some(record) = visible_record(&fleet, &headers, &id).await else {
+        return job_not_found(&id);
+    };
+    if record.job.status.is_terminal() {
+        return Json(record.job).into_response();
+    }
+    // Aborting the worker drops its reservation guard, which releases the
+    // budget. A render already running on a neuron finishes there unseen.
+    let handle = fleet
+        .audio_jobs
+        .workers
+        .lock()
+        .ok()
+        .and_then(|mut w| w.remove(&id));
+    if let Some(h) = handle {
+        h.abort();
+    }
+    let updated = fleet
+        .audio_jobs
+        .update(&id, |r| {
+            r.job.status = JobStatus::Cancelled;
+            r.job.completed_at = Some(now_secs());
+        })
+        .await;
+    metrics::counter!("cortex_audio_jobs_cancelled_total", "model" => record.job.model.clone())
+        .increment(1);
+    tracing::info!(job = %id, "audio job cancelled");
+    match updated {
+        Some(r) => Json(r.job).into_response(),
+        None => job_not_found(&id),
+    }
+}
+
+/// A job the caller may see: owned by the caller's account, or anonymous.
+async fn visible_record(fleet: &CortexState, headers: &HeaderMap, id: &str) -> Option<JobRecord> {
+    let record = fleet.audio_jobs.get(id).await?;
+    let caller = crate::metering::principal_from_headers(headers);
+    match (&record.owner, &caller) {
+        (None, _) => Some(record),
+        (Some(owner), Some(caller)) if owner.account_id == caller.account_id => Some(record),
+        _ => None,
+    }
+}
+
+fn spawn_worker(fleet: Arc<CortexState>, id: String, guard: ReservationGuard) {
+    let jobs_fleet = Arc::clone(&fleet);
+    let task_id = id.clone();
+    let handle = tokio::spawn(async move {
+        run_job(&jobs_fleet, &task_id, guard).await;
+        jobs_fleet.audio_jobs.forget_worker(&task_id);
+    });
+    if let Ok(mut w) = fleet.audio_jobs.workers.lock() {
+        // The task may already have finished (and tried to forget itself)
+        // before this insert; only track it while it is still running.
+        if !handle.is_finished() {
+            w.insert(id, handle);
+        }
+    }
+}
+
+/// Outcome of one dispatch to a neuron.
+enum Dispatch {
+    Done(Bytes, NeuronAudioResult),
+    /// The model is busy; try again after the given wait.
+    Busy(Duration),
+    /// The request itself was refused; retrying will not help.
+    Rejected(Value),
+    /// The node failed or was unreachable; another attempt may succeed.
+    Failed(Value),
+}
+
+async fn run_job(fleet: &Arc<CortexState>, id: &str, guard: ReservationGuard) {
+    let jobs = &fleet.audio_jobs;
+    let Ok(_slot) = Arc::clone(&jobs.slots).acquire_owned().await else {
+        return;
+    };
+    let Some(record) = jobs
+        .update(id, |r| {
+            if !r.job.status.is_terminal() {
+                r.job.status = JobStatus::Running;
+                r.job.started_at = Some(now_secs());
+            }
+        })
+        .await
+    else {
+        return;
+    };
+    // Cancelled while it waited for a slot.
+    if record.job.status != JobStatus::Running {
+        return;
+    }
+    let model = record.job.model.clone();
+    let deadline = Instant::now() + Duration::from_secs(jobs.config.queue_timeout_secs);
+    let mut attempts = 0u32;
+    let mut last_error = error_value(
+        "api_error",
+        "audio_generation_failed",
+        "the audio job could not be dispatched",
+    );
+
+    while attempts < jobs.config.max_attempts.max(1) {
+        let outcome = dispatch_once(fleet, &record).await;
+        match outcome {
+            Dispatch::Done(bytes, result) => {
+                complete(fleet, id, &record, bytes, result, guard).await;
+                return;
+            }
+            Dispatch::Busy(wait) => {
+                if Instant::now() + wait > deadline {
+                    last_error = error_value(
+                        "rate_limit_error",
+                        "queue_timeout",
+                        "the model stayed busy for longer than the job queue timeout",
+                    );
+                    break;
+                }
+                tokio::time::sleep(wait).await;
+            }
+            Dispatch::Rejected(err) => {
+                last_error = err;
+                break;
+            }
+            Dispatch::Failed(err) => {
+                attempts += 1;
+                tracing::warn!(job = %id, model = %model, attempt = attempts, error = %err, "audio dispatch failed");
+                last_error = err;
+                if attempts < jobs.config.max_attempts {
+                    tokio::time::sleep(Duration::from_secs(2)).await;
+                }
+            }
+        }
+    }
+    drop(guard); // releases the reservation: nothing was rendered
+    fail(fleet, id, last_error).await;
+}
+
+async fn dispatch_once(fleet: &Arc<CortexState>, record: &JobRecord) -> Dispatch {
+    let route = match crate::router::resolve(fleet, &record.job.model).await {
+        Ok(r) => r,
+        Err(e) => {
+            let err = error_value(e.broad_type(), e.code(), &e.to_string());
+            return match e.retry_after_secs() {
+                Some(secs) => Dispatch::Busy(Duration::from_secs(secs.max(1))),
+                None if e.http_status() < 500 => Dispatch::Rejected(err),
+                None => Dispatch::Failed(err),
+            };
+        }
+    };
+    crate::handlers::touch_model(fleet, &route.node_name, &route.resolved_model_id).await;
+
+    let mut body = record.request.clone();
+    body.model = route.resolved_model_id.clone();
+    let url = format!("{}{JOBS_PATH}", route.endpoint);
+    let mut req = fleet.audio_jobs.http.post(&url).json(&body);
+    if let Some(owner) = &record.owner {
+        req = req
+            .header(HEADER_ACCOUNT_ID, &owner.account_id)
+            .header(HEADER_KEY_ID, &owner.key_id);
+    }
+    let labels = [
+        ("model", route.resolved_model_id.clone()),
+        ("node", route.node_name.clone()),
+    ];
+    metrics::counter!("cortex_requests_total", &labels).increment(1);
+    if route.cold_start {
+        metrics::counter!("cortex_cold_starts_total", &labels).increment(1);
+    }
+    let start = Instant::now();
+    let resp = match req.send().await {
+        Ok(r) => r,
+        Err(e) => {
+            metrics::counter!("cortex_request_errors_total", &labels).increment(1);
+            return Dispatch::Failed(error_value(
+                "api_error",
+                "upstream_unreachable",
+                &format!("failed to reach the serving node: {e}"),
+            ));
+        }
+    };
+    let status = resp.status();
+    let retry_after = resp
+        .headers()
+        .get(header::RETRY_AFTER)
+        .and_then(|v| v.to_str().ok())
+        .and_then(|v| v.trim().parse::<u64>().ok())
+        .map(|s| Duration::from_secs(s.max(1)));
+    let result_header = resp
+        .headers()
+        .get(RESULT_HEADER)
+        .and_then(|v| v.to_str().ok())
+        .map(str::to_string);
+    let bytes = match resp.bytes().await {
+        Ok(b) => b,
+        Err(e) => {
+            metrics::counter!("cortex_request_errors_total", &labels).increment(1);
+            return Dispatch::Failed(error_value(
+                "api_error",
+                "upstream_read_failed",
+                &format!("failed to read the serving node's response: {e}"),
+            ));
+        }
+    };
+    metrics::histogram!("cortex_audio_render_seconds", &labels)
+        .record(start.elapsed().as_secs_f64());
+
+    if status.is_success() {
+        let parsed = result_header
+            .as_deref()
+            .and_then(|h| serde_json::from_str::<NeuronAudioResult>(h).ok());
+        return match parsed {
+            Some(result) => Dispatch::Done(bytes, result),
+            None => Dispatch::Failed(error_value(
+                "api_error",
+                "upstream_bad_response",
+                "the serving node returned audio without its result metadata",
+            )),
+        };
+    }
+    metrics::counter!("cortex_request_errors_total", &labels).increment(1);
+    if status.as_u16() == 429 || status.as_u16() == 503 {
+        return Dispatch::Busy(retry_after.unwrap_or(DEFAULT_BUSY_WAIT));
+    }
+    let err = serde_json::from_slice::<Value>(&bytes)
+        .ok()
+        .and_then(|v| v.get("error").cloned())
+        .unwrap_or_else(|| {
+            error_value(
+                "api_error",
+                "upstream_error",
+                &format!("the serving node answered {status}"),
+            )
+        });
+    if status.is_client_error() {
+        Dispatch::Rejected(err)
+    } else {
+        Dispatch::Failed(err)
+    }
+}
+
+async fn complete(
+    fleet: &Arc<CortexState>,
+    id: &str,
+    record: &JobRecord,
+    bytes: Bytes,
+    result: NeuronAudioResult,
+    guard: ReservationGuard,
+) {
+    let jobs = &fleet.audio_jobs;
+    // Cancelled while the render was in flight: drop the result, release
+    // the budget, and leave the job cancelled.
+    if jobs.get(id).await.map(|r| r.job.status) == Some(JobStatus::Cancelled) {
+        drop(guard);
+        return;
+    }
+    let size = bytes.len() as u64;
+    match jobs.audio_path(id) {
+        Some(path) => {
+            if let Err(e) = tokio::fs::write(&path, &bytes).await {
+                drop(guard);
+                fail(
+                    fleet,
+                    id,
+                    error_value(
+                        "api_error",
+                        "storage_failed",
+                        &format!("could not store the rendered audio: {e}"),
+                    ),
+                )
+                .await;
+                return;
+            }
+        }
+        None => jobs.keep_in_memory(id, bytes),
+    }
+
+    let units = result.usage.helexa_audio_units;
+    match &record.owner {
+        Some(owner) => {
+            // An owned job always carries the rate it was reserved at.
+            let rate = record
+                .tokens_per_second
+                .unwrap_or(crate::rates::LOCAL_DEFAULT_TOKENS_PER_AUDIO_SECOND);
+            crate::metering::settle_audio_usage(owner, guard, &fleet.served_usage, units, rate)
+        }
+        None => drop(guard),
+    }
+    jobs.update(id, |r| {
+        r.job.status = JobStatus::Succeeded;
+        r.job.completed_at = Some(now_secs());
+        r.job.result = Some(AudioResult {
+            duration: result.duration,
+            sample_rate: result.sample_rate,
+            channels: result.channels,
+            seed: result.seed,
+            format: "wav".into(),
+            bytes: size,
+        });
+        r.job.usage = Some(result.usage.clone());
+    })
+    .await;
+    metrics::counter!("cortex_audio_jobs_succeeded_total", "model" => record.job.model.clone())
+        .increment(1);
+    tracing::info!(job = %id, model = %record.job.model, seconds = result.duration, "audio job succeeded");
+}
+
+async fn fail(fleet: &CortexState, id: &str, error: Value) {
+    let updated = fleet
+        .audio_jobs
+        .update(id, |r| {
+            r.job.status = JobStatus::Failed;
+            r.job.completed_at = Some(now_secs());
+            r.job.error = Some(error.clone());
+        })
+        .await;
+    if let Some(r) = updated {
+        metrics::counter!("cortex_audio_jobs_failed_total", "model" => r.job.model.clone())
+            .increment(1);
+        tracing::warn!(job = %id, model = %r.job.model, error = %error, "audio job failed");
+    }
+}
+
+/// Reserve budget for a job at its model's current rate (#386). Returns the
+/// guard and the rate it reserved at; anonymous jobs are unmetered.
+#[allow(clippy::result_large_err)] // the Err is the #63 envelope; see metering::reserve_or_reject
+async fn reserve(
+    fleet: &Arc<CortexState>,
+    owner: Option<&Principal>,
+    request: &AudioGenerationRequest,
+) -> Result<(ReservationGuard, Option<u64>), OpenAiError> {
+    let Some(p) = owner else {
+        return Ok((
+            ReservationGuard::anonymous(Arc::clone(&fleet.entitlements)),
+            None,
+        ));
+    };
+    let Some(rate) = fleet.rates.audio_tokens_per_second(&request.model).await else {
+        return Err(OpenAiError::service_unavailable(
+            format!("no audio rate is configured for model '{}'", request.model),
+            Some(5),
+        ));
+    };
+    let default = fleet
+        .catalogue
+        .get(&request.model)
+        .and_then(|m| m.audio.as_ref())
+        .map(|l| l.default_duration);
+    let seconds = crate::metering::audio_reservation_seconds(request.duration, default);
+    let guard = crate::metering::reserve_or_reject(
+        Arc::clone(&fleet.entitlements),
+        p,
+        crate::metering::audio_seconds_to_tokens(seconds, rate),
+    )
+    .await?;
+    Ok((guard, Some(rate)))
+}
+
+impl AudioJobs {
+    fn keep_in_memory(&self, id: &str, bytes: Bytes) {
+        if let Ok(mut m) = self.memory_audio.lock() {
+            m.insert(id.to_string(), bytes);
+        }
+    }
+
+    fn memory_audio(&self, id: &str) -> Option<Vec<u8>> {
+        self.memory_audio
+            .lock()
+            .ok()
+            .and_then(|m| m.get(id).map(|b| b.to_vec()))
+    }
+}
+
+fn now_secs() -> u64 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .map(|d| d.as_secs())
+        .unwrap_or(0)
+}
+
+fn status_word(s: JobStatus) -> &'static str {
+    match s {
+        JobStatus::Queued => "queued",
+        JobStatus::Running => "running",
+        JobStatus::Succeeded => "succeeded",
+        JobStatus::Failed => "failed",
+        JobStatus::Cancelled => "cancelled",
+    }
+}
+
+fn error_value(typ: &str, code: &str, message: &str) -> Value {
+    serde_json::json!({"type": typ, "code": code, "message": message})
+}
+
+fn envelope_error(env: &OpenAiError) -> Value {
+    env.body()
+        .get("error")
+        .cloned()
+        .unwrap_or_else(|| error_value("api_error", "budget_rejected", "budget reservation failed"))
+}
+
+fn invalid_params(message: impl Into<String>) -> Response {
+    crate::error::envelope_response(OpenAiError::new(
+        400,
+        "invalid_request_error",
+        "invalid_audio_params",
+        format!("invalid audio generation parameters: {}", message.into()),
+    ))
+}
+
+fn job_not_found(id: &str) -> Response {
+    crate::error::envelope_response(OpenAiError::new(
+        404,
+        "invalid_request_error",
+        "job_not_found",
+        format!("no audio job '{id}'"),
+    ))
+}
@@ -24,6 +24,8 @@ pub fn api_routes() -> Router<Arc<CortexState>> {
        .route("/v1/completions", post(completions))
        .route("/v1/responses", post(responses))
        .route("/v1/images/generations", post(images_generations))
+        .merge(crate::audio_jobs::routes())
+        .route("/v1/systemone", post(systemone))
        .route("/v1/models", get(list_models))
        .route("/v1/messages", post(anthropic_messages))
        .route("/health", get(health))
@@ -349,6 +351,192 @@ async fn images_generations(
        })
 }

+/// `POST /v1/systemone` — typed decisions (TypeSafe Jev wire protocol)
+/// from a decision model such as Laya (#338).
+///
+/// Routing differs from the other inference paths in one way: a
+/// missing or unknown `model` resolves to
+/// [`DECISION_DEFAULT_ALIAS`](cortex_core::catalogue::DECISION_DEFAULT_ALIAS)
+/// instead of 404ing, because Jev clients name their own service there.
+/// A `model` that *does* resolve — an alias, or a model this fleet
+/// knows — routes as usual, so asking a text model for a decision
+/// reaches neuron and gets its `wrong_modality` answer rather than a
+/// silent substitution.
+///
+/// The body is forwarded untouched (its `model` may pin a checkpoint
+/// inside the served family); the routed model travels in
+/// [`DECISION_MODEL_HEADER`](cortex_core::catalogue::DECISION_MODEL_HEADER).
+///
+/// Metered like Jev: input tokens only. The reservation is the exact
+/// upper bound — questions × the sequence cap — and settles to
+/// `usage.input_tokens` from the response.
+async fn systemone(
+    State(fleet): State<Arc<CortexState>>,
+    headers: HeaderMap,
+    body: Bytes,
+) -> Response {
+    use cortex_core::catalogue::{DECISION_DEFAULT_ALIAS, DECISION_MODEL_HEADER};
+
+    log_inbound("systemone", "/v1/systemone", &body);
+    let requested = extract_model(&body);
+    let route = match requested.as_deref() {
+        Some(model) => match router::resolve(&fleet, model).await {
+            Err(router::RouteError::ModelNotFound(_)) => {
+                tracing::debug!(
+                    requested = model,
+                    default = DECISION_DEFAULT_ALIAS,
+                    "unknown decision model; using the default"
+                );
+                router::resolve(&fleet, DECISION_DEFAULT_ALIAS).await
+            }
+            other => other,
+        },
+        None => router::resolve(&fleet, DECISION_DEFAULT_ALIAS).await,
+    };
+    let route = match route {
+        Ok(r) => r,
+        Err(e) => {
+            tracing::warn!(
+                handler = "systemone",
+                requested = requested.as_deref().unwrap_or("<none>"),
+                error = %e,
+                "route resolve failed"
+            );
+            return route_error_response(&e);
+        }
+    };
+
+    touch_model(&fleet, &route.node_name, &route.resolved_model_id).await;
+
+    let labels = [
+        ("model", route.resolved_model_id.clone()),
+        ("node", route.node_name.clone()),
+    ];
+    metrics::counter!("cortex_requests_total", &labels).increment(1);
+    metrics::counter!("cortex_decision_requests_total", &labels).increment(1);
+    if route.cold_start {
+        metrics::counter!("cortex_cold_starts_total", &labels).increment(1);
+    }
+
+    let principal = crate::metering::principal_from_headers(&headers);
+    let bound = crate::metering::decision_reservation_tokens(&body);
+    let guard = match &principal {
+        Some(principal) => {
+            match crate::metering::reserve_or_reject(
+                Arc::clone(&fleet.entitlements),
+                principal,
+                bound,
+            )
+            .await
+            {
+                Ok(guard) => Some(guard),
+                Err(env) => return crate::error::envelope_response(env),
+            }
+        }
+        None => None,
+    };
+
+    let mut forward_headers = strip_hop_headers(&headers);
+    match axum::http::HeaderValue::from_str(&route.resolved_model_id) {
+        Ok(v) => {
+            forward_headers.insert(DECISION_MODEL_HEADER, v);
+        }
+        Err(_) => {
+            // A model id that is not a valid header value cannot have
+            // come from the catalogue or a neuron's model list.
+            return error_response(500, "api_error", "invalid_model_id", "internal error");
+        }
+    }
+
+    let url = format!("{}/v1/systemone", route.endpoint);
+    let start = Instant::now();
+    let upstream = fleet
+        .http_client
+        .post(&url)
+        .headers(forward_headers)
+        .header("content-type", "application/json")
+        .body(body)
+        .send()
+        .await;
+    let upstream = match upstream {
+        Ok(r) => r,
+        Err(e) => {
+            metrics::counter!("cortex_request_errors_total", &labels).increment(1);
+            tracing::warn!(node = %route.node_name, url = %url, error = %e, "systemone proxy failed");
+            return error_response(
+                502,
+                "api_error",
+                "upstream_unreachable",
+                "failed to reach the serving node",
+            );
+        }
+    };
+    let status = upstream.status();
+    let passthrough: Vec<(axum::http::HeaderName, axum::http::HeaderValue)> = upstream
+        .headers()
+        .iter()
+        .filter(|(n, _)| {
+            matches!(
+                n.as_str(),
+                "server-timing" | "x-inference-time-ms" | "retry-after"
+            )
+        })
+        .map(|(n, v)| (n.clone(), v.clone()))
+        .collect();
+    let resp_bytes = match upstream.bytes().await {
+        Ok(b) => b,
+        Err(e) => {
+            metrics::counter!("cortex_request_errors_total", &labels).increment(1);
+            tracing::warn!(node = %route.node_name, error = %e, "systemone response read failed");
+            return error_response(
+                502,
+                "api_error",
+                "upstream_read_failed",
+                "failed to read the serving node's response",
+            );
+        }
+    };
+    metrics::histogram!("cortex_request_duration_seconds", &labels)
+        .record(start.elapsed().as_secs_f64());
+
+    if status.is_success() {
+        let input_tokens = serde_json::from_slice::<Value>(&resp_bytes)
+            .ok()
+            .and_then(|v| v.get("usage")?.get("input_tokens")?.as_u64());
+        if let Some(n) = input_tokens {
+            metrics::counter!("cortex_decision_input_tokens_total", &labels).increment(n);
+        }
+        if let (Some(principal), Some(guard)) = (principal, guard) {
+            // No usage block means the node answered outside the
+            // contract; settle at the reserved bound rather than free.
+            let spent = input_tokens.unwrap_or_else(|| {
+                tracing::warn!(node = %route.node_name, "systemone response carried no usage");
+                bound
+            });
+            crate::metering::record_spend(&principal, spent, 0);
+            fleet
+                .served_usage
+                .add(&principal.account_id, &principal.key_id, spent);
+            guard.settle(spent);
+        }
+    } else {
+        metrics::counter!("cortex_request_errors_total", &labels).increment(1);
+        // guard drops here -> reservation released, no spend recorded.
+    }
+
+    let mut response = axum::response::Response::builder()
+        .status(status)
+        .header(axum::http::header::CONTENT_TYPE, "application/json");
+    for (name, value) in passthrough {
+        response = response.header(name, value);
+    }
+    response
+        .body(axum::body::Body::from(resp_bytes))
+        .unwrap_or_else(|_| {
+            error_response(500, "api_error", "response_build_failed", "internal error")
+        })
+}
+
 /// Forward client headers minus hop-by-hop and host headers, keeping
 /// the cortex-stamped principal headers intact for neuron's fair-share
 /// admission (#54).
@@ -790,6 +978,8 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
                // `cost` stays operator-set and flows from the catalogue.
                limit: None,
                cost: profile.cost.clone(),
+                audio: profile.audio.clone(),
+                audio_tokens_per_second: None,
                // Runtime-detected — will be OR-ed in Pass 2 from neuron data.
                tool_call: false,
                reasoning: false,
@@ -872,6 +1062,8 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
                    capabilities: entry.capabilities.clone(),
                    limit: entry.limit.clone(),
                    cost: None,
+                    audio: None,
+                    audio_tokens_per_second: None,
                    tool_call: entry.tool_call,
                    reasoning: entry.reasoning,
                    max_model_len: None,
@@ -941,6 +1133,8 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
                    capabilities: discovered.get(model_id).cloned().unwrap_or_default(),
                    limit: None,
                    cost: None,
+                    audio: None,
+                    audio_tokens_per_second: None,
                    tool_call: false,
                    reasoning: false,
                    max_model_len: None,
@@ -982,6 +1176,8 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
                capabilities: target_entry.capabilities,
                limit: target_entry.limit.clone(),
                cost: target_entry.cost.clone(),
+                audio: target_entry.audio.clone(),
+                audio_tokens_per_second: None,
                tool_call: target_entry.tool_call,
                reasoning: target_entry.reasoning,
                max_model_len: None,
@@ -999,10 +1195,19 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
    // from each entry's now-settled `limit`, so vLLM-convention clients
    // (Hermes Agent et al.) can read the window without knowing helexa's
    // `limit` schema.
+    //
+    // Audio models also carry their current rate (#386), so a client can
+    // estimate a job's cost before submitting it. An alias is priced as
+    // the model it names.
+    let rates = fleet.rates.snapshot().await;
    let data: Vec<Value> = entries
        .values_mut()
        .map(|e| {
            e.sync_flat_limit();
+            if e.audio.is_some() {
+                let priced_as = catalogue.resolve_alias(&e.id);
+                e.audio_tokens_per_second = rates.audio_tokens_per_second(priced_as);
+            }
            json!(e)
        })
        .collect();
@@ -1268,7 +1473,7 @@ fn context_length_exceeded_response(
 }

 /// Update `last_accessed` timestamp for a model on a node (drives LRU eviction).
-async fn touch_model(fleet: &CortexState, node_name: &str, model_id: &str) {
+pub(crate) async fn touch_model(fleet: &CortexState, node_name: &str, model_id: &str) {
    let mut nodes = fleet.nodes.write().await;
    if let Some(node) = nodes.get_mut(node_name)
        && let Some(entry) = node.models.get_mut(model_id)
@@ -1,4 +1,5 @@
 pub mod anthropic_sse;
+pub mod audio_jobs;
 pub mod auth;
 pub mod entitlements_chain;
 pub mod entitlements_local;
@@ -10,6 +11,7 @@ pub mod metering;
 pub mod metrics;
 pub mod poller;
 pub mod proxy;
+pub mod rates;
 pub mod router;
 pub mod served_usage;
 pub mod state;
@@ -52,6 +54,15 @@ pub async fn run(config: GatewayConfig) -> Result<()> {
        poller::poll_loop(poller_fleet).await;
    });

+    // Per-model rates (#386): keep the cached table in step with upstream.
+    if config.upstream.enabled {
+        rates::spawn_refresh(Arc::clone(&fleet.rates), &config.upstream);
+    }
+
+    // Async audio jobs (#382): reload persisted jobs, re-queue unfinished
+    // ones, start the retention sweep.
+    audio_jobs::start(Arc::clone(&fleet)).await;
+
    // Spawn the evictor (reacts to VRAM pressure events from the router).
    let evictor_fleet = Arc::clone(&fleet);
    tokio::spawn(async move {
@@ -228,6 +228,42 @@ pub fn settle_image_usage(
    guard.settle(tokens);
 }

+/// Audio-seconds to reserve for a job: the requested duration, else the
+/// model's advertised default, else a minute. Settle corrects to the actual
+/// length the neuron reports.
+pub fn audio_reservation_seconds(requested: Option<f64>, model_default: Option<f64>) -> f64 {
+    requested
+        .or(model_default)
+        .filter(|d| d.is_finite() && *d > 0.0)
+        .unwrap_or(60.0)
+}
+
+/// Convert audio-seconds to budget tokens at a per-model rate (#386),
+/// rounding up.
+pub fn audio_seconds_to_tokens(seconds: f64, tokens_per_second: u64) -> u64 {
+    (seconds * tokens_per_second as f64).ceil() as u64
+}
+
+/// Record per-principal audio spend and settle the reservation with the
+/// audio-seconds actually generated.
+pub fn settle_audio_usage(
+    principal: &Principal,
+    guard: ReservationGuard,
+    served_usage: &crate::served_usage::ServedUsage,
+    audio_seconds: f64,
+    tokens_per_second: u64,
+) {
+    let labels = [
+        ("account", principal.account_id.clone()),
+        ("key", principal.key_id.clone()),
+    ];
+    metrics::counter!("cortex_spend_audio_milliseconds_total", &labels)
+        .increment((audio_seconds * 1000.0).round() as u64);
+    let tokens = audio_seconds_to_tokens(audio_seconds, tokens_per_second);
+    served_usage.add(&principal.account_id, &principal.key_id, tokens);
+    guard.settle(tokens);
+}
+
 /// Map a [`BudgetError`] to the #63 envelope. The provider chose the window
 /// semantics; this only translates them to HTTP.
 fn budget_error_to_envelope(err: BudgetError) -> OpenAiError {
@@ -239,6 +275,35 @@ fn budget_error_to_envelope(err: BudgetError) -> OpenAiError {
    }
 }

+/// Sequence cap assumed for a decision request that does not set
+/// `max_len`: the largest default any shipped Laya checkpoint uses.
+const DECISION_DEFAULT_MAX_LEN: u64 = 1024;
+
+/// The ceiling a decision request may raise `max_len` to (Laya's own
+/// server caps the budget at the same value).
+const DECISION_MAX_LEN_CAP: u64 = 8192;
+
+/// Upper-bound input tokens for a `/v1/systemone` request (#338/#340).
+///
+/// A decision model encodes one sequence per question, each at most
+/// `max_len` tokens, and bills input tokens only — so questions ×
+/// `max_len` is a true ceiling, not an estimate. The node reports the
+/// exact count in `usage.input_tokens` and settle corrects to it.
+pub fn decision_reservation_tokens(body: &[u8]) -> u64 {
+    let v: serde_json::Value = serde_json::from_slice(body).unwrap_or_default();
+    let questions = match v.get("questions") {
+        Some(serde_json::Value::Object(m)) => m.len() as u64,
+        Some(serde_json::Value::Array(a)) => a.len() as u64,
+        _ => 1,
+    };
+    let max_len = v
+        .get("max_len")
+        .and_then(serde_json::Value::as_u64)
+        .unwrap_or(DECISION_DEFAULT_MAX_LEN)
+        .min(DECISION_MAX_LEN_CAP);
+    questions.max(1).saturating_mul(max_len.max(1))
+}
+
 /// Upper-bound tokens to reserve for a request (#52): an over-estimate of
 /// the prompt plus the maximum output. `advertised_output` is the model's
 /// `limit.output` (#62), used when the request omits `max_(completion_)tokens`.
@@ -271,6 +336,33 @@ fn estimate_prompt_tokens(body: &[u8]) -> u64 {
 mod tests {
    use super::*;

+    #[test]
+    fn decision_reservation_is_questions_times_sequence_cap() {
+        let body = br#"{"state":"x","questions":{"a":{},"b":{},"c":{}}}"#;
+        assert_eq!(
+            decision_reservation_tokens(body),
+            3 * DECISION_DEFAULT_MAX_LEN
+        );
+        let body = br#"{"state":"x","questions":{"a":{},"b":{}},"max_len":256}"#;
+        assert_eq!(decision_reservation_tokens(body), 512);
+    }
+
+    #[test]
+    fn decision_reservation_caps_a_raised_max_len() {
+        let body = br#"{"questions":{"a":{}},"max_len":1000000}"#;
+        assert_eq!(decision_reservation_tokens(body), DECISION_MAX_LEN_CAP);
+    }
+
+    #[test]
+    fn decision_reservation_never_zero() {
+        assert_eq!(
+            decision_reservation_tokens(b"not json"),
+            DECISION_DEFAULT_MAX_LEN
+        );
+        let body = br#"{"questions":{},"max_len":0}"#;
+        assert_eq!(decision_reservation_tokens(body), 1);
+    }
+
    #[test]
    fn requested_max_output_prefers_max_completion_tokens() {
        let body = br#"{"model":"m","max_completion_tokens":256,"max_tokens":99}"#;
@@ -0,0 +1,208 @@
+//! Per-model rates for work metered in units other than tokens (#386).
+//!
+//! The budget ledger counts tokens. Audio jobs are metered in seconds of
+//! generated audio, converted to ledger tokens at a per-model rate. The
+//! rates are configuration held by the clearing house (helexa-upstream's
+//! `model_rates` table), so pricing changes without a deploy: cortex
+//! fetches the table from `/authz/v1/rates`, caches it, and refreshes it in
+//! the background.
+//!
+//! Resolution: the model's own rate, else the unit's `"*"` default. A model
+//! with neither is not priced, and cortex refuses its metered work rather
+//! than serving it free. Until the first fetch succeeds nothing is priced,
+//! for the same reason.
+//!
+//! A cortex running without upstream has no clearing house to ask. It
+//! prices every audio-second at [`LOCAL_DEFAULT_TOKENS_PER_AUDIO_SECOND`],
+//! which only matters for the operator's own local key caps.
+
+use cortex_core::config::UpstreamClientConfig;
+use serde::Deserialize;
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::sync::RwLock;
+
+/// Unit name for one second of generated audio, as stored upstream.
+pub const AUDIO_SECOND: &str = "audio_second";
+
+/// The default-rate model id.
+pub const ANY_MODEL: &str = "*";
+
+/// Audio rate used when this cortex runs without an upstream authority.
+pub const LOCAL_DEFAULT_TOKENS_PER_AUDIO_SECOND: u64 = 1000;
+
+/// How often the rate table is refreshed from upstream.
+const REFRESH_INTERVAL: Duration = Duration::from_secs(60);
+
+/// Retry interval while no table has loaded yet.
+const INITIAL_RETRY: Duration = Duration::from_secs(5);
+
+/// A snapshot of the rate table: `(unit, model_id) → tokens per unit`.
+#[derive(Debug, Clone, Default)]
+pub struct RateTable {
+    rates: HashMap<(String, String), u64>,
+}
+
+impl RateTable {
+    /// The table a cortex without upstream uses.
+    pub fn local_default() -> Self {
+        let mut t = Self::default();
+        t.set(
+            AUDIO_SECOND,
+            ANY_MODEL,
+            LOCAL_DEFAULT_TOKENS_PER_AUDIO_SECOND,
+        );
+        t
+    }
+
+    pub fn set(&mut self, unit: &str, model_id: &str, tokens_per_unit: u64) {
+        self.rates
+            .insert((unit.to_string(), model_id.to_string()), tokens_per_unit);
+    }
+
+    /// The model's rate for `unit`, else the unit's default, else `None`.
+    pub fn rate(&self, unit: &str, model_id: &str) -> Option<u64> {
+        self.rates
+            .get(&(unit.to_string(), model_id.to_string()))
+            .or_else(|| self.rates.get(&(unit.to_string(), ANY_MODEL.to_string())))
+            .copied()
+    }
+
+    /// Ledger tokens per second of audio for `model_id`.
+    pub fn audio_tokens_per_second(&self, model_id: &str) -> Option<u64> {
+        self.rate(AUDIO_SECOND, model_id)
+    }
+}
+
+/// The live rate table, held in `CortexState`.
+pub struct Rates {
+    table: RwLock<RateTable>,
+}
+
+impl Rates {
+    /// Rates for a cortex with no upstream: the local default.
+    pub fn local() -> Self {
+        Self {
+            table: RwLock::new(RateTable::local_default()),
+        }
+    }
+
+    /// Rates for a cortex with upstream: empty (nothing priced) until the
+    /// first fetch lands.
+    pub fn upstream() -> Self {
+        Self {
+            table: RwLock::new(RateTable::default()),
+        }
+    }
+
+    pub async fn snapshot(&self) -> RateTable {
+        self.table.read().await.clone()
+    }
+
+    pub async fn replace(&self, table: RateTable) {
+        *self.table.write().await = table;
+    }
+
+    pub async fn audio_tokens_per_second(&self, model_id: &str) -> Option<u64> {
+        self.table.read().await.audio_tokens_per_second(model_id)
+    }
+}
+
+#[derive(Deserialize)]
+struct RatesResp {
+    data: Vec<RateRow>,
+}
+
+#[derive(Deserialize)]
+struct RateRow {
+    model_id: String,
+    unit: String,
+    tokens_per_unit: i64,
+}
+
+/// Fetch the full rate table from upstream's `/authz/v1/rates`.
+pub async fn fetch(
+    client: &reqwest::Client,
+    base_url: &str,
+    bearer: &str,
+) -> anyhow::Result<RateTable> {
+    let url = format!("{}/authz/v1/rates", base_url.trim_end_matches('/'));
+    let resp = client
+        .get(&url)
+        .bearer_auth(bearer)
+        .send()
+        .await?
+        .error_for_status()?;
+    let body: RatesResp = resp.json().await?;
+    let mut table = RateTable::default();
+    for row in body.data {
+        // The database forbids non-positive rates; skip one rather than
+        // trust it if that ever changes.
+        if let Ok(tokens) = u64::try_from(row.tokens_per_unit)
+            && tokens > 0
+        {
+            table.set(&row.unit, &row.model_id, tokens);
+        }
+    }
+    Ok(table)
+}
+
+/// Keep `rates` in step with upstream. Retries quickly until the first
+/// table loads, then refreshes every minute. A failed refresh keeps the
+/// last good table.
+pub fn spawn_refresh(rates: Arc<Rates>, cfg: &UpstreamClientConfig) {
+    let client = reqwest::Client::builder()
+        .timeout(Duration::from_secs(cfg.timeout_secs.max(1)))
+        .build()
+        .expect("failed to build rates HTTP client");
+    let url = cfg.url.clone();
+    let bearer = cfg.bearer.clone();
+    tokio::spawn(async move {
+        let mut loaded = false;
+        loop {
+            match fetch(&client, &url, &bearer).await {
+                Ok(table) => {
+                    if !loaded {
+                        tracing::info!("model rates loaded from upstream");
+                    }
+                    loaded = true;
+                    rates.replace(table).await;
+                }
+                Err(e) => tracing::warn!(error = %e, "model rate refresh failed"),
+            }
+            tokio::time::sleep(if loaded {
+                REFRESH_INTERVAL
+            } else {
+                INITIAL_RETRY
+            })
+            .await;
+        }
+    });
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn model_rate_beats_default_and_absent_unit_is_unpriced() {
+        let mut t = RateTable::default();
+        assert_eq!(t.audio_tokens_per_second("m"), None);
+        t.set(AUDIO_SECOND, ANY_MODEL, 1000);
+        assert_eq!(t.audio_tokens_per_second("m"), Some(1000));
+        t.set(AUDIO_SECOND, "m", 250);
+        assert_eq!(t.audio_tokens_per_second("m"), Some(250));
+        assert_eq!(t.audio_tokens_per_second("other"), Some(1000));
+        assert_eq!(t.rate("some_other_unit", "m"), None);
+    }
+
+    #[test]
+    fn local_default_prices_every_audio_model() {
+        let t = RateTable::local_default();
+        assert_eq!(
+            t.audio_tokens_per_second("anything"),
+            Some(LOCAL_DEFAULT_TOKENS_PER_AUDIO_SECOND)
+        );
+    }
+}
@@ -663,6 +663,7 @@ mod tests {
            limit: None,
            cost: None,
            capabilities: vec![],
+            audio: None,
        }
    }

@@ -52,6 +52,11 @@ pub struct CortexState {
    /// rather than an edge — it hid a GGUF that two of three nodes could
    /// describe perfectly well.
    pub capability_probe_attempts: RwLock<HashMap<(String, String), std::time::Instant>>,
+    /// Async audio generation jobs (#382).
+    pub audio_jobs: crate::audio_jobs::AudioJobs,
+    /// Per-model rates for non-token units (#386), from upstream when it
+    /// is enabled, else the local default.
+    pub rates: Arc<crate::rates::Rates>,
 }

 impl CortexState {
@@ -131,6 +136,12 @@ impl CortexState {
            served_usage: Arc::new(crate::served_usage::ServedUsage::new()),
            discovered_capabilities: RwLock::new(HashMap::new()),
            capability_probe_attempts: RwLock::new(HashMap::new()),
+            audio_jobs: crate::audio_jobs::AudioJobs::new(&config.audio_jobs),
+            rates: Arc::new(if config.upstream.enabled {
+                crate::rates::Rates::upstream()
+            } else {
+                crate::rates::Rates::local()
+            }),
        }
    }
 }
@@ -87,6 +87,7 @@ async fn test_alias_resolves_in_chat_completions() {
        models_config: models_path.to_string_lossy().to_string(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };

    let fleet = Arc::new(CortexState::from_config(&config));
@@ -176,6 +177,7 @@ async fn test_aliases_surface_in_v1_models() {
        models_config: models_path.to_string_lossy().to_string(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };

    let fleet = Arc::new(CortexState::from_config(&config));
@@ -268,6 +270,7 @@ async fn test_alias_falls_through_for_unmapped_model() {
        models_config: models_path.to_string_lossy().to_string(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };

    let fleet = Arc::new(CortexState::from_config(&config));
@@ -0,0 +1,822 @@
+//! Async audio generation jobs (#382): submit → poll → content, cancel,
+//! ownership, retries, busy back-off, and restart recovery, against a
+//! scriptable mock neuron.
+
+use axum::Router;
+use axum::extract::Path;
+use axum::http::{HeaderMap, StatusCode};
+use axum::response::{IntoResponse, Response};
+use axum::routing::{get, post};
+use axum::{Json, body::Bytes};
+use cortex_core::audio::{AudioUsage, NeuronAudioResult, RESULT_HEADER};
+use cortex_core::config::{
+    ApiKeyConfig, AudioJobsConfig, EntitlementsConfig, EvictionSettings, EvictionStrategy,
+    GatewayConfig, GatewaySettings, NeuronEndpoint,
+};
+use cortex_core::node::{ModelEntry, ModelStatus};
+use cortex_gateway::state::CortexState;
+use serde_json::{Value, json};
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::time::Duration;
+use tokio::net::TcpListener;
+
+const MODEL: &str = "music-model";
+const WAV: &[u8] = b"RIFF\x24\x00\x00\x00WAVEfmt fake-audio-bytes";
+
+/// How the mock neuron answers `POST /v1/audio/generations`.
+#[derive(Clone, Copy)]
+enum Mode {
+    /// Render immediately.
+    Ok,
+    /// Answer 503 + Retry-After for the first N calls, then render.
+    BusyThenOk(usize),
+    /// Answer 500 for the first N calls, then render.
+    FailThenOk(usize),
+    /// Always 500.
+    AlwaysFail,
+    /// Always 422 wrong_modality.
+    Reject,
+    /// Sleep this long, then render.
+    Slow(Duration),
+    /// Always 503.
+    AlwaysBusy,
+}
+
+struct Mock {
+    url: String,
+    calls: Arc<AtomicUsize>,
+    /// Account header seen on each render call.
+    accounts: Arc<std::sync::Mutex<Vec<Option<String>>>>,
+}
+
+async fn spawn_mock(mode: Mode) -> Mock {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let url = format!("http://{}", listener.local_addr().unwrap());
+    let calls = Arc::new(AtomicUsize::new(0));
+    let accounts = Arc::new(std::sync::Mutex::new(Vec::new()));
+    let endpoint = url.clone();
+    let (c, a) = (Arc::clone(&calls), Arc::clone(&accounts));
+    let app = Router::new()
+        .route(
+            "/models/{id}/endpoint",
+            get(move |Path(_id): Path<String>| {
+                let url = endpoint.clone();
+                async move { Json(json!({"url": url})) }
+            }),
+        )
+        .route(
+            "/v1/audio/generations",
+            post(move |headers: HeaderMap, body: Bytes| {
+                let (c, a) = (Arc::clone(&c), Arc::clone(&a));
+                async move { render(mode, &c, &a, &headers, &body).await }
+            }),
+        );
+    tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
+    Mock {
+        url,
+        calls,
+        accounts,
+    }
+}
+
+async fn render(
+    mode: Mode,
+    calls: &AtomicUsize,
+    accounts: &std::sync::Mutex<Vec<Option<String>>>,
+    headers: &HeaderMap,
+    body: &[u8],
+) -> Response {
+    let n = calls.fetch_add(1, Ordering::SeqCst);
+    accounts.lock().unwrap().push(
+        headers
+            .get(cortex_core::entitlements::HEADER_ACCOUNT_ID)
+            .and_then(|v| v.to_str().ok())
+            .map(str::to_string),
+    );
+    let busy = || {
+        (
+            StatusCode::SERVICE_UNAVAILABLE,
+            [("retry-after", "1")],
+            Json(json!({"error": {"type": "rate_limit_error", "code": "rate_limit_exceeded", "message": "busy"}})),
+        )
+            .into_response()
+    };
+    let fail = || {
+        (
+            StatusCode::INTERNAL_SERVER_ERROR,
+            Json(json!({"error": {"type": "api_error", "code": "inference_failed", "message": "boom"}})),
+        )
+            .into_response()
+    };
+    match mode {
+        Mode::BusyThenOk(k) if n < k => return busy(),
+        Mode::AlwaysBusy => return busy(),
+        Mode::FailThenOk(k) if n < k => return fail(),
+        Mode::AlwaysFail => return fail(),
+        Mode::Reject => {
+            return (
+                StatusCode::UNPROCESSABLE_ENTITY,
+                Json(json!({"error": {"type": "invalid_request_error", "code": "wrong_modality", "message": "not an audio model"}})),
+            )
+                .into_response();
+        }
+        Mode::Slow(d) => tokio::time::sleep(d).await,
+        _ => {}
+    }
+    let req: Value = serde_json::from_slice(body).unwrap();
+    let duration = req["duration"].as_f64().unwrap_or(30.0);
+    let result = NeuronAudioResult {
+        duration,
+        sample_rate: 48_000,
+        channels: 2,
+        seed: req["seed"].as_u64().unwrap_or(7),
+        usage: AudioUsage {
+            helexa_audio_units: duration,
+            helexa_timing: None,
+        },
+    };
+    (
+        StatusCode::OK,
+        [
+            ("content-type", "audio/wav".to_string()),
+            (RESULT_HEADER, serde_json::to_string(&result).unwrap()),
+        ],
+        WAV,
+    )
+        .into_response()
+}
+
+fn config(neuron: &str, dir: &std::path::Path) -> GatewayConfig {
+    GatewayConfig {
+        gateway: GatewaySettings {
+            listen: "127.0.0.1:0".into(),
+            metrics_listen: "127.0.0.1:0".into(),
+        },
+        eviction: EvictionSettings {
+            strategy: EvictionStrategy::Lru,
+            defrag_after_cycles: 0,
+        },
+        neurons: vec![NeuronEndpoint {
+            name: "mock-node".into(),
+            endpoint: neuron.to_string(),
+        }],
+        models_config: "/dev/null".into(),
+        entitlements: EntitlementsConfig {
+            require_auth: false,
+            keys: ["alice", "bob"]
+                .iter()
+                .map(|who| ApiKeyConfig {
+                    key: format!("key-{who}"),
+                    account_id: format!("acct-{who}"),
+                    key_id: None,
+                    hard_cap: None,
+                    window: Default::default(),
+                })
+                .collect(),
+        },
+        upstream: Default::default(),
+        audio_jobs: AudioJobsConfig {
+            dir: dir.to_string_lossy().into_owned(),
+            queue_timeout_secs: 20,
+            ..Default::default()
+        },
+    }
+}
+
+async fn gateway(neuron: &str, dir: &std::path::Path) -> (Arc<CortexState>, String) {
+    let fleet = Arc::new(CortexState::from_config(&config(neuron, dir)));
+    seed_model(&fleet).await;
+    let url = serve(&fleet).await;
+    (fleet, url)
+}
+
+async fn seed_model(fleet: &CortexState) {
+    let mut nodes = fleet.nodes.write().await;
+    let node = nodes.get_mut("mock-node").unwrap();
+    node.healthy = true;
+    node.models.insert(
+        MODEL.into(),
+        ModelEntry {
+            id: MODEL.into(),
+            status: ModelStatus::Loaded,
+            last_accessed: None,
+            vram_estimate_mb: Some(8000),
+            capabilities: vec!["audio".into()],
+            tool_call: false,
+            reasoning: false,
+            limit: None,
+            servable: None,
+            reasoning_budget: Vec::new(),
+        },
+    );
+}
+
+async fn serve(fleet: &Arc<CortexState>) -> String {
+    let app = cortex_gateway::build_app(Arc::clone(fleet));
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
+    format!("http://{addr}")
+}
+
+async fn submit(gw: &str, key: Option<&str>, body: Value) -> reqwest::Response {
+    let mut req = reqwest::Client::new()
+        .post(format!("{gw}/v1/audio/generations"))
+        .json(&body);
+    if let Some(k) = key {
+        req = req.bearer_auth(k);
+    }
+    req.send().await.unwrap()
+}
+
+async fn get_job(gw: &str, key: Option<&str>, id: &str) -> reqwest::Response {
+    let mut req = reqwest::Client::new().get(format!("{gw}/v1/audio/generations/{id}"));
+    if let Some(k) = key {
+        req = req.bearer_auth(k);
+    }
+    req.send().await.unwrap()
+}
+
+/// Poll until the job reaches a terminal status (or panic after `secs`).
+async fn wait_terminal(gw: &str, key: Option<&str>, id: &str, secs: u64) -> Value {
+    let deadline = std::time::Instant::now() + Duration::from_secs(secs);
+    loop {
+        let job: Value = get_job(gw, key, id).await.json().await.unwrap();
+        if matches!(
+            job["status"].as_str(),
+            Some("succeeded" | "failed" | "cancelled")
+        ) {
+            return job;
+        }
+        assert!(
+            std::time::Instant::now() < deadline,
+            "job {id} never finished: {job}"
+        );
+        tokio::time::sleep(Duration::from_millis(50)).await;
+    }
+}
+
+fn song() -> Value {
+    json!({"model": MODEL, "prompt": "warm lofi", "duration": 12.5, "seed": 42})
+}
+
+#[tokio::test]
+async fn submit_poll_fetch_content() {
+    let mock = spawn_mock(Mode::Ok).await;
+    let dir = tempfile::tempdir().unwrap();
+    let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
+
+    let resp = submit(&gw, Some("key-alice"), song()).await;
+    assert_eq!(resp.status(), 202);
+    let job: Value = resp.json().await.unwrap();
+    assert_eq!(job["object"], "audio.generation");
+    assert_eq!(job["status"], "queued");
+    assert_eq!(job["model"], MODEL);
+    assert_eq!(job["prompt"], "warm lofi");
+    assert!(job.get("lyrics").is_none());
+    let id = job["id"].as_str().unwrap().to_string();
+    assert!(id.starts_with("audiogen_"));
+
+    let done = wait_terminal(&gw, Some("key-alice"), &id, 10).await;
+    assert_eq!(done["status"], "succeeded", "{done}");
+    assert_eq!(done["result"]["duration"], 12.5);
+    assert_eq!(done["result"]["sample_rate"], 48_000);
+    assert_eq!(done["result"]["seed"], 42);
+    assert_eq!(done["result"]["bytes"], WAV.len());
+    assert_eq!(done["usage"]["helexa_audio_units"], 12.5);
+
+    let content = reqwest::Client::new()
+        .get(format!("{gw}/v1/audio/generations/{id}/content"))
+        .bearer_auth("key-alice")
+        .send()
+        .await
+        .unwrap();
+    assert_eq!(content.status(), 200);
+    assert_eq!(content.headers()["content-type"], "audio/wav");
+    assert_eq!(content.bytes().await.unwrap().as_ref(), WAV);
+
+    // The neuron saw the submitting account, for fair-share admission.
+    assert_eq!(
+        mock.accounts.lock().unwrap().as_slice(),
+        &[Some("acct-alice".to_string())]
+    );
+    // Record and audio are on disk.
+    assert!(dir.path().join(format!("{id}.json")).exists());
+    assert!(dir.path().join(format!("{id}.wav")).exists());
+}
+
+#[tokio::test]
+async fn content_before_completion_is_409() {
+    let mock = spawn_mock(Mode::Slow(Duration::from_secs(3))).await;
+    let dir = tempfile::tempdir().unwrap();
+    let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
+    let job: Value = submit(&gw, None, song()).await.json().await.unwrap();
+    let id = job["id"].as_str().unwrap();
+    let resp = reqwest::get(format!("{gw}/v1/audio/generations/{id}/content"))
+        .await
+        .unwrap();
+    assert_eq!(resp.status(), 409);
+    let body: Value = resp.json().await.unwrap();
+    assert_eq!(body["error"]["code"], "job_not_complete");
+}
+
+#[tokio::test]
+async fn jobs_are_private_to_their_account() {
+    let mock = spawn_mock(Mode::Ok).await;
+    let dir = tempfile::tempdir().unwrap();
+    let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
+
+    let job: Value = submit(&gw, Some("key-alice"), song())
+        .await
+        .json()
+        .await
+        .unwrap();
+    let id = job["id"].as_str().unwrap().to_string();
+    wait_terminal(&gw, Some("key-alice"), &id, 10).await;
+
+    // Another account, and an anonymous caller, cannot see it.
+    assert_eq!(get_job(&gw, Some("key-bob"), &id).await.status(), 404);
+    assert_eq!(get_job(&gw, None, &id).await.status(), 404);
+    let content = reqwest::Client::new()
+        .get(format!("{gw}/v1/audio/generations/{id}/content"))
+        .bearer_auth("key-bob")
+        .send()
+        .await
+        .unwrap();
+    assert_eq!(content.status(), 404);
+    let cancel = reqwest::Client::new()
+        .delete(format!("{gw}/v1/audio/generations/{id}"))
+        .bearer_auth("key-bob")
+        .send()
+        .await
+        .unwrap();
+    assert_eq!(cancel.status(), 404);
+
+    // Listings are per account.
+    let list = |key: &'static str| {
+        let gw = gw.clone();
+        async move {
+            reqwest::Client::new()
+                .get(format!("{gw}/v1/audio/generations"))
+                .bearer_auth(key)
+                .send()
+                .await
+                .unwrap()
+                .json::<Value>()
+                .await
+                .unwrap()
+        }
+    };
+    let alice = list("key-alice").await;
+    assert_eq!(alice["object"], "list");
+    assert_eq!(alice["data"].as_array().unwrap().len(), 1);
+    assert_eq!(alice["data"][0]["id"], id.as_str());
+    assert_eq!(list("key-bob").await["data"].as_array().unwrap().len(), 0);
+}
+
+#[tokio::test]
+async fn anonymous_jobs_are_reachable_by_id_but_never_listed() {
+    let mock = spawn_mock(Mode::Ok).await;
+    let dir = tempfile::tempdir().unwrap();
+    let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
+    let job: Value = submit(&gw, None, song()).await.json().await.unwrap();
+    let id = job["id"].as_str().unwrap().to_string();
+    let done = wait_terminal(&gw, None, &id, 10).await;
+    assert_eq!(done["status"], "succeeded");
+    let list: Value = reqwest::get(format!("{gw}/v1/audio/generations"))
+        .await
+        .unwrap()
+        .json()
+        .await
+        .unwrap();
+    assert_eq!(list["data"].as_array().unwrap().len(), 0);
+}
+
+#[tokio::test]
+async fn cancel_running_job() {
+    let mock = spawn_mock(Mode::Slow(Duration::from_secs(5))).await;
+    let dir = tempfile::tempdir().unwrap();
+    let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
+    let job: Value = submit(&gw, Some("key-alice"), song())
+        .await
+        .json()
+        .await
+        .unwrap();
+    let id = job["id"].as_str().unwrap().to_string();
+
+    // Wait until it is running on the neuron.
+    let deadline = std::time::Instant::now() + Duration::from_secs(5);
+    while mock.calls.load(Ordering::SeqCst) == 0 {
+        assert!(std::time::Instant::now() < deadline);
+        tokio::time::sleep(Duration::from_millis(20)).await;
+    }
+    let resp = reqwest::Client::new()
+        .delete(format!("{gw}/v1/audio/generations/{id}"))
+        .bearer_auth("key-alice")
+        .send()
+        .await
+        .unwrap();
+    assert_eq!(resp.status(), 200);
+    let cancelled: Value = resp.json().await.unwrap();
+    assert_eq!(cancelled["status"], "cancelled");
+
+    // It stays cancelled: the slow render finishing must not resurrect it.
+    tokio::time::sleep(Duration::from_secs(6)).await;
+    let job: Value = get_job(&gw, Some("key-alice"), &id)
+        .await
+        .json()
+        .await
+        .unwrap();
+    assert_eq!(job["status"], "cancelled");
+    assert!(job.get("result").is_none());
+}
+
+#[tokio::test]
+async fn busy_neuron_is_waited_on_without_spending_attempts() {
+    // Three busy answers exceed max_attempts (2), so this only succeeds if
+    // busy answers don't count as attempts.
+    let mock = spawn_mock(Mode::BusyThenOk(3)).await;
+    let dir = tempfile::tempdir().unwrap();
+    let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
+    let job: Value = submit(&gw, None, song()).await.json().await.unwrap();
+    let id = job["id"].as_str().unwrap().to_string();
+    let done = wait_terminal(&gw, None, &id, 15).await;
+    assert_eq!(done["status"], "succeeded", "{done}");
+    assert_eq!(mock.calls.load(Ordering::SeqCst), 4);
+}
+
+#[tokio::test]
+async fn failing_node_is_retried_then_succeeds() {
+    let mock = spawn_mock(Mode::FailThenOk(1)).await;
+    let dir = tempfile::tempdir().unwrap();
+    let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
+    let job: Value = submit(&gw, None, song()).await.json().await.unwrap();
+    let done = wait_terminal(&gw, None, job["id"].as_str().unwrap(), 15).await;
+    assert_eq!(done["status"], "succeeded", "{done}");
+    assert_eq!(mock.calls.load(Ordering::SeqCst), 2);
+}
+
+#[tokio::test]
+async fn failing_node_exhausts_attempts() {
+    let mock = spawn_mock(Mode::AlwaysFail).await;
+    let dir = tempfile::tempdir().unwrap();
+    let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
+    let job: Value = submit(&gw, None, song()).await.json().await.unwrap();
+    let done = wait_terminal(&gw, None, job["id"].as_str().unwrap(), 15).await;
+    assert_eq!(done["status"], "failed");
+    assert_eq!(done["error"]["code"], "inference_failed");
+    assert_eq!(mock.calls.load(Ordering::SeqCst), 2);
+}
+
+#[tokio::test]
+async fn neuron_rejection_fails_without_retry() {
+    let mock = spawn_mock(Mode::Reject).await;
+    let dir = tempfile::tempdir().unwrap();
+    let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
+    let job: Value = submit(&gw, None, song()).await.json().await.unwrap();
+    let done = wait_terminal(&gw, None, job["id"].as_str().unwrap(), 10).await;
+    assert_eq!(done["status"], "failed");
+    assert_eq!(done["error"]["code"], "wrong_modality");
+    assert_eq!(mock.calls.load(Ordering::SeqCst), 1);
+}
+
+#[tokio::test]
+async fn submit_validation() {
+    let mock = spawn_mock(Mode::Ok).await;
+    let dir = tempfile::tempdir().unwrap();
+    let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
+
+    let resp = submit(&gw, None, json!({"model": MODEL, "prompt": ""})).await;
+    assert_eq!(resp.status(), 400);
+    let body: Value = resp.json().await.unwrap();
+    assert_eq!(body["error"]["code"], "invalid_audio_params");
+
+    let resp = submit(
+        &gw,
+        None,
+        json!({"model": MODEL, "prompt": "x", "response_format": "ogg"}),
+    )
+    .await;
+    assert_eq!(resp.status(), 400);
+
+    let resp = submit(&gw, None, json!({"model": "nope", "prompt": "x"})).await;
+    assert_eq!(resp.status(), 404);
+    let body: Value = resp.json().await.unwrap();
+    assert_eq!(body["error"]["code"], "model_not_found");
+
+    let resp = submit(&gw, None, json!({"prompt": "x"})).await;
+    assert_eq!(resp.status(), 400);
+
+    // Nothing reached the neuron.
+    assert_eq!(mock.calls.load(Ordering::SeqCst), 0);
+}
+
+#[tokio::test]
+async fn unknown_job_is_404() {
+    let mock = spawn_mock(Mode::Ok).await;
+    let dir = tempfile::tempdir().unwrap();
+    let (_fleet, gw) = gateway(&mock.url, dir.path()).await;
+    let resp = get_job(&gw, None, "audiogen_doesnotexist").await;
+    assert_eq!(resp.status(), 404);
+    let body: Value = resp.json().await.unwrap();
+    assert_eq!(body["error"]["code"], "job_not_found");
+}
+
+#[tokio::test]
+async fn jobs_survive_a_restart() {
+    let dir = tempfile::tempdir().unwrap();
+
+    // First cortex: one job finishes, one is stuck behind a busy model.
+    let ok = spawn_mock(Mode::Ok).await;
+    let (_a, gw_a) = gateway(&ok.url, dir.path()).await;
+    let done: Value = submit(&gw_a, Some("key-alice"), song())
+        .await
+        .json()
+        .await
+        .unwrap();
+    let done_id = done["id"].as_str().unwrap().to_string();
+    wait_terminal(&gw_a, Some("key-alice"), &done_id, 10).await;
+
+    let busy = spawn_mock(Mode::AlwaysBusy).await;
+    let (_b, gw_b) = gateway(&busy.url, dir.path()).await;
+    let stuck: Value = submit(&gw_b, Some("key-alice"), song())
+        .await
+        .json()
+        .await
+        .unwrap();
+    let stuck_id = stuck["id"].as_str().unwrap().to_string();
+    let deadline = std::time::Instant::now() + Duration::from_secs(5);
+    while busy.calls.load(Ordering::SeqCst) == 0 {
+        assert!(std::time::Instant::now() < deadline);
+        tokio::time::sleep(Duration::from_millis(20)).await;
+    }
+
+    // A fresh cortex over the same directory, with a healthy neuron.
+    let healthy = spawn_mock(Mode::Ok).await;
+    let fleet = Arc::new(CortexState::from_config(&config(&healthy.url, dir.path())));
+    seed_model(&fleet).await;
+    cortex_gateway::audio_jobs::start(Arc::clone(&fleet)).await;
+    let gw = serve(&fleet).await;
+
+    // The finished job and its audio are still there.
+    let job: Value = get_job(&gw, Some("key-alice"), &done_id)
+        .await
+        .json()
+        .await
+        .unwrap();
+    assert_eq!(job["status"], "succeeded");
+    let content = reqwest::Client::new()
+        .get(format!("{gw}/v1/audio/generations/{done_id}/content"))
+        .bearer_auth("key-alice")
+        .send()
+        .await
+        .unwrap();
+    assert_eq!(content.bytes().await.unwrap().as_ref(), WAV);
+
+    // The unfinished one was re-queued and now completes.
+    let resumed = wait_terminal(&gw, Some("key-alice"), &stuck_id, 10).await;
+    assert_eq!(resumed["status"], "succeeded", "{resumed}");
+}
+
+#[tokio::test]
+async fn expired_jobs_are_swept() {
+    let mock = spawn_mock(Mode::Ok).await;
+    let dir = tempfile::tempdir().unwrap();
+    let mut cfg = config(&mock.url, dir.path());
+    cfg.audio_jobs.retention_hours = 0;
+    let fleet = Arc::new(CortexState::from_config(&cfg));
+    seed_model(&fleet).await;
+    let gw = serve(&fleet).await;
+
+    let job: Value = submit(&gw, None, song()).await.json().await.unwrap();
+    let id = job["id"].as_str().unwrap().to_string();
+    wait_terminal(&gw, None, &id, 10).await;
+    // completed_at has one-second resolution; let it fall behind the cutoff.
+    tokio::time::sleep(Duration::from_millis(1100)).await;
+
+    cortex_gateway::audio_jobs::sweep_expired(&fleet).await;
+    assert_eq!(get_job(&gw, None, &id).await.status(), 404);
+    assert!(!dir.path().join(format!("{id}.json")).exists());
+    assert!(!dir.path().join(format!("{id}.wav")).exists());
+}
+
+// ── Per-model rates (#386) ───────────────────────────────────────────
+
+/// A gateway whose keys are capped at `cap` tokens and whose audio rate
+/// table is `rates` (model id → tokens per second; "*" is the default).
+async fn priced_gateway(
+    neuron: &str,
+    dir: &std::path::Path,
+    cap: Option<u64>,
+    rates: &[(&str, u64)],
+) -> (Arc<CortexState>, String) {
+    let mut cfg = config(neuron, dir);
+    for k in &mut cfg.entitlements.keys {
+        k.hard_cap = cap;
+    }
+    let fleet = Arc::new(CortexState::from_config(&cfg));
+    let mut table = cortex_gateway::rates::RateTable::default();
+    for (model, rate) in rates {
+        table.set(cortex_gateway::rates::AUDIO_SECOND, model, *rate);
+    }
+    fleet.rates.replace(table).await;
+    seed_model(&fleet).await;
+    let url = serve(&fleet).await;
+    (fleet, url)
+}
+
+async fn spent(fleet: &CortexState, who: &str) -> u64 {
+    let principal = cortex_core::entitlements::Principal {
+        account_id: format!("acct-{who}"),
+        key_id: format!("acct-{who}"),
+    };
+    fleet
+        .entitlements
+        .snapshot(&principal)
+        .await
+        .map(|s| s.spent)
+        .unwrap_or(0)
+}
+
+#[tokio::test]
+async fn reservation_uses_the_models_rate() {
+    // A 12.5 s song against a 10,000-token cap: refused at 1000 tokens/s
+    // (12,500), accepted at 500 tokens/s (6,250).
+    let mock = spawn_mock(Mode::Ok).await;
+    let dir = tempfile::tempdir().unwrap();
+    let (_f, gw) = priced_gateway(&mock.url, dir.path(), Some(10_000), &[("*", 1000)]).await;
+    let resp = submit(&gw, Some("key-alice"), song()).await;
+    assert_eq!(resp.status(), 429, "{:?}", resp.text().await);
+
+    let dir = tempfile::tempdir().unwrap();
+    let (_f, gw) = priced_gateway(
+        &mock.url,
+        dir.path(),
+        Some(10_000),
+        &[("*", 1000), (MODEL, 500)],
+    )
+    .await;
+    assert_eq!(submit(&gw, Some("key-alice"), song()).await.status(), 202);
+}
+
+#[tokio::test]
+async fn unpriced_model_is_refused_for_metered_callers_only() {
+    let mock = spawn_mock(Mode::Ok).await;
+    let dir = tempfile::tempdir().unwrap();
+    // An empty table: what a cortex with upstream sees before its first fetch.
+    let (_f, gw) = priced_gateway(&mock.url, dir.path(), None, &[]).await;
+    let resp = submit(&gw, Some("key-alice"), song()).await;
+    assert_eq!(resp.status(), 503);
+    let body: Value = resp.json().await.unwrap();
+    assert!(
+        body["error"]["message"]
+            .as_str()
+            .unwrap()
+            .contains("no audio rate"),
+        "{body}"
+    );
+    // Anonymous work is not metered, so it does not need a price.
+    assert_eq!(submit(&gw, None, song()).await.status(), 202);
+}
+
+#[tokio::test]
+async fn settles_at_the_rate_reserved_even_if_it_changes_mid_job() {
+    let mock = spawn_mock(Mode::Slow(Duration::from_millis(800))).await;
+    let dir = tempfile::tempdir().unwrap();
+    let (fleet, gw) = priced_gateway(&mock.url, dir.path(), None, &[(MODEL, 300)]).await;
+    let job: Value = submit(&gw, Some("key-alice"), song())
+        .await
+        .json()
+        .await
+        .unwrap();
+    // Reprice while the render is in flight.
+    let mut table = cortex_gateway::rates::RateTable::default();
+    table.set(cortex_gateway::rates::AUDIO_SECOND, MODEL, 1);
+    fleet.rates.replace(table).await;
+
+    let done = wait_terminal(&gw, Some("key-alice"), job["id"].as_str().unwrap(), 10).await;
+    assert_eq!(done["status"], "succeeded");
+    // 12.5 s × 300 tokens/s, settled asynchronously.
+    let deadline = std::time::Instant::now() + Duration::from_secs(5);
+    loop {
+        let s = spent(&fleet, "alice").await;
+        if s == 3750 {
+            break;
+        }
+        assert!(std::time::Instant::now() < deadline, "spent {s}, want 3750");
+        tokio::time::sleep(Duration::from_millis(20)).await;
+    }
+}
+
+#[tokio::test]
+async fn rates_are_fetched_from_upstream_with_the_client_bearer() {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let url = format!("http://{}", listener.local_addr().unwrap());
+    let app = Router::new().route(
+        "/authz/v1/rates",
+        get(|headers: HeaderMap| async move {
+            if headers.get("authorization").and_then(|v| v.to_str().ok())
+                != Some("Bearer op-secret")
+            {
+                return StatusCode::UNAUTHORIZED.into_response();
+            }
+            Json(json!({"data": [
+                {"model_id": "*", "unit": "audio_second", "tokens_per_unit": 1000},
+                {"model_id": MODEL, "unit": "audio_second", "tokens_per_unit": 40},
+                {"model_id": "bad", "unit": "audio_second", "tokens_per_unit": 0}
+            ]}))
+            .into_response()
+        }),
+    );
+    tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
+
+    let client = reqwest::Client::new();
+    let table = cortex_gateway::rates::fetch(&client, &url, "op-secret")
+        .await
+        .unwrap();
+    assert_eq!(table.audio_tokens_per_second(MODEL), Some(40));
+    assert_eq!(table.audio_tokens_per_second("other"), Some(1000));
+    // A non-positive rate is skipped, falling back to the default.
+    assert_eq!(table.audio_tokens_per_second("bad"), Some(1000));
+    assert!(
+        cortex_gateway::rates::fetch(&client, &url, "wrong")
+            .await
+            .is_err()
+    );
+}
+
+#[tokio::test]
+async fn models_endpoint_advertises_audio_limits_and_rate() {
+    let mock = spawn_mock(Mode::Ok).await;
+    let dir = tempfile::tempdir().unwrap();
+    let catalogue = dir.path().join("models.toml");
+    std::fs::write(
+        &catalogue,
+        format!(
+            r#"
+[[models]]
+id = "{MODEL}"
+harness = "candle"
+quant = "q8_0"
+
+[models.audio]
+min_duration = 10.0
+max_duration = 600.0
+default_duration = 30.0
+vocals = true
+sample_rate = 48000
+channels = 2
+max_steps = 20
+
+[aliases]
+"helexa/music" = "{MODEL}"
+"#
+        ),
+    )
+    .unwrap();
+    let mut cfg = config(&mock.url, dir.path());
+    cfg.models_config = catalogue.to_string_lossy().into_owned();
+    let fleet = Arc::new(CortexState::from_config(&cfg));
+    let mut table = cortex_gateway::rates::RateTable::default();
+    table.set(cortex_gateway::rates::AUDIO_SECOND, MODEL, 77);
+    fleet.rates.replace(table).await;
+    seed_model(&fleet).await;
+    let gw = serve(&fleet).await;
+
+    let models: Value = reqwest::get(format!("{gw}/v1/models"))
+        .await
+        .unwrap()
+        .json()
+        .await
+        .unwrap();
+    let find = |id: &str| {
+        models["data"]
+            .as_array()
+            .unwrap()
+            .iter()
+            .find(|m| m["id"] == id)
+            .cloned()
+            .unwrap_or_else(|| panic!("{id} missing from {models}"))
+    };
+    for id in [MODEL, "helexa/music"] {
+        let m = find(id);
+        assert_eq!(m["audio"]["max_duration"], 600.0, "{m}");
+        assert_eq!(m["audio"]["vocals"], true);
+        assert_eq!(m["audio_tokens_per_second"], 77, "{m}");
+    }
+
+    // The catalogue's limits are enforced at submit.
+    let resp = submit(
+        &gw,
+        None,
+        json!({"model": MODEL, "prompt": "x", "duration": 5.0}),
+    )
+    .await;
+    assert_eq!(resp.status(), 400);
+    // An alias is accepted and resolved to the model.
+    let resp = submit(&gw, None, json!({"model": "helexa/music", "prompt": "x"})).await;
+    assert_eq!(resp.status(), 202);
+    let job: Value = resp.json().await.unwrap();
+    assert_eq!(job["model"], MODEL);
+}
@@ -106,6 +106,7 @@ async fn spawn_gateway(neuron_url: &str, entitlements: EntitlementsConfig) -> St
        models_config: "/dev/null".into(),
        entitlements,
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };

    let fleet = Arc::new(CortexState::from_config(&config));
@@ -82,6 +82,7 @@ async fn spawn_gateway(neuron_url: &str, key: ApiKeyConfig) -> (Arc<CortexState>
            keys: vec![key],
        },
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };
    let fleet = Arc::new(CortexState::from_config(&config));
    {
@@ -498,6 +498,7 @@ pub async fn spawn_gateway_with_state(mock_url: &str) -> (Arc<CortexState>, Stri
        models_config: "/dev/null".into(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };

    let fleet = Arc::new(CortexState::from_config(&config));
@@ -559,6 +560,7 @@ pub async fn two_neuron_fleet(endpoint_a: &str, endpoint_b: &str) -> Arc<CortexS
        models_config: "/dev/null".into(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };
    Arc::new(CortexState::from_config(&config))
 }
@@ -90,6 +90,7 @@ async fn error_response_no_healthy_nodes() {
        models_config: "/dev/null".into(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };

    let fleet = Arc::new(cortex_gateway::state::CortexState::from_config(&config));
@@ -82,6 +82,7 @@ fn make_fleet_with_catalogue(
        models_config: path.to_string_lossy().into_owned(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };
    (Arc::new(CortexState::from_config(&config)), path)
 }
@@ -144,6 +145,7 @@ fn make_fleet(endpoint: &str, defrag_after: u32) -> Arc<CortexState> {
        models_config: "/dev/null".into(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };
    Arc::new(CortexState::from_config(&config))
 }
@@ -88,6 +88,7 @@ async fn fleet_with(big_healthy: bool, big_devices: usize) -> Arc<CortexState> {
        models_config: cat.to_string_lossy().into_owned(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };
    let fleet = Arc::new(CortexState::from_config(&config));
    {
@@ -165,6 +165,7 @@ async fn spawn_keyed_gateway(neuron_url: &str, hard_cap: u64) -> String {
            }],
        },
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };
    let fleet = Arc::new(CortexState::from_config(&config));
    {
@@ -89,6 +89,7 @@ async fn two_neuron_fleet(endpoint_a: &str, endpoint_b: &str) -> Arc<CortexState
        models_config: "/dev/null".into(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };
    Arc::new(CortexState::from_config(&config))
 }
@@ -54,6 +54,7 @@ async fn spawn_metered_gateway(neuron_url: &str) -> (Arc<CortexState>, String) {
            }],
        },
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };

    let fleet = Arc::new(CortexState::from_config(&config));
@@ -162,6 +163,7 @@ async fn anonymous_request_records_no_spend() {
        models_config: "/dev/null".into(),
        entitlements: EntitlementsConfig::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };
    let fleet = Arc::new(CortexState::from_config(&config));
    {
@@ -204,6 +204,7 @@ async fn test_capacity_gauges_exported_from_health_poll() {
        models_config: "/dev/null".into(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };
    let fleet = std::sync::Arc::new(cortex_gateway::state::CortexState::from_config(&config));
    cortex_gateway::poller::poll_once(&fleet).await;
@@ -67,6 +67,7 @@ harness = "candle"
        models_config: cat_path.to_string_lossy().into_owned(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };

    let fleet = Arc::new(CortexState::from_config(&config));
@@ -59,6 +59,7 @@ capabilities = ["text"]
        models_config: cat_path.to_string_lossy().into_owned(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };

    let fleet = Arc::new(CortexState::from_config(&config));
@@ -33,6 +33,7 @@ async fn test_poller_discovers_models() {
        models_config: "/dev/null".into(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };

    let fleet = Arc::new(CortexState::from_config(&config));
@@ -86,6 +87,7 @@ async fn test_poller_updates_gateway_models_endpoint() {
        models_config: "/dev/null".into(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };

    let fleet = Arc::new(CortexState::from_config(&config));
@@ -159,6 +161,7 @@ async fn test_models_endpoint_unions_capabilities_across_nodes() {
        models_config: "/dev/null".into(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };

    let fleet = Arc::new(CortexState::from_config(&config));
@@ -223,6 +226,7 @@ async fn test_poller_marks_unreachable_node_unhealthy() {
        models_config: "/dev/null".into(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };

    let fleet = Arc::new(CortexState::from_config(&config));
@@ -278,6 +282,7 @@ async fn test_poller_removes_stale_models() {
        models_config: "/dev/null".into(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };

    let fleet = Arc::new(CortexState::from_config(&config));
@@ -310,6 +315,7 @@ async fn test_poller_removes_stale_models() {
        models_config: "/dev/null".into(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };

    let fleet2 = Arc::new(CortexState::from_config(&config2));
@@ -397,6 +403,7 @@ async fn test_poller_captures_activation_from_health() {
        models_config: "/dev/null".into(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };

    let fleet = Arc::new(CortexState::from_config(&config));
@@ -443,6 +450,7 @@ async fn test_poller_parses_recovering_status() {
        models_config: "/dev/null".into(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };

    let fleet = Arc::new(CortexState::from_config(&config));
@@ -76,6 +76,7 @@ async fn spawn_gateway(neuron: &str, context: usize) -> String {
        models_config: "/dev/null".into(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };
    let fleet = Arc::new(CortexState::from_config(&config));
    {
@@ -119,6 +119,7 @@ async fn test_no_healthy_nodes() {
        models_config: "/dev/null".into(),
        entitlements: Default::default(),
        upstream: Default::default(),
+        audio_jobs: Default::default(),
    };
    let fleet = std::sync::Arc::new(cortex_gateway::state::CortexState::from_config(&config));

@@ -0,0 +1,275 @@
+//! `/v1/systemone` decision proxy tests (#338).
+//!
+//! The mock neuron echoes what it received — the `model` field and the
+//! routed-model header — so each test can assert both halves of the
+//! contract: *where* cortex sent the request, and that it did not
+//! rewrite the body on the way.
+
+use axum::Router;
+use axum::http::HeaderMap;
+use axum::response::{IntoResponse, Json};
+use axum::routing::{get, post};
+use cortex_core::catalogue::{DECISION_DEFAULT_ALIAS, DECISION_MODEL_HEADER};
+use cortex_core::config::{
+    ApiKeyConfig, EntitlementsConfig, EvictionSettings, EvictionStrategy, GatewayConfig,
+    GatewaySettings, NeuronEndpoint,
+};
+use cortex_core::entitlements::CapWindow;
+use cortex_core::node::{ModelEntry, ModelStatus};
+use cortex_gateway::state::CortexState;
+use serde_json::{Value, json};
+use std::sync::Arc;
+use tokio::net::TcpListener;
+
+const DECISION_MODEL: &str = "convaiinnovations/laya";
+const TEXT_MODEL: &str = "Qwen/Qwen3-1.7B";
+
+/// Mock neuron: `/v1/systemone` answers with a Jev-shaped body that
+/// echoes the forwarded `model` and routed-model header, and reports
+/// `usage.input_tokens` = 41 per question.
+async fn spawn_decision_mock_neuron() -> String {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    let base_url = format!("http://{addr}");
+    let inference_url = base_url.clone();
+    let app = Router::new()
+        .route(
+            "/models/{*rest}",
+            get(move || {
+                let url = inference_url.clone();
+                async move { Json(json!({"url": url})) }
+            }),
+        )
+        .route(
+            "/v1/systemone",
+            post(|headers: HeaderMap, Json(body): Json<Value>| async move {
+                let routed = headers
+                    .get(DECISION_MODEL_HEADER)
+                    .and_then(|v| v.to_str().ok())
+                    .map(str::to_string);
+                let questions = body
+                    .get("questions")
+                    .and_then(Value::as_object)
+                    .map(|m| m.len() as u64)
+                    .unwrap_or(0);
+                (
+                    [("server-timing", "inference;dur=12.50")],
+                    Json(json!({
+                        "model": "laya-rl-agent",
+                        "answers": {},
+                        "usage": {"input_tokens": 41 * questions, "output_tokens": 0},
+                        "echo_model": body.get("model").cloned().unwrap_or(Value::Null),
+                        "echo_routed": routed,
+                    })),
+                )
+                    .into_response()
+            }),
+        );
+    tokio::spawn(async move {
+        axum::serve(listener, app).await.unwrap();
+    });
+    base_url
+}
+
+fn entry(id: &str, capabilities: &[&str]) -> ModelEntry {
+    ModelEntry {
+        id: id.into(),
+        status: ModelStatus::Loaded,
+        last_accessed: None,
+        vram_estimate_mb: Some(1000),
+        capabilities: capabilities.iter().map(|c| c.to_string()).collect(),
+        tool_call: false,
+        reasoning: false,
+        limit: None,
+        servable: None,
+        reasoning_budget: Vec::new(),
+    }
+}
+
+/// A gateway fronting the mock with one decision and one text model
+/// loaded. `default_alias` controls whether `helexa/one` is configured;
+/// `hard_cap` switches on keyed auth with that balance.
+async fn spawn_gateway(neuron_url: &str, default_alias: bool, hard_cap: Option<u64>) -> String {
+    let entitlements = match hard_cap {
+        Some(cap) => EntitlementsConfig {
+            require_auth: true,
+            keys: vec![ApiKeyConfig {
+                key: "sk-dec".into(),
+                account_id: "acct-dec".into(),
+                key_id: Some("key-dec".into()),
+                hard_cap: Some(cap),
+                window: CapWindow::Balance,
+            }],
+        },
+        None => EntitlementsConfig::default(),
+    };
+    let config = GatewayConfig {
+        gateway: GatewaySettings {
+            listen: "127.0.0.1:0".into(),
+            metrics_listen: "127.0.0.1:0".into(),
+        },
+        eviction: EvictionSettings {
+            strategy: EvictionStrategy::Lru,
+            defrag_after_cycles: 0,
+        },
+        neurons: vec![NeuronEndpoint {
+            name: "mock-node".into(),
+            endpoint: neuron_url.to_string(),
+        }],
+        models_config: "/dev/null".into(),
+        entitlements,
+        upstream: Default::default(),
+        audio_jobs: Default::default(),
+    };
+    let mut state = CortexState::from_config(&config);
+    if default_alias {
+        state
+            .catalogue
+            .aliases
+            .insert(DECISION_DEFAULT_ALIAS.into(), DECISION_MODEL.into());
+    }
+    let fleet = Arc::new(state);
+    {
+        let mut nodes = fleet.nodes.write().await;
+        let node = nodes.get_mut("mock-node").unwrap();
+        node.healthy = true;
+        node.models
+            .insert(DECISION_MODEL.into(), entry(DECISION_MODEL, &["decision"]));
+        node.models
+            .insert(TEXT_MODEL.into(), entry(TEXT_MODEL, &[]));
+    }
+    let app = cortex_gateway::build_app(Arc::clone(&fleet));
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    tokio::spawn(async move {
+        axum::serve(listener, app).await.unwrap();
+    });
+    format!("http://{addr}")
+}
+
+async fn send(gateway: &str, body: Value) -> (u16, HeaderMap, Value) {
+    let resp = reqwest::Client::new()
+        .post(format!("{gateway}/v1/systemone"))
+        .json(&body)
+        .send()
+        .await
+        .expect("request should reach the gateway");
+    let status = resp.status().as_u16();
+    let mut headers = HeaderMap::new();
+    for (k, v) in resp.headers() {
+        headers.insert(k.clone(), v.clone());
+    }
+    (status, headers, resp.json().await.unwrap())
+}
+
+fn request(model: Option<&str>) -> Value {
+    let mut body = json!({
+        "state": "I was charged twice this month",
+        "questions": {"refund": {"type": "noul", "instructions": "Refund request?"}},
+    });
+    if let Some(m) = model {
+        body["model"] = json!(m);
+    }
+    body
+}
+
+#[tokio::test]
+async fn test_explicit_decision_model_routes_and_body_is_untouched() {
+    let gateway = spawn_gateway(&spawn_decision_mock_neuron().await, true, None).await;
+    let (status, headers, body) = send(&gateway, request(Some(DECISION_MODEL))).await;
+    assert_eq!(status, 200, "{body}");
+    assert_eq!(body["echo_routed"], DECISION_MODEL);
+    assert_eq!(body["echo_model"], DECISION_MODEL);
+    assert_eq!(body["usage"]["input_tokens"], 41);
+    // Upstream timing headers survive the hop.
+    assert_eq!(
+        headers.get("server-timing").and_then(|v| v.to_str().ok()),
+        Some("inference;dur=12.50")
+    );
+}
+
+#[tokio::test]
+async fn test_foreign_model_name_falls_back_to_default_alias() {
+    // A Jev client names its own service. The request must still be
+    // served — routed via helexa/one — and the client's value must reach
+    // the node unchanged (it may be a checkpoint pin the node honours).
+    let gateway = spawn_gateway(&spawn_decision_mock_neuron().await, true, None).await;
+    let (status, _, body) = send(&gateway, request(Some("jev-1"))).await;
+    assert_eq!(status, 200, "{body}");
+    assert_eq!(body["echo_routed"], DECISION_MODEL);
+    assert_eq!(body["echo_model"], "jev-1");
+}
+
+#[tokio::test]
+async fn test_missing_model_falls_back_to_default_alias() {
+    let gateway = spawn_gateway(&spawn_decision_mock_neuron().await, true, None).await;
+    let (status, _, body) = send(&gateway, request(None)).await;
+    assert_eq!(status, 200, "{body}");
+    assert_eq!(body["echo_routed"], DECISION_MODEL);
+    assert_eq!(body["echo_model"], Value::Null);
+}
+
+#[tokio::test]
+async fn test_default_alias_itself_resolves() {
+    let gateway = spawn_gateway(&spawn_decision_mock_neuron().await, true, None).await;
+    let (status, _, body) = send(&gateway, request(Some(DECISION_DEFAULT_ALIAS))).await;
+    assert_eq!(status, 200, "{body}");
+    assert_eq!(body["echo_routed"], DECISION_MODEL);
+}
+
+#[tokio::test]
+async fn test_known_non_decision_model_is_not_substituted() {
+    // A model the fleet does serve routes to itself; the node owns the
+    // wrong_modality answer. Silently swapping in the decision model
+    // would hide the client's mistake.
+    let gateway = spawn_gateway(&spawn_decision_mock_neuron().await, true, None).await;
+    let (status, _, body) = send(&gateway, request(Some(TEXT_MODEL))).await;
+    assert_eq!(status, 200, "{body}");
+    assert_eq!(body["echo_routed"], TEXT_MODEL);
+}
+
+#[tokio::test]
+async fn test_unknown_model_without_default_alias_is_404() {
+    let gateway = spawn_gateway(&spawn_decision_mock_neuron().await, false, None).await;
+    let (status, _, body) = send(&gateway, request(Some("jev-1"))).await;
+    assert_eq!(status, 404, "{body}");
+    assert_eq!(body["error"]["code"], "model_not_found");
+}
+
+#[tokio::test]
+async fn test_budget_below_reservation_rejected_before_dispatch() {
+    // One question reserves one full sequence (1024 tokens by default);
+    // a 100-token balance cannot cover it.
+    let gateway = spawn_gateway(&spawn_decision_mock_neuron().await, true, Some(100)).await;
+    let resp = reqwest::Client::new()
+        .post(format!("{gateway}/v1/systemone"))
+        .bearer_auth("sk-dec")
+        .json(&request(Some(DECISION_MODEL)))
+        .send()
+        .await
+        .unwrap();
+    assert_eq!(resp.status(), 429);
+    let body: Value = resp.json().await.unwrap();
+    assert_eq!(body["error"]["code"], "insufficient_quota");
+}
+
+#[tokio::test]
+async fn test_spend_settles_to_reported_input_tokens() {
+    // Balance for exactly two one-question requests at the reserved
+    // bound (2 × 1024). Settling to the reported 41 tokens leaves room
+    // for many more; settling at the bound would exhaust it on the third.
+    let gateway = spawn_gateway(&spawn_decision_mock_neuron().await, true, Some(2048)).await;
+    let client = reqwest::Client::new();
+    for i in 0..5 {
+        let resp = client
+            .post(format!("{gateway}/v1/systemone"))
+            .bearer_auth("sk-dec")
+            .json(&request(Some(DECISION_MODEL)))
+            .send()
+            .await
+            .unwrap();
+        assert_eq!(resp.status(), 200, "request {i} should be affordable");
+        // Settlement is spawned; let it land before the next reservation.
+        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
+    }
+}
@@ -153,17 +153,29 @@ pub trait Scenario: Send + Sync {
    fn prompt_size(&self) -> u32;

    /// Whether this scenario should run against the given model. The
-    /// default runs against everything except image-generation models
-    /// (a chat request at an image model is a 422 `wrong_modality`);
-    /// modality-specific scenarios override this.
+    /// default runs against chat models only: an image-generation model
+    /// or a decision model (#333) answers a chat request with a 422
+    /// `wrong_modality`. Modality-specific scenarios override this.
    fn applies_to(&self, model: &ModelInfo) -> bool {
-        !model.capabilities.iter().any(|c| c == "image")
+        serves_chat(model)
    }

    /// Issue one shaped request and measure it.
    async fn run(&self, ctx: &RunCtx) -> Result<ScenarioMetrics>;
 }

+/// Capabilities that mark a model as serving an endpoint other than
+/// chat completions.
+const NON_CHAT_MODALITIES: [&str; 3] = ["image", "decision", "audio"];
+
+/// Whether a model answers `/v1/chat/completions`.
+pub fn serves_chat(model: &ModelInfo) -> bool {
+    !model
+        .capabilities
+        .iter()
+        .any(|c| NON_CHAT_MODALITIES.contains(&c.as_str()))
+}
+
 /// Fixed-seed image-latency scenario (#203): one 9-step generation at a
 /// configured square size against `/v1/images/generations`, recording
 /// the server-side phase timing (`usage.helexa_timing`) and the metered
@@ -883,6 +895,35 @@ async fn stream_and_measure_inner(
 mod tests {
    use super::*;

+    fn model(capabilities: &[&str]) -> ModelInfo {
+        serde_json::from_value(serde_json::json!({
+            "id": "m",
+            "harness": "candle",
+            "status": "loaded",
+            "devices": [0],
+            "vram_used_mb": null,
+            "capabilities": capabilities,
+        }))
+        .expect("model info")
+    }
+
+    /// Chat scenarios skip every model that answers chat with
+    /// `wrong_modality`: image generators and decision models alike.
+    #[test]
+    fn chat_scenarios_skip_image_and_decision_models() {
+        assert!(serves_chat(&model(&[])));
+        assert!(serves_chat(&model(&["text", "vision"])));
+        assert!(!serves_chat(&model(&["image"])));
+        assert!(!serves_chat(&model(&["decision"])));
+        assert!(!serves_chat(&model(&["audio"])));
+        let image = ImageLatencyScenario {
+            id: "image:512".into(),
+            side_px: 512,
+        };
+        assert!(image.applies_to(&model(&["image"])));
+        assert!(!image.applies_to(&model(&["decision"])));
+    }
+
    /// A burst must report server-side phase timing, or `queue_wait_ms`
    /// (TTFT minus server prefill, #85) cannot be computed for the very
    /// cells it was designed for, and the prefill-tok/s panel stays
@@ -0,0 +1,173 @@
+//! Async audio generation jobs through the router (#382).
+//!
+//! Submitting a job is ordinary capacity dispatch: the router picks a cortex
+//! that serves the model and forwards the body, and that cortex owns the job
+//! from then on. Every later call names only a job id, which carries no hint
+//! of its cortex, so those calls are asked of each cortex in turn and the
+//! first answer that is not `404` wins. Listing merges every cortex's list.
+//!
+//! Ownership is enforced by the cortex holding the job; the router forwards
+//! the caller's bearer unchanged, exactly as for inference.
+
+use crate::dispatch;
+use crate::error::envelope_response;
+use crate::state::RouterState;
+use axum::Router;
+use axum::body::{Body, Bytes};
+use axum::extract::{Path, State};
+use axum::http::{HeaderMap, Method, StatusCode};
+use axum::response::{IntoResponse, Json, Response};
+use axum::routing::get;
+use cortex_core::audio::{AudioJob, AudioJobList};
+use cortex_core::error_envelope::OpenAiError;
+use std::sync::Arc;
+
+const PATH: &str = "/v1/audio/generations";
+
+pub fn routes() -> Router<Arc<RouterState>> {
+    Router::new()
+        .route(PATH, get(list).post(submit))
+        .route("/v1/audio/generations/{id}", get(get_job).delete(cancel))
+        .route("/v1/audio/generations/{id}/content", get(content))
+}
+
+async fn submit(
+    State(state): State<Arc<RouterState>>,
+    headers: HeaderMap,
+    body: Bytes,
+) -> Response {
+    dispatch::dispatch(&state, PATH, headers, body).await
+}
+
+async fn get_job(
+    State(state): State<Arc<RouterState>>,
+    headers: HeaderMap,
+    Path(id): Path<String>,
+) -> Response {
+    find(&state, Method::GET, &format!("{PATH}/{id}"), &headers, &id).await
+}
+
+async fn cancel(
+    State(state): State<Arc<RouterState>>,
+    headers: HeaderMap,
+    Path(id): Path<String>,
+) -> Response {
+    find(
+        &state,
+        Method::DELETE,
+        &format!("{PATH}/{id}"),
+        &headers,
+        &id,
+    )
+    .await
+}
+
+async fn content(
+    State(state): State<Arc<RouterState>>,
+    headers: HeaderMap,
+    Path(id): Path<String>,
+) -> Response {
+    find(
+        &state,
+        Method::GET,
+        &format!("{PATH}/{id}/content"),
+        &headers,
+        &id,
+    )
+    .await
+}
+
+/// Ask each cortex for the job until one answers with something other than
+/// `404`, and stream that answer back verbatim.
+async fn find(
+    state: &RouterState,
+    method: Method,
+    path: &str,
+    headers: &HeaderMap,
+    id: &str,
+) -> Response {
+    for ep in &state.cortexes {
+        let Some(client) = state.client_for(&ep.name) else {
+            continue;
+        };
+        let url = format!("{}{path}", ep.endpoint);
+        let resp = match client
+            .request(method.clone(), &url)
+            .headers(forward_headers(headers))
+            .send()
+            .await
+        {
+            Ok(r) => r,
+            Err(e) => {
+                tracing::warn!(cortex = %ep.name, url = %url, error = %e, "cortex unreachable");
+                continue;
+            }
+        };
+        if resp.status() == reqwest::StatusCode::NOT_FOUND {
+            continue;
+        }
+        let status =
+            StatusCode::from_u16(resp.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
+        let mut out = Response::builder().status(status);
+        for (k, v) in resp.headers() {
+            if k == "transfer-encoding" || k == "connection" {
+                continue;
+            }
+            out = out.header(k, v);
+        }
+        return out
+            .body(Body::from_stream(resp.bytes_stream()))
+            .unwrap_or_else(|_| StatusCode::BAD_GATEWAY.into_response());
+    }
+    envelope_response(OpenAiError::new(
+        404,
+        "invalid_request_error",
+        "job_not_found",
+        format!("no audio job '{id}'"),
+    ))
+}
+
+/// Merge every reachable cortex's job list, newest first.
+async fn list(State(state): State<Arc<RouterState>>, headers: HeaderMap) -> Response {
+    let mut data: Vec<AudioJob> = Vec::new();
+    for ep in &state.cortexes {
+        let Some(client) = state.client_for(&ep.name) else {
+            continue;
+        };
+        let url = format!("{}{PATH}", ep.endpoint);
+        let page = client
+            .get(&url)
+            .headers(forward_headers(&headers))
+            .send()
+            .await
+            .ok()
+            .filter(|r| r.status().is_success());
+        let Some(page) = page else {
+            tracing::warn!(cortex = %ep.name, "audio job list unavailable from cortex");
+            continue;
+        };
+        if let Ok(list) = page.json::<AudioJobList>().await {
+            data.extend(list.data);
+        }
+    }
+    data.sort_by(|a, b| b.created_at.cmp(&a.created_at).then(b.id.cmp(&a.id)));
+    Json(AudioJobList {
+        object: "list".into(),
+        data,
+    })
+    .into_response()
+}
+
+fn forward_headers(headers: &HeaderMap) -> HeaderMap {
+    let mut out = HeaderMap::new();
+    for (k, v) in headers {
+        if matches!(
+            k.as_str(),
+            "host" | "content-length" | "connection" | "transfer-encoding"
+        ) {
+            continue;
+        }
+        out.append(k.clone(), v.clone());
+    }
+    out
+}
@@ -168,6 +168,8 @@ mod tests {
            capabilities: vec![],
            limit: None,
            cost: None,
+            audio: None,
+            audio_tokens_per_second: None,
            tool_call: false,
            reasoning: false,
            max_model_len: None,
@@ -149,25 +149,81 @@ pub async fn dispatch(

    let candidates = match select_cortexes(state, &model).await {
        Selection::Candidates(c) => c,
-        Selection::UnknownModel => {
-            return envelope_response(
-                OpenAiError::new(
-                    404,
-                    "invalid_request_error",
-                    "model_not_found",
-                    format!("no operator serves model '{model}'"),
-                )
-                .with_param("model"),
-            );
-        }
-        Selection::NoReachableCapacity => {
-            return envelope_response(OpenAiError::service_unavailable(
-                format!("model '{model}' is temporarily unavailable on all operators"),
-                Some(RETRY_AFTER_SECS),
-            ));
-        }
+        other => return selection_error(other, &model),
    };
+    forward(state, path, headers, body, &model, candidates).await
+}

+/// Proxy a `/v1/systemone` decision request (#338).
+///
+/// Same as [`dispatch`], except that a missing `model`, or one no
+/// operator serves, selects capacity for
+/// [`DECISION_DEFAULT_ALIAS`](cortex_core::catalogue::DECISION_DEFAULT_ALIAS)
+/// instead of 404ing: decision clients written against the Jev API name
+/// their own service in `model`. The body is forwarded with the client's
+/// `model` intact — cortex applies the same fallback and the serving node
+/// may read it as a checkpoint pin.
+pub async fn dispatch_decision(
+    state: &RouterState,
+    path: &str,
+    headers: HeaderMap,
+    body: Bytes,
+) -> Response {
+    use cortex_core::catalogue::DECISION_DEFAULT_ALIAS;
+
+    let requested = extract_model(&body);
+    let (mut model, body) = match requested.as_deref() {
+        Some(m) => match state.aliases.get(m) {
+            Some(real) => {
+                tracing::debug!(alias = m, model = %real, "resolving tier alias");
+                (real.clone(), rewrite_model_in_body(&body, real))
+            }
+            None => (m.to_string(), body),
+        },
+        None => (DECISION_DEFAULT_ALIAS.to_string(), body),
+    };
+    let mut selection = select_cortexes(state, &model).await;
+    if selection == Selection::UnknownModel
+        && let Some(default) = state.aliases.get(DECISION_DEFAULT_ALIAS)
+    {
+        tracing::debug!(requested = %model, default = %default, "unknown decision model; using the default");
+        model = default.clone();
+        selection = select_cortexes(state, &model).await;
+    }
+    match selection {
+        Selection::Candidates(c) => forward(state, path, headers, body, &model, c).await,
+        other => selection_error(other, &model),
+    }
+}
+
+/// The #63 envelope for a selection that produced no candidates.
+fn selection_error(selection: Selection, model: &str) -> Response {
+    match selection {
+        Selection::Candidates(_) | Selection::UnknownModel => envelope_response(
+            OpenAiError::new(
+                404,
+                "invalid_request_error",
+                "model_not_found",
+                format!("no operator serves model '{model}'"),
+            )
+            .with_param("model"),
+        ),
+        Selection::NoReachableCapacity => envelope_response(OpenAiError::service_unavailable(
+            format!("model '{model}' is temporarily unavailable on all operators"),
+            Some(RETRY_AFTER_SECS),
+        )),
+    }
+}
+
+/// Try `candidates` in order, failing over only on transport errors.
+async fn forward(
+    state: &RouterState,
+    path: &str,
+    headers: HeaderMap,
+    body: Bytes,
+    model: &str,
+    candidates: Vec<CortexEndpoint>,
+) -> Response {
    // Try candidates in order, failing over only on transport errors. A
    // genuine HTTP response (any status — including cortex's #63 429/400)
    // is returned verbatim and never retried away.
@@ -14,6 +14,8 @@ pub fn api_routes() -> Router<Arc<RouterState>> {
        .route("/v1/chat/completions", post(chat_completions))
        .route("/v1/completions", post(completions))
        .route("/v1/images/generations", post(images_generations))
+        .merge(crate::audio::routes())
+        .route("/v1/systemone", post(systemone))
        .route("/v1/responses", post(responses))
        .route("/v1/messages", post(messages))
        .route("/v1/models", get(list_models))
@@ -63,6 +65,17 @@ async fn images_generations(
    dispatch::dispatch(&state, "/v1/images/generations", headers, body).await
 }

+/// `POST /v1/systemone` (#338): typed decisions (TypeSafe Jev wire
+/// protocol) from a decision model. Unlike the other paths an unknown or
+/// missing `model` is not a 404 — see [`dispatch::dispatch_decision`].
+async fn systemone(
+    State(state): State<Arc<RouterState>>,
+    headers: HeaderMap,
+    body: Bytes,
+) -> Response {
+    dispatch::dispatch_decision(&state, "/v1/systemone", headers, body).await
+}
+
 async fn messages(
    State(state): State<Arc<RouterState>>,
    headers: HeaderMap,
@@ -11,6 +11,7 @@
 //! on capacity (epic #69). A background [`poller`] keeps a live
 //! per-cortex topology (#72) that the dispatcher (#73) will route on.

+pub mod audio;
 pub mod catalogue;
 pub mod config;
 pub mod dispatch;
@@ -0,0 +1,187 @@
+//! Audio job lookups through the router (#382): a job id carries no hint of
+//! its cortex, so job calls are asked of each cortex until one answers with
+//! something other than 404, and listings merge every cortex's jobs.
+
+use axum::extract::Path;
+use axum::http::{HeaderMap, StatusCode};
+use axum::response::{IntoResponse, Response};
+use axum::routing::get;
+use axum::{Json, Router};
+use helexa_router::config::{CortexEndpoint, RouterConfig};
+use helexa_router::state::RouterState;
+use serde_json::{Value, json};
+use std::sync::Arc;
+use tokio::net::TcpListener;
+
+/// A cortex that holds exactly one job, `job_id`, created at `created`.
+async fn spawn_cortex(job_id: &'static str, created: u64) -> String {
+    let job = move || {
+        json!({"id": job_id, "object": "audio.generation", "status": "succeeded",
+               "model": "m", "created_at": created})
+    };
+    let app = Router::new()
+        .route(
+            "/v1/audio/generations",
+            get(move |headers: HeaderMap| async move {
+                // Only the bearer's own jobs are listed.
+                if headers.get("authorization").is_none() {
+                    return Json(json!({"object": "list", "data": []})).into_response();
+                }
+                Json(json!({"object": "list", "data": [job()]})).into_response()
+            }),
+        )
+        .route(
+            "/v1/audio/generations/{id}",
+            get(move |Path(id): Path<String>| async move {
+                if id == job_id {
+                    Json(job()).into_response()
+                } else {
+                    not_found()
+                }
+            })
+            .delete(move |Path(id): Path<String>| async move {
+                if id == job_id {
+                    let mut j = job();
+                    j["status"] = json!("cancelled");
+                    Json(j).into_response()
+                } else {
+                    not_found()
+                }
+            }),
+        )
+        .route(
+            "/v1/audio/generations/{id}/content",
+            get(move |Path(id): Path<String>| async move {
+                if id == job_id {
+                    (
+                        StatusCode::OK,
+                        [("content-type", "audio/wav")],
+                        format!("audio-of-{job_id}"),
+                    )
+                        .into_response()
+                } else {
+                    not_found()
+                }
+            }),
+        );
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
+    format!("http://{addr}")
+}
+
+fn not_found() -> Response {
+    (
+        StatusCode::NOT_FOUND,
+        Json(json!({"error": {"type": "invalid_request_error", "code": "job_not_found", "message": "no"}})),
+    )
+        .into_response()
+}
+
+async fn router_over(cortexes: Vec<(&str, String)>) -> String {
+    let cfg = RouterConfig {
+        cortexes: cortexes
+            .into_iter()
+            .map(|(name, endpoint)| CortexEndpoint {
+                name: name.into(),
+                endpoint,
+                region: None,
+                tls_ca: None,
+            })
+            .collect(),
+        ..Default::default()
+    };
+    let state = Arc::new(RouterState::from_config(&cfg));
+    let app = helexa_router::handlers::api_routes().with_state(state);
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
+    format!("http://{addr}")
+}
+
+#[tokio::test]
+async fn job_calls_find_the_cortex_that_holds_the_job() {
+    let a = spawn_cortex("audiogen_a", 100).await;
+    let b = spawn_cortex("audiogen_b", 200).await;
+    // An unreachable cortex in the list must be skipped, not fatal.
+    let router = router_over(vec![
+        ("dead", "http://127.0.0.1:1".into()),
+        ("a", a),
+        ("b", b),
+    ])
+    .await;
+    let http = reqwest::Client::new();
+
+    let job: Value = http
+        .get(format!("{router}/v1/audio/generations/audiogen_b"))
+        .send()
+        .await
+        .unwrap()
+        .json()
+        .await
+        .unwrap();
+    assert_eq!(job["id"], "audiogen_b");
+
+    let content = http
+        .get(format!("{router}/v1/audio/generations/audiogen_b/content"))
+        .send()
+        .await
+        .unwrap();
+    assert_eq!(content.status(), 200);
+    assert_eq!(content.headers()["content-type"], "audio/wav");
+    assert_eq!(content.text().await.unwrap(), "audio-of-audiogen_b");
+
+    let cancelled: Value = http
+        .delete(format!("{router}/v1/audio/generations/audiogen_a"))
+        .send()
+        .await
+        .unwrap()
+        .json()
+        .await
+        .unwrap();
+    assert_eq!(cancelled["status"], "cancelled");
+
+    let missing = http
+        .get(format!("{router}/v1/audio/generations/audiogen_nope"))
+        .send()
+        .await
+        .unwrap();
+    assert_eq!(missing.status(), 404);
+    let body: Value = missing.json().await.unwrap();
+    assert_eq!(body["error"]["code"], "job_not_found");
+}
+
+#[tokio::test]
+async fn listing_merges_cortexes_newest_first_and_forwards_the_bearer() {
+    let a = spawn_cortex("audiogen_a", 100).await;
+    let b = spawn_cortex("audiogen_b", 200).await;
+    let router = router_over(vec![("a", a), ("b", b)]).await;
+    let http = reqwest::Client::new();
+
+    let list: Value = http
+        .get(format!("{router}/v1/audio/generations"))
+        .bearer_auth("k")
+        .send()
+        .await
+        .unwrap()
+        .json()
+        .await
+        .unwrap();
+    let ids: Vec<&str> = list["data"]
+        .as_array()
+        .unwrap()
+        .iter()
+        .map(|j| j["id"].as_str().unwrap())
+        .collect();
+    assert_eq!(ids, ["audiogen_b", "audiogen_a"]);
+
+    let anon: Value = http
+        .get(format!("{router}/v1/audio/generations"))
+        .send()
+        .await
+        .unwrap()
+        .json()
+        .await
+        .unwrap();
+    assert_eq!(anon["data"].as_array().unwrap().len(), 0);
+}
@@ -14,7 +14,7 @@ use axum::routing::post;
 use axum::{Json, Router};
 use cortex_core::node::CortexModelEntry;
 use helexa_router::config::{CortexEndpoint, RouterConfig};
-use helexa_router::dispatch::{Selection, dispatch, select_cortexes};
+use helexa_router::dispatch::{Selection, dispatch, dispatch_decision, select_cortexes};
 use helexa_router::state::{CortexTopology, RouterState};
 use serde_json::{Value, json};
 use std::collections::HashMap;
@@ -37,6 +37,8 @@ fn model_entry(loaded: bool, feasible: bool) -> CortexModelEntry {
        capabilities: vec![],
        limit: None,
        cost: None,
+        audio: None,
+        audio_tokens_per_second: None,
        tool_call: false,
        reasoning: false,
        max_model_len: None,
@@ -87,6 +89,7 @@ async fn mock_handler(State(m): State<MockCortex>, headers: HeaderMap, body: Byt
 async fn spawn_cortex(mock: MockCortex) -> String {
    let app = Router::new()
        .route("/v1/chat/completions", post(mock_handler))
+        .route("/v1/systemone", post(mock_handler))
        .with_state(mock);
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
@@ -349,3 +352,75 @@ async fn ranking_prefers_loaded_then_region_then_headroom() {
    let names: Vec<&str> = order.iter().map(|e| e.name.as_str()).collect();
    assert_eq!(names, vec!["warm-eu", "warm-us", "cold-eu"]);
 }
+
+// ── /v1/systemone default-model fallback (#338) ──────────────────────
+
+/// Router state with c1 serving MODEL, and `helexa/one` pointing at it
+/// when `with_default` is set.
+async fn decision_state(with_default: bool) -> RouterState {
+    let url = spawn_cortex(ok_cortex("c1")).await;
+    let mut state = state_with(vec![ep("c1", &url, None)], None);
+    if with_default {
+        state.aliases.insert(
+            cortex_core::catalogue::DECISION_DEFAULT_ALIAS.to_string(),
+            MODEL.to_string(),
+        );
+    }
+    set_topology(&state, "c1", true, true, true, 2).await;
+    state
+}
+
+#[tokio::test]
+async fn decision_foreign_model_selects_default_and_keeps_body_model() {
+    // A Jev client names its own service. Capacity is chosen for the
+    // default decision model, and the client's `model` is forwarded
+    // untouched for cortex (and the node) to interpret.
+    let state = decision_state(true).await;
+    let body = Bytes::from_static(b"{\"model\":\"jev-1\",\"questions\":{}}");
+    let resp = dispatch_decision(&state, "/v1/systemone", HeaderMap::new(), body).await;
+    let (status, body) = body_json(resp).await;
+    assert_eq!(status, StatusCode::OK);
+    assert_eq!(body["served_by"], "c1");
+    assert_eq!(body["model_seen"], "jev-1");
+}
+
+#[tokio::test]
+async fn decision_missing_model_selects_default() {
+    let state = decision_state(true).await;
+    let body = Bytes::from_static(b"{\"questions\":{}}");
+    let resp = dispatch_decision(&state, "/v1/systemone", HeaderMap::new(), body).await;
+    let (status, body) = body_json(resp).await;
+    assert_eq!(status, StatusCode::OK);
+    assert_eq!(body["served_by"], "c1");
+    assert_eq!(body["model_seen"], "");
+}
+
+#[tokio::test]
+async fn decision_default_alias_is_rewritten_like_any_alias() {
+    let state = decision_state(true).await;
+    let body = Bytes::from_static(b"{\"model\":\"helexa/one\",\"questions\":{}}");
+    let resp = dispatch_decision(&state, "/v1/systemone", HeaderMap::new(), body).await;
+    let (status, body) = body_json(resp).await;
+    assert_eq!(status, StatusCode::OK);
+    assert_eq!(body["model_seen"], MODEL);
+}
+
+#[tokio::test]
+async fn decision_unknown_model_without_default_is_404() {
+    let state = decision_state(false).await;
+    let body = Bytes::from_static(b"{\"model\":\"jev-1\",\"questions\":{}}");
+    let resp = dispatch_decision(&state, "/v1/systemone", HeaderMap::new(), body).await;
+    let (status, body) = body_json(resp).await;
+    assert_eq!(status, StatusCode::NOT_FOUND);
+    assert_eq!(body["error"]["code"], "model_not_found");
+}
+
+#[tokio::test]
+async fn chat_unknown_model_does_not_take_the_decision_fallback() {
+    // The fallback is scoped to the decision endpoint.
+    let state = decision_state(true).await;
+    let body = Bytes::from_static(b"{\"model\":\"jev-1\",\"stream\":false}");
+    let resp = dispatch(&state, "/v1/chat/completions", HeaderMap::new(), body).await;
+    let (status, _) = body_json(resp).await;
+    assert_eq!(status, StatusCode::NOT_FOUND);
+}
@@ -0,0 +1,32 @@
+-- ── Per-model rates for non-token units (#386) ───────────────────────
+--
+-- The ledger counts tokens. Work metered in another unit — audio-seconds
+-- for music generation today — is converted into ledger tokens at a
+-- per-model rate kept here, so an operator can price a model, or change
+-- its price, without a deploy. cortex fetches the table from
+-- `/authz/v1/rates` and reserves and settles at the rate it read.
+--
+-- `model_id = '*'` is the default for a unit: it applies to every model
+-- without a row of its own. A unit with no row at all, for the model or
+-- for '*', is not priced, and cortex refuses that work rather than serving
+-- it free.
+CREATE TABLE model_rates (
+    model_id         TEXT NOT NULL,
+    -- 'audio_second' — one second of generated audio.
+    unit             TEXT NOT NULL,
+    -- Ledger tokens charged per unit.
+    tokens_per_unit  BIGINT NOT NULL,
+    description      TEXT,
+    -- Audit: who changed it and when (NULL updated_by = shipped default).
+    updated_at       TIMESTAMPTZ NOT NULL DEFAULT now(),
+    updated_by       UUID REFERENCES users(id) ON DELETE SET NULL,
+
+    PRIMARY KEY (model_id, unit),
+    CONSTRAINT model_rates_unit_known CHECK (unit IN ('audio_second')),
+    CONSTRAINT model_rates_positive CHECK (tokens_per_unit > 0)
+);
+
+-- Placeholder default while music generation is being built; set real
+-- per-model prices with an UPDATE or per-model INSERT.
+INSERT INTO model_rates (model_id, unit, tokens_per_unit, description)
+VALUES ('*', 'audio_second', 1000, 'default audio rate (placeholder)');
@@ -21,7 +21,7 @@ use axum::extract::{Request, State};
 use axum::http::{StatusCode, header};
 use axum::middleware::Next;
 use axum::response::{IntoResponse, Response};
-use axum::routing::post;
+use axum::routing::{get, post};
 use axum::{Extension, Json, Router};
 use cortex_core::error_envelope::OpenAiError;
 use serde::{Deserialize, Serialize};
@@ -42,6 +42,7 @@ pub fn router(state: &AppState) -> Router<AppState> {
        .route("/authz/v1/release", post(release))
        .route("/authz/v1/snapshot", post(snapshot))
        .route("/authz/v1/served-usage", post(served_usage))
+        .route("/authz/v1/rates", get(rates))
        .layer(axum::middleware::from_fn_with_state(
            state.clone(),
            client_auth,
@@ -338,3 +339,23 @@ fn bad_request(msg: &str) -> Response {
        msg,
    ))
 }
+
+// ── rates (#386) ────────────────────────────────────────────────────
+
+#[derive(Serialize)]
+struct RatesResp {
+    data: Vec<crate::rates::ModelRate>,
+}
+
+/// `GET /authz/v1/rates` — every per-model rate, for cortex to cache and
+/// resolve locally. A database error is a 500: cortex keeps its last good
+/// table and refuses unpriced work rather than guessing.
+async fn rates(State(state): State<AppState>) -> Response {
+    match crate::rates::list(&state.pool).await {
+        Ok(data) => Json(RatesResp { data }).into_response(),
+        Err(e) => {
+            tracing::error!(error = %e, "rates query failed");
+            (StatusCode::INTERNAL_SERVER_ERROR, "rates unavailable").into_response()
+        }
+    }
+}
@@ -21,6 +21,7 @@ pub mod email;
 pub mod error;
 pub mod handlers;
 pub mod ledger;
+pub mod rates;
 pub mod reconcile;
 pub mod state;
 pub mod topup;
@@ -0,0 +1,35 @@
+//! Per-model rates for work metered in units other than tokens (#386).
+//!
+//! The ledger only counts tokens; a rate converts another unit (one second
+//! of generated audio) into ledger tokens. Rates live in `model_rates` so
+//! they can change without a deploy. cortex reads the whole table from
+//! `/authz/v1/rates` and resolves a model's rate itself: the model's own
+//! row, else the unit's `'*'` default, else the work is not priced.
+
+use serde::Serialize;
+use sqlx::postgres::PgPool;
+
+/// One row of `model_rates`, as served to cortex.
+#[derive(Debug, Clone, Serialize)]
+pub struct ModelRate {
+    pub model_id: String,
+    pub unit: String,
+    pub tokens_per_unit: i64,
+}
+
+/// Every configured rate.
+pub async fn list(pool: &PgPool) -> Result<Vec<ModelRate>, sqlx::Error> {
+    let rows = sqlx::query(
+        "SELECT model_id, unit, tokens_per_unit FROM model_rates ORDER BY unit, model_id",
+    )
+    .fetch_all(pool)
+    .await?;
+    Ok(rows
+        .into_iter()
+        .map(|r| ModelRate {
+            model_id: sqlx::Row::get(&r, "model_id"),
+            unit: sqlx::Row::get(&r, "unit"),
+            tokens_per_unit: sqlx::Row::get(&r, "tokens_per_unit"),
+        })
+        .collect())
+}
@@ -0,0 +1,73 @@
+//! `model_rates` (#386) against a real PostgreSQL. Gated on
+//! `UPSTREAM_TEST_DATABASE_URL`, like the ledger tests: unset, every test
+//! logs a skip and passes.
+
+use helexa_upstream::db::connect_and_migrate;
+use helexa_upstream::rates;
+use sqlx::postgres::PgPool;
+
+async fn pool_or_skip(test: &str) -> Option<PgPool> {
+    let Ok(url) = std::env::var("UPSTREAM_TEST_DATABASE_URL") else {
+        eprintln!("skipping {test}: UPSTREAM_TEST_DATABASE_URL not set");
+        return None;
+    };
+    Some(
+        connect_and_migrate(&url, 4)
+            .await
+            .expect("connect + migrate"),
+    )
+}
+
+#[tokio::test]
+async fn default_audio_rate_is_seeded_and_model_rows_are_listed() {
+    let Some(pool) = pool_or_skip("default_audio_rate_is_seeded").await else {
+        return;
+    };
+    let model = format!("test/model-{}", uuid::Uuid::new_v4());
+    sqlx::query(
+        "INSERT INTO model_rates (model_id, unit, tokens_per_unit) VALUES ($1, 'audio_second', 2500)",
+    )
+    .bind(&model)
+    .execute(&pool)
+    .await
+    .unwrap();
+
+    let all = rates::list(&pool).await.unwrap();
+    assert!(
+        all.iter()
+            .any(|r| r.model_id == "*" && r.unit == "audio_second" && r.tokens_per_unit > 0),
+        "seeded default missing: {all:?}"
+    );
+    assert!(
+        all.iter()
+            .any(|r| r.model_id == model && r.tokens_per_unit == 2500)
+    );
+
+    sqlx::query("DELETE FROM model_rates WHERE model_id = $1")
+        .bind(&model)
+        .execute(&pool)
+        .await
+        .unwrap();
+}
+
+#[tokio::test]
+async fn rates_must_be_positive_and_of_a_known_unit() {
+    let Some(pool) = pool_or_skip("rates_must_be_positive").await else {
+        return;
+    };
+    let insert = |unit: &'static str, tokens: i64| {
+        let pool = pool.clone();
+        async move {
+            sqlx::query(
+                "INSERT INTO model_rates (model_id, unit, tokens_per_unit) VALUES ('test/bad', $1, $2)",
+            )
+            .bind(unit)
+            .bind(tokens)
+            .execute(&pool)
+            .await
+        }
+    };
+    assert!(insert("audio_second", 0).await.is_err());
+    assert!(insert("audio_second", -5).await.is_err());
+    assert!(insert("furlongs", 10).await.is_err());
+}
@@ -76,6 +76,12 @@ toml.workspace = true
 # transitively via candle-core.
 rayon = "1"

+# Unicode tables for the decision-model language detector (#339), which
+# must classify characters exactly as Python's `str.isalpha` / `re` and
+# `unicodedata.combining` do. Both already in the tree via tokenizers.
+unicode-normalization = "0.1"
+unicode-properties = "0.1"
+
 # candle for in-process inference. CUDA support is gated behind the
 # crate's `cuda` feature (default off) so the workspace builds on
 # non-CUDA hosts and CI runners.
@@ -95,6 +101,9 @@ cudarc = { version = "0.19", optional = true, default-features = false, features
 # compiling the `half` crate at conflicting versions.
 half = { version = "2.5", optional = true }
 tokenizers = { version = "0.22", default-features = false, features = ["onig"] }
+# Reads ACE-Step's `silence_latent.pt` (a torch zip archive of one raw
+# tensor). Already in the graph via candle; stored entries need no codecs.
+zip = { version = "8", default-features = false }
 hf-hub = { version = "0.4", features = ["tokio"] }
 # Jinja-compatible template renderer for the model's chat template
 # (standalone `chat_template.jinja` or `tokenizer_config.json::chat_template`).
@@ -55,6 +55,8 @@ pub fn neuron_routes() -> Router<Arc<NeuronState>> {
        .route("/models/{model_id}/capabilities", get(model_capabilities))
        .route("/v1/chat/completions", post(chat_completions))
        .route("/v1/images/generations", post(images_generations))
+        .route("/v1/audio/generations", post(audio_generations))
+        .route("/v1/systemone", post(systemone))
        .route("/v1/responses", post(responses))
 }

@@ -669,6 +671,138 @@ async fn images_generations(
    }
 }

+/// `POST /v1/audio/generations` (#382): render one audio generation and
+/// return it. This is the neuron half of the async job API — cortex owns
+/// the job and calls this once per job, waiting for the whole render. The
+/// body is the audio (`audio/wav`); the result metadata travels in the
+/// [`cortex_core::audio::RESULT_HEADER`] header so the bytes need no
+/// base64 envelope.
+async fn audio_generations(
+    State(state): State<Arc<NeuronState>>,
+    headers: axum::http::HeaderMap,
+    ApiJson(req): ApiJson<cortex_core::audio::AudioGenerationRequest>,
+) -> axum::response::Response {
+    use cortex_core::audio::RESULT_HEADER;
+    use cortex_core::error_envelope::OpenAiError;
+
+    let Some(candle) = state.candle.as_ref().map(Arc::clone) else {
+        return (
+            StatusCode::SERVICE_UNAVAILABLE,
+            Json(json!({"error": "candle harness not enabled on this neuron"})),
+        )
+            .into_response();
+    };
+    if let Err(msg) = req.validate_shape() {
+        return envelope_response(OpenAiError::new(
+            400,
+            "invalid_request_error",
+            "invalid_audio_params",
+            format!("invalid audio generation parameters: {msg}"),
+        ));
+    }
+    let principal = principal_key(&headers);
+    match candle.audio_generation(&req, principal.as_deref()).await {
+        Ok((wav, result)) => {
+            let meta = match serde_json::to_string(&result) {
+                Ok(m) => m,
+                Err(e) => {
+                    return inference_error_response(InferenceError::Other(anyhow::anyhow!(
+                        "serialise audio result: {e}"
+                    )));
+                }
+            };
+            let mut resp = axum::response::Response::new(axum::body::Body::from(wav));
+            let h = resp.headers_mut();
+            h.insert(
+                axum::http::header::CONTENT_TYPE,
+                axum::http::HeaderValue::from_static("audio/wav"),
+            );
+            if let Ok(v) = axum::http::HeaderValue::from_str(&meta) {
+                h.insert(RESULT_HEADER, v);
+            }
+            resp
+        }
+        Err(e) => inference_error_response(e),
+    }
+}
+
+/// `POST /v1/systemone` (#338): typed decisions over the TypeSafe Jev
+/// wire protocol, from a loaded decision model (Laya). Parsing and the
+/// transport-level limits happen here; routing, validation, admission,
+/// scoring and shaping in `CandleHarness::systemone`.
+///
+/// Errors carry the #63 envelope *and* the reference server's top-level
+/// `detail` string, so OpenAI-style and Jev/laya clients both read them.
+async fn systemone(
+    State(state): State<Arc<NeuronState>>,
+    headers: axum::http::HeaderMap,
+    body: axum::body::Bytes,
+) -> axum::response::Response {
+    use crate::harness::candle::SystemOneError;
+    use cortex_core::decisions::{SystemOneRequest, limits};
+
+    let Some(candle) = state.candle.as_ref().map(Arc::clone) else {
+        return decision_error_response(503, "candle harness not enabled on this neuron");
+    };
+    if body.len() > limits::MAX_BODY_BYTES {
+        return decision_error_response(413, "request body too large");
+    }
+    let request = match SystemOneRequest::parse(&body, limits::DEFAULT_MAX_TOKEN_BUDGET) {
+        Ok(r) => r,
+        Err(e) => return decision_error_response(e.status, &e.detail),
+    };
+    let served = headers
+        .get(cortex_core::catalogue::DECISION_MODEL_HEADER)
+        .and_then(|v| v.to_str().ok())
+        .map(str::to_string);
+    let principal = principal_key(&headers);
+
+    let started = std::time::Instant::now();
+    match candle
+        .systemone(served.as_deref(), request, principal.as_deref())
+        .await
+    {
+        Ok(response) => {
+            let ms = started.elapsed().as_secs_f64() * 1e3;
+            let mut out = Json(response).into_response();
+            for (name, value) in [
+                ("server-timing", format!("inference;dur={ms:.2}")),
+                ("x-inference-time-ms", format!("{ms:.2}")),
+            ] {
+                if let Ok(v) = axum::http::HeaderValue::from_str(&value) {
+                    out.headers_mut().insert(name, v);
+                }
+            }
+            out
+        }
+        Err(SystemOneError::Request(e)) => decision_error_response(e.status, &e.detail),
+        Err(SystemOneError::Inference(e)) => {
+            tracing::warn!(error = %e, "systemone request failed");
+            inference_error_response(e)
+        }
+    }
+}
+
+/// A rejected decision request: the #63 envelope, plus the `detail` field
+/// the reference (FastAPI) server answers with.
+fn decision_error_response(status: u16, detail: &str) -> axum::response::Response {
+    use cortex_core::error_envelope::OpenAiError;
+    let (kind, code) = match status {
+        404 => ("invalid_request_error", "model_not_found"),
+        413 => ("invalid_request_error", "request_too_large"),
+        422 => ("invalid_request_error", "invalid_decision_request"),
+        400..=499 => ("invalid_request_error", "invalid_request"),
+        _ => ("api_error", "service_unavailable"),
+    };
+    let env = OpenAiError::new(status, kind, code, detail.to_string());
+    let mut body = env.body();
+    if let Some(obj) = body.as_object_mut() {
+        obj.insert("detail".into(), json!(detail));
+    }
+    let status = StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_REQUEST);
+    (status, Json(body)).into_response()
+}
+
 fn finish_reason_from_str(s: &str) -> crate::wire::FinishReason {
    use crate::wire::FinishReason;
    match s {
@@ -769,6 +903,12 @@ fn inference_error_response(err: InferenceError) -> axum::response::Response {
        // model. Not retryable against this model — the client should
        // pick a model whose `capabilities` include the endpoint's
        // modality.
+        InferenceError::InvalidAudioParams { detail } => OpenAiError::new(
+            400,
+            "invalid_request_error",
+            "invalid_audio_params",
+            format!("invalid audio generation parameters: {detail}"),
+        ),
        InferenceError::InvalidImageParams { detail } => OpenAiError::new(
            400,
            "invalid_request_error",
@@ -0,0 +1,831 @@
+//! Condition encoders for ACE-Step 1.5.
+//!
+//! This module implements the conditioning pipeline that processes text descriptions,
+//! lyrics, and timbre (reference audio) into condition tensors consumed by the DiT
+//! diffusion model.
+//!
+//! The main entry point is [`ConditionEncoder`], which orchestrates:
+//! - A simple linear projection for text encoder hidden states.
+//! - A [`LyricEncoder`] (8-layer bidirectional transformer) for lyric embeddings.
+//! - A [`TimbreEncoder`] (4-layer bidirectional transformer with CLS pooling) for
+//!   reference audio embeddings.
+//!
+//! All encoder layers share the same architecture: pre-norm (RmsNorm) bidirectional
+//! GQA attention followed by a gated SiLU MLP, matching the Qwen3 block design.
+
+use candle_core::{DType, Device, Module, Result, Tensor};
+use candle_nn::{Activation, VarBuilder};
+use candle_transformers::models::with_tracing::{Linear, RmsNorm, linear_b, linear_no_bias};
+use candle_transformers::utils::repeat_kv;
+
+use super::AceStepConfig;
+
+// ---------------------------------------------------------------------------
+// RotaryEmbedding (standalone for encoder, same algorithm as the DiT variant)
+// ---------------------------------------------------------------------------
+
+/// Precomputed sin/cos tables for rotary position embeddings.
+///
+/// The encoder uses its own instance because it may have a different head dimension
+/// or maximum sequence length than the DiT backbone.
+#[derive(Debug, Clone)]
+pub struct RotaryEmbedding {
+    sin: Tensor,
+    cos: Tensor,
+}
+
+impl RotaryEmbedding {
+    /// Build the sin/cos lookup tables up to `max_seq_len` positions.
+    pub fn new(
+        dtype: DType,
+        head_dim: usize,
+        max_seq_len: usize,
+        rope_theta: f64,
+        dev: &Device,
+    ) -> Result<Self> {
+        let inv_freq: Vec<_> = (0..head_dim)
+            .step_by(2)
+            .map(|i| 1f32 / rope_theta.powf(i as f64 / head_dim as f64) as f32)
+            .collect();
+        let inv_freq_len = inv_freq.len();
+        let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(DType::F32)?;
+        let t = Tensor::arange(0u32, max_seq_len as u32, dev)?
+            .to_dtype(DType::F32)?
+            .reshape((max_seq_len, 1))?;
+        let freqs = t.matmul(&inv_freq)?;
+        Ok(Self {
+            sin: freqs.sin()?.to_dtype(dtype)?,
+            cos: freqs.cos()?.to_dtype(dtype)?,
+        })
+    }
+
+    /// Apply RoPE to Q and K tensors of shape `(B, H, L, D)`.
+    pub fn apply(&self, q: &Tensor, k: &Tensor, offset: usize) -> Result<(Tensor, Tensor)> {
+        let (_, _, seq_len, _) = q.dims4()?;
+        let cos = self.cos.narrow(0, offset, seq_len)?;
+        let sin = self.sin.narrow(0, offset, seq_len)?;
+        let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?;
+        let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?;
+        Ok((q_embed, k_embed))
+    }
+}
+
+// ---------------------------------------------------------------------------
+// EncoderMLP
+// ---------------------------------------------------------------------------
+
+/// Gated SiLU MLP identical to the Qwen3 feed-forward block.
+///
+/// ```text
+/// output = down_proj(silu(gate_proj(x)) * up_proj(x))
+/// ```
+#[derive(Debug, Clone)]
+pub struct EncoderMLP {
+    gate_proj: Linear,
+    up_proj: Linear,
+    down_proj: Linear,
+    act_fn: Activation,
+}
+
+impl EncoderMLP {
+    pub fn new(cfg: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let hidden = cfg.encoder_hidden_size();
+        let intermediate = cfg.encoder_intermediate_size();
+        Ok(Self {
+            gate_proj: linear_no_bias(hidden, intermediate, vb.pp("gate_proj"))?,
+            up_proj: linear_no_bias(hidden, intermediate, vb.pp("up_proj"))?,
+            down_proj: linear_no_bias(intermediate, hidden, vb.pp("down_proj"))?,
+            act_fn: cfg.hidden_act,
+        })
+    }
+}
+
+impl Module for EncoderMLP {
+    fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let gate = x.apply(&self.gate_proj)?.apply(&self.act_fn)?;
+        let up = x.apply(&self.up_proj)?;
+        (gate * up)?.apply(&self.down_proj)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// EncoderAttention
+// ---------------------------------------------------------------------------
+
+/// Bidirectional grouped-query attention for encoder layers.
+///
+/// Uses per-head RmsNorm on Q and K (Qwen3 style) and rotary position
+/// embeddings. No KV cache is needed because the encoder runs a single
+/// forward pass over the full sequence.
+#[derive(Debug, Clone)]
+pub struct EncoderAttention {
+    q_proj: Linear,
+    k_proj: Linear,
+    v_proj: Linear,
+    o_proj: Linear,
+    q_norm: RmsNorm,
+    k_norm: RmsNorm,
+    num_heads: usize,
+    num_kv_heads: usize,
+    num_kv_groups: usize,
+    head_dim: usize,
+    hidden_size: usize,
+}
+
+impl EncoderAttention {
+    pub fn new(cfg: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let head_dim = cfg.head_dim;
+        let num_heads = cfg.encoder_num_attention_heads();
+        let num_kv_heads = cfg.encoder_num_key_value_heads();
+        let num_kv_groups = num_heads / num_kv_heads;
+        let hidden = cfg.encoder_hidden_size();
+
+        let q_proj = linear_b(
+            hidden,
+            num_heads * head_dim,
+            cfg.attention_bias,
+            vb.pp("q_proj"),
+        )?;
+        let k_proj = linear_b(
+            hidden,
+            num_kv_heads * head_dim,
+            cfg.attention_bias,
+            vb.pp("k_proj"),
+        )?;
+        let v_proj = linear_b(
+            hidden,
+            num_kv_heads * head_dim,
+            cfg.attention_bias,
+            vb.pp("v_proj"),
+        )?;
+        let o_proj = linear_b(
+            num_heads * head_dim,
+            hidden,
+            cfg.attention_bias,
+            vb.pp("o_proj"),
+        )?;
+
+        let q_norm = RmsNorm::new(head_dim, cfg.rms_norm_eps, vb.pp("q_norm"))?;
+        let k_norm = RmsNorm::new(head_dim, cfg.rms_norm_eps, vb.pp("k_norm"))?;
+
+        let hidden_size = head_dim * num_heads;
+
+        Ok(Self {
+            q_proj,
+            k_proj,
+            v_proj,
+            o_proj,
+            q_norm,
+            k_norm,
+            num_heads,
+            num_kv_heads,
+            num_kv_groups,
+            head_dim,
+            hidden_size,
+        })
+    }
+
+    /// Run bidirectional self-attention.
+    ///
+    /// * `x` - input tensor of shape `(B, L, D)`.
+    /// * `mask` - optional additive attention mask `(B, 1, L, L)` where `-inf`
+    ///   blocks attention. Pass `None` for fully bidirectional attention.
+    /// * `rotary` - precomputed rotary embedding tables.
+    pub fn forward(
+        &self,
+        x: &Tensor,
+        mask: Option<&Tensor>,
+        rotary: &RotaryEmbedding,
+    ) -> Result<Tensor> {
+        let (b, l, _) = x.dims3()?;
+
+        // Project
+        let q = self.q_proj.forward(x)?;
+        let k = self.k_proj.forward(x)?;
+        let v = self.v_proj.forward(x)?;
+
+        // Reshape to (B, H, L, D)
+        let q = q
+            .reshape((b, l, self.num_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let k = k
+            .reshape((b, l, self.num_kv_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let v = v
+            .reshape((b, l, self.num_kv_heads, self.head_dim))?
+            .transpose(1, 2)?;
+
+        // Per-head RmsNorm
+        let q = self.q_norm.forward(&q.flatten(0, 2)?)?.reshape((
+            b,
+            self.num_heads,
+            l,
+            self.head_dim,
+        ))?;
+        let k = self.k_norm.forward(&k.flatten(0, 2)?)?.reshape((
+            b,
+            self.num_kv_heads,
+            l,
+            self.head_dim,
+        ))?;
+
+        // Rotary position embeddings
+        let (q, k) = rotary.apply(&q, &k, 0)?;
+
+        // GQA key/value expansion
+        let k = repeat_kv(k, self.num_kv_groups)?.contiguous()?;
+        let v = repeat_kv(v, self.num_kv_groups)?.contiguous()?;
+
+        // Scaled dot-product attention
+        let scale = 1.0 / (self.head_dim as f64).sqrt();
+        let mut scores = (q.matmul(&k.transpose(2, 3)?)? * scale)?;
+        if let Some(m) = mask {
+            scores = scores.broadcast_add(m)?;
+        }
+        let probs = candle_nn::ops::softmax_last_dim(&scores)?;
+        let ctx = probs.matmul(&v)?;
+
+        // Reshape back and output projection
+        ctx.transpose(1, 2)?
+            .reshape((b, l, self.hidden_size))?
+            .apply(&self.o_proj)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// EncoderLayer
+// ---------------------------------------------------------------------------
+
+/// A single bidirectional pre-norm transformer layer.
+///
+/// ```text
+/// x = x + self_attn(input_layernorm(x))
+/// x = x + mlp(post_attention_layernorm(x))
+/// ```
+///
+/// Layers alternate between sliding-window and full attention based on
+/// `attention_type`, which controls the mask passed to the attention block.
+#[derive(Debug, Clone)]
+pub struct EncoderLayer {
+    self_attn: EncoderAttention,
+    input_layernorm: RmsNorm,
+    post_attention_layernorm: RmsNorm,
+    mlp: EncoderMLP,
+    attention_type: String,
+}
+
+impl EncoderLayer {
+    pub fn new(cfg: &AceStepConfig, attention_type: &str, vb: VarBuilder) -> Result<Self> {
+        let hidden = cfg.encoder_hidden_size();
+        Ok(Self {
+            self_attn: EncoderAttention::new(cfg, vb.pp("self_attn"))?,
+            input_layernorm: RmsNorm::new(hidden, cfg.rms_norm_eps, vb.pp("input_layernorm"))?,
+            post_attention_layernorm: RmsNorm::new(
+                hidden,
+                cfg.rms_norm_eps,
+                vb.pp("post_attention_layernorm"),
+            )?,
+            mlp: EncoderMLP::new(cfg, vb.pp("mlp"))?,
+            attention_type: attention_type.to_string(),
+        })
+    }
+
+    /// Returns the attention type for this layer (`"sliding_attention"` or
+    /// `"full_attention"`), used to select the appropriate mask.
+    pub fn attention_type(&self) -> &str {
+        &self.attention_type
+    }
+
+    /// Forward pass with pre-norm residual connections.
+    pub fn forward(
+        &self,
+        x: &Tensor,
+        mask: Option<&Tensor>,
+        rotary: &RotaryEmbedding,
+    ) -> Result<Tensor> {
+        // Self-attention with residual
+        let residual = x.clone();
+        let h = self.input_layernorm.forward(x)?;
+        let h = self.self_attn.forward(&h, mask, rotary)?;
+        let x = (residual + h)?;
+
+        // MLP with residual
+        let residual = x.clone();
+        let h = self.post_attention_layernorm.forward(&x)?;
+        let h = h.apply(&self.mlp)?;
+        residual + h
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Bidirectional attention mask helpers
+// ---------------------------------------------------------------------------
+
+/// Build a bidirectional sliding-window mask.
+///
+/// Positions outside the window receive `-inf`; positions inside receive `0`.
+/// The returned tensor has shape `(1, 1, seq_len, seq_len)`.
+fn sliding_window_mask(
+    seq_len: usize,
+    window: usize,
+    dtype: DType,
+    dev: &Device,
+) -> Result<Tensor> {
+    let minf = f32::NEG_INFINITY;
+    let mask: Vec<f32> = (0..seq_len)
+        .flat_map(|i| {
+            (0..seq_len).map(move |j| {
+                if (i as i64 - j as i64).unsigned_abs() as usize <= window {
+                    0.0
+                } else {
+                    minf
+                }
+            })
+        })
+        .collect();
+    Tensor::from_slice(&mask, (1, 1, seq_len, seq_len), dev)?.to_dtype(dtype)
+}
+
+/// Build an attention mask from a padding mask, optionally combined with a
+/// sliding window constraint.
+///
+/// * `attention_mask` - `(B, L)` with `1` for valid tokens, `0` for padding.
+/// * `sliding_window` - if `Some(w)`, restricts attention to a local window.
+///
+/// Returns `(B, 1, L, L)` additive mask.
+fn build_bidirectional_mask(
+    attention_mask: &Tensor,
+    seq_len: usize,
+    sliding_window: Option<usize>,
+    dtype: DType,
+    dev: &Device,
+) -> Result<Tensor> {
+    let (b, _) = attention_mask.dims2()?;
+
+    // Padding mask: convert (B, L) 1/0 mask → (B, 1, 1, L) with 0.0 for valid, -inf for padding.
+    let pad_mask = attention_mask
+        .reshape((b, 1, 1, seq_len))?
+        .to_dtype(DType::F32)?;
+    // Use where_cond to avoid 0 * -inf = NaN
+    let zeros = Tensor::zeros_like(&pad_mask)?;
+    let neginf = Tensor::full(f32::NEG_INFINITY, pad_mask.shape(), dev)?.to_dtype(DType::F32)?;
+    let pad_mask_bool = pad_mask.gt(&zeros)?; // true for valid (mask==1)
+    let pad_mask = pad_mask_bool.where_cond(&zeros, &neginf)?; // 0.0 for valid, -inf for padding
+
+    let mask = match sliding_window {
+        Some(w) => {
+            let sw = sliding_window_mask(seq_len, w, DType::F32, dev)?;
+            // Combine: both must allow the position
+            pad_mask.broadcast_add(&sw)?
+        }
+        None => pad_mask
+            .broadcast_as((b, 1, seq_len, seq_len))?
+            .contiguous()?,
+    };
+    mask.to_dtype(dtype)
+}
+
+// ---------------------------------------------------------------------------
+// LyricEncoder
+// ---------------------------------------------------------------------------
+
+/// Bidirectional transformer encoder for lyric embeddings.
+///
+/// Takes pre-embedded lyric token features (from a text encoder) and refines
+/// them through 8 layers of bidirectional self-attention.
+///
+/// Weight prefix: `lyric_encoder.*`
+#[derive(Debug, Clone)]
+pub struct LyricEncoder {
+    embed_tokens: Linear,
+    layers: Vec<EncoderLayer>,
+    norm: RmsNorm,
+    rotary_emb: RotaryEmbedding,
+    sliding_window: Option<usize>,
+    dtype: DType,
+    device: Device,
+}
+
+impl LyricEncoder {
+    pub fn new(cfg: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let enc_hidden = cfg.encoder_hidden_size();
+
+        // Linear projection from text embedding dim to encoder hidden size (with bias).
+        let embed_tokens = linear_b(cfg.text_hidden_dim, enc_hidden, true, vb.pp("embed_tokens"))?;
+
+        let num_layers = cfg.num_lyric_encoder_hidden_layers;
+        let mut layers = Vec::with_capacity(num_layers);
+        let vb_l = vb.pp("layers");
+        for i in 0..num_layers {
+            let attn_type = if i % 2 == 0 {
+                "sliding_attention"
+            } else {
+                "full_attention"
+            };
+            layers.push(EncoderLayer::new(cfg, attn_type, vb_l.pp(i))?);
+        }
+
+        let norm = RmsNorm::new(enc_hidden, cfg.rms_norm_eps, vb.pp("norm"))?;
+
+        let rotary_emb = RotaryEmbedding::new(
+            vb.dtype(),
+            cfg.head_dim,
+            cfg.max_position_embeddings,
+            cfg.rope_theta,
+            vb.device(),
+        )?;
+
+        let sliding_window = if cfg.use_sliding_window {
+            Some(cfg.sliding_window)
+        } else {
+            None
+        };
+
+        Ok(Self {
+            embed_tokens,
+            layers,
+            norm,
+            rotary_emb,
+            sliding_window,
+            dtype: vb.dtype(),
+            device: vb.device().clone(),
+        })
+    }
+
+    /// Encode lyric embeddings.
+    ///
+    /// * `inputs_embeds` - `(B, T, text_hidden_dim)` pre-embedded lyric tokens.
+    /// * `attention_mask` - `(B, T)` with `1` for valid tokens, `0` for padding.
+    ///
+    /// Returns `(B, T, encoder_hidden_size)`.
+    pub fn forward(&self, inputs_embeds: &Tensor, attention_mask: &Tensor) -> Result<Tensor> {
+        let (_, seq_len, _) = inputs_embeds.dims3()?;
+
+        let mut hidden = self.embed_tokens.forward(inputs_embeds)?;
+
+        for layer in &self.layers {
+            let sw = match layer.attention_type() {
+                "sliding_attention" => self.sliding_window,
+                _ => None,
+            };
+            let mask =
+                build_bidirectional_mask(attention_mask, seq_len, sw, self.dtype, &self.device)?;
+            hidden = layer.forward(&hidden, Some(&mask), &self.rotary_emb)?;
+        }
+
+        self.norm.forward(&hidden)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// TimbreEncoder
+// ---------------------------------------------------------------------------
+
+/// Bidirectional transformer encoder for timbre (reference audio) embeddings.
+///
+/// Uses a learned CLS token prepended to the input sequence. After processing,
+/// the CLS token representation (position 0) is extracted as the timbre summary.
+///
+/// Weight prefix: `timbre_encoder.*`
+#[derive(Debug, Clone)]
+pub struct TimbreEncoder {
+    embed_tokens: Linear,
+    layers: Vec<EncoderLayer>,
+    norm: RmsNorm,
+    rotary_emb: RotaryEmbedding,
+    /// CLS token prepended to timbre sequence (XL models only).
+    special_token: Tensor,
+    /// Whether to prepend the CLS special_token (true for XL, false for 2B).
+    use_cls: bool,
+    sliding_window: Option<usize>,
+    dtype: DType,
+    device: Device,
+}
+
+impl TimbreEncoder {
+    pub fn new(cfg: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let enc_hidden = cfg.encoder_hidden_size();
+
+        // Linear projection from timbre embedding dim to encoder hidden size (with bias).
+        let embed_tokens = linear_b(
+            cfg.timbre_hidden_dim,
+            enc_hidden,
+            true,
+            vb.pp("embed_tokens"),
+        )?;
+
+        let num_layers = cfg.num_timbre_encoder_hidden_layers;
+        let mut layers = Vec::with_capacity(num_layers);
+        let vb_l = vb.pp("layers");
+        for i in 0..num_layers {
+            let attn_type = if i % 2 == 0 {
+                "sliding_attention"
+            } else {
+                "full_attention"
+            };
+            layers.push(EncoderLayer::new(cfg, attn_type, vb_l.pp(i))?);
+        }
+
+        let norm = RmsNorm::new(enc_hidden, cfg.rms_norm_eps, vb.pp("norm"))?;
+
+        // XL models prepend a CLS special token to the timbre sequence.
+        // Base/2B models have this commented out in Python.
+        let special_token = vb.get((1, 1, enc_hidden), "special_token")?;
+
+        let rotary_emb = RotaryEmbedding::new(
+            vb.dtype(),
+            cfg.head_dim,
+            cfg.max_position_embeddings,
+            cfg.rope_theta,
+            vb.device(),
+        )?;
+
+        let sliding_window = if cfg.use_sliding_window {
+            Some(cfg.sliding_window)
+        } else {
+            None
+        };
+
+        // XL models have encoder_hidden_size != hidden_size and use the CLS token.
+        let use_cls = cfg.encoder_hidden_size.is_some();
+
+        Ok(Self {
+            embed_tokens,
+            layers,
+            norm,
+            rotary_emb,
+            special_token,
+            use_cls,
+            sliding_window,
+            dtype: vb.dtype(),
+            device: vb.device().clone(),
+        })
+    }
+
+    /// Unpack packed timbre embeddings into per-batch format.
+    ///
+    /// Matches the Python `unpack_timbre_embeddings` exactly:
+    /// * `timbre_embs_packed` - `(N, D)` embeddings from each reference audio segment.
+    /// * `refer_audio_order_mask` - `(N,)` 1D integer tensor mapping each segment
+    ///   to its batch index (0..B-1).
+    ///
+    /// Returns `(timbre_embs, timbre_mask)`:
+    /// - `timbre_embs`: `(B, max_count, D)` gathered embeddings.
+    /// - `timbre_mask`: `(B, max_count)` integer mask (1=valid, 0=padding).
+    fn unpack_timbre_embeddings(
+        &self,
+        timbre_embs_packed: &Tensor,
+        refer_audio_order_mask: &Tensor,
+    ) -> Result<(Tensor, Tensor)> {
+        let (n, d) = timbre_embs_packed.dims2()?;
+        let mask_1d = refer_audio_order_mask.flatten_all()?;
+
+        // B = max(order_mask) + 1
+        let b = (mask_1d.max(0)?.to_scalar::<i64>()? + 1) as usize;
+
+        // Count elements per batch and find max_count
+        let mask_vec = mask_1d.to_vec1::<i64>()?;
+        let mut counts = vec![0usize; b];
+        for &idx in &mask_vec {
+            counts[idx as usize] += 1;
+        }
+        let max_count = *counts.iter().max().unwrap_or(&1);
+
+        // Build the output by placing each embedding at the right (batch, position)
+        let mut position_in_batch = vec![0usize; b];
+        let mut result = vec![0f32; b * max_count * d];
+        let mut mask_out = vec![0f32; b * max_count];
+
+        let packed_data = timbre_embs_packed.to_dtype(DType::F32)?.to_vec2::<f32>()?;
+        for i in 0..n {
+            let batch_idx = mask_vec[i] as usize;
+            let pos = position_in_batch[batch_idx];
+            position_in_batch[batch_idx] += 1;
+            let dst_offset = (batch_idx * max_count + pos) * d;
+            result[dst_offset..dst_offset + d].copy_from_slice(&packed_data[i]);
+            mask_out[batch_idx * max_count + pos] = 1.0;
+        }
+
+        let timbre_embs =
+            Tensor::from_vec(result, (b, max_count, d), &self.device)?.to_dtype(self.dtype)?;
+        let timbre_mask =
+            Tensor::from_vec(mask_out, (b, max_count), &self.device)?.to_dtype(self.dtype)?;
+
+        Ok((timbre_embs, timbre_mask))
+    }
+
+    /// Encode timbre reference audio segments.
+    ///
+    /// * `refer_audio_packed` - `(N, T, timbre_hidden_dim)` packed reference
+    ///   audio features for all segments across the batch.
+    /// * `refer_audio_order_mask` - `(N,)` 1D integer tensor mapping each
+    ///   segment to its batch index (0..B-1).
+    ///
+    /// Returns `(timbre_embs, timbre_mask)`:
+    /// - `timbre_embs`: `(B, max_count, encoder_hidden_size)`.
+    /// - `timbre_mask`: `(B, max_count)` integer mask.
+    pub fn forward(
+        &self,
+        refer_audio_packed: &Tensor,
+        refer_audio_order_mask: &Tensor,
+    ) -> Result<(Tensor, Tensor)> {
+        let (n, _t, _) = refer_audio_packed.dims3()?;
+
+        // Project input features: (N, T, timbre_dim) -> (N, T, hidden_size)
+        let mut x = self.embed_tokens.forward(refer_audio_packed)?;
+
+        // XL models prepend a CLS special_token for timbre aggregation.
+        if self.use_cls {
+            let cls = self
+                .special_token
+                .broadcast_as((n, 1, x.dim(2)?))?
+                .contiguous()?;
+            x = Tensor::cat(&[&cls, &x], 1)?;
+        }
+
+        let seq_len = x.dim(1)?;
+
+        // Build attention mask
+        let ones = Tensor::ones((n, seq_len), self.dtype, &self.device)?;
+
+        for layer in &self.layers {
+            let sw = match layer.attention_type() {
+                "sliding_attention" => self.sliding_window,
+                _ => None,
+            };
+            let mask = build_bidirectional_mask(&ones, seq_len, sw, self.dtype, &self.device)?;
+            x = layer.forward(&x, Some(&mask), &self.rotary_emb)?;
+        }
+
+        x = self.norm.forward(&x)?;
+
+        // Extract position 0 as the timbre summary: (N, T, D) -> (N, D)
+        let embeds = x.narrow(1, 0, 1)?.squeeze(1)?;
+
+        self.unpack_timbre_embeddings(&embeds, refer_audio_order_mask)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// pack_sequences
+// ---------------------------------------------------------------------------
+
+/// Pack two (hidden, mask) pairs: concatenate, then reorder so valid tokens
+/// (mask=1) come before padding tokens (mask=0) in each batch element.
+///
+/// Uses a CPU-side stable gather (valid tokens first in original order, then
+/// padding tokens) to match Python's `argsort(descending=True, stable=True)`.
+/// Length computation uses I64 to avoid BF16 precision loss for long sequences.
+///
+/// * `hidden1` - `(B, L1, D)`
+/// * `hidden2` - `(B, L2, D)`
+/// * `mask1` - `(B, L1)` float with 1.0=valid, 0.0=padding
+/// * `mask2` - `(B, L2)` float
+///
+/// Returns `(hidden, mask)` with shapes `(B, L1+L2, D)` and `(B, L1+L2)`.
+pub fn pack_sequences(
+    hidden1: &Tensor,
+    hidden2: &Tensor,
+    mask1: &Tensor,
+    mask2: &Tensor,
+) -> Result<(Tensor, Tensor)> {
+    let hidden_cat = Tensor::cat(&[hidden1, hidden2], 1)?;
+    let mask_cat = Tensor::cat(&[mask1, mask2], 1)?;
+
+    let (b, l, d) = hidden_cat.dims3()?;
+
+    // Build stable sort indices on CPU: valid positions first (in order),
+    // then padding positions (in order). This is equivalent to Python's
+    // stable descending argsort on a binary mask.
+    let mask_f32 = mask_cat
+        .to_dtype(candle_core::DType::F32)?
+        .to_device(&candle_core::Device::Cpu)?;
+    let mask_data = mask_f32.to_vec2::<f32>()?;
+
+    let mut sort_indices = vec![0u32; b * l];
+    let mut valid_counts = vec![0usize; b];
+    for bi in 0..b {
+        let mut valid = Vec::new();
+        let mut invalid = Vec::new();
+        for (j, &val) in mask_data[bi].iter().enumerate().take(l) {
+            if val > 0.5 {
+                valid.push(j as u32);
+            } else {
+                invalid.push(j as u32);
+            }
+        }
+        valid_counts[bi] = valid.len();
+        let offset = bi * l;
+        for (k, &idx) in valid.iter().chain(invalid.iter()).enumerate() {
+            sort_indices[offset + k] = idx;
+        }
+    }
+
+    let sort_idx = Tensor::from_vec(sort_indices, (b, l), hidden_cat.device())?;
+    let idx_expanded = sort_idx
+        .unsqueeze(2)?
+        .broadcast_as((b, l, d))?
+        .contiguous()?;
+    let hidden_sorted = hidden_cat.gather(&idx_expanded, 1)?;
+
+    // Build new mask using integer counts (avoids BF16 precision loss)
+    let mut new_mask_data = vec![0f32; b * l];
+    for bi in 0..b {
+        for j in 0..valid_counts[bi] {
+            new_mask_data[bi * l + j] = 1.0;
+        }
+    }
+    let new_mask =
+        Tensor::from_vec(new_mask_data, (b, l), hidden_cat.device())?.to_dtype(mask1.dtype())?;
+
+    Ok((hidden_sorted, new_mask))
+}
+
+// ---------------------------------------------------------------------------
+// ConditionEncoder
+// ---------------------------------------------------------------------------
+
+/// Top-level condition encoder that produces the cross-attention context for
+/// the DiT diffusion backbone.
+///
+/// Combines three conditioning signals:
+/// 1. **Text** - style/prompt hidden states projected to encoder dimension.
+/// 2. **Lyrics** - lyric token embeddings processed by a bidirectional
+///    transformer.
+/// 3. **Timbre** - reference audio segments encoded with CLS-token pooling.
+///
+/// The outputs are packed into a single sequence with an accompanying mask.
+///
+/// Weight prefix: `encoder.*`
+#[derive(Debug, Clone)]
+pub struct ConditionEncoder {
+    text_projector: Linear,
+    lyric_encoder: LyricEncoder,
+    timbre_encoder: TimbreEncoder,
+}
+
+impl ConditionEncoder {
+    pub fn new(cfg: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let enc_hidden = cfg.encoder_hidden_size();
+
+        // Text projector: no bias
+        let text_projector =
+            linear_no_bias(cfg.text_hidden_dim, enc_hidden, vb.pp("text_projector"))?;
+
+        let lyric_encoder = LyricEncoder::new(cfg, vb.pp("lyric_encoder"))?;
+        let timbre_encoder = TimbreEncoder::new(cfg, vb.pp("timbre_encoder"))?;
+
+        Ok(Self {
+            text_projector,
+            lyric_encoder,
+            timbre_encoder,
+        })
+    }
+
+    /// Encode all conditioning inputs into a single packed sequence.
+    ///
+    /// * `text_hidden_states` - `(B, T_text, text_hidden_dim)` from the text encoder.
+    /// * `text_attention_mask` - `(B, T_text)` mask for text tokens.
+    /// * `lyric_hidden_states` - `(B, T_lyric, text_hidden_dim)` pre-embedded lyrics.
+    /// * `lyric_attention_mask` - `(B, T_lyric)` mask for lyric tokens.
+    /// * `refer_audio_packed` - `(N, T_audio, timbre_hidden_dim)` packed reference
+    ///   audio features.
+    /// * `refer_audio_order_mask` - `(N,)` 1D integer tensor mapping each
+    ///   segment to its batch index.
+    ///
+    /// Returns `(enc_hidden, enc_mask)`:
+    /// - `enc_hidden`: `(B, T_lyric + max_ref + T_text, encoder_hidden_size)`.
+    /// - `enc_mask`: `(B, T_lyric + max_ref + T_text)` float mask.
+    pub fn forward(
+        &self,
+        text_hidden_states: &Tensor,
+        text_attention_mask: &Tensor,
+        lyric_hidden_states: &Tensor,
+        lyric_attention_mask: &Tensor,
+        refer_audio_packed: &Tensor,
+        refer_audio_order_mask: &Tensor,
+    ) -> Result<(Tensor, Tensor)> {
+        // Project text hidden states
+        let text = self.text_projector.forward(text_hidden_states)?;
+
+        // Encode lyrics
+        let lyric = self
+            .lyric_encoder
+            .forward(lyric_hidden_states, lyric_attention_mask)?;
+
+        // Encode timbre
+        let (timbre, timbre_mask) = self
+            .timbre_encoder
+            .forward(refer_audio_packed, refer_audio_order_mask)?;
+
+        // Pack lyric + timbre
+        let (enc_hidden, enc_mask) =
+            pack_sequences(&lyric, &timbre, lyric_attention_mask, &timbre_mask)?;
+
+        // Pack (lyric+timbre) + text
+        let (enc_hidden, enc_mask) =
+            pack_sequences(&enc_hidden, &text, &enc_mask, text_attention_mask)?;
+
+        Ok((enc_hidden, enc_mask))
+    }
+}
@@ -0,0 +1,358 @@
+//! ACE-Step 1.5 music generation model.
+//!
+//! Vendored from the candle port in huggingface/candle#3451
+//! (n0madic/candle@294086f4, MIT OR Apache-2.0), which matches the Python
+//! reference to about 2e-5. Changes here: crate paths, a local copy of
+//! candle's `qwen3` with the two accessors the port needs (`embed_tokens`,
+//! a public `clear_kv_cache`), and lint fixes. Drop the copy once the PR
+//! lands upstream and the candle pin moves past it.
+//!
+//! ACE-Step is a music generation foundation model based on a diffusion transformer (DiT)
+//! architecture with text conditioning for lyrics and style prompts. Supports two
+//! inference modes:
+//!
+//! - **DiT-only**: text encoder → DiT denoising → VAE decode (base/sft models)
+//! - **LM+DiT**: Qwen3 LM generates metadata + audio codes → DiT conditioned on
+//!   codes → VAE decode (turbo models, higher quality)
+//!
+//! ## Modules
+//!
+//! - [`pipeline`] — `AceStepPipeline` high-level API (text2music, cover)
+//! - [`model`] — `AceStepConditionGenerationModel` (condition encoder + DiT + tokenizer)
+//! - [`lm`] — `LmPipeline` for Qwen3-based audio code generation
+//! - [`dit`] — `AceStepDiTModel` diffusion transformer
+//! - [`condition`] — `ConditionEncoder` (text/lyric/timbre encoding)
+//! - [`tokenizer`] — `AudioTokenizer` / `AudioTokenDetokenizer` / `ResidualFSQ`
+//! - [`sampling`] — timestep schedules, Euler ODE/SDE, APG guidance
+//! - [`vae`] — `AutoencoderOobleck` / `OobleckDecoder` for audio ↔ latent
+//!
+//! - [HuggingFace Models](https://huggingface.co/ACE-Step)
+//! - [GitHub Repository](https://github.com/ACE-Step/ACE-Step-1.5)
+
+pub mod condition;
+pub mod dit;
+pub mod lm;
+pub mod model;
+pub mod pipeline;
+pub mod qwen3;
+pub mod sampling;
+pub mod tokenizer;
+pub mod vae;
+
+use candle_nn::Activation;
+
+// -- AceStepConfig default value functions -----------------------------------
+
+fn default_hidden_size() -> usize {
+    2048
+}
+fn default_intermediate_size() -> usize {
+    6144
+}
+fn default_num_hidden_layers() -> usize {
+    24
+}
+fn default_num_attention_heads() -> usize {
+    16
+}
+fn default_num_key_value_heads() -> usize {
+    8
+}
+fn default_head_dim() -> usize {
+    128
+}
+fn default_hidden_act() -> Activation {
+    Activation::Silu
+}
+fn default_rms_norm_eps() -> f64 {
+    1e-6
+}
+fn default_rope_theta() -> f64 {
+    1_000_000.0
+}
+fn default_use_sliding_window() -> bool {
+    true
+}
+fn default_sliding_window() -> usize {
+    128
+}
+fn default_in_channels() -> usize {
+    192
+}
+fn default_patch_size() -> usize {
+    2
+}
+fn default_text_hidden_dim() -> usize {
+    1024
+}
+fn default_audio_acoustic_hidden_dim() -> usize {
+    64
+}
+fn default_timbre_hidden_dim() -> usize {
+    64
+}
+fn default_pool_window_size() -> usize {
+    5
+}
+fn default_vocab_size() -> usize {
+    64003
+}
+fn default_max_position_embeddings() -> usize {
+    32768
+}
+fn default_fsq_dim() -> usize {
+    2048
+}
+fn default_fsq_input_levels() -> Vec<usize> {
+    vec![8, 8, 8, 5, 5, 5]
+}
+fn default_fsq_input_num_quantizers() -> usize {
+    1
+}
+fn default_attention_bias() -> bool {
+    false
+}
+fn default_attention_dropout() -> f64 {
+    0.0
+}
+fn default_num_lyric_encoder_hidden_layers() -> usize {
+    8
+}
+fn default_num_timbre_encoder_hidden_layers() -> usize {
+    4
+}
+fn default_num_attention_pooler_hidden_layers() -> usize {
+    2
+}
+fn default_num_audio_decoder_hidden_layers() -> usize {
+    24
+}
+fn default_layer_types() -> Vec<String> {
+    (0..24)
+        .map(|i| {
+            if i % 2 == 0 {
+                "sliding_attention".to_string()
+            } else {
+                "full_attention".to_string()
+            }
+        })
+        .collect()
+}
+
+/// Main configuration for the ACE-Step 1.5 DiT-based music generation model.
+///
+/// Deserializable from HuggingFace `config.json`. Supports both the base (3.5B)
+/// and XL variants via optional encoder dimension overrides.
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct AceStepConfig {
+    // -- DiT dimensions ------------------------------------------------------
+    #[serde(default = "default_hidden_size")]
+    pub hidden_size: usize,
+    #[serde(default = "default_intermediate_size")]
+    pub intermediate_size: usize,
+    #[serde(default = "default_num_hidden_layers")]
+    pub num_hidden_layers: usize,
+    #[serde(default = "default_num_attention_heads")]
+    pub num_attention_heads: usize,
+    #[serde(default = "default_num_key_value_heads")]
+    pub num_key_value_heads: usize,
+    #[serde(default = "default_head_dim")]
+    pub head_dim: usize,
+
+    // -- Encoder dimensions (XL only, optional) ------------------------------
+    #[serde(default)]
+    pub encoder_hidden_size: Option<usize>,
+    #[serde(default)]
+    pub encoder_intermediate_size: Option<usize>,
+    #[serde(default)]
+    pub encoder_num_attention_heads: Option<usize>,
+    #[serde(default)]
+    pub encoder_num_key_value_heads: Option<usize>,
+
+    // -- Common parameters ---------------------------------------------------
+    #[serde(default = "default_hidden_act")]
+    pub hidden_act: Activation,
+    #[serde(default = "default_rms_norm_eps")]
+    pub rms_norm_eps: f64,
+    #[serde(default = "default_rope_theta")]
+    pub rope_theta: f64,
+    #[serde(default = "default_use_sliding_window")]
+    pub use_sliding_window: bool,
+    #[serde(default = "default_sliding_window")]
+    pub sliding_window: usize,
+    #[serde(default = "default_in_channels")]
+    pub in_channels: usize,
+    #[serde(default = "default_patch_size")]
+    pub patch_size: usize,
+    #[serde(default = "default_text_hidden_dim")]
+    pub text_hidden_dim: usize,
+    #[serde(default = "default_audio_acoustic_hidden_dim")]
+    pub audio_acoustic_hidden_dim: usize,
+    #[serde(default = "default_timbre_hidden_dim")]
+    pub timbre_hidden_dim: usize,
+    #[serde(default = "default_pool_window_size")]
+    pub pool_window_size: usize,
+    #[serde(default = "default_vocab_size")]
+    pub vocab_size: usize,
+    #[serde(default = "default_max_position_embeddings")]
+    pub max_position_embeddings: usize,
+    #[serde(default = "default_fsq_dim")]
+    pub fsq_dim: usize,
+    #[serde(default = "default_fsq_input_levels")]
+    pub fsq_input_levels: Vec<usize>,
+    #[serde(default = "default_fsq_input_num_quantizers")]
+    pub fsq_input_num_quantizers: usize,
+    #[serde(default = "default_attention_bias")]
+    pub attention_bias: bool,
+    #[serde(default = "default_attention_dropout")]
+    pub attention_dropout: f64,
+
+    // -- Sub-model layer counts ----------------------------------------------
+    #[serde(default = "default_num_lyric_encoder_hidden_layers")]
+    pub num_lyric_encoder_hidden_layers: usize,
+    #[serde(default = "default_num_timbre_encoder_hidden_layers")]
+    pub num_timbre_encoder_hidden_layers: usize,
+    #[serde(default = "default_num_attention_pooler_hidden_layers")]
+    pub num_attention_pooler_hidden_layers: usize,
+    #[serde(default = "default_num_audio_decoder_hidden_layers")]
+    pub num_audio_decoder_hidden_layers: usize,
+
+    // -- Layer type schedule -------------------------------------------------
+    #[serde(default = "default_layer_types")]
+    pub layer_types: Vec<String>,
+
+    // -- Turbo flag ----------------------------------------------------------
+    #[serde(default)]
+    pub is_turbo: bool,
+}
+
+impl AceStepConfig {
+    /// Returns the encoder hidden size, falling back to the main `hidden_size`.
+    pub fn encoder_hidden_size(&self) -> usize {
+        self.encoder_hidden_size.unwrap_or(self.hidden_size)
+    }
+
+    /// Returns the encoder intermediate size, falling back to the main `intermediate_size`.
+    pub fn encoder_intermediate_size(&self) -> usize {
+        self.encoder_intermediate_size
+            .unwrap_or(self.intermediate_size)
+    }
+
+    /// Returns the encoder attention head count, falling back to `num_attention_heads`.
+    pub fn encoder_num_attention_heads(&self) -> usize {
+        self.encoder_num_attention_heads
+            .unwrap_or(self.num_attention_heads)
+    }
+
+    /// Returns the encoder key-value head count, falling back to `num_key_value_heads`.
+    pub fn encoder_num_key_value_heads(&self) -> usize {
+        self.encoder_num_key_value_heads
+            .unwrap_or(self.num_key_value_heads)
+    }
+}
+
+impl Default for AceStepConfig {
+    fn default() -> Self {
+        Self {
+            hidden_size: 2048,
+            intermediate_size: 6144,
+            num_hidden_layers: 24,
+            num_attention_heads: 16,
+            num_key_value_heads: 8,
+            head_dim: 128,
+            encoder_hidden_size: None,
+            encoder_intermediate_size: None,
+            encoder_num_attention_heads: None,
+            encoder_num_key_value_heads: None,
+            hidden_act: Activation::Silu,
+            rms_norm_eps: 1e-6,
+            rope_theta: 1_000_000.0,
+            use_sliding_window: true,
+            sliding_window: 128,
+            in_channels: 192,
+            patch_size: 2,
+            text_hidden_dim: 1024,
+            audio_acoustic_hidden_dim: 64,
+            timbre_hidden_dim: 64,
+            pool_window_size: 5,
+            vocab_size: 64003,
+            max_position_embeddings: 32768,
+            fsq_dim: 2048,
+            fsq_input_levels: vec![8, 8, 8, 5, 5, 5],
+            fsq_input_num_quantizers: 1,
+            attention_bias: false,
+            attention_dropout: 0.0,
+            num_lyric_encoder_hidden_layers: 8,
+            num_timbre_encoder_hidden_layers: 4,
+            num_attention_pooler_hidden_layers: 2,
+            num_audio_decoder_hidden_layers: 24,
+            layer_types: default_layer_types(),
+            is_turbo: false,
+        }
+    }
+}
+
+// -- VaeConfig default value functions ---------------------------------------
+
+fn default_vae_encoder_hidden_size() -> usize {
+    128
+}
+fn default_vae_decoder_channels() -> usize {
+    128
+}
+fn default_vae_decoder_input_channels() -> usize {
+    64
+}
+fn default_vae_audio_channels() -> usize {
+    2
+}
+fn default_vae_downsampling_ratios() -> Vec<usize> {
+    vec![2, 4, 4, 6, 10]
+}
+fn default_vae_channel_multiples() -> Vec<usize> {
+    vec![1, 2, 4, 8, 16]
+}
+fn default_vae_sampling_rate() -> usize {
+    48000
+}
+
+/// Configuration for the ACE-Step VAE (variational autoencoder) that converts
+/// between waveform audio and the latent space consumed by the DiT.
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct VaeConfig {
+    #[serde(default = "default_vae_encoder_hidden_size")]
+    pub encoder_hidden_size: usize,
+    #[serde(default = "default_vae_decoder_channels")]
+    pub decoder_channels: usize,
+    #[serde(default = "default_vae_decoder_input_channels")]
+    pub decoder_input_channels: usize,
+    #[serde(default = "default_vae_audio_channels")]
+    pub audio_channels: usize,
+    #[serde(default = "default_vae_downsampling_ratios")]
+    pub downsampling_ratios: Vec<usize>,
+    #[serde(default = "default_vae_channel_multiples")]
+    pub channel_multiples: Vec<usize>,
+    #[serde(default = "default_vae_sampling_rate")]
+    pub sampling_rate: usize,
+}
+
+impl VaeConfig {
+    /// Computes the total hop length as the product of all downsampling ratios.
+    pub fn hop_length(&self) -> usize {
+        self.downsampling_ratios.iter().product()
+    }
+}
+
+impl Default for VaeConfig {
+    fn default() -> Self {
+        Self {
+            encoder_hidden_size: 128,
+            decoder_channels: 128,
+            decoder_input_channels: 64,
+            audio_channels: 2,
+            downsampling_ratios: vec![2, 4, 4, 6, 10],
+            channel_multiples: vec![1, 2, 4, 8, 16],
+            sampling_rate: 48000,
+        }
+    }
+}
@@ -0,0 +1,520 @@
+//! Top-level ACE-Step conditional generation model.
+//!
+//! Mirrors the subset of Python's `AceStepConditionGenerationModel` needed for
+//! **single-sample DiT-only Text2Music** with base/sft models. Owns the
+//! condition encoder, DiT decoder, audio tokenizer/detokenizer, and
+//! null-condition embedding. Provides `prepare_condition` and
+//! `generate_latents` as the primary API, operating in latent space.
+//!
+//! ## Current scope vs Python
+//!
+//! Implemented:
+//! - Text2Music with base/sft (continuous schedule + APG guidance)
+//! - Text2Music with turbo (discrete schedule, no CFG)
+//! - Real tokenizer masks, encoder_attention_mask propagation
+//! - Cover mode (`is_covers`, tokenize/detokenize round-trip in prepare_condition)
+//! - `precomputed_lm_hints_25Hz` / `audio_codes` (LM-generated audio codes)
+//! - Audio cover strength blending (switch condition mid-loop with KV cache reset)
+//! - Custom timestep schedules for turbo models (mapped to VALID_TIMESTEPS)
+//!
+//! Not yet implemented (Python `generate_audio` features):
+//! - Batched inference (batch > 1)
+//! - LoRA support
+
+use candle_core::{DType, Device, Result, Tensor};
+use candle_nn::VarBuilder;
+
+use super::AceStepConfig;
+use super::condition::ConditionEncoder;
+use super::dit::AceStepDiTModel;
+use super::sampling;
+use super::tokenizer::{AudioTokenDetokenizer, AudioTokenizer};
+
+/// Inputs for the condition encoder.
+pub struct ConditionInputs<'a> {
+    /// Text encoder hidden states `(B, T_text, text_hidden_dim)`.
+    pub text_hidden_states: &'a Tensor,
+    /// Text attention mask `(B, T_text)`.
+    pub text_attention_mask: &'a Tensor,
+    /// Lyric token embeddings `(B, T_lyric, text_hidden_dim)`.
+    pub lyric_hidden_states: &'a Tensor,
+    /// Lyric attention mask `(B, T_lyric)`.
+    pub lyric_attention_mask: &'a Tensor,
+    /// Packed reference audio features `(N, T_audio, timbre_hidden_dim)`.
+    pub refer_audio_packed: &'a Tensor,
+    /// 1D order mask mapping segments to batch indices `(N,)`.
+    pub refer_audio_order_mask: &'a Tensor,
+    /// Source latents `(B, T, acoustic_dim)` — silence for text2music.
+    pub src_latents: &'a Tensor,
+    /// Chunk masks `(B, T, acoustic_dim)` — ones for full generation.
+    pub chunk_masks: &'a Tensor,
+    /// Silence latent for tokenizer padding `(1, T_max, acoustic_dim)`.
+    pub silence_latent: &'a Tensor,
+    /// Boolean tensor `(B,)` — true for cover mode items.
+    pub is_covers: &'a Tensor,
+    /// Precomputed LM hints at 25Hz `(B, T, acoustic_dim)`. Skips tokenize/detokenize.
+    pub precomputed_lm_hints_25hz: Option<&'a Tensor>,
+    /// Audio codes from LM `(B, N, 1)`. Skips tokenize, uses quantizer → detokenize.
+    pub audio_codes: Option<&'a Tensor>,
+    /// Text hidden states for the non-cover condition (used when `audio_cover_strength < 1.0`).
+    pub non_cover_text_hidden_states: Option<&'a Tensor>,
+    /// Attention mask for non-cover text.
+    pub non_cover_text_attention_mask: Option<&'a Tensor>,
+}
+
+/// Prepared conditioning for the DiT decoder.
+pub struct PreparedCondition {
+    pub encoder_hidden_states: Tensor,
+    pub encoder_attention_mask: Tensor,
+    pub context_latents: Tensor,
+}
+
+/// Options for latent generation.
+pub struct GenerateOptions {
+    pub seed: Option<u64>,
+    pub num_steps: usize,
+    pub guidance_scale: f64,
+    pub shift: f64,
+    pub is_turbo: bool,
+    /// ODE (deterministic) or SDE (stochastic) sampling.
+    pub infer_method: sampling::InferMethod,
+    /// Timestep range in which CFG/APG guidance is applied.
+    /// Outside this range, only the conditional prediction is used.
+    /// Default: `(0.0, 1.0)` — apply guidance at all timesteps.
+    pub cfg_interval: (f64, f64),
+    /// Custom timestep schedule for turbo models (1–20 values).
+    /// Each value is mapped to the nearest valid timestep. When `None`,
+    /// the shift-based schedule is used.
+    pub timesteps: Option<Vec<f64>>,
+    /// Fraction of denoising steps that use the cover condition (0.0–1.0).
+    /// At `cover_steps = round(num_steps * audio_cover_strength)` the loop
+    /// switches to the non-cover condition and clears the KV cache.
+    /// Default: `1.0` (no blending — cover condition used for all steps).
+    pub audio_cover_strength: f64,
+    /// Starting noise `(1, frames, acoustic_dim)`, replacing the seeded
+    /// draw. Lets a parity test start from the reference's saved noise.
+    pub initial_noise: Option<Tensor>,
+    /// Wavelet correction after each turbo step; on by default, as upstream.
+    pub dcw: Option<sampling::Dcw>,
+}
+
+impl Default for GenerateOptions {
+    fn default() -> Self {
+        Self {
+            seed: None,
+            num_steps: 50,
+            guidance_scale: 5.0,
+            shift: 1.0,
+            is_turbo: false,
+            infer_method: sampling::InferMethod::Ode,
+            cfg_interval: (0.0, 1.0),
+            timesteps: None,
+            audio_cover_strength: 1.0,
+            initial_noise: None,
+            dcw: Some(sampling::Dcw::default()),
+        }
+    }
+}
+
+/// Top-level ACE-Step model that owns all sub-modules and provides
+/// `prepare_condition` + `generate_latents` as the main API.
+pub struct AceStepConditionGenerationModel {
+    config: AceStepConfig,
+    encoder: ConditionEncoder,
+    decoder: AceStepDiTModel,
+    pub tokenizer: AudioTokenizer,
+    pub detokenizer: AudioTokenDetokenizer,
+    null_condition_emb: Tensor,
+}
+
+impl AceStepConditionGenerationModel {
+    /// Load all sub-modules from a single safetensors checkpoint.
+    pub fn new(config: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let encoder = ConditionEncoder::new(config, vb.pp("encoder"))?;
+        let decoder = AceStepDiTModel::new(config, vb.pp("decoder"))?;
+        let tokenizer = AudioTokenizer::new(config, vb.pp("tokenizer"))?;
+        let detokenizer = AudioTokenDetokenizer::new(config, vb.pp("detokenizer"))?;
+        // XL models use encoder_hidden_size for null_condition_emb (it operates
+        // in encoder space, not decoder space).
+        let null_condition_emb =
+            vb.get((1, 1, config.encoder_hidden_size()), "null_condition_emb")?;
+
+        Ok(Self {
+            config: config.clone(),
+            encoder,
+            decoder,
+            tokenizer,
+            detokenizer,
+            null_condition_emb,
+        })
+    }
+
+    /// Tokenize audio latents: pad to `pool_window_size` multiple, reshape, quantize.
+    ///
+    /// Returns `(quantized, indices, downsampled_mask)`.
+    fn tokenize(
+        &self,
+        hidden_states: &Tensor,
+        silence_latent: &Tensor,
+        attention_mask: &Tensor,
+    ) -> Result<(Tensor, Tensor, Tensor)> {
+        let (b, t, d) = hidden_states.dims3()?;
+        let pws = self.tokenizer.pool_window_size();
+
+        // Pad to multiple of pool_window_size using silence_latent
+        let (hidden_states, attention_mask) = if t % pws != 0 {
+            let pad_len = pws - (t % pws);
+            let pad = silence_latent
+                .narrow(1, 0, pad_len)?
+                .broadcast_as((b, pad_len, d))?
+                .contiguous()?;
+            let hs = Tensor::cat(&[hidden_states, &pad], 1)?;
+            let mask_pad = Tensor::zeros(
+                (b, pad_len),
+                attention_mask.dtype(),
+                attention_mask.device(),
+            )?;
+            let am = Tensor::cat(&[attention_mask, &mask_pad], 1)?;
+            (hs, am)
+        } else {
+            (hidden_states.clone(), attention_mask.clone())
+        };
+
+        let t_padded = hidden_states.dim(1)?;
+        let num_patches = t_padded / pws;
+
+        // Reshape: (B, T, D) -> (B, T/pws, pws, D)
+        let reshaped = hidden_states.reshape((b, num_patches, pws, d))?;
+
+        // Downsample attention_mask: reshape (B, T) -> (B, T/pws, pws) then max
+        let mask_reshaped = attention_mask.reshape((b, num_patches, pws))?;
+        let downsampled_mask = mask_reshaped.max(2)?;
+
+        let (quantized, indices) = self.tokenizer.forward(&reshaped)?;
+        Ok((quantized, indices, downsampled_mask))
+    }
+
+    /// Detokenize quantized representations back to acoustic space.
+    fn detokenize(&self, quantized: &Tensor) -> Result<Tensor> {
+        self.detokenizer.forward(quantized)
+    }
+
+    /// Prepare conditioning: encode text/lyrics/timbre into packed encoder
+    /// hidden states, and build context latents from src_latents + chunk_masks.
+    ///
+    /// When cover mode is active (`is_covers > 0`), source latents are replaced
+    /// with tokenized-then-detokenized LM hints.
+    pub fn prepare_condition(&self, inputs: &ConditionInputs) -> Result<PreparedCondition> {
+        let (encoder_hidden_states, encoder_attention_mask) = self.encoder.forward(
+            inputs.text_hidden_states,
+            inputs.text_attention_mask,
+            inputs.lyric_hidden_states,
+            inputs.lyric_attention_mask,
+            inputs.refer_audio_packed,
+            inputs.refer_audio_order_mask,
+        )?;
+
+        // Cover mode: compute LM hints and replace src_latents where is_covers > 0.
+        // Skip entirely when is_covers is all-false (text2music) — no tokenize/detokenize needed.
+        let latent_len = inputs.src_latents.dim(1)?;
+
+        let any_covers = inputs
+            .is_covers
+            .gt(0u8)?
+            .to_dtype(candle_core::DType::U32)?
+            .sum_all()?
+            .to_scalar::<u32>()?
+            > 0;
+
+        let src_latents = if any_covers {
+            let lm_hints_25hz = if let Some(precomputed) = inputs.precomputed_lm_hints_25hz {
+                precomputed.narrow(1, 0, latent_len)?
+            } else if let Some(audio_codes) = inputs.audio_codes {
+                let quantized = self.tokenizer.get_output_from_indices(audio_codes)?;
+                let detokenized = self.detokenize(&quantized)?;
+                let det_len = detokenized.dim(1)?;
+                if det_len >= latent_len {
+                    detokenized.narrow(1, 0, latent_len)?
+                } else {
+                    let pad = inputs.silence_latent.narrow(1, 0, latent_len - det_len)?;
+                    Tensor::cat(&[&detokenized, &pad], 1)?
+                }
+            } else {
+                let b = inputs.src_latents.dim(0)?;
+                let attention_mask = Tensor::ones(
+                    (b, latent_len),
+                    inputs.src_latents.dtype(),
+                    inputs.src_latents.device(),
+                )?;
+                let (quantized, _indices, _mask) =
+                    self.tokenize(inputs.src_latents, inputs.silence_latent, &attention_mask)?;
+                let detokenized = self.detokenize(&quantized)?;
+                detokenized.narrow(1, 0, latent_len)?
+            };
+
+            let is_covers_3d = inputs
+                .is_covers
+                .gt(0u8)?
+                .unsqueeze(1)?
+                .unsqueeze(2)?
+                .broadcast_as(inputs.src_latents.shape())?;
+            is_covers_3d.where_cond(&lm_hints_25hz, inputs.src_latents)?
+        } else {
+            inputs.src_latents.clone()
+        };
+
+        // Context = [src_latents, chunk_masks] along feature dim
+        let context_latents =
+            Tensor::cat(&[&src_latents, inputs.chunk_masks], candle_core::D::Minus1)?;
+
+        Ok(PreparedCondition {
+            encoder_hidden_states,
+            encoder_attention_mask,
+            context_latents,
+        })
+    }
+
+    /// Clear cross-attention KV caches. Call between generations with different
+    /// encoder states (different prompts).
+    pub fn clear_kv_cache(&mut self) {
+        self.decoder.clear_kv_cache();
+    }
+
+    /// Generate audio latents from prepared conditions.
+    ///
+    /// Returns the denoised latent `(B, T, acoustic_dim)` ready for VAE decode.
+    /// Automatically caches cross-attention K/V on the first denoising step and
+    /// reuses them for subsequent steps (encoder states are constant).
+    ///
+    /// When `non_cover_condition` is provided and `opts.audio_cover_strength < 1.0`,
+    /// the loop switches from cover to non-cover condition partway through and
+    /// clears the KV cache at the transition.
+    pub fn generate_latents(
+        &mut self,
+        condition: &PreparedCondition,
+        non_cover_condition: Option<&PreparedCondition>,
+        opts: &GenerateOptions,
+        device: &Device,
+        dtype: DType,
+    ) -> Result<Tensor> {
+        // Clear KV cache so it gets rebuilt for this generation's encoder states.
+        self.decoder.clear_kv_cache();
+
+        let latent_frames = condition.context_latents.dim(1)?;
+        let acoustic_dim = self.config.audio_acoustic_hidden_dim;
+
+        // Drawn on the host from a seeded generator rather than with the
+        // device RNG: candle cannot seed its CPU RNG at all, and a CUDA
+        // draw differs from a CPU one. A host draw makes the same seed give
+        // the same song on every device.
+        let noise = match &opts.initial_noise {
+            Some(n) => n.to_device(device)?.to_dtype(dtype)?,
+            None => seeded_normal(opts.seed, (1, latent_frames, acoustic_dim))?
+                .to_device(device)?
+                .to_dtype(dtype)?,
+        };
+
+        let mut xt = noise;
+
+        if opts.is_turbo {
+            // Turbo: discrete schedule, no CFG, final step = x0 prediction
+            let t_schedule = match &opts.timesteps {
+                Some(ts) => sampling::get_custom_turbo_schedule(ts)?,
+                None => sampling::get_turbo_schedule(opts.shift),
+            };
+            let num_steps = t_schedule.len();
+            let cover_steps = (num_steps as f64 * opts.audio_cover_strength).round() as usize;
+            let mut active = if cover_steps == 0 {
+                if let Some(nc) = non_cover_condition {
+                    nc
+                } else {
+                    condition
+                }
+            } else {
+                condition
+            };
+
+            for step in 0..num_steps {
+                // Switch to non-cover condition at the transition point
+                if let Some(nc) = non_cover_condition
+                    && step == cover_steps
+                    && step > 0
+                {
+                    active = nc;
+                    self.decoder.clear_kv_cache();
+                }
+
+                let t_curr = t_schedule[step];
+                let t_tensor =
+                    Tensor::from_vec(vec![t_curr as f32], (1,), device)?.to_dtype(dtype)?;
+
+                let vt = self.decoder.forward(
+                    &xt,
+                    &t_tensor,
+                    &t_tensor,
+                    None,
+                    &active.encoder_hidden_states,
+                    Some(&active.encoder_attention_mask),
+                    &active.context_latents,
+                )?;
+
+                let xt_before = xt.clone();
+                if step == num_steps - 1 {
+                    xt = sampling::get_x0_from_noise(&xt, &vt, t_curr)?;
+                } else {
+                    let t_next = t_schedule[step + 1];
+                    match opts.infer_method {
+                        sampling::InferMethod::Ode => {
+                            xt = sampling::euler_step(&xt, &vt, t_curr - t_next)?;
+                        }
+                        sampling::InferMethod::Sde => {
+                            xt = sampling::sde_step(&xt, &vt, t_curr, t_next)?;
+                        }
+                    }
+                }
+                // DCW against the clean sample this step predicted,
+                // after every step including the final x0 (as upstream).
+                if let Some(dcw) = &opts.dcw {
+                    let denoised = (&xt_before - (&vt * t_curr)?)?;
+                    xt = dcw.apply(&xt, &denoised, t_curr)?;
+                }
+            }
+        } else {
+            // Base/SFT: continuous schedule with APG guidance
+            let schedule = sampling::get_schedule(opts.num_steps, opts.shift);
+            let mut momentum_buffer = sampling::MomentumBuffer::new(-0.75);
+            let do_cfg = opts.guidance_scale > 1.0;
+            let cover_steps = (opts.num_steps as f64 * opts.audio_cover_strength).round() as usize;
+
+            // Build doubled tensors for CFG from cover condition
+            let build_cfg_tensors =
+                |cond: &PreparedCondition, null_emb: &Tensor| -> Result<(Tensor, Tensor, Tensor)> {
+                    let null_enc = null_emb
+                        .broadcast_as(cond.encoder_hidden_states.shape())?
+                        .contiguous()?;
+                    if do_cfg {
+                        let enc = Tensor::cat(&[&cond.encoder_hidden_states, &null_enc], 0)?;
+                        let mask = Tensor::cat(
+                            &[&cond.encoder_attention_mask, &cond.encoder_attention_mask],
+                            0,
+                        )?;
+                        let ctx = Tensor::cat(&[&cond.context_latents, &cond.context_latents], 0)?;
+                        Ok((enc, mask, ctx))
+                    } else {
+                        Ok((
+                            cond.encoder_hidden_states.clone(),
+                            cond.encoder_attention_mask.clone(),
+                            cond.context_latents.clone(),
+                        ))
+                    }
+                };
+
+            let initial_condition = if cover_steps == 0 {
+                non_cover_condition.unwrap_or(condition)
+            } else {
+                condition
+            };
+            let (mut enc_cond, mut enc_mask_cond, mut ctx_doubled) =
+                build_cfg_tensors(initial_condition, &self.null_condition_emb)?;
+
+            let (cfg_start, cfg_end) = opts.cfg_interval;
+
+            for step in 0..opts.num_steps {
+                // Switch to non-cover condition at the transition point
+                if let Some(nc) = non_cover_condition
+                    && step == cover_steps
+                    && step > 0
+                {
+                    self.decoder.clear_kv_cache();
+                    let rebuilt = build_cfg_tensors(nc, &self.null_condition_emb)?;
+                    enc_cond = rebuilt.0;
+                    enc_mask_cond = rebuilt.1;
+                    ctx_doubled = rebuilt.2;
+                }
+
+                let t_curr = schedule[step];
+                let t_prev = schedule[step + 1];
+
+                let x_input = if do_cfg {
+                    Tensor::cat(&[&xt, &xt], 0)?
+                } else {
+                    xt.clone()
+                };
+                let batch_size = x_input.dim(0)?;
+                let t_tensor =
+                    Tensor::from_vec(vec![t_curr as f32; batch_size], (batch_size,), device)?
+                        .to_dtype(dtype)?;
+
+                let vt = self.decoder.forward(
+                    &x_input,
+                    &t_tensor,
+                    &t_tensor,
+                    None,
+                    &enc_cond,
+                    Some(&enc_mask_cond),
+                    &ctx_doubled,
+                )?;
+
+                let vt = if do_cfg {
+                    let v_cond = vt.narrow(0, 0, 1)?;
+                    let v_uncond = vt.narrow(0, 1, 1)?;
+                    let apply_guidance = t_curr >= cfg_start && t_curr <= cfg_end;
+                    if apply_guidance {
+                        sampling::apg_forward(
+                            &v_cond,
+                            &v_uncond,
+                            opts.guidance_scale,
+                            &mut momentum_buffer,
+                            2.5,
+                        )?
+                    } else {
+                        v_cond
+                    }
+                } else {
+                    vt
+                };
+
+                match opts.infer_method {
+                    sampling::InferMethod::Ode => {
+                        xt = sampling::euler_step(&xt, &vt, t_curr - t_prev)?;
+                    }
+                    sampling::InferMethod::Sde => {
+                        xt = sampling::sde_step(&xt, &vt, t_curr, t_prev)?;
+                    }
+                }
+            }
+        }
+
+        Ok(xt)
+    }
+}
+
+/// Standard-normal noise from a seeded host generator (SplitMix64 +
+/// Box-Muller). An unseeded draw takes its seed from the process's
+/// random hasher keys.
+pub fn seeded_normal(seed: Option<u64>, shape: (usize, usize, usize)) -> Result<Tensor> {
+    let mut state = seed.unwrap_or_else(|| {
+        use std::hash::BuildHasher;
+        std::collections::hash_map::RandomState::new().hash_one(0u8)
+    });
+    let mut next = move || {
+        state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
+        let mut z = state;
+        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
+        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
+        z ^ (z >> 31)
+    };
+    // 53 random bits mapped into (0, 1]: never zero, so ln() is finite.
+    let mut unit = move || ((next() >> 11) as f64 + 1.0) / (1u64 << 53) as f64;
+    let n = shape.0 * shape.1 * shape.2;
+    let mut out = Vec::with_capacity(n + 1);
+    while out.len() < n {
+        let (u1, u2) = (unit(), unit());
+        let r = (-2.0 * u1.ln()).sqrt();
+        let theta = 2.0 * std::f64::consts::PI * u2;
+        out.push((r * theta.cos()) as f32);
+        out.push((r * theta.sin()) as f32);
+    }
+    out.truncate(n);
+    Tensor::from_vec(out, shape, &Device::Cpu)
+}
@@ -0,0 +1,696 @@
+//! High-level ACE-Step pipeline with text encoder inside.
+//!
+//! Vendored with the rest of this module from the candle example in
+//! huggingface/candle#3451 (`candle-examples/examples/ace-step/pipeline.rs`).
+//!
+//! Accepts raw text prompts — handles tokenization, text encoding,
+//! conditioning, DiT denoising, and VAE decoding internally.
+
+use super::qwen3;
+use super::{
+    AceStepConfig, VaeConfig,
+    lm::{LmConfig, LmPipeline, TokenIds},
+    model::{AceStepConditionGenerationModel, ConditionInputs, GenerateOptions},
+    sampling,
+    vae::{AutoencoderOobleck, OobleckDecoder},
+};
+use anyhow::Result;
+use candle_core::{DType, Device, Tensor};
+use candle_nn::VarBuilder;
+
+/// Metadata from LM CoT output, used to enrich the DiT text prompt.
+///
+/// When the LM generates chain-of-thought metadata (BPM, keyscale, etc.),
+/// these values should be fed back into the DiT caption — matching the Python
+/// `_dict_to_meta_string` / `_extract_caption_and_language` flow.
+#[derive(Default)]
+pub struct LmMetadata {
+    pub bpm: Option<String>,
+    pub timesignature: Option<String>,
+    pub keyscale: Option<String>,
+    /// LM-generated caption — replaces the user prompt in the DiT text input.
+    pub caption: Option<String>,
+    /// Language code for lyrics (default: "en").
+    pub language: Option<String>,
+}
+
+/// Generation parameters.
+pub struct GenerationParams {
+    pub duration_secs: f64,
+    pub seed: Option<u64>,
+    pub num_steps: Option<usize>,
+    pub guidance_scale: f64,
+    pub shift: f64,
+    pub is_turbo: bool,
+    /// Precomputed LM hints at 25Hz `(1, T, acoustic_dim)`.
+    pub lm_hints_25hz: Option<Tensor>,
+    /// LM-generated metadata to enrich the DiT text prompt.
+    pub lm_metadata: Option<LmMetadata>,
+    /// Fraction of denoising steps using cover condition (0.0–1.0).
+    pub audio_cover_strength: f64,
+    /// Timestep range for CFG/APG guidance.
+    pub cfg_interval: (f64, f64),
+    /// ODE (deterministic) or SDE (stochastic) sampling.
+    pub infer_method: sampling::InferMethod,
+    /// Latent frames per VAE-decoder chunk. `None` uses the built-in default
+    /// (128 frames ≈ 5.1 s of audio). Lower values cut peak decode memory at
+    /// the cost of more compute overhead; raise when memory is ample.
+    pub vae_chunk_frames: Option<usize>,
+    /// Overlap in latent frames on each side of a VAE-decoder chunk. `None`
+    /// uses the default (16). Must satisfy `chunk > 2 * overlap`.
+    pub vae_chunk_overlap: Option<usize>,
+    /// Starting noise for the DiT, replacing the seeded draw (parity tests).
+    pub initial_noise: Option<Tensor>,
+    /// Vocal language code (e.g. `"en"`); the planner's, when it ran.
+    pub language: Option<String>,
+}
+
+impl Default for GenerationParams {
+    fn default() -> Self {
+        Self {
+            duration_secs: 10.0,
+            seed: None,
+            num_steps: None,
+            guidance_scale: 5.0,
+            shift: 1.0,
+            is_turbo: false,
+            lm_hints_25hz: None,
+            lm_metadata: None,
+            audio_cover_strength: 1.0,
+            cfg_interval: (0.0, 1.0),
+            infer_method: sampling::InferMethod::Ode,
+            vae_chunk_frames: None,
+            vae_chunk_overlap: None,
+            initial_noise: None,
+            language: None,
+        }
+    }
+}
+
+/// Stereo audio output.
+pub struct AudioOutput {
+    /// Stereo waveform `(2, T)` in `[-1, 1]`, 48 kHz.
+    pub audio: Tensor,
+    pub sample_rate: u32,
+    /// The denoised latents `(1, T, acoustic_dim)` the VAE decoded.
+    pub latents: Tensor,
+}
+
+/// High-level pipeline: text prompt → stereo audio.
+pub struct AceStepPipeline {
+    model: AceStepConditionGenerationModel,
+    text_encoder: qwen3::Model,
+    tokenizer: tokenizers::Tokenizer,
+    vae_decoder: OobleckDecoder,
+    vae_encoder: Option<AutoencoderOobleck>,
+    lm: Option<LmState>,
+    silence_latent: Tensor,
+    dit_config: AceStepConfig,
+    vae_config: VaeConfig,
+    device: Device,
+    dtype: DType,
+}
+
+struct LmState {
+    pipeline: LmPipeline,
+    tokenizer: tokenizers::Tokenizer,
+}
+
+impl AceStepPipeline {
+    /// Load all components.
+    ///
+    /// - `dit_vb` / `vae_vb` / `text_vb`: VarBuilders for each model
+    /// - `text_config`: Qwen3 config (for the embedding model, NOT the LM)
+    /// - `tokenizer_path`: `tokenizer.json` for the text encoder
+    /// - `silence_latent`: `(1, T, acoustic_dim)` pre-loaded
+    /// - `load_encoder`: load VAE encoder for cover mode
+    #[allow(clippy::too_many_arguments)]
+    pub fn new(
+        dit_config: &AceStepConfig,
+        vae_config: &VaeConfig,
+        text_config: &qwen3::Config,
+        dit_vb: VarBuilder,
+        vae_vb: VarBuilder,
+        text_vb: VarBuilder,
+        tokenizer_path: &std::path::Path,
+        silence_latent: Tensor,
+        load_encoder: bool,
+    ) -> Result<Self> {
+        let device = dit_vb.device().clone();
+        let dtype = dit_vb.dtype();
+
+        let text_encoder = qwen3::Model::new(text_config, text_vb)?;
+        let tokenizer = tokenizers::Tokenizer::from_file(tokenizer_path)
+            .map_err(|e| anyhow::anyhow!("tokenizer: {e}"))?;
+        let model = AceStepConditionGenerationModel::new(dit_config, dit_vb)?;
+        let vae_decoder = OobleckDecoder::new(vae_config, vae_vb.pp("decoder"))?;
+        let vae_encoder = if load_encoder {
+            Some(AutoencoderOobleck::new(vae_config, vae_vb)?)
+        } else {
+            None
+        };
+
+        Ok(Self {
+            model,
+            text_encoder,
+            tokenizer,
+            vae_decoder,
+            vae_encoder,
+            lm: None,
+            silence_latent,
+            dit_config: dit_config.clone(),
+            vae_config: vae_config.clone(),
+            device,
+            dtype,
+        })
+    }
+
+    /// Number of latent frames for a given duration.
+    pub fn latent_frames(&self, duration_secs: f64) -> usize {
+        let hop = self.vae_config.hop_length();
+        (duration_secs * self.vae_config.sampling_rate as f64 / hop as f64).ceil() as usize
+    }
+
+    /// Generate music from a text prompt and optional lyrics.
+    pub fn text2music(
+        &mut self,
+        prompt: &str,
+        lyrics: &str,
+        params: &GenerationParams,
+    ) -> Result<AudioOutput> {
+        let (text_hs, text_mask, lyric_hs, lyric_mask) = self.encode_text(
+            prompt,
+            lyrics,
+            params.duration_secs,
+            params.lm_metadata.as_ref(),
+            params.language.as_deref(),
+        )?;
+
+        let latent_frames = self.latent_frames(params.duration_secs);
+        let timbre_frames = 750.min(self.silence_latent.dim(1)?);
+
+        self.generate_inner(
+            &text_hs,
+            &text_mask,
+            &lyric_hs,
+            &lyric_mask,
+            &self.silence_latent.narrow(1, 0, latent_frames)?,
+            &self.silence_latent.narrow(1, 0, timbre_frames)?,
+            params,
+            latent_frames,
+        )
+    }
+
+    /// Generate a cover conditioned on reference audio `(1, 2, T)` at 48kHz.
+    ///
+    /// - `reference_audio`: structural source (melody/rhythm)
+    /// - `timbre_audio`: optional separate timbre source; if `None`, uses
+    ///   `reference_audio` for both structure and timbre (Python separates these
+    ///   as `src_audio` vs `reference_audio`)
+    ///
+    /// When `params.lm_hints_25hz` is set, uses the LM-derived hints with
+    /// `is_covers=true` (tokenize/detokenize path in prepare_condition).
+    /// Otherwise, uses raw reference latents directly as context (no lossy
+    /// FSQ round-trip).
+    pub fn cover(
+        &mut self,
+        prompt: &str,
+        lyrics: &str,
+        reference_audio: &Tensor,
+        timbre_audio: Option<&Tensor>,
+        params: &GenerationParams,
+    ) -> Result<AudioOutput> {
+        let (text_hs, text_mask, lyric_hs, lyric_mask) = self.encode_text(
+            prompt,
+            lyrics,
+            params.duration_secs,
+            params.lm_metadata.as_ref(),
+            params.language.as_deref(),
+        )?;
+
+        let vae = self
+            .vae_encoder
+            .as_ref()
+            .ok_or_else(|| anyhow::anyhow!("VAE encoder not loaded (pass load_encoder=true)"))?;
+        let ref_latent = vae.encode_mean(reference_audio)?.transpose(1, 2)?;
+
+        let latent_frames = self.latent_frames(params.duration_secs);
+        let ref_len = ref_latent.dim(1)?;
+
+        let src_latents = if ref_len >= latent_frames {
+            ref_latent.narrow(1, 0, latent_frames)?
+        } else {
+            let pad = self.silence_latent.narrow(1, 0, latent_frames - ref_len)?;
+            Tensor::cat(&[&ref_latent, &pad], 1)?
+        };
+
+        // Timbre: use separate timbre audio if provided, otherwise reference audio
+        let refer_audio_packed = if let Some(timbre) = timbre_audio {
+            let timbre_latent = vae.encode_mean(timbre)?.transpose(1, 2)?;
+            let timbre_len = timbre_latent.dim(1)?;
+            let timbre_frames = 750.min(timbre_len);
+            timbre_latent.narrow(1, 0, timbre_frames)?
+        } else {
+            let timbre_frames = 750.min(ref_len);
+            ref_latent.narrow(1, 0, timbre_frames)?
+        };
+
+        // Only enable is_covers (tokenize/detokenize path) when LM hints
+        // are provided. Without LM hints, raw reference latents go directly
+        // into context — the lossy FSQ round-trip on raw VAE latents produces
+        // poor results.
+        self.generate_inner(
+            &text_hs,
+            &text_mask,
+            &lyric_hs,
+            &lyric_mask,
+            &src_latents,
+            &refer_audio_packed,
+            params,
+            latent_frames,
+        )
+    }
+
+    /// Selectively re-generate a region of existing audio while preserving the rest.
+    ///
+    /// - `source_audio`: original stereo waveform `(1, 2, T)` at 48kHz
+    /// - `start_secs` / `end_secs`: time range to regenerate (seconds)
+    ///
+    /// The preserved regions keep the original audio latents; the repaint region
+    /// is replaced with silence and the chunk_mask signals the DiT to generate new
+    /// content there.
+    pub fn repaint(
+        &mut self,
+        prompt: &str,
+        lyrics: &str,
+        source_audio: &Tensor,
+        start_secs: f64,
+        end_secs: f64,
+        params: &GenerationParams,
+    ) -> Result<AudioOutput> {
+        let (text_hs, text_mask, lyric_hs, lyric_mask) = self.encode_text(
+            prompt,
+            lyrics,
+            params.duration_secs,
+            params.lm_metadata.as_ref(),
+            params.language.as_deref(),
+        )?;
+
+        let vae = self
+            .vae_encoder
+            .as_ref()
+            .ok_or_else(|| anyhow::anyhow!("VAE encoder not loaded (pass load_encoder=true)"))?;
+        let src_latent = vae.encode_mean(source_audio)?.transpose(1, 2)?; // (1, T_lat, D)
+
+        let latent_frames = self.latent_frames(params.duration_secs);
+        let hop = self.vae_config.hop_length() as f64;
+        let sr = self.vae_config.sampling_rate as f64;
+
+        // Convert time range to latent frame indices
+        let start_frame = ((start_secs * sr / hop).floor() as usize).min(latent_frames);
+        let end_frame = ((end_secs * sr / hop).ceil() as usize).min(latent_frames);
+
+        anyhow::ensure!(
+            end_frame > start_frame,
+            "repaint region is empty: {start_secs}s..{end_secs}s → frames {start_frame}..{end_frame}"
+        );
+
+        // Crop/pad source latents to target length
+        let src_len = src_latent.dim(1)?;
+        let mut src_latents = if src_len >= latent_frames {
+            src_latent.narrow(1, 0, latent_frames)?
+        } else {
+            let pad = self.silence_latent.narrow(1, 0, latent_frames - src_len)?;
+            Tensor::cat(&[&src_latent, &pad], 1)?
+        };
+
+        // Replace repaint region with silence in src_latents
+        let silence_region = self
+            .silence_latent
+            .narrow(1, start_frame, end_frame - start_frame)?;
+        // Build: [original[:start], silence[start:end], original[end:]]
+        let parts: Vec<Tensor> = if start_frame > 0 && end_frame < latent_frames {
+            vec![
+                src_latents.narrow(1, 0, start_frame)?,
+                silence_region,
+                src_latents.narrow(1, end_frame, latent_frames - end_frame)?,
+            ]
+        } else if start_frame == 0 {
+            vec![
+                silence_region,
+                src_latents.narrow(1, end_frame, latent_frames - end_frame)?,
+            ]
+        } else {
+            vec![src_latents.narrow(1, 0, start_frame)?, silence_region]
+        };
+        src_latents = Tensor::cat(&parts.iter().collect::<Vec<_>>(), 1)?;
+
+        // Timbre from original audio (first 750 frames)
+        let timbre_frames = 750.min(src_len);
+        let refer_audio_packed = src_latent.narrow(1, 0, timbre_frames)?;
+
+        // Build chunk_mask: 0 = preserve, 1 = regenerate
+        let acoustic_dim = self.dit_config.audio_acoustic_hidden_dim;
+        let chunk_mask_dim = self.dit_config.in_channels - 2 * acoustic_dim;
+        let mut mask_data = vec![0.0f32; latent_frames * chunk_mask_dim];
+        for t in start_frame..end_frame {
+            for c in 0..chunk_mask_dim {
+                mask_data[t * chunk_mask_dim + c] = 1.0;
+            }
+        }
+        let chunk_masks =
+            Tensor::from_vec(mask_data, (1, latent_frames, chunk_mask_dim), &self.device)?
+                .to_dtype(self.dtype)?;
+
+        self.generate_inner_with_mask(
+            &text_hs,
+            &text_mask,
+            &lyric_hs,
+            &lyric_mask,
+            &src_latents,
+            &refer_audio_packed,
+            &chunk_masks,
+            params,
+            latent_frames,
+        )
+    }
+
+    /// Unload the LM model to free GPU/RAM memory before DiT denoising.
+    ///
+    /// The LM (~1.2 GB for 0.6B, ~3.5 GB for 1.7B) is only needed for
+    /// `generate_lm_hints()`. Dropping it before the DiT pass frees memory
+    /// for the denoising loop.
+    pub fn unload_lm(&mut self) {
+        self.lm = None;
+    }
+
+    /// Load the 5Hz LM model for audio code generation.
+    ///
+    /// Call once before `generate_lm_hints()`. The LM weights and tokenizer
+    /// are loaded from a HuggingFace repo (e.g. `ACE-Step/acestep-5Hz-lm-0.6B`).
+    pub fn load_lm(
+        &mut self,
+        lm_config: &qwen3::Config,
+        lm_vb: VarBuilder,
+        lm_tokenizer_path: &std::path::Path,
+    ) -> Result<()> {
+        let lm_model = qwen3::ModelForCausalLM::new(lm_config, lm_vb)?;
+        let lm_tokenizer = tokenizers::Tokenizer::from_file(lm_tokenizer_path)
+            .map_err(|e| anyhow::anyhow!("LM tokenizer: {e}"))?;
+        let token_ids = TokenIds::discover(|s| {
+            let enc = lm_tokenizer
+                .encode(s, false)
+                .map_err(|e| candle_core::Error::Msg(format!("encode {s:?}: {e}")))?;
+            Ok(enc.get_ids().to_vec())
+        })?;
+        let pipeline = LmPipeline::new(lm_model, token_ids, self.device.clone());
+        self.lm = Some(LmState {
+            pipeline,
+            tokenizer: lm_tokenizer,
+        });
+        Ok(())
+    }
+
+    /// Generate LM audio codes and convert to 25Hz latent hints.
+    ///
+    /// Returns `(lm_hints_tensor, metadata)`. The hints can be passed as
+    /// `GenerationParams.lm_hints_25hz` to `text2music()` or `cover()`.
+    pub fn generate_lm_hints(
+        &mut self,
+        prompt: &str,
+        lyrics: &str,
+        duration_secs: f64,
+        lm_config: &LmConfig,
+    ) -> Result<(Tensor, std::collections::BTreeMap<String, String>)> {
+        let lm = self
+            .lm
+            .as_mut()
+            .ok_or_else(|| anyhow::anyhow!("LM not loaded — call load_lm() first"))?;
+
+        let prompt_text = LmPipeline::format_prompt(prompt, lyrics);
+        let prompt_enc = lm
+            .tokenizer
+            .encode(prompt_text.as_str(), false)
+            .map_err(|e| anyhow::anyhow!("LM encode: {e}"))?;
+
+        let decode_fn = |t: u32| -> candle_core::Result<String> {
+            lm.tokenizer
+                .decode(&[t], false)
+                .map_err(|e| candle_core::Error::Msg(format!("{e}")))
+        };
+        let encode_fn = |s: &str| -> candle_core::Result<Vec<u32>> {
+            let e = lm
+                .tokenizer
+                .encode(s, false)
+                .map_err(|e| candle_core::Error::Msg(format!("{e}")))?;
+            Ok(e.get_ids().to_vec())
+        };
+
+        let output = lm.pipeline.generate(
+            prompt_enc.get_ids(),
+            duration_secs,
+            lm_config,
+            &decode_fn,
+            &encode_fn,
+        )?;
+
+        anyhow::ensure!(
+            !output.audio_codes.is_empty(),
+            "LM generated no audio codes"
+        );
+        let hints = self.audio_codes_to_latents(&output.audio_codes)?;
+        Ok((hints, output.metadata))
+    }
+
+    /// Convert LM-generated audio codes to 25Hz latent hints.
+    pub fn audio_codes_to_latents(&self, codes: &[i64]) -> Result<Tensor> {
+        let code_tensor = Tensor::new(codes, &self.device)?
+            .unsqueeze(0)?
+            .unsqueeze(2)?;
+        let quantized = self.model.tokenizer.get_output_from_indices(&code_tensor)?;
+        let quantized = quantized.to_dtype(self.dtype)?;
+        Ok(self.model.detokenizer.forward(&quantized)?)
+    }
+
+    // ---- Internals ----
+
+    /// The text-encoder and lyric prompts, exactly as upstream builds them:
+    /// the caption (the planner's, when it ran) with its metas, and the
+    /// lyrics under their language. `language` defaults to `"unknown"`,
+    /// upstream's own default when the caller names none.
+    pub fn format_prompts(
+        prompt: &str,
+        lyrics: &str,
+        duration_secs: f64,
+        metadata: Option<&LmMetadata>,
+        language: Option<&str>,
+    ) -> (String, String) {
+        let caption = metadata
+            .and_then(|m| m.caption.as_deref())
+            .unwrap_or(prompt);
+        let bpm = metadata.and_then(|m| m.bpm.as_deref()).unwrap_or("N/A");
+        let timesig = metadata
+            .and_then(|m| m.timesignature.as_deref())
+            .unwrap_or("N/A");
+        let keyscale = metadata
+            .and_then(|m| m.keyscale.as_deref())
+            .unwrap_or("N/A");
+        let caption_text = format!(
+            "# Instruction\nFill the audio semantic mask based on the given conditions:\n\n\
+             # Caption\n{caption}\n\n# Metas\n\
+             - bpm: {bpm}\n- timesignature: {timesig}\n- keyscale: {keyscale}\n- duration: {} seconds\n\
+             <|endoftext|>\n",
+            duration_secs as u32
+        );
+        let language = metadata
+            .and_then(|m| m.language.as_deref())
+            .or(language)
+            .unwrap_or("unknown");
+        let lyrics = if lyrics.is_empty() {
+            "[Instrumental]"
+        } else {
+            lyrics
+        };
+        let lyric_text = format!("# Languages\n{language}\n\n# Lyric\n{lyrics}<|endoftext|>");
+        (caption_text, lyric_text)
+    }
+
+    /// Token ids of both prompts, as the encoders will see them.
+    pub fn prompt_token_ids(
+        &self,
+        caption_text: &str,
+        lyric_text: &str,
+    ) -> Result<(Vec<u32>, Vec<u32>)> {
+        let text = self
+            .tokenizer
+            .encode(caption_text, true)
+            .map_err(|e| anyhow::anyhow!("tokenization: {e}"))?;
+        let lyric = self
+            .tokenizer
+            .encode(lyric_text, true)
+            .map_err(|e| anyhow::anyhow!("lyric tokenization: {e}"))?;
+        Ok((text.get_ids().to_vec(), lyric.get_ids().to_vec()))
+    }
+
+    fn encode_text(
+        &mut self,
+        prompt: &str,
+        lyrics: &str,
+        duration_secs: f64,
+        metadata: Option<&LmMetadata>,
+        language: Option<&str>,
+    ) -> Result<(Tensor, Tensor, Tensor, Tensor)> {
+        let (caption_text, lyric_text) =
+            Self::format_prompts(prompt, lyrics, duration_secs, metadata, language);
+        let encoding = self
+            .tokenizer
+            .encode(caption_text.as_str(), true)
+            .map_err(|e| anyhow::anyhow!("tokenization: {e}"))?;
+        let prompt_tokens = Tensor::new(encoding.get_ids(), &self.device)?.unsqueeze(0)?;
+        let text_hidden_states = self.text_encoder.forward(&prompt_tokens, 0)?;
+        self.text_encoder.clear_kv_cache();
+
+        let lyric_encoding = self
+            .tokenizer
+            .encode(lyric_text.as_str(), true)
+            .map_err(|e| anyhow::anyhow!("lyric tokenization: {e}"))?;
+        let lyric_tokens = Tensor::new(lyric_encoding.get_ids(), &self.device)?.unsqueeze(0)?;
+        let lyric_hidden_states = self.text_encoder.embed_tokens(&lyric_tokens)?;
+
+        let text_mask = Tensor::new(encoding.get_attention_mask(), &self.device)?
+            .unsqueeze(0)?
+            .to_dtype(self.dtype)?;
+        let lyric_mask = Tensor::new(lyric_encoding.get_attention_mask(), &self.device)?
+            .unsqueeze(0)?
+            .to_dtype(self.dtype)?;
+
+        Ok((
+            text_hidden_states,
+            text_mask,
+            lyric_hidden_states,
+            lyric_mask,
+        ))
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    fn generate_inner(
+        &mut self,
+        text_hidden_states: &Tensor,
+        text_attention_mask: &Tensor,
+        lyric_hidden_states: &Tensor,
+        lyric_attention_mask: &Tensor,
+        src_latents: &Tensor,
+        refer_audio_packed: &Tensor,
+        params: &GenerationParams,
+        latent_frames: usize,
+    ) -> Result<AudioOutput> {
+        let acoustic_dim = self.dit_config.audio_acoustic_hidden_dim;
+        let chunk_mask_dim = self.dit_config.in_channels - 2 * acoustic_dim;
+        let chunk_masks =
+            Tensor::ones((1, latent_frames, chunk_mask_dim), self.dtype, &self.device)?;
+
+        self.generate_inner_with_mask(
+            text_hidden_states,
+            text_attention_mask,
+            lyric_hidden_states,
+            lyric_attention_mask,
+            src_latents,
+            refer_audio_packed,
+            &chunk_masks,
+            params,
+            latent_frames,
+        )
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    fn generate_inner_with_mask(
+        &mut self,
+        text_hidden_states: &Tensor,
+        text_attention_mask: &Tensor,
+        lyric_hidden_states: &Tensor,
+        lyric_attention_mask: &Tensor,
+        src_latents: &Tensor,
+        refer_audio_packed: &Tensor,
+        chunk_masks: &Tensor,
+        params: &GenerationParams,
+        latent_frames: usize,
+    ) -> Result<AudioOutput> {
+        let is_cover = params.lm_hints_25hz.is_some();
+
+        let condition = self.model.prepare_condition(&ConditionInputs {
+            text_hidden_states,
+            text_attention_mask,
+            lyric_hidden_states,
+            lyric_attention_mask,
+            refer_audio_packed,
+            refer_audio_order_mask: &Tensor::zeros((1,), DType::I64, &self.device)?,
+            src_latents,
+            chunk_masks,
+            silence_latent: &self.silence_latent.narrow(1, 0, latent_frames)?,
+            is_covers: &Tensor::from_vec(vec![if is_cover { 1u8 } else { 0u8 }], 1, &self.device)?,
+            precomputed_lm_hints_25hz: params.lm_hints_25hz.as_ref(),
+            audio_codes: None,
+            non_cover_text_hidden_states: None,
+            non_cover_text_attention_mask: None,
+        })?;
+
+        let num_steps = params
+            .num_steps
+            .unwrap_or(if params.is_turbo { 8 } else { 50 });
+
+        let xt = self.model.generate_latents(
+            &condition,
+            None,
+            &GenerateOptions {
+                seed: params.seed,
+                num_steps,
+                guidance_scale: params.guidance_scale,
+                shift: params.shift,
+                is_turbo: params.is_turbo,
+                audio_cover_strength: params.audio_cover_strength,
+                cfg_interval: params.cfg_interval,
+                infer_method: params.infer_method,
+                initial_noise: params.initial_noise.clone(),
+                ..Default::default()
+            },
+            &self.device,
+            self.dtype,
+        )?;
+
+        let audio = self.vae_decoder.tiled_decode(
+            &xt.transpose(1, 2)?,
+            params.vae_chunk_frames,
+            params.vae_chunk_overlap,
+        )?;
+        let audio = audio.squeeze(0)?;
+
+        Ok(AudioOutput {
+            audio,
+            sample_rate: self.vae_config.sampling_rate as u32,
+            latents: xt,
+        })
+    }
+}
+
+/// Load silence_latent from PyTorch `.pt` file → `(1, T, acoustic_dim)`.
+pub fn load_silence_latent(
+    path: &std::path::Path,
+    acoustic_dim: usize,
+    dtype: DType,
+    device: &Device,
+) -> Result<Tensor> {
+    let file = std::fs::File::open(path)?;
+    let mut archive = zip::ZipArchive::new(std::io::BufReader::new(file))?;
+    let mut raw = Vec::new();
+    std::io::Read::read_to_end(&mut archive.by_name("silence_latent/data/0")?, &mut raw)?;
+    let n_elements = raw.len() / 4;
+    let t_len = n_elements / acoustic_dim;
+    let cpu = Device::Cpu;
+    Ok(
+        Tensor::from_raw_buffer(&raw, DType::F32, &[1, acoustic_dim, t_len], &cpu)?
+            .transpose(1, 2)?
+            .contiguous()?
+            .to_dtype(dtype)?
+            .to_device(device)?,
+    )
+}
@@ -0,0 +1,391 @@
+use candle_core::{DType, Device, Module, Result, Tensor};
+use candle_nn::{Activation, VarBuilder, kv_cache::ConcatKvCache};
+use candle_transformers::{
+    models::with_tracing::{Linear, RmsNorm, linear_b, linear_no_bias},
+    utils::repeat_kv,
+};
+use std::sync::Arc;
+
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct Config {
+    pub vocab_size: usize,
+    pub hidden_size: usize,
+    pub intermediate_size: usize,
+    pub num_hidden_layers: usize,
+    pub num_attention_heads: usize,
+    pub head_dim: usize,
+    pub attention_bias: bool,
+    pub num_key_value_heads: usize,
+    pub max_position_embeddings: usize,
+    pub sliding_window: Option<usize>,
+    pub max_window_layers: usize,
+    pub tie_word_embeddings: bool,
+    pub rope_theta: f64,
+    pub rms_norm_eps: f64,
+    pub use_sliding_window: bool,
+    pub hidden_act: Activation,
+}
+
+#[derive(Debug, Clone)]
+pub(crate) struct Qwen3RotaryEmbedding {
+    sin: Tensor,
+    cos: Tensor,
+}
+
+impl Qwen3RotaryEmbedding {
+    pub(crate) fn new(dtype: DType, cfg: &Config, dev: &Device) -> Result<Self> {
+        let dim = cfg.head_dim;
+        let max_seq_len = cfg.max_position_embeddings;
+        let inv_freq: Vec<_> = (0..dim)
+            .step_by(2)
+            .map(|i| 1f32 / cfg.rope_theta.powf(i as f64 / dim as f64) as f32)
+            .collect();
+        let inv_freq_len = inv_freq.len();
+        let inv_freq = Tensor::from_vec(inv_freq, (1, inv_freq_len), dev)?.to_dtype(DType::F32)?;
+        let t = Tensor::arange(0u32, max_seq_len as u32, dev)?
+            .to_dtype(DType::F32)?
+            .reshape((max_seq_len, 1))?;
+        let freqs = t.matmul(&inv_freq)?;
+        Ok(Self {
+            sin: freqs.sin()?.to_dtype(dtype)?,
+            cos: freqs.cos()?.to_dtype(dtype)?,
+        })
+    }
+
+    /// Apply RoPE (q, k shape: B x H x L x D)
+    pub(crate) fn apply(&self, q: &Tensor, k: &Tensor, offset: usize) -> Result<(Tensor, Tensor)> {
+        let (_, _, seq_len, _) = q.dims4()?;
+        let cos = self.cos.narrow(0, offset, seq_len)?;
+        let sin = self.sin.narrow(0, offset, seq_len)?;
+        let q_embed = candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?;
+        let k_embed = candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?;
+        Ok((q_embed, k_embed))
+    }
+}
+
+#[derive(Debug, Clone)]
+pub(crate) struct Qwen3MLP {
+    gate_proj: Linear,
+    up_proj: Linear,
+    down_proj: Linear,
+    act_fn: Activation,
+}
+
+impl Qwen3MLP {
+    pub(crate) fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
+        Ok(Self {
+            gate_proj: linear_no_bias(cfg.hidden_size, cfg.intermediate_size, vb.pp("gate_proj"))?,
+            up_proj: linear_no_bias(cfg.hidden_size, cfg.intermediate_size, vb.pp("up_proj"))?,
+            down_proj: linear_no_bias(cfg.intermediate_size, cfg.hidden_size, vb.pp("down_proj"))?,
+            act_fn: cfg.hidden_act,
+        })
+    }
+}
+
+impl Module for Qwen3MLP {
+    fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let lhs = x.apply(&self.gate_proj)?.apply(&self.act_fn)?;
+        let rhs = x.apply(&self.up_proj)?;
+        (lhs * rhs)?.apply(&self.down_proj)
+    }
+}
+
+#[derive(Debug, Clone)]
+pub(crate) struct Qwen3Attention {
+    // projections
+    q_proj: Linear,
+    k_proj: Linear,
+    v_proj: Linear,
+    o_proj: Linear,
+    // norms
+    q_norm: RmsNorm,
+    k_norm: RmsNorm,
+    // hyper params
+    num_heads: usize,
+    num_kv_heads: usize,
+    num_kv_groups: usize,
+    head_dim: usize,
+    hidden_size: usize,
+    // utils
+    rotary_emb: Arc<Qwen3RotaryEmbedding>,
+    kv_cache: ConcatKvCache,
+}
+
+impl Qwen3Attention {
+    pub(crate) fn new(
+        cfg: &Config,
+        rotary_emb: Arc<Qwen3RotaryEmbedding>,
+        vb: VarBuilder,
+    ) -> Result<Self> {
+        if cfg.use_sliding_window {
+            candle_core::bail!("sliding window is not supported")
+        }
+
+        let head_dim = cfg.head_dim;
+        let num_heads = cfg.num_attention_heads;
+        let num_kv_heads = cfg.num_key_value_heads;
+        let num_kv_groups = num_heads / num_kv_heads;
+
+        let q_proj = linear_b(
+            cfg.hidden_size,
+            num_heads * head_dim,
+            cfg.attention_bias,
+            vb.pp("q_proj"),
+        )?;
+        let k_proj = linear_b(
+            cfg.hidden_size,
+            num_kv_heads * head_dim,
+            cfg.attention_bias,
+            vb.pp("k_proj"),
+        )?;
+        let v_proj = linear_b(
+            cfg.hidden_size,
+            num_kv_heads * head_dim,
+            cfg.attention_bias,
+            vb.pp("v_proj"),
+        )?;
+        let o_proj = linear_b(
+            num_heads * head_dim,
+            cfg.hidden_size,
+            cfg.attention_bias,
+            vb.pp("o_proj"),
+        )?;
+
+        let q_norm = RmsNorm::new(head_dim, cfg.rms_norm_eps, vb.pp("q_norm"))?;
+        let k_norm = RmsNorm::new(head_dim, cfg.rms_norm_eps, vb.pp("k_norm"))?;
+
+        // Necessary because the hidden_size in the config isn't always accurate
+        let hidden_size = head_dim * cfg.num_attention_heads;
+
+        // dim=2 because we concatenate along the sequence dimension
+        // For tensors of shape [batch, heads, seq, head_dim]
+        let kv_cache = ConcatKvCache::new(2);
+
+        Ok(Self {
+            q_proj,
+            k_proj,
+            v_proj,
+            o_proj,
+            q_norm,
+            k_norm,
+            num_heads,
+            num_kv_heads,
+            num_kv_groups,
+            head_dim,
+            hidden_size,
+            rotary_emb,
+            kv_cache,
+        })
+    }
+
+    pub(crate) fn forward(
+        &mut self,
+        x: &Tensor,
+        attn_mask: Option<&Tensor>,
+        offset: usize,
+    ) -> Result<Tensor> {
+        let (b, l, _) = x.dims3()?;
+
+        // 1. Proj
+        let q = self.q_proj.forward(x)?;
+        let k = self.k_proj.forward(x)?;
+        let v = self.v_proj.forward(x)?;
+
+        // 2. Reshape: (B, L, H, D) -> (B, H, L, D)
+        let q = q
+            .reshape((b, l, self.num_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let k = k
+            .reshape((b, l, self.num_kv_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let v = v
+            .reshape((b, l, self.num_kv_heads, self.head_dim))?
+            .transpose(1, 2)?;
+
+        // 3. Per‑head RMSNorm
+        let q_flat = q.flatten(0, 2)?; // (B*H, L, D) -> (BHL, D) after transpose later
+        let k_flat = k.flatten(0, 2)?;
+        let q_flat = self.q_norm.forward(&q_flat)?;
+        let k_flat = self.k_norm.forward(&k_flat)?;
+        let q = q_flat.reshape((b, self.num_heads, l, self.head_dim))?;
+        let k = k_flat.reshape((b, self.num_kv_heads, l, self.head_dim))?;
+
+        // 4. RoPE
+        let (q, k) = self.rotary_emb.apply(&q, &k, offset)?;
+
+        // 5. Accumulate KV cache
+        let (k, v) = self.kv_cache.append(&k, &v)?;
+
+        // 6. GQA repeat_kv
+        let k = repeat_kv(k, self.num_kv_groups)?.contiguous()?;
+        let v = repeat_kv(v, self.num_kv_groups)?.contiguous()?;
+
+        // 7. Attention score
+        let scale = 1.0 / (self.head_dim as f64).sqrt();
+        let mut scores = (q.matmul(&k.transpose(2, 3)?)? * scale)?;
+        if let Some(m) = attn_mask {
+            scores = scores.broadcast_add(m)?;
+        }
+        let probs = candle_nn::ops::softmax_last_dim(&scores)?;
+        let ctx = probs.matmul(&v)?; // (B, H, L, D)
+
+        // 8. Output proj
+        ctx.transpose(1, 2)?
+            .reshape((b, l, self.hidden_size))?
+            .apply(&self.o_proj)
+    }
+
+    pub(crate) fn clear_kv_cache(&mut self) {
+        self.kv_cache.reset();
+    }
+}
+
+#[derive(Debug, Clone)]
+struct DecoderLayer {
+    self_attn: Qwen3Attention,
+    mlp: Qwen3MLP,
+    ln1: RmsNorm,
+    ln2: RmsNorm,
+}
+
+impl DecoderLayer {
+    fn new(cfg: &Config, rotary: Arc<Qwen3RotaryEmbedding>, vb: VarBuilder) -> Result<Self> {
+        let self_attn = Qwen3Attention::new(cfg, rotary, vb.pp("self_attn"))?;
+        let mlp = Qwen3MLP::new(cfg, vb.pp("mlp"))?;
+        let ln1 = RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?;
+        let ln2 = RmsNorm::new(
+            cfg.hidden_size,
+            cfg.rms_norm_eps,
+            vb.pp("post_attention_layernorm"),
+        )?;
+        Ok(Self {
+            self_attn,
+            mlp,
+            ln1,
+            ln2,
+        })
+    }
+
+    fn forward(&mut self, x: &Tensor, mask: Option<&Tensor>, offset: usize) -> Result<Tensor> {
+        let h = self.ln1.forward(x)?;
+        let h = self.self_attn.forward(&h, mask, offset)?;
+        let x = (x + h)?;
+        let h2 = self.ln2.forward(&x)?;
+        let h2 = h2.apply(&self.mlp)?;
+        x + h2
+    }
+
+    fn clear_kv_cache(&mut self) {
+        self.self_attn.clear_kv_cache();
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct Model {
+    embed_tokens: candle_nn::Embedding,
+    layers: Vec<DecoderLayer>,
+    norm: RmsNorm,
+    device: Device,
+    dtype: DType,
+}
+
+impl Model {
+    pub fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
+        let embed_tokens =
+            candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("model.embed_tokens"))?;
+        let rotary = Arc::new(Qwen3RotaryEmbedding::new(vb.dtype(), cfg, vb.device())?);
+        let mut layers = Vec::with_capacity(cfg.num_hidden_layers);
+        let vb_l = vb.pp("model.layers");
+        for i in 0..cfg.num_hidden_layers {
+            layers.push(DecoderLayer::new(cfg, rotary.clone(), vb_l.pp(i))?);
+        }
+        Ok(Self {
+            embed_tokens,
+            layers,
+            norm: RmsNorm::new(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("model.norm"))?,
+            device: vb.device().clone(),
+            dtype: vb.dtype(),
+        })
+    }
+
+    pub fn clear_kv_cache(&mut self) {
+        for l in &mut self.layers {
+            l.clear_kv_cache();
+        }
+    }
+
+    fn causal_mask(
+        &self,
+        b: usize,
+        tgt: usize,
+        offset: usize,
+        sw: Option<usize>,
+    ) -> Result<Tensor> {
+        let minf = f32::NEG_INFINITY;
+        let mask: Vec<_> = (0..tgt)
+            .flat_map(|i| {
+                (0..(tgt + offset)).map(move |j| {
+                    let past_ok = j <= i + offset;
+                    let sw_ok = match sw {
+                        Some(w) => (i + offset) as i64 - j as i64 <= w as i64,
+                        None => true,
+                    };
+                    if past_ok && sw_ok { 0. } else { minf }
+                })
+            })
+            .collect();
+        Tensor::from_slice(&mask, (b, 1, tgt, tgt + offset), &self.device)?.to_dtype(self.dtype)
+    }
+
+    pub fn forward(&mut self, input: &Tensor, offset: usize) -> Result<Tensor> {
+        let (b, l) = input.dims2()?;
+        let mut h = self.embed_tokens.forward(input)?;
+
+        let causal = if l == 1 {
+            None
+        } else {
+            Some(self.causal_mask(b, l, offset, None)?)
+        };
+
+        for layer in &mut self.layers {
+            h = layer.forward(&h, causal.as_ref(), offset)?;
+        }
+        self.norm.forward(&h)
+    }
+
+    /// Embed token IDs without running through transformer layers.
+    /// Returns raw token embeddings of shape `(B, L, hidden_size)`.
+    pub fn embed_tokens(&self, input: &Tensor) -> Result<Tensor> {
+        self.embed_tokens.forward(input)
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct ModelForCausalLM {
+    base: Model,
+    lm_head: Linear,
+}
+
+impl ModelForCausalLM {
+    pub fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
+        let base = Model::new(cfg, vb.clone())?;
+        let lm_head = if cfg.tie_word_embeddings {
+            Linear::from_weights(base.embed_tokens.embeddings().clone(), None)
+        } else {
+            linear_no_bias(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))?
+        };
+        Ok(Self { base, lm_head })
+    }
+
+    pub fn forward(&mut self, input: &Tensor, offset: usize) -> Result<Tensor> {
+        let (_, l) = input.dims2()?;
+        self.base
+            .forward(input, offset)?
+            .narrow(1, l - 1, 1)?
+            .apply(&self.lm_head)
+    }
+
+    pub fn clear_kv_cache(&mut self) {
+        self.base.clear_kv_cache();
+    }
+}
@@ -0,0 +1,466 @@
+//! Flow matching sampler for ACE-Step diffusion model.
+//!
+//! Implements the Euler ODE solver for rectified flow / flow matching,
+//! with Adaptive Prompt Guidance (APG) support.
+
+use candle_core::{Result, Tensor};
+
+/// All unique timesteps from shift=1,2,3 with fix_nfe=8.
+/// Turbo models were distilled on these exact values — custom timesteps are
+/// mapped to the nearest entry in this table.
+const VALID_TIMESTEPS: [f64; 20] = [
+    1.0,
+    0.9545454545454546,
+    0.9333333333333333,
+    0.9,
+    0.875,
+    0.8571428571428571,
+    0.8333333333333334,
+    0.7692307692307693,
+    0.75,
+    0.6666666666666666,
+    0.6428571428571429,
+    0.625,
+    0.5454545454545454,
+    0.5,
+    0.4,
+    0.375,
+    0.3,
+    0.25,
+    0.2222222222222222,
+    0.125,
+];
+
+/// Generate a linear timestep schedule from 1.0 to 0.0.
+///
+/// Optionally applies a time shift transformation:
+/// `t_shifted = shift * t / (1 + (shift - 1) * t)`
+pub fn get_schedule(num_steps: usize, shift: f64) -> Vec<f64> {
+    if num_steps == 0 {
+        return vec![0.0];
+    }
+    let mut schedule = Vec::with_capacity(num_steps + 1);
+    for i in 0..=num_steps {
+        let t = 1.0 - (i as f64 / num_steps as f64);
+        let t = if (shift - 1.0).abs() > 1e-6 {
+            shift * t / (1.0 + (shift - 1.0) * t)
+        } else {
+            t
+        };
+        schedule.push(t);
+    }
+    schedule
+}
+
+/// Pre-defined discrete timestep schedules for turbo models.
+/// These are the only valid schedules — turbo models were distilled on these
+/// exact timesteps and will produce poor results with continuous schedules.
+///
+/// Each schedule has `fix_nfe` steps (default 8) plus an implicit final t=0.
+pub fn get_turbo_schedule(shift: f64) -> Vec<f64> {
+    if (shift - 1.0).abs() < 0.5 {
+        // shift ≈ 1
+        vec![1.0, 0.875, 0.75, 0.625, 0.5, 0.375, 0.25, 0.125]
+    } else if (shift - 2.0).abs() < 0.5 {
+        // shift ≈ 2
+        vec![
+            1.0,
+            0.9333333333333333,
+            0.8571428571428571,
+            0.7692307692307693,
+            0.6666666666666666,
+            0.5454545454545454,
+            0.4,
+            0.2222222222222222,
+        ]
+    } else {
+        // shift ≈ 3 (default for turbo)
+        vec![
+            1.0,
+            0.9545454545454546,
+            0.9,
+            0.8333333333333334,
+            0.75,
+            0.6428571428571429,
+            0.5,
+            0.3,
+        ]
+    }
+}
+
+/// Build a custom turbo schedule from user-provided timestep values.
+///
+/// Each value is mapped to the nearest entry in `VALID_TIMESTEPS`. Trailing
+/// zeros are stripped, and the resulting schedule must have 1–20 entries.
+pub fn get_custom_turbo_schedule(timesteps: &[f64]) -> Result<Vec<f64>> {
+    let end = timesteps
+        .iter()
+        .rposition(|&t| t != 0.0)
+        .map(|i| i + 1)
+        .unwrap_or(0);
+    let timesteps = &timesteps[..end];
+    if timesteps.is_empty() || timesteps.len() > 20 {
+        candle_core::bail!(
+            "custom timesteps must have 1-20 non-zero values, got {}",
+            timesteps.len()
+        );
+    }
+    for &t in timesteps {
+        if !t.is_finite() {
+            candle_core::bail!("custom timesteps contain non-finite value: {t}");
+        }
+    }
+    Ok(timesteps
+        .iter()
+        .map(|&t| {
+            *VALID_TIMESTEPS
+                .iter()
+                .min_by(|&&a, &&b| {
+                    (a - t)
+                        .abs()
+                        .partial_cmp(&(b - t).abs())
+                        .unwrap_or(std::cmp::Ordering::Equal)
+                })
+                .unwrap()
+        })
+        .collect())
+}
+
+/// Get x0 (clean sample) from noisy sample and velocity prediction.
+/// Flow matching: x_t = t * noise + (1 - t) * x_0, v = noise - x_0
+/// Therefore: x_0 = x_t - v * t
+pub fn get_x0_from_noise(xt: &Tensor, vt: &Tensor, t: f64) -> Result<Tensor> {
+    xt - (vt * t)?
+}
+
+/// Euler ODE step for flow matching.
+///
+/// Given current sample `x_t` and velocity prediction `v_t`:
+/// `x_{t-dt} = x_t - v_t * dt`
+pub fn euler_step(x: &Tensor, v: &Tensor, dt: f64) -> Result<Tensor> {
+    x - (v * dt)?
+}
+
+/// SDE step: predict clean sample, then re-add noise at the next timestep.
+///
+/// `x_0 = get_x0_from_noise(x_t, v_t, t_curr)`
+/// `x_{t_next} = (1 - t_next) * x_0 + t_next * noise`
+pub fn sde_step(xt: &Tensor, vt: &Tensor, t_curr: f64, t_next: f64) -> Result<Tensor> {
+    let x0 = get_x0_from_noise(xt, vt, t_curr)?;
+    let noise = Tensor::randn_like(&x0, 0., 1.)?;
+    let clean_part = (&x0 * (1.0 - t_next))?;
+    let noise_part = (noise * t_next)?;
+    clean_part + noise_part
+}
+
+/// Inference method for the denoising loop.
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum InferMethod {
+    /// Ordinary Differential Equation — deterministic Euler integration.
+    Ode,
+    /// Stochastic Differential Equation — predicts x0 then re-noises.
+    Sde,
+}
+
+/// Momentum buffer for APG (Adaptive Prompt Guidance).
+pub struct MomentumBuffer {
+    momentum: f64,
+    running_average: Option<Tensor>,
+}
+
+impl MomentumBuffer {
+    pub fn new(momentum: f64) -> Self {
+        Self {
+            momentum,
+            running_average: None,
+        }
+    }
+
+    /// Update running average: `running_avg = update_value + momentum * running_avg`
+    pub fn update(&mut self, update_value: &Tensor) -> Result<()> {
+        let new_average = match &self.running_average {
+            Some(ra) => (ra * self.momentum)?,
+            None => Tensor::zeros_like(update_value)?,
+        };
+        self.running_average = Some((update_value + new_average)?);
+        Ok(())
+    }
+
+    pub fn get(&self) -> Option<&Tensor> {
+        self.running_average.as_ref()
+    }
+}
+
+/// Project v0 onto v1 and decompose into parallel and orthogonal components.
+///
+/// Operates along `dim` (typically dim=1 for sequence dimension).
+/// Uses F64 precision internally for numerical stability (matching Python).
+/// On Metal, computation is moved to CPU to avoid MPS numeric issues.
+fn project(v0: &Tensor, v1: &Tensor, dim: usize) -> Result<(Tensor, Tensor)> {
+    let orig_device = v0.device().clone();
+    let orig_dtype = v0.dtype();
+
+    // Use CPU for Metal (MPS numeric issues); stay on device for CUDA.
+    let compute_device = if orig_device.is_metal() {
+        candle_core::Device::Cpu
+    } else {
+        orig_device.clone()
+    };
+    let v0 = v0
+        .to_device(&compute_device)?
+        .to_dtype(candle_core::DType::F64)?;
+    let v1 = v1
+        .to_device(&compute_device)?
+        .to_dtype(candle_core::DType::F64)?;
+
+    // L2 normalize v1 along dim
+    let v1_norm = v1.sqr()?.sum_keepdim(dim)?.sqrt()?;
+    let v1_normalized = v1.broadcast_div(&(v1_norm + 1e-8)?)?;
+
+    // Project: v0_parallel = (v0 · v1_hat) * v1_hat
+    let dot = (&v0 * &v1_normalized)?.sum_keepdim(dim)?;
+    let v0_parallel = dot.broadcast_mul(&v1_normalized)?;
+    let v0_orthogonal = (&v0 - &v0_parallel)?;
+
+    Ok((
+        v0_parallel.to_dtype(orig_dtype)?.to_device(&orig_device)?,
+        v0_orthogonal
+            .to_dtype(orig_dtype)?
+            .to_device(&orig_device)?,
+    ))
+}
+
+/// Adaptive Prompt Guidance (APG).
+///
+/// Decomposes the cond-uncond difference into parallel/orthogonal components
+/// relative to the conditional prediction, applies norm thresholding, and
+/// uses momentum for temporal smoothing.
+///
+/// Reference: ACE-Step's `apg_forward` in `apg_guidance.py`.
+pub fn apg_forward(
+    pred_cond: &Tensor,
+    pred_uncond: &Tensor,
+    guidance_scale: f64,
+    momentum_buffer: &mut MomentumBuffer,
+    norm_threshold: f64,
+) -> Result<Tensor> {
+    // dims=[1] in Python — operates along the sequence dimension
+    let dim = 1usize;
+    let orig_device = pred_cond.device().clone();
+    let orig_dtype = pred_cond.dtype();
+
+    // On Metal, move to CPU for numerical stability (Python does this for MPS).
+    // On CUDA/CPU, stay on the original device to avoid transfer overhead.
+    let compute_device = if orig_device.is_metal() {
+        candle_core::Device::Cpu
+    } else {
+        orig_device.clone()
+    };
+    let pred_cond_cpu = pred_cond.to_device(&compute_device)?;
+    let pred_uncond_cpu = pred_uncond.to_device(&compute_device)?;
+
+    let diff = (&pred_cond_cpu - &pred_uncond_cpu)?;
+
+    // Apply momentum smoothing
+    momentum_buffer.update(&diff)?;
+    let diff = momentum_buffer.get().unwrap().clone();
+
+    // Norm thresholding: clamp L2 norm along dim to prevent oversaturation
+    let diff = if norm_threshold > 0.0 {
+        let diff_norm = diff.sqr()?.sum_keepdim(dim)?.sqrt()?;
+        let ones = Tensor::ones_like(&diff_norm)?;
+        let scale = (ones * norm_threshold)?.broadcast_div(&(&diff_norm + 1e-8)?)?;
+        let scale = scale.minimum(&Tensor::ones_like(&scale)?)?;
+        diff.broadcast_mul(&scale)?
+    } else {
+        diff
+    };
+
+    // Decompose into parallel and orthogonal components
+    let (_diff_parallel, diff_orthogonal) = project(&diff, &pred_cond_cpu, dim)?;
+
+    // eta=0.0 by default: use only orthogonal component
+    let result = (pred_cond_cpu + diff_orthogonal * (guidance_scale - 1.0))?;
+    result.to_dtype(orig_dtype)?.to_device(&orig_device)
+}
+
+/// Differential Correction in Wavelet domain ("double" mode, Haar), which
+/// upstream applies by default after every turbo step: split the latent and
+/// the predicted clean sample into low and high bands along time, and push
+/// each band of the latent away from the clean sample's. The low band's
+/// strength decays with `t`, the high band's grows as `t → 0`.
+#[derive(Debug, Clone, Copy)]
+pub struct Dcw {
+    pub low_scaler: f64,
+    pub high_scaler: f64,
+}
+
+impl Default for Dcw {
+    /// Upstream's defaults for turbo (`dcw_scaler`, `dcw_high_scaler`).
+    fn default() -> Self {
+        Self {
+            low_scaler: 0.05,
+            high_scaler: 0.02,
+        }
+    }
+}
+
+impl Dcw {
+    /// Correct `x_next` (after the sampler step at `t_curr`) against
+    /// `denoised = x_t - v·t_curr`. Shapes are `(B, T, C)`.
+    pub fn apply(&self, x_next: &Tensor, denoised: &Tensor, t_curr: f64) -> Result<Tensor> {
+        let low = t_curr * self.low_scaler;
+        let high = (1.0 - t_curr) * self.high_scaler;
+        if low == 0.0 && high == 0.0 {
+            return Ok(x_next.clone());
+        }
+        let t_out = x_next.dim(1)?;
+        let (xl, xh) = haar_dwt_t(x_next)?;
+        let (yl, yh) = haar_dwt_t(denoised)?;
+        let xl = (&xl + ((&xl - &yl)? * low)?)?;
+        let xh = (&xh + ((&xh - &yh)? * high)?)?;
+        haar_idwt_t(&xl, &xh, t_out)
+    }
+}
+
+/// One-level Haar DWT along dim 1 of `(B, T, C)`, zero-padding an odd `T`
+/// as `pytorch_wavelets` does in `mode="zero"`. Returns `(low, high)`, each
+/// `(B, ceil(T/2), C)`.
+fn haar_dwt_t(x: &Tensor) -> Result<(Tensor, Tensor)> {
+    let (b, t, c) = x.dims3()?;
+    let x = if t % 2 == 1 {
+        let pad = Tensor::zeros((b, 1, c), x.dtype(), x.device())?;
+        Tensor::cat(&[x, &pad], 1)?
+    } else {
+        x.clone()
+    };
+    let half = x.dim(1)? / 2;
+    let pairs = x.reshape((b, half, 2, c))?;
+    let even = pairs.narrow(2, 0, 1)?.squeeze(2)?;
+    let odd = pairs.narrow(2, 1, 1)?.squeeze(2)?;
+    let k = std::f64::consts::FRAC_1_SQRT_2;
+    Ok((((&even + &odd)? * k)?, ((&even - &odd)? * k)?))
+}
+
+/// Inverse of [`haar_dwt_t`], trimmed back to `t_out` frames.
+fn haar_idwt_t(low: &Tensor, high: &Tensor, t_out: usize) -> Result<Tensor> {
+    let (b, half, c) = low.dims3()?;
+    let k = std::f64::consts::FRAC_1_SQRT_2;
+    let even = ((low + high)? * k)?;
+    let odd = ((low - high)? * k)?;
+    let x = Tensor::stack(&[&even, &odd], 2)?.reshape((b, half * 2, c))?;
+    x.narrow(1, 0, t_out)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_get_custom_turbo_schedule_basic() -> Result<()> {
+        // Exact valid timesteps should pass through unchanged
+        let ts = vec![1.0, 0.875, 0.75, 0.5, 0.25];
+        let result = get_custom_turbo_schedule(&ts)?;
+        assert_eq!(result, ts);
+        Ok(())
+    }
+
+    #[test]
+    fn test_get_custom_turbo_schedule_nearest_mapping() -> Result<()> {
+        // Values should be mapped to nearest valid timestep
+        let ts = vec![0.99, 0.51, 0.13];
+        let result = get_custom_turbo_schedule(&ts)?;
+        assert_eq!(result, vec![1.0, 0.5, 0.125]);
+        Ok(())
+    }
+
+    #[test]
+    fn test_get_custom_turbo_schedule_strips_trailing_zeros() -> Result<()> {
+        let ts = vec![1.0, 0.5, 0.0, 0.0];
+        let result = get_custom_turbo_schedule(&ts)?;
+        assert_eq!(result.len(), 2);
+        Ok(())
+    }
+
+    #[test]
+    fn test_get_custom_turbo_schedule_empty_after_strip() {
+        let ts = vec![0.0, 0.0];
+        assert!(get_custom_turbo_schedule(&ts).is_err());
+    }
+
+    #[test]
+    fn test_get_custom_turbo_schedule_too_long() {
+        let ts = vec![1.0; 21];
+        assert!(get_custom_turbo_schedule(&ts).is_err());
+    }
+
+    #[test]
+    fn test_get_custom_turbo_schedule_single() -> Result<()> {
+        let ts = vec![0.87];
+        let result = get_custom_turbo_schedule(&ts)?;
+        assert_eq!(result, vec![0.875]);
+        Ok(())
+    }
+
+    #[test]
+    fn test_turbo_schedule_shift1_matches_custom() -> Result<()> {
+        // The shift=1 hardcoded schedule should match custom input of same values
+        let hardcoded = get_turbo_schedule(1.0);
+        let custom = get_custom_turbo_schedule(&hardcoded)?;
+        assert_eq!(hardcoded, custom);
+        Ok(())
+    }
+
+    #[test]
+    fn test_turbo_schedule_shift3_matches_custom() -> Result<()> {
+        let hardcoded = get_turbo_schedule(3.0);
+        let custom = get_custom_turbo_schedule(&hardcoded)?;
+        assert_eq!(hardcoded, custom);
+        Ok(())
+    }
+
+    #[test]
+    fn haar_round_trips_and_dcw_matches_the_reference_formula() {
+        use candle_core::Device;
+        let d = Device::Cpu;
+        // Odd T exercises the zero pad and the trim.
+        let x = Tensor::arange(0f32, 21.0, &d)
+            .unwrap()
+            .reshape((1, 7, 3))
+            .unwrap();
+        let (l, h) = super::haar_dwt_t(&x).unwrap();
+        assert_eq!(l.dims3().unwrap(), (1, 4, 3));
+        let back = super::haar_idwt_t(&l, &h, 7).unwrap();
+        let diff = (back - &x)
+            .unwrap()
+            .abs()
+            .unwrap()
+            .max_all()
+            .unwrap()
+            .to_scalar::<f32>()
+            .unwrap();
+        assert!(diff < 1e-5, "{diff}");
+
+        // Against the band formula directly, for a 2-frame signal:
+        // L = (a+b)/√2, H = (a-b)/√2; corrected bands, then inverted.
+        let xn = Tensor::new(&[[[3f32], [1.0]]], &d).unwrap();
+        let dn = Tensor::new(&[[[1f32], [1.0]]], &d).unwrap();
+        let t = 0.5;
+        let out = Dcw::default()
+            .apply(&xn, &dn, t)
+            .unwrap()
+            .flatten_all()
+            .unwrap()
+            .to_vec1::<f32>()
+            .unwrap();
+        let r = std::f64::consts::FRAC_1_SQRT_2;
+        let (xl, xh, yl, yh) = (4.0 * r, 2.0 * r, 2.0 * r, 0.0);
+        let (s_lo, s_hi) = (t * 0.05, (1.0 - t) * 0.02);
+        let (cl, ch) = (xl + s_lo * (xl - yl), xh + s_hi * (xh - yh));
+        let want = [((cl + ch) * r) as f32, ((cl - ch) * r) as f32];
+        assert!(
+            (out[0] - want[0]).abs() < 1e-5 && (out[1] - want[1]).abs() < 1e-5,
+            "{out:?} vs {want:?}"
+        );
+    }
+}
@@ -0,0 +1,656 @@
+//! Audio tokenizer and detokenizer for ACE-Step.
+//!
+//! Implements Finite Scalar Quantization (FSQ) based audio tokenization.
+//! These components are used in the cover/audio-code pipeline (LM + DiT hybrid
+//! mode) but are **not exercised** by the Text2Music path.
+//!
+//! **Validation status** (against Python `vector_quantize_pytorch.ResidualFSQ`):
+//! - `AudioTokenDetokenizer`: bit-exact (diff_max < 0.00001)
+//! - `ResidualFSQ` (quantizer): FSQ bound/round matches; small rounding
+//!   boundary differences (diff_max ~0.018) are expected for discrete
+//!   quantization due to float precision at level boundaries.
+//! - `AttentionPooler` / `AudioTokenizer`: structurally correct, minor
+//!   float precision differences propagate through FSQ rounding.
+
+use candle_core::{D, DType, Module, Result, Tensor};
+use candle_nn::{Activation, VarBuilder};
+use candle_transformers::models::with_tracing::{Linear, RmsNorm, linear, linear_no_bias};
+use candle_transformers::utils::repeat_kv;
+use std::sync::Arc;
+
+use super::AceStepConfig;
+use super::dit::RotaryEmbedding;
+
+// ---------------------------------------------------------------------------
+// Encoder MLP (same SiLU-gated pattern as Qwen3MLP)
+// ---------------------------------------------------------------------------
+#[derive(Debug, Clone)]
+#[allow(clippy::upper_case_acronyms)]
+struct MLP {
+    gate_proj: Linear,
+    up_proj: Linear,
+    down_proj: Linear,
+    act_fn: Activation,
+}
+
+impl MLP {
+    fn new(
+        hidden_size: usize,
+        intermediate_size: usize,
+        act: Activation,
+        vb: VarBuilder,
+    ) -> Result<Self> {
+        Ok(Self {
+            gate_proj: linear_no_bias(hidden_size, intermediate_size, vb.pp("gate_proj"))?,
+            up_proj: linear_no_bias(hidden_size, intermediate_size, vb.pp("up_proj"))?,
+            down_proj: linear_no_bias(intermediate_size, hidden_size, vb.pp("down_proj"))?,
+            act_fn: act,
+        })
+    }
+}
+
+impl Module for MLP {
+    fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let lhs = x.apply(&self.gate_proj)?.apply(&self.act_fn)?;
+        let rhs = x.apply(&self.up_proj)?;
+        (lhs * rhs)?.apply(&self.down_proj)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Encoder Attention (bidirectional, GQA, with RoPE)
+// ---------------------------------------------------------------------------
+#[derive(Debug, Clone)]
+struct TokenizerAttention {
+    q_proj: Linear,
+    k_proj: Linear,
+    v_proj: Linear,
+    o_proj: Linear,
+    q_norm: RmsNorm,
+    k_norm: RmsNorm,
+    num_heads: usize,
+    num_kv_heads: usize,
+    num_kv_groups: usize,
+    head_dim: usize,
+    hidden_size: usize,
+}
+
+impl TokenizerAttention {
+    fn new(cfg: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let enc_hs = cfg.encoder_hidden_size();
+        let num_heads = cfg.encoder_num_attention_heads();
+        let num_kv_heads = cfg.encoder_num_key_value_heads();
+        let head_dim = cfg.head_dim;
+        let hidden_size = num_heads * head_dim;
+
+        Ok(Self {
+            q_proj: linear_no_bias(enc_hs, num_heads * head_dim, vb.pp("q_proj"))?,
+            k_proj: linear_no_bias(enc_hs, num_kv_heads * head_dim, vb.pp("k_proj"))?,
+            v_proj: linear_no_bias(enc_hs, num_kv_heads * head_dim, vb.pp("v_proj"))?,
+            o_proj: linear_no_bias(num_heads * head_dim, enc_hs, vb.pp("o_proj"))?,
+            q_norm: RmsNorm::new(head_dim, cfg.rms_norm_eps, vb.pp("q_norm"))?,
+            k_norm: RmsNorm::new(head_dim, cfg.rms_norm_eps, vb.pp("k_norm"))?,
+            num_heads,
+            num_kv_heads,
+            num_kv_groups: num_heads / num_kv_heads,
+            head_dim,
+            hidden_size,
+        })
+    }
+
+    fn forward(
+        &self,
+        x: &Tensor,
+        attn_mask: Option<&Tensor>,
+        rotary: &RotaryEmbedding,
+    ) -> Result<Tensor> {
+        let (b, l, _) = x.dims3()?;
+
+        let q = self.q_proj.forward(x)?;
+        let k = self.k_proj.forward(x)?;
+        let v = self.v_proj.forward(x)?;
+
+        let q = q
+            .reshape((b, l, self.num_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let k = k
+            .reshape((b, l, self.num_kv_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let v = v
+            .reshape((b, l, self.num_kv_heads, self.head_dim))?
+            .transpose(1, 2)?;
+
+        // Per-head RMSNorm
+        let q = self.q_norm.forward(&q.flatten(0, 1)?)?.reshape(q.shape())?;
+        let k = self.k_norm.forward(&k.flatten(0, 1)?)?.reshape(k.shape())?;
+
+        // RoPE
+        let (q, k) = rotary.apply(&q, &k, 0)?;
+
+        // GQA
+        let k = repeat_kv(k, self.num_kv_groups)?.contiguous()?;
+        let v = repeat_kv(v, self.num_kv_groups)?.contiguous()?;
+
+        // Attention
+        let scale = 1.0 / (self.head_dim as f64).sqrt();
+        let mut scores = (q.matmul(&k.transpose(2, 3)?)? * scale)?;
+        if let Some(m) = attn_mask {
+            scores = scores.broadcast_add(m)?;
+        }
+        let probs = candle_nn::ops::softmax_last_dim(&scores)?;
+        let ctx = probs.matmul(&v)?;
+
+        ctx.transpose(1, 2)?
+            .reshape((b, l, self.hidden_size))?
+            .apply(&self.o_proj)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Encoder Layer
+// ---------------------------------------------------------------------------
+#[derive(Debug, Clone)]
+struct EncoderLayer {
+    self_attn: TokenizerAttention,
+    input_layernorm: RmsNorm,
+    post_attention_layernorm: RmsNorm,
+    mlp: MLP,
+}
+
+impl EncoderLayer {
+    fn new(cfg: &AceStepConfig, _layer_idx: usize, vb: VarBuilder) -> Result<Self> {
+        let enc_hs = cfg.encoder_hidden_size();
+        Ok(Self {
+            self_attn: TokenizerAttention::new(cfg, vb.pp("self_attn"))?,
+            input_layernorm: RmsNorm::new(enc_hs, cfg.rms_norm_eps, vb.pp("input_layernorm"))?,
+            post_attention_layernorm: RmsNorm::new(
+                enc_hs,
+                cfg.rms_norm_eps,
+                vb.pp("post_attention_layernorm"),
+            )?,
+            mlp: MLP::new(
+                enc_hs,
+                cfg.encoder_intermediate_size(),
+                cfg.hidden_act,
+                vb.pp("mlp"),
+            )?,
+        })
+    }
+
+    fn forward(
+        &self,
+        x: &Tensor,
+        attn_mask: Option<&Tensor>,
+        rotary: &RotaryEmbedding,
+    ) -> Result<Tensor> {
+        let residual = x.clone();
+        let h = self.input_layernorm.forward(x)?;
+        let h = self.self_attn.forward(&h, attn_mask, rotary)?;
+        let x = (residual + h)?;
+
+        let residual = x.clone();
+        let h = self.post_attention_layernorm.forward(&x)?;
+        let h = h.apply(&self.mlp)?;
+        residual + h
+    }
+}
+
+// ---------------------------------------------------------------------------
+// AttentionPooler - CLS-token based attention pooling
+// ---------------------------------------------------------------------------
+#[derive(Debug, Clone)]
+pub struct AttentionPooler {
+    embed_tokens: Linear,
+    layers: Vec<EncoderLayer>,
+    norm: RmsNorm,
+    special_token: Tensor,
+    rotary_emb: Arc<RotaryEmbedding>,
+}
+
+impl AttentionPooler {
+    pub fn new(cfg: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let enc_hs = cfg.encoder_hidden_size();
+        let mut layers = Vec::with_capacity(cfg.num_attention_pooler_hidden_layers);
+        let vb_l = vb.pp("layers");
+        for i in 0..cfg.num_attention_pooler_hidden_layers {
+            layers.push(EncoderLayer::new(cfg, i, vb_l.pp(i))?);
+        }
+        let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb.device())?);
+        Ok(Self {
+            embed_tokens: linear(enc_hs, enc_hs, vb.pp("embed_tokens"))?,
+            layers,
+            norm: RmsNorm::new(enc_hs, cfg.rms_norm_eps, vb.pp("norm"))?,
+            special_token: vb.get((1, 1, enc_hs), "special_token")?,
+            rotary_emb,
+        })
+    }
+
+    /// Pool patches into sequence-level representations.
+    /// Input: (B, T, pool_window_size, D) → Output: (B, T, D)
+    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let (b, t, p, _d) = x.dims4()?;
+        let x = x.apply(&self.embed_tokens)?;
+
+        // Prepend CLS token to each patch group
+        let cls = self
+            .special_token
+            .broadcast_as((b * t, 1, x.dim(D::Minus1)?))?;
+        let x = x.reshape((b * t, p, ()))?;
+        let x = Tensor::cat(&[&cls, &x], 1)?;
+
+        let mut hidden = x;
+        for layer in &self.layers {
+            hidden = layer.forward(&hidden, None, &self.rotary_emb)?;
+        }
+        let hidden = self.norm.forward(&hidden)?;
+
+        // Extract CLS output (position 0)
+        hidden.narrow(1, 0, 1)?.squeeze(1)?.reshape((b, t, ()))
+    }
+}
+
+// ---------------------------------------------------------------------------
+// ResidualFSQ - Finite Scalar Quantization
+// ---------------------------------------------------------------------------
+
+/// Residual Finite Scalar Quantization.
+///
+/// Matches `vector_quantize_pytorch.ResidualFSQ` with a single top-level
+/// `project_in` (dim → codebook_dim) and `project_out` (codebook_dim → dim).
+/// Each FSQ layer independently quantizes each codebook dimension to its
+/// corresponding number of levels using `tanh` bounding and rounding.
+#[derive(Debug, Clone)]
+pub struct ResidualFSQ {
+    project_in: Linear,
+    project_out: Linear,
+    /// Precomputed: (levels - 1) * (1 + eps) / 2 for each codebook dim
+    half_l: Vec<f32>,
+    /// Precomputed: levels // 2
+    half_width: Vec<f32>,
+    /// Precomputed: 0.5 for even levels, 0.0 for odd
+    offset: Vec<f32>,
+    /// Precomputed: atanh(offset / half_l) shift
+    shift: Vec<f32>,
+    /// Precomputed basis for index calculation: [1, L0, L0*L1, ...]
+    basis: Vec<i64>,
+    /// Quantization levels per codebook dimension (e.g. [8, 8, 8, 5, 5, 5]).
+    levels: Vec<usize>,
+    /// Model working dtype (inferred from VarBuilder at construction time).
+    model_dtype: DType,
+}
+
+impl ResidualFSQ {
+    pub fn new(cfg: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let model_dtype = vb.dtype();
+        let dim = cfg.fsq_dim;
+        let levels = &cfg.fsq_input_levels;
+        let codebook_dim = levels.len();
+
+        // Single project_in/project_out at the ResidualFSQ level
+        let project_in = linear(dim, codebook_dim, vb.pp("project_in"))?;
+        let project_out = linear(codebook_dim, dim, vb.pp("project_out"))?;
+
+        let eps = 1e-3f32;
+        let half_l: Vec<f32> = levels
+            .iter()
+            .map(|&l| (l as f32 - 1.0) * (1.0 + eps) / 2.0)
+            .collect();
+        let half_width: Vec<f32> = levels.iter().map(|&l| (l / 2) as f32).collect();
+        let offset: Vec<f32> = levels
+            .iter()
+            .map(|&l| if l % 2 == 0 { 0.5 } else { 0.0 })
+            .collect();
+        let shift: Vec<f32> = half_l
+            .iter()
+            .zip(&offset)
+            .map(|(&hl, &o)| (o / hl).atanh())
+            .collect();
+
+        let mut basis = vec![1i64];
+        for i in 0..levels.len() - 1 {
+            basis.push(basis[i] * levels[i] as i64);
+        }
+
+        Ok(Self {
+            project_in,
+            project_out,
+            half_l,
+            half_width,
+            offset,
+            shift,
+            basis,
+            levels: levels.to_vec(),
+            model_dtype,
+        })
+    }
+
+    /// FSQ bound + round: `round(tanh(z + shift) * half_l - offset) / half_width`
+    fn fsq_quantize(&self, z: &Tensor) -> Result<Tensor> {
+        let dev = z.device();
+        let half_l = Tensor::new(&self.half_l[..], dev)?
+            .unsqueeze(0)?
+            .unsqueeze(0)?;
+        let offset = Tensor::new(&self.offset[..], dev)?
+            .unsqueeze(0)?
+            .unsqueeze(0)?;
+        let shift = Tensor::new(&self.shift[..], dev)?
+            .unsqueeze(0)?
+            .unsqueeze(0)?;
+        let half_width = Tensor::new(&self.half_width[..], dev)?
+            .unsqueeze(0)?
+            .unsqueeze(0)?;
+
+        // bound: tanh(z + shift) * half_l - offset, then round, then / half_width
+        let bounded = z
+            .broadcast_add(&shift)?
+            .tanh()?
+            .broadcast_mul(&half_l)?
+            .broadcast_sub(&offset)?;
+        let rounded = bounded.round()?;
+        rounded.broadcast_div(&half_width)
+    }
+
+    /// Quantize continuous representations into discrete tokens.
+    /// Input: `(B, N, dim)`. Returns `(quantized_out, indices)` where
+    /// `quantized_out` has the same shape as input and `indices` is `(B, N, num_quantizers)`.
+    pub fn forward(&self, x: &Tensor) -> Result<(Tensor, Tensor)> {
+        // project_in: (B, N, dim) -> (B, N, codebook_dim)
+        let projected = x.apply(&self.project_in)?;
+
+        // FSQ quantize
+        let codes = self.fsq_quantize(&projected)?; // (B, N, codebook_dim)
+
+        // Compute indices from codes: scale_and_shift then dot with basis
+        let dev = x.device();
+        let half_width = Tensor::new(&self.half_width[..], dev)?
+            .unsqueeze(0)?
+            .unsqueeze(0)?;
+        let offset = Tensor::new(&self.offset[..], dev)?
+            .unsqueeze(0)?
+            .unsqueeze(0)?;
+        let basis = Tensor::new(&self.basis[..], dev)?
+            .unsqueeze(0)?
+            .unsqueeze(0)?;
+        // scale_and_shift: codes * half_width + offset
+        let shifted = codes.broadcast_mul(&half_width)?.broadcast_add(&offset)?;
+        let indices = shifted
+            .to_dtype(DType::F32)?
+            .broadcast_mul(&basis.to_dtype(DType::F32)?)?
+            .sum(2)?
+            .round()?
+            .to_dtype(DType::I64)?
+            .unsqueeze(2)?; // (B, N, 1)
+
+        // project_out: (B, N, codebook_dim) -> (B, N, dim)
+        let quantized_out = codes.apply(&self.project_out)?;
+
+        Ok((quantized_out, indices))
+    }
+
+    /// Reconstruct from quantized codes (not indices).
+    /// Input: `(B, N, codebook_dim)` → Output: `(B, N, dim)`.
+    pub fn get_output_from_codes(&self, codes: &Tensor) -> Result<Tensor> {
+        codes.apply(&self.project_out)
+    }
+
+    /// Reconstruct from packed indices.
+    /// Input: `(B, N, 1)` → Output: `(B, N, dim)`.
+    ///
+    /// Inverse of the forward path: decompose flat index into per-level values
+    /// via basis, then invert scale_and_shift, then `project_out`.
+    pub fn get_output_from_indices(&self, indices: &Tensor) -> Result<Tensor> {
+        let dev = indices.device();
+        let indices = indices.squeeze(D::Minus1)?.to_dtype(DType::F32)?;
+        let codebook_dim = self.levels.len();
+
+        // Decompose: level_val[i] = floor(index / basis[i]) % levels[i]
+        // F32 has enough precision for indices up to 64000 (max codebook size).
+        let mut code_dims = Vec::with_capacity(codebook_dim);
+        for i in 0..codebook_dim {
+            let b = self.basis[i] as f32;
+            let l = self.levels[i] as f32;
+            let dim_val = indices.affine(1.0 / b as f64, 0.0)?.floor()?;
+            // modulo: dim_val - floor(dim_val / l) * l
+            let dim_val = (&dim_val
+                - dim_val
+                    .affine(1.0 / l as f64, 0.0)?
+                    .floor()?
+                    .affine(l as f64, 0.0)?)?;
+            code_dims.push(dim_val);
+        }
+        let shifted = Tensor::stack(&code_dims, D::Minus1)?;
+
+        // Inverse scale_and_shift (Python convention):
+        // forward: shifted = codes * half_width + half_width  →  shifted ∈ [0, levels-1]
+        // inverse: codes = (shifted - half_width) / half_width
+        let hw = Tensor::new(&self.half_width[..], dev)?
+            .unsqueeze(0)?
+            .unsqueeze(0)?;
+        let codes = shifted.broadcast_sub(&hw)?.broadcast_div(&hw)?;
+
+        codes.to_dtype(self.model_dtype)?.apply(&self.project_out)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use candle_core::Device;
+
+    /// Build a minimal ResidualFSQ with identity projections for testing.
+    fn test_fsq() -> Result<ResidualFSQ> {
+        let dev = &Device::Cpu;
+        let levels = vec![8, 8, 8, 5, 5, 5];
+        let codebook_dim = levels.len();
+
+        // Identity-like project_in/project_out (codebook_dim x codebook_dim)
+        let eye = Tensor::eye(codebook_dim, DType::F32, dev)?;
+        let zeros = Tensor::zeros(codebook_dim, DType::F32, dev)?;
+
+        let eps = 1e-3f32;
+        let half_l: Vec<f32> = levels
+            .iter()
+            .map(|&l| (l as f32 - 1.0) * (1.0 + eps) / 2.0)
+            .collect();
+        let half_width: Vec<f32> = levels.iter().map(|&l| (l / 2) as f32).collect();
+        let offset: Vec<f32> = levels
+            .iter()
+            .map(|&l| if l % 2 == 0 { 0.5 } else { 0.0 })
+            .collect();
+        let shift: Vec<f32> = half_l
+            .iter()
+            .zip(&offset)
+            .map(|(&hl, &o)| (o / hl).atanh())
+            .collect();
+        let mut basis = vec![1i64];
+        for i in 0..levels.len() - 1 {
+            basis.push(basis[i] * levels[i] as i64);
+        }
+
+        Ok(ResidualFSQ {
+            project_in: Linear::from_weights(eye.clone(), Some(zeros.clone())),
+            project_out: Linear::from_weights(eye, Some(zeros)),
+            half_l,
+            half_width,
+            offset,
+            shift,
+            basis,
+            levels,
+            model_dtype: DType::F32,
+        })
+    }
+
+    #[test]
+    fn test_get_output_from_indices_roundtrip() -> Result<()> {
+        // Test that manually-packed Python-convention indices correctly
+        // reconstruct codes via get_output_from_indices.
+        let fsq = test_fsq()?;
+        let dev = &Device::Cpu;
+        // levels = [8,8,8,5,5,5], half_width = [4,4,4,2,2,2]
+        // basis = [1, 8, 64, 512, 2560, 12800]
+        //
+        // Choose codes: [0.25, -0.5, 0.75, 0.5, -1.0, 0.0]
+        // Python scale_and_shift: shifted = codes * half_width + half_width
+        //   = [0.25*4+4, -0.5*4+4, 0.75*4+4, 0.5*2+2, -1.0*2+2, 0.0*2+2]
+        //   = [5, 2, 7, 3, 0, 2]
+        // index = 5*1 + 2*8 + 7*64 + 3*512 + 0*2560 + 2*12800
+        //       = 5 + 16 + 448 + 1536 + 0 + 25600 = 27605
+        let expected_codes: Vec<f32> = vec![0.25, -0.5, 0.75, 0.5, -1.0, 0.0];
+        let index: i64 = 27605;
+
+        let indices = Tensor::new(&[[index]], dev)?.unsqueeze(2)?; // (1, 1, 1)
+        let reconstructed = fsq.get_output_from_indices(&indices)?;
+        let vals = reconstructed.squeeze(0)?.squeeze(0)?.to_vec1::<f32>()?;
+
+        for (i, (&got, &want)) in vals.iter().zip(expected_codes.iter()).enumerate() {
+            assert!((got - want).abs() < 1e-5, "dim {i}: got={got}, want={want}");
+        }
+        Ok(())
+    }
+
+    #[test]
+    fn test_get_output_from_indices_known_values() -> Result<()> {
+        let fsq = test_fsq()?;
+        let dev = &Device::Cpu;
+
+        // Index 0 should decompose to level_vals = [0, 0, 0, 0, 0, 0]
+        // Python convention: codes = (level_vals - half_width) / half_width
+        // For levels [8,8,8,5,5,5]: half_width=[4,4,4,2,2,2]
+        // codes = [(0-4)/4, (0-4)/4, (0-4)/4, (0-2)/2, (0-2)/2, (0-2)/2]
+        //       = [-1, -1, -1, -1, -1, -1]
+        let indices = Tensor::new(&[[0i64]], dev)?.unsqueeze(2)?; // (1, 1, 1)
+        let result = fsq.get_output_from_indices(&indices)?;
+        let vals = result.squeeze(0)?.squeeze(0)?.to_vec1::<f32>()?;
+        assert!((vals[0] - (-1.0)).abs() < 1e-5, "val[0]={}", vals[0]);
+        assert!((vals[3] - (-1.0)).abs() < 1e-5, "val[3]={}", vals[3]);
+
+        // Index that corresponds to center: half_width values for each level
+        // level_vals = [4, 4, 4, 2, 2, 2]
+        // packed index = 4*1 + 4*8 + 4*64 + 2*512 + 2*2560 + 2*12800
+        //             = 4 + 32 + 256 + 1024 + 5120 + 25600 = 32036
+        // codes = [0, 0, 0, 0, 0, 0]
+        let indices = Tensor::new(&[[32036i64]], dev)?.unsqueeze(2)?;
+        let result = fsq.get_output_from_indices(&indices)?;
+        let vals = result.squeeze(0)?.squeeze(0)?.to_vec1::<f32>()?;
+        for (i, &v) in vals.iter().enumerate() {
+            assert!(v.abs() < 1e-5, "center val[{i}]={v}, expected 0");
+        }
+        Ok(())
+    }
+}
+
+// ---------------------------------------------------------------------------
+// AudioTokenizer
+// ---------------------------------------------------------------------------
+#[derive(Debug, Clone)]
+pub struct AudioTokenizer {
+    audio_acoustic_proj: Linear,
+    attention_pooler: AttentionPooler,
+    quantizer: ResidualFSQ,
+    pool_window_size: usize,
+}
+
+impl AudioTokenizer {
+    pub fn new(cfg: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let enc_hs = cfg.encoder_hidden_size();
+        Ok(Self {
+            audio_acoustic_proj: linear(
+                cfg.audio_acoustic_hidden_dim,
+                enc_hs,
+                vb.pp("audio_acoustic_proj"),
+            )?,
+            attention_pooler: AttentionPooler::new(cfg, vb.pp("attention_pooler"))?,
+            quantizer: ResidualFSQ::new(cfg, vb.pp("quantizer"))?,
+            pool_window_size: cfg.pool_window_size,
+        })
+    }
+
+    pub fn pool_window_size(&self) -> usize {
+        self.pool_window_size
+    }
+
+    /// Reconstruct from packed indices via the quantizer.
+    /// Input: `(B, N, 1)` → Output: `(B, N, dim)`.
+    pub fn get_output_from_indices(&self, indices: &Tensor) -> Result<Tensor> {
+        self.quantizer.get_output_from_indices(indices)
+    }
+
+    /// Tokenize audio features into quantized tokens.
+    /// Input: (B, T_patches, pool_window_size, acoustic_dim)
+    /// Output: (quantized, indices)
+    pub fn forward(&self, x: &Tensor) -> Result<(Tensor, Tensor)> {
+        let x = x.apply(&self.audio_acoustic_proj)?;
+        let pooled = self.attention_pooler.forward(&x)?;
+        self.quantizer.forward(&pooled)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// AudioTokenDetokenizer
+// ---------------------------------------------------------------------------
+#[derive(Debug, Clone)]
+pub struct AudioTokenDetokenizer {
+    embed_tokens: Linear,
+    layers: Vec<EncoderLayer>,
+    norm: RmsNorm,
+    proj_out: Linear,
+    special_tokens: Tensor,
+    pool_window_size: usize,
+    rotary_emb: Arc<RotaryEmbedding>,
+}
+
+impl AudioTokenDetokenizer {
+    pub fn new(cfg: &AceStepConfig, vb: VarBuilder) -> Result<Self> {
+        let enc_hs = cfg.encoder_hidden_size();
+        let num_layers = cfg.num_attention_pooler_hidden_layers;
+
+        let mut layers = Vec::with_capacity(num_layers);
+        let vb_l = vb.pp("layers");
+        for i in 0..num_layers {
+            layers.push(EncoderLayer::new(cfg, i, vb_l.pp(i))?);
+        }
+        let rotary_emb = Arc::new(RotaryEmbedding::new(vb.dtype(), cfg, vb.device())?);
+
+        Ok(Self {
+            embed_tokens: linear(enc_hs, enc_hs, vb.pp("embed_tokens"))?,
+            layers,
+            norm: RmsNorm::new(enc_hs, cfg.rms_norm_eps, vb.pp("norm"))?,
+            proj_out: linear(enc_hs, cfg.audio_acoustic_hidden_dim, vb.pp("proj_out"))?,
+            special_tokens: vb.get((1, cfg.pool_window_size, enc_hs), "special_tokens")?,
+            pool_window_size: cfg.pool_window_size,
+            rotary_emb,
+        })
+    }
+
+    /// Detokenize quantized tokens back to continuous acoustic representations.
+    /// Input: (B, T, D) → Output: (B, T * pool_window_size, acoustic_dim)
+    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let (b, t, _d) = x.dims3()?;
+        let x = x.apply(&self.embed_tokens)?;
+
+        // Expand each token into pool_window_size patches
+        let x = x
+            .unsqueeze(2)?
+            .broadcast_as((b, t, self.pool_window_size, x.dim(D::Minus1)?))?;
+        let special = self.special_tokens.broadcast_as((
+            b,
+            t,
+            self.pool_window_size,
+            self.special_tokens.dim(D::Minus1)?,
+        ))?;
+        let x = (x.contiguous()? + special.contiguous()?)?;
+
+        // Reshape for processing: (B*T, pool_window_size, hidden)
+        let x = x.reshape((b * t, self.pool_window_size, ()))?;
+
+        let mut hidden = x;
+        for layer in &self.layers {
+            hidden = layer.forward(&hidden, None, &self.rotary_emb)?;
+        }
+        let hidden = self.norm.forward(&hidden)?;
+        let hidden = hidden.apply(&self.proj_out)?;
+
+        // Reshape back: (B*T, P, acoustic_dim) → (B, T*P, acoustic_dim)
+        hidden.reshape((b, t * self.pool_window_size, ()))
+    }
+}
@@ -0,0 +1,721 @@
+//! AutoencoderOobleck VAE for the ACE-Step music generation model.
+//!
+//! This module implements the Oobleck audio VAE that encodes 48kHz stereo audio
+//! into a compact latent representation and decodes it back. The architecture
+//! uses Snake1d activations, weight-normalized convolutions, and a progressive
+//! channel structure with strided up/downsampling.
+//!
+//! The encoder compresses (B, 2, T_audio) to (B, 128, T_latent) which is then
+//! split into mean and log-variance (each 64 channels) for the diagonal Gaussian.
+//! The decoder reconstructs (B, 64, T_latent) back to (B, 2, T_audio) where
+//! T_audio is approximately T_latent * 1920 (the product of downsampling ratios).
+//!
+//! Reference: diffusers `AutoencoderOobleck`
+
+use candle_core::{D, Module, Result, Tensor};
+use candle_nn::{Conv1d, Conv1dConfig, ConvTranspose1d, ConvTranspose1dConfig, VarBuilder};
+
+use super::VaeConfig;
+use candle_transformers::models::encodec;
+
+// ---------------------------------------------------------------------------
+// Snake1d activation
+// ---------------------------------------------------------------------------
+
+/// Learnable Snake activation with alpha and beta parameters stored in log-scale.
+///
+/// Computes: `x + (exp(beta) + 1e-9).recip() * sin(exp(alpha) * x)^2`
+///
+/// Both `alpha` and `beta` have shape `(1, channels, 1)` and are exponentiated
+/// before use to ensure positivity.
+#[derive(Debug, Clone)]
+pub struct Snake1d {
+    alpha: Tensor,
+    beta: Tensor,
+}
+
+impl Snake1d {
+    pub fn new(channels: usize, vb: VarBuilder) -> Result<Self> {
+        let alpha = vb.get((1, channels, 1), "alpha")?;
+        let beta = vb.get((1, channels, 1), "beta")?;
+        Ok(Self { alpha, beta })
+    }
+}
+
+impl Module for Snake1d {
+    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let alpha = self.alpha.exp()?;
+        let beta = self.beta.exp()?;
+        let sin = alpha.broadcast_mul(xs)?.sin()?;
+        let sin_sq = (&sin * &sin)?;
+        xs + ((&beta + 1e-9)?.recip()?.broadcast_mul(&sin_sq)?)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// OobleckResidualUnit
+// ---------------------------------------------------------------------------
+
+/// Residual unit with dilated convolutions and Snake1d activations.
+///
+/// Structure: `snake1 -> conv1(k=7, dilation=d) -> snake2 -> conv2(k=1) + skip`
+///
+/// All convolutions use weight normalization. The skip connection trims the
+/// input if the output is shorter due to dilation-induced size differences.
+#[derive(Debug, Clone)]
+pub struct OobleckResidualUnit {
+    snake1: Snake1d,
+    conv1: Conv1d,
+    snake2: Snake1d,
+    conv2: Conv1d,
+}
+
+impl OobleckResidualUnit {
+    pub fn new(dim: usize, dilation: usize, vb: VarBuilder) -> Result<Self> {
+        let pad = ((7 - 1) * dilation) / 2;
+        let snake1 = Snake1d::new(dim, vb.pp("snake1"))?;
+        let cfg1 = Conv1dConfig {
+            dilation,
+            padding: pad,
+            ..Default::default()
+        };
+        let conv1 = encodec::conv1d_weight_norm(dim, dim, 7, cfg1, vb.pp("conv1"))?;
+        let snake2 = Snake1d::new(dim, vb.pp("snake2"))?;
+        let conv2 = encodec::conv1d_weight_norm(dim, dim, 1, Default::default(), vb.pp("conv2"))?;
+        Ok(Self {
+            snake1,
+            conv1,
+            snake2,
+            conv2,
+        })
+    }
+}
+
+impl Module for OobleckResidualUnit {
+    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let ys = xs
+            .apply(&self.snake1)?
+            .apply(&self.conv1)?
+            .apply(&self.snake2)?
+            .apply(&self.conv2)?;
+        let xs_len = xs.dim(D::Minus1)?;
+        let ys_len = ys.dim(D::Minus1)?;
+        if xs_len >= ys_len {
+            let pad = (xs_len - ys_len) / 2;
+            if pad > 0 {
+                &ys + xs.narrow(D::Minus1, pad, ys_len)?
+            } else {
+                ys + xs
+            }
+        } else {
+            let pad = (ys_len - xs_len) / 2;
+            &ys.narrow(D::Minus1, pad, xs_len)? + xs
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+// OobleckEncoderBlock
+// ---------------------------------------------------------------------------
+
+/// Encoder block: 3 residual units (dilations 1, 3, 9) followed by Snake1d
+/// activation and a strided downsampling convolution.
+#[derive(Debug, Clone)]
+pub struct OobleckEncoderBlock {
+    res_unit1: OobleckResidualUnit,
+    res_unit2: OobleckResidualUnit,
+    res_unit3: OobleckResidualUnit,
+    snake1: Snake1d,
+    conv1: Conv1d,
+}
+
+impl OobleckEncoderBlock {
+    pub fn new(input_dim: usize, output_dim: usize, stride: usize, vb: VarBuilder) -> Result<Self> {
+        let res_unit1 = OobleckResidualUnit::new(input_dim, 1, vb.pp("res_unit1"))?;
+        let res_unit2 = OobleckResidualUnit::new(input_dim, 3, vb.pp("res_unit2"))?;
+        let res_unit3 = OobleckResidualUnit::new(input_dim, 9, vb.pp("res_unit3"))?;
+        let snake1 = Snake1d::new(input_dim, vb.pp("snake1"))?;
+        let cfg = Conv1dConfig {
+            stride,
+            padding: stride.div_ceil(2),
+            ..Default::default()
+        };
+        let conv1 =
+            encodec::conv1d_weight_norm(input_dim, output_dim, 2 * stride, cfg, vb.pp("conv1"))?;
+        Ok(Self {
+            res_unit1,
+            res_unit2,
+            res_unit3,
+            snake1,
+            conv1,
+        })
+    }
+}
+
+impl Module for OobleckEncoderBlock {
+    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        xs.apply(&self.res_unit1)?
+            .apply(&self.res_unit2)?
+            .apply(&self.res_unit3)?
+            .apply(&self.snake1)?
+            .apply(&self.conv1)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// OobleckDecoderBlock
+// ---------------------------------------------------------------------------
+
+/// Decoder block: Snake1d activation and transposed convolution for upsampling,
+/// followed by 3 residual units (dilations 1, 3, 9).
+#[derive(Debug, Clone)]
+pub struct OobleckDecoderBlock {
+    snake1: Snake1d,
+    conv_t1: ConvTranspose1d,
+    res_unit1: OobleckResidualUnit,
+    res_unit2: OobleckResidualUnit,
+    res_unit3: OobleckResidualUnit,
+}
+
+impl OobleckDecoderBlock {
+    pub fn new(input_dim: usize, output_dim: usize, stride: usize, vb: VarBuilder) -> Result<Self> {
+        let snake1 = Snake1d::new(input_dim, vb.pp("snake1"))?;
+        let cfg = ConvTranspose1dConfig {
+            stride,
+            padding: stride.div_ceil(2),
+            ..Default::default()
+        };
+        let conv_t1 = encodec::conv_transpose1d_weight_norm(
+            input_dim,
+            output_dim,
+            2 * stride,
+            true,
+            cfg,
+            vb.pp("conv_t1"),
+        )?;
+        let res_unit1 = OobleckResidualUnit::new(output_dim, 1, vb.pp("res_unit1"))?;
+        let res_unit2 = OobleckResidualUnit::new(output_dim, 3, vb.pp("res_unit2"))?;
+        let res_unit3 = OobleckResidualUnit::new(output_dim, 9, vb.pp("res_unit3"))?;
+        Ok(Self {
+            snake1,
+            conv_t1,
+            res_unit1,
+            res_unit2,
+            res_unit3,
+        })
+    }
+}
+
+impl Module for OobleckDecoderBlock {
+    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        xs.apply(&self.snake1)?
+            .apply(&self.conv_t1)?
+            .apply(&self.res_unit1)?
+            .apply(&self.res_unit2)?
+            .apply(&self.res_unit3)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// OobleckEncoder
+// ---------------------------------------------------------------------------
+
+/// Oobleck encoder that compresses stereo audio to a latent representation.
+///
+/// Channel progression (default config):
+/// `2 -> 128 -> 128 -> 256 -> 512 -> 1024 -> 2048 -> 128`
+#[derive(Debug, Clone)]
+pub struct OobleckEncoder {
+    conv1: Conv1d,
+    blocks: Vec<OobleckEncoderBlock>,
+    snake1: Snake1d,
+    conv2: Conv1d,
+}
+
+impl OobleckEncoder {
+    pub fn new(config: &VaeConfig, vb: VarBuilder) -> Result<Self> {
+        let enc_hs = config.encoder_hidden_size;
+
+        // Prepend 1 to channel_multiples: [1, 1, 2, 4, 8, 16]
+        let mut multiples = vec![1usize];
+        multiples.extend_from_slice(&config.channel_multiples);
+
+        let cfg1 = Conv1dConfig {
+            padding: 3,
+            ..Default::default()
+        };
+        let conv1 =
+            encodec::conv1d_weight_norm(config.audio_channels, enc_hs, 7, cfg1, vb.pp("conv1"))?;
+
+        let mut blocks = Vec::with_capacity(config.downsampling_ratios.len());
+        for (i, &stride) in config.downsampling_ratios.iter().enumerate() {
+            let in_channels = enc_hs * multiples[i];
+            let out_channels = enc_hs * multiples[i + 1];
+            let block =
+                OobleckEncoderBlock::new(in_channels, out_channels, stride, vb.pp("block").pp(i))?;
+            blocks.push(block);
+        }
+
+        let final_channels = enc_hs * multiples[multiples.len() - 1];
+        let snake1 = Snake1d::new(final_channels, vb.pp("snake1"))?;
+
+        let cfg2 = Conv1dConfig {
+            padding: 1,
+            ..Default::default()
+        };
+        let conv2 = encodec::conv1d_weight_norm(final_channels, enc_hs, 3, cfg2, vb.pp("conv2"))?;
+
+        Ok(Self {
+            conv1,
+            blocks,
+            snake1,
+            conv2,
+        })
+    }
+}
+
+impl Module for OobleckEncoder {
+    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let mut xs = xs.apply(&self.conv1)?;
+        for block in &self.blocks {
+            xs = xs.apply(block)?;
+        }
+        xs.apply(&self.snake1)?.apply(&self.conv2)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// OobleckDecoder
+// ---------------------------------------------------------------------------
+
+/// Oobleck decoder that reconstructs stereo audio from a latent representation.
+///
+/// Channel progression (default config):
+/// `64 -> 2048 -> 1024 -> 512 -> 256 -> 128 -> 2`
+#[derive(Debug, Clone)]
+pub struct OobleckDecoder {
+    conv1: Conv1d,
+    blocks: Vec<OobleckDecoderBlock>,
+    snake1: Snake1d,
+    conv2: Conv1d,
+}
+
+impl OobleckDecoder {
+    pub fn new(config: &VaeConfig, vb: VarBuilder) -> Result<Self> {
+        let dec_ch = config.decoder_channels;
+
+        // Prepend 1 to channel_multiples: [1, 1, 2, 4, 8, 16]
+        let mut multiples = vec![1usize];
+        multiples.extend_from_slice(&config.channel_multiples);
+        let n = multiples.len(); // 6
+
+        // Upsampling ratios are reversed downsampling: [10, 6, 4, 4, 2]
+        let upsampling_ratios: Vec<usize> =
+            config.downsampling_ratios.iter().copied().rev().collect();
+
+        let cfg1 = Conv1dConfig {
+            padding: 3,
+            ..Default::default()
+        };
+        // conv1: decoder_input_channels -> dec_ch * mult[-1] (e.g. 64 -> 2048)
+        let first_channels = dec_ch * multiples[n - 1];
+        let conv1 = encodec::conv1d_weight_norm(
+            config.decoder_input_channels,
+            first_channels,
+            7,
+            cfg1,
+            vb.pp("conv1"),
+        )?;
+
+        let mut blocks = Vec::with_capacity(upsampling_ratios.len());
+        for (i, &stride) in upsampling_ratios.iter().enumerate() {
+            // Channels go from mult[N-1-i] to mult[N-2-i]
+            let in_channels = dec_ch * multiples[n - 1 - i];
+            let out_channels = dec_ch * multiples[n - 2 - i];
+            let block =
+                OobleckDecoderBlock::new(in_channels, out_channels, stride, vb.pp("block").pp(i))?;
+            blocks.push(block);
+        }
+
+        let snake1 = Snake1d::new(dec_ch, vb.pp("snake1"))?;
+
+        let cfg2 = Conv1dConfig {
+            padding: 3,
+            ..Default::default()
+        };
+        // Final conv: dec_ch -> audio_channels, NO bias
+        let conv2 = encodec::conv1d_weight_norm_no_bias(
+            dec_ch,
+            config.audio_channels,
+            7,
+            cfg2,
+            vb.pp("conv2"),
+        )?;
+
+        Ok(Self {
+            conv1,
+            blocks,
+            snake1,
+            conv2,
+        })
+    }
+
+    /// Tiled decoding for long latents that may not fit in memory.
+    ///
+    /// Splits the latent along the time axis into overlapping chunks, decodes
+    /// each chunk independently via `forward`, trims the overlap regions in
+    /// audio space, and concatenates the cores. The decoder is fully
+    /// convolutional with local padding, so cores away from chunk boundaries
+    /// are numerically identical to a single full-length decode (up to the
+    /// decoder's finite receptive field, absorbed by the overlap).
+    ///
+    /// For short latents (`T_latent <= chunk_latent_frames`) this falls back
+    /// to a plain `forward` call.
+    ///
+    /// - `chunk_latent_frames`: latent frames per chunk (default 128 ≈ 5.1s of
+    ///   audio at 48 kHz with a 1920× upsample ratio).
+    /// - `overlap_latent_frames`: overlap on each side in latent frames
+    ///   (default 16). Must satisfy `chunk > 2 * overlap`.
+    ///
+    /// Input: `(B, decoder_input_channels, T_latent)`.
+    /// Output: `(B, audio_channels, T_latent * upsample_ratio)`.
+    pub fn tiled_decode(
+        &self,
+        z: &Tensor,
+        chunk_latent_frames: Option<usize>,
+        overlap_latent_frames: Option<usize>,
+    ) -> Result<Tensor> {
+        let chunk_size = chunk_latent_frames.unwrap_or(128);
+        let overlap = overlap_latent_frames.unwrap_or(16);
+        let total_frames = z.dim(2)?;
+
+        if total_frames <= chunk_size {
+            return self.forward(z);
+        }
+
+        if chunk_size <= 2 * overlap {
+            candle_core::bail!(
+                "tiled_decode: chunk_latent_frames ({chunk_size}) must be > 2 * overlap_latent_frames ({overlap})"
+            );
+        }
+
+        let stride = chunk_size - 2 * overlap;
+        let num_steps = total_frames.div_ceil(stride);
+        let mut upsample_factor: Option<usize> = None;
+        let mut audio_chunks = Vec::with_capacity(num_steps);
+
+        for i in 0..num_steps {
+            let core_start = i * stride;
+            let core_end = (core_start + stride).min(total_frames);
+            if core_start >= core_end {
+                break;
+            }
+
+            // Expand the window with overlap on both sides, clipped at the edges.
+            let win_start = core_start.saturating_sub(overlap);
+            let win_end = (core_end + overlap).min(total_frames);
+
+            let latent_chunk = z.narrow(2, win_start, win_end - win_start)?;
+            let audio_chunk = self.forward(&latent_chunk)?;
+
+            // Determine the upsample factor from the first chunk. The decoder
+            // is strided-conv-transpose only, so audio_len is an integer
+            // multiple of latent_len.
+            let up = match upsample_factor {
+                Some(up) => up,
+                None => {
+                    let audio_len = audio_chunk.dim(2)?;
+                    let latent_len = latent_chunk.dim(2)?;
+                    if latent_len == 0 || audio_len % latent_len != 0 {
+                        candle_core::bail!(
+                            "tiled_decode: audio length {audio_len} is not a multiple of latent length {latent_len}"
+                        );
+                    }
+                    let up = audio_len / latent_len;
+                    upsample_factor = Some(up);
+                    up
+                }
+            };
+
+            // Trim the overlap in audio space to recover only the core.
+            let trim_start = (core_start - win_start) * up;
+            let trim_end = (win_end - core_end) * up;
+            let audio_len = audio_chunk.dim(2)?;
+            let core_len = audio_len - trim_start - trim_end;
+            let audio_core = audio_chunk.narrow(2, trim_start, core_len)?;
+            audio_chunks.push(audio_core);
+        }
+
+        Tensor::cat(&audio_chunks.iter().collect::<Vec<_>>(), 2)
+    }
+}
+
+impl Module for OobleckDecoder {
+    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let mut xs = xs.apply(&self.conv1)?;
+        for block in &self.blocks {
+            xs = xs.apply(block)?;
+        }
+        xs.apply(&self.snake1)?.apply(&self.conv2)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// AutoencoderOobleck
+// ---------------------------------------------------------------------------
+
+/// AutoencoderOobleck: full VAE combining encoder and decoder for audio.
+///
+/// Encodes stereo 48kHz audio `(B, 2, T)` to latent `(B, 64, T/1920)` and
+/// decodes back. Uses a diagonal Gaussian latent distribution with
+/// reparameterization sampling.
+#[derive(Debug, Clone)]
+pub struct AutoencoderOobleck {
+    encoder: OobleckEncoder,
+    decoder: OobleckDecoder,
+    config: VaeConfig,
+}
+
+impl AutoencoderOobleck {
+    pub fn new(config: &VaeConfig, vb: VarBuilder) -> Result<Self> {
+        let encoder = OobleckEncoder::new(config, vb.pp("encoder"))?;
+        let decoder = OobleckDecoder::new(config, vb.pp("decoder"))?;
+        Ok(Self {
+            encoder,
+            decoder,
+            config: config.clone(),
+        })
+    }
+
+    /// Encode audio to a latent sample using the diagonal Gaussian.
+    ///
+    /// The encoder output `(B, encoder_hidden_size, T_latent)` is split along
+    /// the channel dimension into mean and log-variance (each with
+    /// `decoder_input_channels` channels), then sampled via reparameterization:
+    /// `z = mean + exp(0.5 * logvar) * noise`.
+    pub fn encode(&self, xs: &Tensor) -> Result<Tensor> {
+        let h = self.encoder.forward(xs)?;
+        let latent_dim = self.config.decoder_input_channels;
+        let mean = h.narrow(1, 0, latent_dim)?;
+        let logvar = h.narrow(1, latent_dim, latent_dim)?;
+        let std = (logvar * 0.5)?.exp()?;
+        let noise = std.randn_like(0., 1.)?;
+        &mean + std.mul(&noise)?
+    }
+
+    /// Encode audio to the latent mean (no sampling noise).
+    ///
+    /// For deterministic encoding (e.g. when exact reconstruction matters),
+    /// this returns only the mean of the diagonal Gaussian, without
+    /// reparameterization noise.
+    pub fn encode_mean(&self, xs: &Tensor) -> Result<Tensor> {
+        let h = self.encoder.forward(xs)?;
+        let latent_dim = self.config.decoder_input_channels;
+        h.narrow(1, 0, latent_dim)
+    }
+
+    /// Decode a latent representation back to audio waveform.
+    ///
+    /// Input shape: `(B, decoder_input_channels, T_latent)`
+    /// Output shape: `(B, audio_channels, T_audio)`
+    pub fn decode(&self, z: &Tensor) -> Result<Tensor> {
+        self.decoder.forward(z)
+    }
+
+    /// Tiled decoding for long latents; see [`OobleckDecoder::tiled_decode`].
+    pub fn tiled_decode(
+        &self,
+        z: &Tensor,
+        chunk_latent_frames: Option<usize>,
+        overlap_latent_frames: Option<usize>,
+    ) -> Result<Tensor> {
+        self.decoder
+            .tiled_decode(z, chunk_latent_frames, overlap_latent_frames)
+    }
+
+    /// Tiled encoding for long audio that may not fit in GPU memory.
+    ///
+    /// Splits the audio into overlapping chunks, encodes each independently,
+    /// trims the overlap regions, and concatenates. Uses the overlap-discard
+    /// strategy to avoid boundary artifacts.
+    ///
+    /// - `chunk_samples`: audio samples per chunk (default: 48000 * 30 = 30s)
+    /// - `overlap_samples`: overlap in audio samples (default: 48000 * 2 = 2s)
+    ///
+    /// Input: `(B, channels, T_audio)`, output: `(B, latent_dim, T_latent)`
+    pub fn tiled_encode(
+        &self,
+        xs: &Tensor,
+        chunk_samples: Option<usize>,
+        overlap_samples: Option<usize>,
+    ) -> Result<Tensor> {
+        let chunk_size = chunk_samples.unwrap_or(48000 * 30);
+        let overlap = overlap_samples.unwrap_or(48000 * 2);
+        let total_samples = xs.dim(2)?;
+
+        // Short audio: encode directly
+        if total_samples <= chunk_size {
+            return self.encode(xs);
+        }
+
+        if chunk_size <= 2 * overlap {
+            candle_core::bail!(
+                "tiled_encode: chunk_size ({chunk_size}) must be > 2 * overlap ({overlap})"
+            );
+        }
+        let stride = chunk_size - 2 * overlap;
+        let num_steps = total_samples.div_ceil(stride);
+        let mut downsample_factor: Option<f64> = None;
+        let mut latent_chunks = Vec::with_capacity(num_steps);
+
+        for i in 0..num_steps {
+            let core_start = i * stride;
+            let core_end = (core_start + stride).min(total_samples);
+
+            // Window with overlap
+            let win_start = core_start.saturating_sub(overlap);
+            let win_end = (core_end + overlap).min(total_samples);
+
+            let audio_chunk = xs.narrow(2, win_start, win_end - win_start)?;
+            let latent_chunk = self.encode(&audio_chunk)?;
+
+            // Determine downsample factor from first chunk
+            let df = match downsample_factor {
+                Some(df) => df,
+                None => {
+                    let df = audio_chunk.dim(2)? as f64 / latent_chunk.dim(2)? as f64;
+                    downsample_factor = Some(df);
+                    df
+                }
+            };
+
+            // Trim overlap in latent space
+            let added_start = core_start - win_start;
+            let trim_start = (added_start as f64 / df).round() as usize;
+
+            let added_end = win_end - core_end;
+            let trim_end = (added_end as f64 / df).round() as usize;
+
+            let latent_len = latent_chunk.dim(2)?;
+            let end_idx = if trim_end > 0 {
+                latent_len - trim_end
+            } else {
+                latent_len
+            };
+            let core_len = end_idx - trim_start;
+            let latent_core = latent_chunk.narrow(2, trim_start, core_len)?;
+            latent_chunks.push(latent_core);
+        }
+
+        Tensor::cat(&latent_chunks.iter().collect::<Vec<_>>(), 2)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use candle_core::{DType, Device};
+    use candle_nn::{VarBuilder, VarMap};
+
+    /// Build a small OobleckDecoder with random weights on CPU. The
+    /// architecture is identical to the production decoder; we just shrink the
+    /// channel count so the test runs quickly.
+    fn tiny_decoder() -> Result<(OobleckDecoder, Device)> {
+        let device = Device::Cpu;
+        let varmap = VarMap::new();
+        let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
+        let config = VaeConfig {
+            decoder_channels: 8,
+            decoder_input_channels: 4,
+            audio_channels: 2,
+            // Five blocks, as in production, but 2x each (hop 32 instead of
+            // 1920): the tiling logic is the same, and the full-size hop made
+            // this test peak at 4.6 GB, which got the CI test binary killed.
+            downsampling_ratios: vec![2, 2, 2, 2, 2],
+            channel_multiples: vec![1, 1, 1, 1, 1],
+            ..VaeConfig::default()
+        };
+        let decoder = OobleckDecoder::new(&config, vb)?;
+        Ok((decoder, device))
+    }
+
+    /// Elementwise compare two tensors tolerant to matching NaN/Inf values.
+    /// The decoder has random weights with no normalization, so activations
+    /// can explode past fp32 range — but `forward` and `tiled_decode` do the
+    /// same computation on the same slice, so whatever NaN/Inf pattern
+    /// `forward` produces, `tiled_decode` must produce in the same places.
+    fn tensors_agree(a: &Tensor, b: &Tensor, atol: f32) -> Result<bool> {
+        let a = a.flatten_all()?.to_vec1::<f32>()?;
+        let b = b.flatten_all()?.to_vec1::<f32>()?;
+        assert_eq!(a.len(), b.len(), "shape mismatch");
+        for (x, y) in a.iter().zip(b.iter()) {
+            let ok = match (x.is_nan(), y.is_nan()) {
+                (true, true) => true,
+                (false, false) => {
+                    if x.is_infinite() || y.is_infinite() {
+                        x == y
+                    } else {
+                        let scale = x.abs().max(y.abs()).max(1.0);
+                        (x - y).abs() <= atol * scale
+                    }
+                }
+                _ => false,
+            };
+            if !ok {
+                eprintln!("mismatch: forward={x} tiled={y}");
+                return Ok(false);
+            }
+        }
+        Ok(true)
+    }
+
+    #[test]
+    fn tiled_decode_short_input_delegates_to_forward() -> Result<()> {
+        let (decoder, device) = tiny_decoder()?;
+        // total < chunk → tiled_decode must short-circuit to forward.
+        let z = Tensor::randn(0f32, 1.0, (1, 4, 64), &device)?;
+
+        let y_direct = decoder.forward(&z)?;
+        let y_tiled = decoder.tiled_decode(&z, Some(128), Some(16))?;
+
+        assert_eq!(y_tiled.dims(), y_direct.dims());
+        assert!(tensors_agree(&y_direct, &y_tiled, 1e-5)?);
+        Ok(())
+    }
+
+    #[test]
+    fn tiled_decode_matches_forward_across_chunks() -> Result<()> {
+        let (decoder, device) = tiny_decoder()?;
+        // Force multiple chunk boundaries.
+        let t_latent = 300;
+        let z = Tensor::randn(0f32, 1.0, (1, 4, t_latent), &device)?;
+
+        let y_direct = decoder.forward(&z)?;
+        let y_tiled = decoder.tiled_decode(&z, Some(64), Some(16))?;
+
+        assert_eq!(
+            y_tiled.dims(),
+            y_direct.dims(),
+            "tiled output shape must equal full-forward shape"
+        );
+        // Receptive field fits inside the 16-frame overlap, so cores should
+        // match within fp32 rounding slack. Values can be extreme due to
+        // random untrained weights; compare relative.
+        assert!(tensors_agree(&y_direct, &y_tiled, 1e-3)?);
+        Ok(())
+    }
+
+    #[test]
+    fn tiled_decode_rejects_bad_overlap() -> Result<()> {
+        let (decoder, device) = tiny_decoder()?;
+        let z = Tensor::randn(0f32, 1.0, (1, 4, 200), &device)?;
+
+        // chunk <= 2 * overlap must fail.
+        let err = decoder.tiled_decode(&z, Some(32), Some(16));
+        assert!(
+            err.is_err(),
+            "expected tiled_decode to reject chunk<=2*overlap"
+        );
+        Ok(())
+    }
+}
@@ -17,6 +17,13 @@
 //! masks; and `attn_norm` is required on every layer except layer 0,
 //! which ModernBERT builds as an identity.
 //!
+//! Precision follows PyTorch autocast, which is how the reference runs on
+//! a GPU: the linear layers (and so the attention inputs) run in the
+//! load dtype, while the residual stream, the embeddings and every
+//! LayerNorm stay in f32. Casting the residual stream to bf16 as well
+//! drifts the released English checkpoint's logits by ~0.25 over its 28
+//! layers; keeping it in f32 is what autocast does and costs little.
+//!
 //! Masked positions use a large *finite* negative rather than `-inf`.
 //! A padded query near the end of a short row can find every key in
 //! its sliding window padded; with `-inf` that row's softmax is NaN,
@@ -33,6 +40,26 @@ use super::config::{EncoderConfig, LayerKind};
 /// Additive mask value for a position that must not be attended.
 pub(crate) const MASKED: f32 = -1e9;

+/// Apply `lin` in its weight's dtype, casting `x` in if needed. The
+/// result stays in that dtype; callers cast back to f32 where it joins
+/// the residual stream.
+pub(crate) fn project(lin: &Linear, x: &Tensor) -> candle_core::Result<Tensor> {
+    x.to_dtype(lin.weight().dtype())?.apply(lin)
+}
+
+/// A bias-free LayerNorm that takes candle's fused kernel.
+///
+/// `candle_nn::LayerNorm` dispatches to the fused `layer_norm` op only
+/// when it has a bias; without one it falls back to six elementwise
+/// kernels (mean, subtract, square, mean, divide, scale), which was a
+/// quarter of the encoder's GPU time. ModernBERT's norms have no bias,
+/// so they get an explicit zero one: same result, one kernel.
+fn fused_layer_norm_no_bias(size: usize, eps: f64, vb: VarBuilder) -> Result<LayerNorm> {
+    let weight = vb.get(size, "weight")?;
+    let bias = weight.zeros_like()?;
+    Ok(LayerNorm::new(weight, bias, eps))
+}
+
 /// Precomputed rotary tables for one base, `(max_pos, head_dim / 2)`.
 struct Rotary {
    cos: Tensor,
@@ -78,37 +105,165 @@ struct Attention {
 }

 impl Attention {
-    fn forward(&self, xs: &Tensor, rotary: &Rotary, mask: &Tensor) -> Result<Tensor> {
+    fn forward(&self, xs: &Tensor, rotary: &Rotary, mask: Option<&Tensor>) -> Result<Tensor> {
        let (b, l, d) = xs.dims3()?;
-        let qkv = xs
-            .apply(&self.wqkv)?
+        let qkv = project(&self.wqkv, xs)?
            .reshape((b, l, 3, self.n_heads, self.head_dim))?
            .permute((2, 0, 3, 1, 4))?;
+        // Rotary in the projection dtype, as under autocast; the scores
+        // are then taken in f32 (see `scaled_attention`).
        let q = qkv.get(0)?;
        let k = qkv.get(1)?;
        let v = qkv.get(2)?.contiguous()?;
        let (q, k) = rotary.apply(&q, &k, l)?;
        let ctx = scaled_attention(&q, &k, &v, mask, self.head_dim)?;
        let ctx = ctx.transpose(1, 2)?.reshape((b, l, d))?;
-        Ok(ctx.apply(&self.wo)?)
+        Ok(ctx.apply(&self.wo)?.to_dtype(DType::F32)?)
+    }
+}
+
+/// A batch packed without padding for the flash-attention path: every
+/// row's tokens back to back, `(total, …)`. Linear layers then do no work
+/// for padding, and flash varlen keeps rows from attending to each other.
+#[cfg(feature = "flash-attn")]
+pub struct Packed {
+    /// `(b + 1,)` u32 cumulative row lengths, as flash varlen takes them.
+    pub cu_seqlens: Tensor,
+    /// Longest row.
+    pub max_len: usize,
+    /// `(total,)` u32: each token's position within its own row.
+    pub positions: Tensor,
+    /// Where each row starts in the packed order.
+    pub offsets: Vec<usize>,
+}
+
+#[cfg(feature = "flash-attn")]
+impl Packed {
+    pub fn new(lens: &[usize], device: &Device) -> Result<Self> {
+        let mut cu = Vec::with_capacity(lens.len() + 1);
+        let mut offsets = Vec::with_capacity(lens.len());
+        let mut positions = Vec::with_capacity(lens.iter().sum());
+        let mut total = 0u32;
+        cu.push(0u32);
+        for &n in lens {
+            offsets.push(total as usize);
+            positions.extend(0..n as u32);
+            total += n as u32;
+            cu.push(total);
+        }
+        let n = cu.len();
+        Ok(Self {
+            cu_seqlens: Tensor::from_vec(cu, n, device)?,
+            max_len: lens.iter().copied().max().unwrap_or(0),
+            positions: Tensor::from_vec(positions, total as usize, device)?,
+            offsets,
+        })
+    }
+}
+
+/// Bidirectional flash attention over a packed batch. `q`, `k`, `v` are
+/// `(total, h, dh)`; `half_window` limits each query to keys within that
+/// distance on either side (ModernBERT's sliding layers), `None` attends
+/// the whole row. Scores are accumulated in f32 inside the kernel.
+#[cfg(feature = "flash-attn")]
+pub(crate) fn flash_attention(
+    q: &Tensor,
+    k: &Tensor,
+    v: &Tensor,
+    packed: &Packed,
+    half_window: Option<usize>,
+) -> Result<Tensor> {
+    let head_dim = q.dim(2)?;
+    let scale = (head_dim as f64).powf(-0.5) as f32;
+    Ok(candle_flash_attn::flash_attn_varlen_windowed(
+        q,
+        k,
+        v,
+        &packed.cu_seqlens,
+        &packed.cu_seqlens,
+        packed.max_len,
+        packed.max_len,
+        scale,
+        half_window,
+        half_window,
+    )?)
+}
+
+/// Split a packed `(total, 3·d)` qkv projection into contiguous
+/// `(total, h, dh)` q, k and v.
+#[cfg(feature = "flash-attn")]
+pub(crate) fn split_qkv(
+    qkv: &Tensor,
+    n_heads: usize,
+    head_dim: usize,
+) -> Result<(Tensor, Tensor, Tensor)> {
+    let t = qkv.dim(0)?;
+    let qkv = qkv.reshape((t, 3, n_heads, head_dim))?;
+    let part = |i: usize| -> Result<Tensor> { Ok(qkv.narrow(1, i, 1)?.squeeze(1)?.contiguous()?) };
+    Ok((part(0)?, part(1)?, part(2)?))
+}
+
+#[cfg(feature = "flash-attn")]
+impl Rotary {
+    /// Rotate packed `(total, h, dh)` q and k by each token's position.
+    fn apply_packed(&self, q: &Tensor, k: &Tensor, positions: &Tensor) -> Result<(Tensor, Tensor)> {
+        let cos = self.cos.index_select(positions, 0)?;
+        let sin = self.sin.index_select(positions, 0)?;
+        let rot = |x: &Tensor| -> Result<Tensor> {
+            Ok(candle_nn::rotary_emb::rope_thd(&x.unsqueeze(0)?, &cos, &sin)?.squeeze(0)?)
+        };
+        Ok((rot(q)?, rot(k)?))
+    }
+}
+
+#[cfg(feature = "flash-attn")]
+impl Attention {
+    /// `xs` is the packed f32 residual stream `(total, d)`.
+    fn forward_packed(
+        &self,
+        xs: &Tensor,
+        rotary: &Rotary,
+        packed: &Packed,
+        half_window: usize,
+    ) -> Result<Tensor> {
+        let (t, d) = xs.dims2()?;
+        let (q, k, v) = split_qkv(&project(&self.wqkv, xs)?, self.n_heads, self.head_dim)?;
+        let (q, k) = rotary.apply_packed(&q, &k, &packed.positions)?;
+        let window = match self.kind {
+            LayerKind::Sliding => Some(half_window),
+            LayerKind::Global => None,
+        };
+        let ctx = flash_attention(&q, &k, &v, packed, window)?.reshape((t, d))?;
+        Ok(ctx.apply(&self.wo)?.to_dtype(DType::F32)?)
    }
 }

 /// `softmax(q kᵀ / √d + mask) v`, scores and softmax in f32.
 ///
 /// `q`, `k`, `v` are `(b, h, l, d)`; `mask` is an f32 additive mask
-/// broadcastable to `(b, h, l, l)`.
+/// broadcastable to `(b, h, l, l)`, or `None` when nothing is masked —
+/// which saves a bandwidth-bound pass over the `(b, h, l, l)` scores.
 pub(crate) fn scaled_attention(
    q: &Tensor,
    k: &Tensor,
    v: &Tensor,
-    mask: &Tensor,
+    mask: Option<&Tensor>,
    head_dim: usize,
 ) -> Result<Tensor> {
-    let dtype = q.dtype();
+    let dtype = v.dtype();
    let scale = (head_dim as f64).powf(-0.5);
-    let scores = (q.matmul(&k.t()?)? * scale)?.to_dtype(DType::F32)?;
-    let probs = candle_nn::ops::softmax_last_dim(&scores.broadcast_add(mask)?)?;
+    // Scores in f32 end to end: a bf16 `q kᵀ` is rounded to 8 mantissa
+    // bits before the softmax ever sees it, which on this model's
+    // logit-sized scores is most of the drift from the reference's GPU
+    // path (whose SDPA kernels keep scores in f32).
+    // The scale goes on q, not on the `(b, h, l, l)` scores: the same
+    // product, on a tensor `l / d` times smaller.
+    let scores = (q.to_dtype(DType::F32)? * scale)?.matmul(&k.to_dtype(DType::F32)?.t()?)?;
+    let scores = match mask {
+        Some(mask) => scores.broadcast_add(mask)?,
+        None => scores,
+    };
+    let probs = candle_nn::ops::softmax_last_dim(&scores)?;
    Ok(probs.to_dtype(dtype)?.matmul(v)?)
 }

@@ -118,10 +273,13 @@ struct Mlp {
 }

 impl Module for Mlp {
+    /// f32 in, f32 out; the projections run in the weights' dtype.
    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
-        let xs = xs.apply(&self.wi)?;
+        let xs = project(&self.wi, xs)?;
        let parts = xs.chunk(2, D::Minus1)?;
-        (parts[0].gelu_erf()? * &parts[1])?.apply(&self.wo)
+        (parts[0].gelu_erf()? * &parts[1])?
+            .apply(&self.wo)?
+            .to_dtype(DType::F32)
    }
 }

@@ -132,13 +290,16 @@ struct Layer {
    mlp: Mlp,
 }

-/// The two additive masks one forward needs, both f32.
+/// The two additive masks one forward needs, both f32. Either is `None`
+/// when it would mask nothing: `global` for a batch with no padding (every
+/// single-question request), `sliding` additionally when no row is longer
+/// than the window reaches.
 pub struct Masks {
    /// `(b, 1, 1, l)`: padded keys masked. Global layers and the
    /// decision head use this one.
-    pub global: Tensor,
+    pub global: Option<Tensor>,
    /// `(b, 1, l, l)`: padded keys, plus keys outside the window.
-    pub sliding: Tensor,
+    pub sliding: Option<Tensor>,
 }

 impl Masks {
@@ -151,11 +312,21 @@ impl Masks {
        device: &Device,
    ) -> Result<Self> {
        let b = lens.len();
-        let mut pad = Vec::with_capacity(b * seq_len);
-        for &n in lens {
-            pad.extend((0..seq_len).map(|j| if j < n { 0f32 } else { MASKED }));
+        let padded = lens.iter().any(|&n| n < seq_len);
+        let windowed = seq_len > half_window + 1;
+        let global = if padded {
+            let mut pad = Vec::with_capacity(b * seq_len);
+            for &n in lens {
+                pad.extend((0..seq_len).map(|j| if j < n { 0f32 } else { MASKED }));
+            }
+            Some(Tensor::from_vec(pad, (b, 1, 1, seq_len), device)?)
+        } else {
+            None
+        };
+        if !windowed {
+            let sliding = global.clone();
+            return Ok(Self { global, sliding });
        }
-        let global = Tensor::from_vec(pad, (b, 1, 1, seq_len), device)?;
        // Built on the device from two aranges, not on the host: at
        // 1024 tokens a host-built window is a 4 MB upload per forward.
        let pos = Tensor::arange(0u32, seq_len as u32, device)?.to_dtype(DType::F32)?;
@@ -169,8 +340,14 @@ impl Masks {
        let window = outside
            .where_cond(&masked, &zeros)?
            .reshape((1, 1, seq_len, seq_len))?;
-        let sliding = global.broadcast_add(&window)?;
-        Ok(Self { global, sliding })
+        let sliding = match &global {
+            Some(g) => g.broadcast_add(&window)?,
+            None => window,
+        };
+        Ok(Self {
+            global,
+            sliding: Some(sliding),
+        })
    }
 }

@@ -186,23 +363,30 @@ pub struct ModernBert {

 impl ModernBert {
    /// Load from a VarBuilder rooted at the encoder's prefix (`encoder`
-    /// in a Laya checkpoint, `model` in a stock ModernBERT one).
+    /// in a Laya checkpoint, `model` in a stock ModernBERT one). Linear
+    /// weights load in the VarBuilder's dtype; embeddings and norms in
+    /// f32.
    pub fn load(vb: VarBuilder, cfg: &EncoderConfig) -> Result<Self> {
        let d = cfg.hidden_size;
-        let embeddings =
-            candle_nn::embedding(cfg.vocab_size, d, vb.pp("embeddings").pp("tok_embeddings"))
-                .context("encoder.embeddings.tok_embeddings")?;
-        let embed_norm = candle_nn::layer_norm_no_bias(d, cfg.norm_eps, vb.pp("embeddings.norm"))?;
+        let f32_vb = vb.clone().set_dtype(DType::F32);
+        let embeddings = candle_nn::embedding(
+            cfg.vocab_size,
+            d,
+            f32_vb.pp("embeddings").pp("tok_embeddings"),
+        )
+        .context("encoder.embeddings.tok_embeddings")?;
+        let embed_norm = fused_layer_norm_no_bias(d, cfg.norm_eps, f32_vb.pp("embeddings.norm"))?;
        let mut layers = Vec::with_capacity(cfg.num_hidden_layers);
        for (i, kind) in cfg.layer_kinds.iter().copied().enumerate() {
            let lvb = vb.pp(format!("layers.{i}"));
+            let norm_vb = f32_vb.pp(format!("layers.{i}"));
            // ModernBERT builds layer 0's attn_norm as an identity and
            // saves nothing for it; every other layer must have one.
            let attn_norm = if lvb.contains_tensor("attn_norm.weight") {
-                Some(candle_nn::layer_norm_no_bias(
+                Some(fused_layer_norm_no_bias(
                    d,
                    cfg.norm_eps,
-                    lvb.pp("attn_norm"),
+                    norm_vb.pp("attn_norm"),
                )?)
            } else {
                ensure!(
@@ -219,7 +403,7 @@ impl ModernBert {
                head_dim: cfg.head_dim(),
                kind,
            };
-            let mlp_norm = candle_nn::layer_norm_no_bias(d, cfg.norm_eps, lvb.pp("mlp_norm"))?;
+            let mlp_norm = fused_layer_norm_no_bias(d, cfg.norm_eps, norm_vb.pp("mlp_norm"))?;
            let mlp = Mlp {
                wi: candle_nn::linear_no_bias(d, 2 * cfg.intermediate_size, lvb.pp("mlp.Wi"))?,
                wo: candle_nn::linear_no_bias(cfg.intermediate_size, d, lvb.pp("mlp.Wo"))?,
@@ -231,16 +415,15 @@ impl ModernBert {
                mlp,
            });
        }
-        let final_norm = candle_nn::layer_norm_no_bias(d, cfg.norm_eps, vb.pp("final_norm"))?;
+        let final_norm = fused_layer_norm_no_bias(d, cfg.norm_eps, f32_vb.pp("final_norm"))?;
        let device = vb.device();
-        let dtype = vb.dtype();
        Ok(Self {
            embeddings,
            embed_norm,
            layers,
            final_norm,
-            global_rotary: Rotary::new(cfg, cfg.global_rope_theta, dtype, device)?,
-            sliding_rotary: Rotary::new(cfg, cfg.local_rope_theta, dtype, device)?,
+            global_rotary: Rotary::new(cfg, cfg.global_rope_theta, vb.dtype(), device)?,
+            sliding_rotary: Rotary::new(cfg, cfg.local_rope_theta, vb.dtype(), device)?,
            half_window: cfg.local_attention / 2,
        })
    }
@@ -249,7 +432,8 @@ impl ModernBert {
        self.half_window
    }

-    /// `ids` is `(b, l)` u32; returns the final hidden states `(b, l, d)`.
+    /// `ids` is `(b, l)` u32; returns the final hidden states `(b, l, d)`
+    /// in f32.
    pub fn forward(&self, ids: &Tensor, masks: &Masks) -> Result<Tensor> {
        let mut xs = ids.apply(&self.embeddings)?.apply(&self.embed_norm)?;
        for layer in &self.layers {
@@ -258,8 +442,8 @@ impl ModernBert {
                None => xs.clone(),
            };
            let (rotary, mask) = match layer.attn.kind {
-                LayerKind::Global => (&self.global_rotary, &masks.global),
-                LayerKind::Sliding => (&self.sliding_rotary, &masks.sliding),
+                LayerKind::Global => (&self.global_rotary, masks.global.as_ref()),
+                LayerKind::Sliding => (&self.sliding_rotary, masks.sliding.as_ref()),
            };
            xs = (&xs + layer.attn.forward(&normed, rotary, mask)?)?;
            let mlp_out = xs.apply(&layer.mlp_norm)?.apply(&layer.mlp)?;
@@ -267,4 +451,28 @@ impl ModernBert {
        }
        Ok(xs.apply(&self.final_norm)?)
    }
+
+    /// The flash path: `ids` is the packed `(total,)` u32 token stream;
+    /// returns the packed final hidden states `(total, d)` in f32.
+    #[cfg(feature = "flash-attn")]
+    pub fn forward_packed(&self, ids: &Tensor, packed: &Packed) -> Result<Tensor> {
+        let mut xs = ids.apply(&self.embeddings)?.apply(&self.embed_norm)?;
+        for layer in &self.layers {
+            let normed = match &layer.attn_norm {
+                Some(norm) => xs.apply(norm)?,
+                None => xs.clone(),
+            };
+            let rotary = match layer.attn.kind {
+                LayerKind::Global => &self.global_rotary,
+                LayerKind::Sliding => &self.sliding_rotary,
+            };
+            xs = (&xs
+                + layer
+                    .attn
+                    .forward_packed(&normed, rotary, packed, self.half_window)?)?;
+            let mlp_out = xs.apply(&layer.mlp_norm)?.apply(&layer.mlp)?;
+            xs = (xs + mlp_out)?;
+        }
+        Ok(xs.apply(&self.final_norm)?)
+    }
 }
@@ -15,10 +15,10 @@
 //!    summary features of the (uncalibrated) answer distribution.

 use anyhow::Result;
-use candle_core::{DType, Module, Tensor};
+use candle_core::{DType, Tensor};
 use candle_nn::{Embedding, LayerNorm, Linear, VarBuilder};

-use super::backbone::scaled_attention;
+use super::backbone::{project, scaled_attention};

 /// PyTorch's `LayerNorm` default.
 const TORCH_LN_EPS: f64 = 1e-5;
@@ -39,11 +39,12 @@ impl HeadLayer {
        // `nn.TransformerEncoder` is built with `nhead = max(1, d // 64)`
        // and `dim_feedforward = 4 * d`.
        let n_heads = (d / 64).max(1);
+        let f32_vb = vb.clone().set_dtype(DType::F32);
        Ok(Self {
-            norm1: candle_nn::layer_norm(d, TORCH_LN_EPS, vb.pp("norm1"))?,
+            norm1: candle_nn::layer_norm(d, TORCH_LN_EPS, f32_vb.pp("norm1"))?,
            in_proj: linear_torch_mha_in_proj(d, vb.pp("self_attn"))?,
            out_proj: candle_nn::linear(d, d, vb.pp("self_attn.out_proj"))?,
-            norm2: candle_nn::layer_norm(d, TORCH_LN_EPS, vb.pp("norm2"))?,
+            norm2: candle_nn::layer_norm(d, TORCH_LN_EPS, f32_vb.pp("norm2"))?,
            linear1: candle_nn::linear(d, 4 * d, vb.pp("linear1"))?,
            linear2: candle_nn::linear(4 * d, d, vb.pp("linear2"))?,
            n_heads,
@@ -51,11 +52,11 @@ impl HeadLayer {
        })
    }

-    fn forward(&self, xs: &Tensor, key_mask: &Tensor) -> Result<Tensor> {
+    /// `xs` is the f32 residual stream; projections run in the weights'
+    /// dtype and join it back in f32, as under autocast.
+    fn forward(&self, xs: &Tensor, key_mask: Option<&Tensor>) -> Result<Tensor> {
        let (b, l, d) = xs.dims3()?;
-        let qkv = xs
-            .apply(&self.norm1)?
-            .apply(&self.in_proj)?
+        let qkv = project(&self.in_proj, &xs.apply(&self.norm1)?)?
            .reshape((b, l, 3, self.n_heads, self.head_dim))?
            .permute((2, 0, 3, 1, 4))?;
        let q = qkv.get(0)?.contiguous()?;
@@ -64,12 +65,27 @@ impl HeadLayer {
        let ctx = scaled_attention(&q, &k, &v, key_mask, self.head_dim)?
            .transpose(1, 2)?
            .reshape((b, l, d))?;
-        let xs = (xs + ctx.apply(&self.out_proj)?)?;
-        let ff = xs
-            .apply(&self.norm2)?
-            .apply(&self.linear1)?
+        let xs = (xs + ctx.apply(&self.out_proj)?.to_dtype(DType::F32)?)?;
+        let ff = project(&self.linear1, &xs.apply(&self.norm2)?)?
            .relu()?
-            .apply(&self.linear2)?;
+            .apply(&self.linear2)?
+            .to_dtype(DType::F32)?;
+        Ok((xs + ff)?)
+    }
+
+    /// The flash path over a packed `(total, d)` f32 stream.
+    #[cfg(feature = "flash-attn")]
+    fn forward_packed(&self, xs: &Tensor, packed: &super::backbone::Packed) -> Result<Tensor> {
+        use super::backbone::{flash_attention, split_qkv};
+        let (t, d) = xs.dims2()?;
+        let qkv = project(&self.in_proj, &xs.apply(&self.norm1)?)?;
+        let (q, k, v) = split_qkv(&qkv, self.n_heads, self.head_dim)?;
+        let ctx = flash_attention(&q, &k, &v, packed, None)?.reshape((t, d))?;
+        let xs = (xs + ctx.apply(&self.out_proj)?.to_dtype(DType::F32)?)?;
+        let ff = project(&self.linear1, &xs.apply(&self.norm2)?)?
+            .relu()?
+            .apply(&self.linear2)?
+            .to_dtype(DType::F32)?;
        Ok((xs + ff)?)
    }
 }
@@ -95,14 +111,17 @@ pub struct DecisionHead {

 impl DecisionHead {
    /// `vb` is the checkpoint root; the head's tensors are top-level.
+    /// Linear weights load in `vb`'s dtype; the type embedding and norms
+    /// in f32.
    pub fn load(vb: VarBuilder, d: usize, head_layers: usize, n_act: usize) -> Result<Self> {
        let layers = (0..head_layers)
            .map(|i| HeadLayer::load(vb.pp(format!("head.layers.{i}")), d))
            .collect::<Result<Vec<_>>>()?;
+        let f32_vb = vb.clone().set_dtype(DType::F32);
        Ok(Self {
-            type_emb: candle_nn::embedding(3, d, vb.pp("type_emb"))?,
+            type_emb: candle_nn::embedding(3, d, f32_vb.pp("type_emb"))?,
            layers,
-            scorer_norm: candle_nn::layer_norm(d, TORCH_LN_EPS, vb.pp("scorer.0"))?,
+            scorer_norm: candle_nn::layer_norm(d, TORCH_LN_EPS, f32_vb.pp("scorer.0"))?,
            scorer_fc: candle_nn::linear(d, d, vb.pp("scorer.1"))?,
            scorer_out: candle_nn::linear(d, 1, vb.pp("scorer.3"))?,
            act_fc: candle_nn::linear(d + 4, 256, vb.pp("act_head.0"))?,
@@ -112,7 +131,12 @@ impl DecisionHead {

    /// Type embedding plus the transformer layers. `h` is the encoder
    /// output `(b, l, d)`, `qtype` is `(b,)` u32.
-    pub fn contextualise(&self, h: &Tensor, qtype: &Tensor, key_mask: &Tensor) -> Result<Tensor> {
+    pub fn contextualise(
+        &self,
+        h: &Tensor,
+        qtype: &Tensor,
+        key_mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
        let mut h = h.broadcast_add(&qtype.apply(&self.type_emb)?.unsqueeze(1)?)?;
        for layer in &self.layers {
            h = layer.forward(&h, key_mask)?;
@@ -120,11 +144,25 @@ impl DecisionHead {
        Ok(h)
    }

+    /// The flash path: `h` is the packed encoder output `(total, d)` and
+    /// `qtype` the per-token question type `(total,)` u32.
+    #[cfg(feature = "flash-attn")]
+    pub fn contextualise_packed(
+        &self,
+        h: &Tensor,
+        qtype: &Tensor,
+        packed: &super::backbone::Packed,
+    ) -> Result<Tensor> {
+        let mut h = (h + qtype.apply(&self.type_emb)?)?;
+        for layer in &self.layers {
+            h = layer.forward_packed(&h, packed)?;
+        }
+        Ok(h)
+    }
+
    /// Score gathered marker states `(b, k, d)` → `(b, k)` f32 logits.
    pub fn score(&self, markers: &Tensor) -> Result<Tensor> {
-        let x = markers
-            .apply(&self.scorer_norm)?
-            .apply(&self.scorer_fc)?
+        let x = project(&self.scorer_fc, &markers.apply(&self.scorer_norm)?)?
            .gelu_erf()?
            .apply(&self.scorer_out)?;
        Ok(x.squeeze(2)?.to_dtype(DType::F32)?)
@@ -134,7 +172,9 @@ impl DecisionHead {
    /// `(b, d)` and the four answer-distribution features `(b, 4)`.
    pub fn act(&self, pooled: &Tensor, feats: &Tensor) -> Result<Tensor> {
        let x = Tensor::cat(&[pooled, &feats.to_dtype(pooled.dtype())?], 1)?;
-        let x = self.act_fc.forward(&x)?.gelu_erf()?.apply(&self.act_out)?;
+        let x = project(&self.act_fc, &x)?
+            .gelu_erf()?
+            .apply(&self.act_out)?;
        Ok(x.to_dtype(DType::F32)?)
    }
 }
@@ -16,6 +16,18 @@ pub struct LayaModel {
    head: DecisionHead,
    config: CheckpointConfig,
    device: Device,
+    /// Precision of the linear layers (the load dtype).
+    dtype: DType,
+}
+
+/// A validated batch, flattened: what both forward paths start from.
+struct Flat {
+    lens: Vec<usize>,
+    qtypes: Vec<u32>,
+    /// Each row's marker positions, padded to `k_max` with 0.
+    marker_idx: Vec<u32>,
+    k_max: usize,
+    seq_len: usize,
 }

 impl LayaModel {
@@ -37,6 +49,7 @@ impl LayaModel {
        config: CheckpointConfig,
        device: &Device,
    ) -> Result<Self> {
+        let dtype = vb.dtype();
        let encoder = ModernBert::load(vb.pp("encoder"), &config.encoder)?;
        let head = DecisionHead::load(
            vb.clone(),
@@ -49,6 +62,7 @@ impl LayaModel {
            head,
            config,
            device: device.clone(),
+            dtype,
        })
    }

@@ -56,10 +70,29 @@ impl LayaModel {
        &self.config
    }

-    /// One forward over every sequence in `batch`. Rows are right-padded
-    /// to the longest, as the reference collates them; padding changes
-    /// nothing about an unpadded row's result.
+    /// Precision the linear layers were loaded in.
+    pub fn dtype(&self) -> DType {
+        self.dtype
+    }
+
+    /// One forward over every sequence in `batch`.
+    ///
+    /// On a CUDA device in reduced precision, in a build with the
+    /// `flash-attn` feature, the batch is packed without padding and
+    /// attention runs in flash varlen kernels. Otherwise rows are
+    /// right-padded to the longest, as the reference collates them, and
+    /// attention is eager; padding changes nothing about an unpadded
+    /// row's result.
    pub fn forward(&self, batch: &DecisionBatch) -> Result<Vec<DecisionRow>> {
+        let flat = self.validate(batch)?;
+        #[cfg(feature = "flash-attn")]
+        if self.device.is_cuda() && self.dtype != DType::F32 && flash_enabled() {
+            return self.forward_packed(batch, &flat);
+        }
+        self.forward_padded(batch, &flat)
+    }
+
+    fn validate(&self, batch: &DecisionBatch) -> Result<Flat> {
        let seqs = &batch.seqs;
        ensure!(!seqs.is_empty(), "empty decision batch");
        let max_pos = self.config.encoder.max_position_embeddings;
@@ -70,18 +103,14 @@ impl LayaModel {
            "sequence of {seq_len} tokens exceeds max_position_embeddings {max_pos}"
        );
        let vocab = self.config.encoder.vocab_size as u32;
-        let pad = self.config.encoder.pad_token_id;
-        let b = seqs.len();
-
-        let mut ids = Vec::with_capacity(b * seq_len);
-        let mut lens = Vec::with_capacity(b);
-        let mut qtypes = Vec::with_capacity(b);
        let k_max = seqs.iter().map(|s| s.markers.len()).max().unwrap_or(0);
        ensure!(
            k_max > 0,
            "decision batch has a sequence with no option markers"
        );
-        let mut marker_idx = Vec::with_capacity(b * k_max);
+        let mut lens = Vec::with_capacity(seqs.len());
+        let mut qtypes = Vec::with_capacity(seqs.len());
+        let mut marker_idx = Vec::with_capacity(seqs.len() * k_max);
        for s in seqs {
            let n = s.input_ids.len();
            ensure!(n > 0, "decision batch has an empty sequence");
@@ -93,40 +122,104 @@ impl LayaModel {
            if let Some(&bad) = s.markers.iter().find(|&&m| m >= n) {
                anyhow::bail!("marker position {bad} outside a {n}-token sequence");
            }
-            ids.extend_from_slice(&s.input_ids);
-            ids.extend(std::iter::repeat_n(pad, seq_len - n));
            lens.push(n);
            qtypes.push(s.qtype as u32);
            marker_idx.extend(s.markers.iter().map(|&m| m as u32));
            // Absent slots gather position 0; their scores are discarded.
            marker_idx.extend(std::iter::repeat_n(0u32, k_max - s.markers.len()));
        }
+        Ok(Flat {
+            lens,
+            qtypes,
+            marker_idx,
+            k_max,
+            seq_len,
+        })
+    }

+    /// The eager path: rows right-padded to `(b, seq_len)`.
+    fn forward_padded(&self, batch: &DecisionBatch, flat: &Flat) -> Result<Vec<DecisionRow>> {
+        let b = batch.seqs.len();
+        let seq_len = flat.seq_len;
+        let pad = self.config.encoder.pad_token_id;
+        let mut ids = Vec::with_capacity(b * seq_len);
+        for s in &batch.seqs {
+            ids.extend_from_slice(&s.input_ids);
+            ids.extend(std::iter::repeat_n(pad, seq_len - s.input_ids.len()));
+        }
        let ids = Tensor::from_vec(ids, (b, seq_len), &self.device)?;
-        let masks = Masks::new(&lens, seq_len, self.encoder.half_window(), &self.device)?;
-        let qtype = Tensor::from_vec(qtypes, b, &self.device)?;
+        let masks = Masks::new(
+            &flat.lens,
+            seq_len,
+            self.encoder.half_window(),
+            &self.device,
+        )?;
+        let qtype = Tensor::from_vec(flat.qtypes.clone(), b, &self.device)?;

        let h = self.encoder.forward(&ids, &masks)?;
-        let h = self.head.contextualise(&h, &qtype, &masks.global)?;
+        let h = self.head.contextualise(&h, &qtype, masks.global.as_ref())?;
        let d = self.config.encoder.hidden_size;
-        let idx = Tensor::from_vec(marker_idx, (b, k_max, 1), &self.device)?
-            .broadcast_as((b, k_max, d))?
+        let idx = Tensor::from_vec(flat.marker_idx.clone(), (b, flat.k_max, 1), &self.device)?
+            .broadcast_as((b, flat.k_max, d))?
            .contiguous()?;
        let marker_h = h.gather(&idx, 1)?;
-        let logits: Vec<Vec<f32>> = self.head.score(&marker_h)?.to_vec2()?;
+        let pooled = h.narrow(1, 0, 1)?.squeeze(1)?;
+        self.finish(batch, &marker_h, &pooled)
+    }

+    /// The flash path: tokens packed without padding, `(total, d)`.
+    #[cfg(feature = "flash-attn")]
+    fn forward_packed(&self, batch: &DecisionBatch, flat: &Flat) -> Result<Vec<DecisionRow>> {
+        use super::backbone::Packed;
+        let b = batch.seqs.len();
+        let packed = Packed::new(&flat.lens, &self.device)?;
+        let total: usize = flat.lens.iter().sum();
+        let mut ids = Vec::with_capacity(total);
+        let mut qtype_tok = Vec::with_capacity(total);
+        for (s, &q) in batch.seqs.iter().zip(&flat.qtypes) {
+            ids.extend_from_slice(&s.input_ids);
+            qtype_tok.extend(std::iter::repeat_n(q, s.input_ids.len()));
+        }
+        let ids = Tensor::from_vec(ids, total, &self.device)?;
+        let qtype_tok = Tensor::from_vec(qtype_tok, total, &self.device)?;
+
+        let h = self.encoder.forward_packed(&ids, &packed)?;
+        let h = self.head.contextualise_packed(&h, &qtype_tok, &packed)?;
+        let d = self.config.encoder.hidden_size;
+        // Marker positions, and each row's first token, in packed order.
+        let global: Vec<u32> = flat
+            .marker_idx
+            .chunks(flat.k_max)
+            .zip(&packed.offsets)
+            .flat_map(|(ms, &off)| ms.iter().map(move |&m| off as u32 + m))
+            .collect();
+        let global = Tensor::from_vec(global, b * flat.k_max, &self.device)?;
+        let marker_h = h.index_select(&global, 0)?.reshape((b, flat.k_max, d))?;
+        let firsts: Vec<u32> = packed.offsets.iter().map(|&o| o as u32).collect();
+        let pooled = h.index_select(&Tensor::from_vec(firsts, b, &self.device)?, 0)?;
+        self.finish(batch, &marker_h, &pooled)
+    }
+
+    /// Score the gathered marker states `(b, k_max, d)` and run the act
+    /// head over the pooled first positions `(b, d)`.
+    fn finish(
+        &self,
+        batch: &DecisionBatch,
+        marker_h: &Tensor,
+        pooled: &Tensor,
+    ) -> Result<Vec<DecisionRow>> {
+        let b = batch.seqs.len();
+        let logits: Vec<Vec<f32>> = self.head.score(marker_h)?.to_vec2()?;
        let mut feats = Vec::with_capacity(b * 4);
        let mut row_logits = Vec::with_capacity(b);
-        for (s, full) in seqs.iter().zip(logits) {
+        for (s, full) in batch.seqs.iter().zip(logits) {
            let row: Vec<f32> = full[..s.markers.len()].to_vec();
            feats.extend_from_slice(&act_features(&row));
            row_logits.push(row);
        }
        let feats = Tensor::from_vec(feats, (b, 4), &self.device)?;
-        let pooled = h.narrow(1, 0, 1)?.squeeze(1)?;
-        let act_logits = self.head.act(&pooled, &feats)?;
+        let act_logits = self.head.act(pooled, &feats)?;
        let act = candle_nn::ops::softmax_last_dim(&act_logits)?.to_vec2::<f32>()?;
-
        Ok(row_logits
            .into_iter()
            .zip(act)
@@ -135,6 +228,16 @@ impl LayaModel {
    }
 }

+/// Runtime kill switch for the flash path: `NEURON_FLASH_ATTN=0` (or
+/// `false`) forces eager attention without a rebuild — the same switch
+/// the text architectures read. Read once.
+#[cfg(feature = "flash-attn")]
+fn flash_enabled() -> bool {
+    static ENABLED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
+    *ENABLED
+        .get_or_init(|| !std::env::var("NEURON_FLASH_ATTN").is_ok_and(|v| v == "0" || v == "false"))
+}
+
 /// The act head's four inputs, from one row's uncalibrated logits:
 /// top-1 probability, top-1 minus top-2, entropy normalised by
 /// `ln(max(k, 2))`, and `k / 255`.
@@ -53,6 +53,36 @@ fn batch_of(seqs: &[TinySeq]) -> DecisionBatch {
    }
 }

+/// Device and precision for the weight-backed tests: `LAYA_DEVICE`
+/// (`cpu`, or `cuda:N` in a build with the `cuda` feature) and
+/// `LAYA_DTYPE` (`f32`, `bf16` or `f16`). Defaults: CPU, f32.
+fn test_device() -> (Device, DType) {
+    let device = match std::env::var("LAYA_DEVICE").as_deref() {
+        Ok("cpu") | Err(_) => Device::Cpu,
+        Ok(s) => {
+            let ordinal = s
+                .strip_prefix("cuda:")
+                .and_then(|n| n.parse().ok())
+                .unwrap_or_else(|| panic!("LAYA_DEVICE={s}: expected cpu or cuda:N"));
+            Device::new_cuda(ordinal).unwrap_or_else(|e| panic!("LAYA_DEVICE={s}: {e}"))
+        }
+    };
+    let dtype = match std::env::var("LAYA_DTYPE").as_deref() {
+        Ok("f32") | Err(_) => DType::F32,
+        Ok("bf16") => DType::BF16,
+        Ok("f16") => DType::F16,
+        Ok(s) => panic!("LAYA_DTYPE={s}: expected f32, bf16 or f16"),
+    };
+    (device, dtype)
+}
+
+/// Logit tolerance against the f32 reference for a given precision.
+/// Reduced precision is judged mainly on argmax agreement; the bound
+/// only catches a broken kernel, not rounding.
+fn logit_tolerance(dtype: DType) -> f32 {
+    if dtype == DType::F32 { 1e-3 } else { 1.0 }
+}
+
 fn max_abs_diff(a: &[f32], b: &[f32]) -> f32 {
    assert_eq!(a.len(), b.len(), "length mismatch");
    a.iter()
@@ -154,7 +184,8 @@ struct RefItem {
 }

 /// Every case in `reference.json`, scored by the checkpoint the
-/// reference routed it to, in f32 on the CPU.
+/// reference routed it to, on the device and precision
+/// [`test_device`] selects (f32 on the CPU by default).
 ///
 /// `LAYA_REFERENCE_DIR` is a snapshot of `convaiinnovations/laya` at the
 /// revision recorded in the fixture header: the root holds the English
@@ -166,11 +197,24 @@ fn released_checkpoints_match_reference() {
        panic!("LAYA_REFERENCE_DIR is not set");
    };
    let root = PathBuf::from(root);
+    let (device, dtype) = test_device();
+    let tolerance = logit_tolerance(dtype);
    let reference: Reference = serde_json::from_str(
        &std::fs::read_to_string(testdata().join("reference.json")).expect("read reference.json"),
    )
    .expect("parse reference.json");

+    // The reference's own GPU run (bf16 autocast), for judging reduced
+    // precision against the drift upstream itself serves with.
+    let torch_bf16: Option<Reference> =
+        std::fs::read_to_string(testdata().join("reference_cuda_bf16.json"))
+            .ok()
+            .map(|t| serde_json::from_str(&t).expect("parse reference_cuda_bf16.json"));
+    // Per checkpoint: summed centred drift of ours and of the reference's
+    // bf16 run, their maxima, and the row count.
+    let mut drift: std::collections::BTreeMap<String, (f32, f32, f32, f32, usize)> =
+        Default::default();
+
    let mut models = std::collections::BTreeMap::new();
    let mut stats: std::collections::BTreeMap<String, (usize, f32, f32, usize)> =
        Default::default();
@@ -187,8 +231,7 @@ fn released_checkpoints_match_reference() {
                "english" => root.clone(),
                other => root.join(other),
            };
-            LayaModel::load(&dir, &Device::Cpu, DType::F32)
-                .unwrap_or_else(|e| panic!("load {ckpt}: {e:#}"))
+            LayaModel::load(&dir, &device, dtype).unwrap_or_else(|e| panic!("load {ckpt}: {e:#}"))
        });
        let batch = DecisionBatch {
            seqs: case
@@ -208,6 +251,21 @@ fn released_checkpoints_match_reference() {
        let t = timing.entry(ckpt.clone()).or_default();
        t.0 += started.elapsed().as_secs_f64() * 1e3;
        t.1 += batch.seqs.iter().map(|s| s.input_ids.len()).sum::<usize>();
+        let torch_case = torch_bf16
+            .as_ref()
+            .and_then(|r| r.cases.iter().find(|c| c.name == case.name));
+        for (i, (got, want)) in rows.iter().zip(&case.items).enumerate() {
+            let d = drift.entry(ckpt.clone()).or_default();
+            let ours = centred_diff(&got.logits, &want.logits);
+            d.0 += ours;
+            d.2 = d.2.max(ours);
+            if let Some(t) = torch_case.and_then(|c| c.items.get(i)) {
+                let theirs = centred_diff(&t.logits, &want.logits);
+                d.1 += theirs;
+                d.3 = d.3.max(theirs);
+            }
+            d.4 += 1;
+        }
        let entry = stats.entry(ckpt.clone()).or_default();
        for (got, want) in rows.iter().zip(&case.items) {
            let dl = max_abs_diff(&got.logits, &want.logits);
@@ -224,23 +282,169 @@ fn released_checkpoints_match_reference() {
            if argmax(&got.logits) == argmax(&want.logits) {
                entry.3 += 1;
            }
+            if dtype != DType::F32 && std::env::var_os("LAYA_VERBOSE").is_some() {
+                eprintln!("  {:<40} {ckpt:<16} |Δlogit| {dl:.3}", case.name);
+            }
            assert!(
-                dl < 1e-3,
+                dl < tolerance,
                "{} ({ckpt}): logit diff {dl:e}\n got {:?}\nwant {:?}",
                case.name,
                got.logits,
                want.logits
            );
-            assert!(da < 1e-3, "{} ({ckpt}): act diff {da:e}", case.name);
+            assert!(da < tolerance, "{} ({ckpt}): act diff {da:e}", case.name);
        }
    }
    for (ckpt, (n, dl, da, agree)) in &stats {
        let (ms, tokens) = timing[ckpt];
        eprintln!(
-            "{ckpt}: {n} questions, max |Δlogit| {dl:.2e}, max |Δact| {da:.2e}, argmax {agree}/{n}; \
+            "[{device:?} {dtype:?}] {ckpt}: {n} questions, max |Δlogit| {dl:.2e}, max |Δact| {da:.2e}, argmax {agree}/{n}; \
             {tokens} tokens in {ms:.0} ms ({:.2} ms/token)",
            ms / tokens as f64
        );
        assert_eq!(agree, n, "{ckpt}: argmax disagreement");
    }
+    // Reduced precision: centred drift (softmax ignores a constant shift)
+    // against f32, next to the reference's own bf16 drift. Per-case drift
+    // is chaotic, so the gate is the mean: no worse than twice upstream's,
+    // plus 0.01 — about one bf16 ulp at these logits' magnitude — so a
+    // checkpoint with a single fixture case whose reference drift happens
+    // to round to nothing is not held to zero.
+    for (ckpt, (ours, theirs, ours_max, theirs_max, n)) in &drift {
+        let (ours, theirs) = (ours / *n as f32, theirs / *n as f32);
+        eprintln!(
+            "[{device:?} {dtype:?}] {ckpt}: centred drift mean {ours:.4} (reference bf16 {theirs:.4}), \
+             max {ours_max:.3} (reference bf16 {theirs_max:.3})"
+        );
+        if dtype != DType::F32 && torch_bf16.is_some() {
+            assert!(
+                ours <= 2.0 * theirs + 0.01,
+                "{ckpt}: mean drift {ours:.4} exceeds twice the reference's own bf16 drift {theirs:.4}"
+            );
+        }
+    }
+}
+
+/// Largest difference between two logit rows after removing each row's
+/// mean — the part of a difference a softmax can see.
+fn centred_diff(a: &[f32], b: &[f32]) -> f32 {
+    let mean = |v: &[f32]| v.iter().sum::<f32>() / v.len() as f32;
+    let (ma, mb) = (mean(a), mean(b));
+    a.iter()
+        .zip(b)
+        .map(|(x, y)| ((x - ma) - (y - mb)).abs())
+        .fold(0.0, f32::max)
+}
+
+/// Forward latency for the four workloads #337 benchmarks against the
+/// PyTorch reference, plus batch throughput, on the English checkpoint.
+///
+/// Sequences are real ones from `reference.json`: `choice_basic` is the
+/// short state (41 tokens), `multi_question` five questions over it, and
+/// `state_right_truncated` a full 512-token sequence (repeated for the
+/// five-question and throughput batches). The timed span is the whole
+/// `forward` call — host-to-device copy, forward, and the CPU readback
+/// that ends it — so it is directly comparable to the reference's
+/// end-to-end `predict()` minus tokenisation.
+///
+/// Run with `cargo test --release -p neuron [--features cuda] --
+/// --ignored --nocapture laya_forward_timing`, with `LAYA_REFERENCE_DIR`
+/// and optionally `LAYA_DEVICE` / `LAYA_DTYPE` set.
+#[test]
+#[ignore = "benchmark: needs the released weights (LAYA_REFERENCE_DIR)"]
+fn laya_forward_timing() {
+    let root = PathBuf::from(std::env::var("LAYA_REFERENCE_DIR").expect("LAYA_REFERENCE_DIR"));
+    let (device, dtype) = test_device();
+    let reference: Reference = serde_json::from_str(
+        &std::fs::read_to_string(testdata().join("reference.json")).expect("read reference.json"),
+    )
+    .expect("parse reference.json");
+    let seqs_of = |name: &str| -> Vec<DecisionSeq> {
+        reference
+            .cases
+            .iter()
+            .find(|c| c.name == name)
+            .unwrap_or_else(|| panic!("no case {name}"))
+            .items
+            .iter()
+            .map(|it| DecisionSeq {
+                input_ids: it.input_ids.clone(),
+                markers: it.markers.clone(),
+                qtype: it.qtype,
+            })
+            .collect()
+    };
+    let long = seqs_of("state_right_truncated").remove(0);
+    let workloads: Vec<(String, DecisionBatch)> = vec![
+        (
+            "short x1q".into(),
+            DecisionBatch {
+                seqs: seqs_of("choice_basic"),
+            },
+        ),
+        (
+            "short x5q".into(),
+            DecisionBatch {
+                seqs: seqs_of("multi_question"),
+            },
+        ),
+        (
+            "512tok x1q".into(),
+            DecisionBatch {
+                seqs: vec![long.clone()],
+            },
+        ),
+        (
+            "512tok x5q".into(),
+            DecisionBatch {
+                seqs: vec![long.clone(); 5],
+            },
+        ),
+        (
+            "512tok x16".into(),
+            DecisionBatch {
+                seqs: vec![long.clone(); 16],
+            },
+        ),
+        (
+            "512tok x64".into(),
+            DecisionBatch {
+                seqs: vec![long.clone(); 64],
+            },
+        ),
+    ];
+
+    let model = LayaModel::load(&root, &device, dtype).expect("load English checkpoint");
+    // `LAYA_TIMING_ONLY` keeps the workloads whose name contains it, and
+    // `LAYA_TIMING_ITERS` sets the timed iterations — for profiling one
+    // workload under a tracer.
+    let only = std::env::var("LAYA_TIMING_ONLY").ok();
+    let iters: usize = std::env::var("LAYA_TIMING_ITERS")
+        .ok()
+        .and_then(|v| v.parse().ok())
+        .unwrap_or(20);
+    for (name, batch) in &workloads {
+        if only.as_deref().is_some_and(|o| !name.contains(o)) {
+            continue;
+        }
+        for _ in 0..3 {
+            model.forward(batch).expect("warm-up forward");
+        }
+        let mut ms: Vec<f64> = (0..iters)
+            .map(|_| {
+                let started = std::time::Instant::now();
+                model.forward(batch).expect("forward");
+                started.elapsed().as_secs_f64() * 1e3
+            })
+            .collect();
+        ms.sort_by(f64::total_cmp);
+        let median = ms[ms.len() / 2];
+        let tokens: usize = batch.seqs.iter().map(|s| s.input_ids.len()).sum();
+        eprintln!(
+            "[{device:?} {dtype:?}] {name:>11}: median {median:7.1} ms, min {:7.1} ms, \
+             {tokens} tokens, {:.0} seq/s",
+            ms[0],
+            batch.seqs.len() as f64 / (median / 1e3)
+        );
+    }
 }
@@ -20,6 +20,7 @@
 //! TP-aware analogues live in `harness/tp/tp_<family>.rs` and follow
 //! the pattern set by `tp_qwen3.rs`.

+pub mod acestep;
 pub mod laya;
 pub mod phase_probe;
 pub mod qwen3_5;
@@ -0,0 +1,543 @@
+//! Music generation engine (#383): ACE-Step 1.5 behind the device worker.
+//!
+//! The ACE-Step repo carries every component as a subfolder of one
+//! snapshot — the turbo DiT (with its `silence_latent.pt`), the song
+//! planner LM, the Qwen3 text encoder and the VAE — so the loader takes
+//! the snapshot root and finds them by name.
+//!
+//! Same discipline as the image engine (`image.rs`): the pipeline lives in
+//! the worker's audio slab, one `Job::GenerateAudio` runs a whole song, and
+//! the reply is CPU-side samples. WAV encoding happens on the async side.
+
+use super::arch::acestep::pipeline::{AceStepPipeline, GenerationParams, LmMetadata};
+use super::arch::acestep::{AceStepConfig, VaeConfig, lm::LmConfig, qwen3, sampling};
+use anyhow::{Context, Result};
+use candle_core::{DType, Device};
+use candle_nn::VarBuilder;
+use std::path::{Path, PathBuf};
+use std::time::Instant;
+
+/// The turbo DiT's subfolder.
+pub const DIT_DIR: &str = "acestep-v15-turbo";
+/// The text encoder's subfolder.
+pub const TEXT_ENCODER_DIR: &str = "Qwen3-Embedding-0.6B";
+/// The VAE's subfolder.
+pub const VAE_DIR: &str = "vae";
+/// Planner LMs in order of preference; the first the snapshot has is used.
+pub const PLANNER_DIRS: [&str; 2] = ["acestep-5Hz-lm-1.7B", "acestep-5Hz-lm-0.6B"];
+
+/// What ACE-Step 1.5 turbo serves (#383). Upstream allows 10 minutes; the
+/// planner's code budget (3072 tokens at 5 Hz) covers that. No `steps`
+/// knob: the turbo schedule is a fixed eight steps.
+pub const ACESTEP_LIMITS: cortex_core::audio::AudioLimits = cortex_core::audio::AudioLimits {
+    min_duration: 10.0,
+    max_duration: 600.0,
+    default_duration: 30.0,
+    vocals: true,
+    sample_rate: 48_000,
+    channels: 2,
+    max_steps: None,
+};
+
+/// The turbo schedule's default step count.
+pub const DEFAULT_STEPS: usize = 8;
+/// Turbo models are trained with timestep shift 3.
+const TURBO_SHIFT: f64 = 3.0;
+/// Upstream's default output peak, in dBFS.
+const PEAK_DB: f64 = -1.0;
+
+/// Files every ACE-Step snapshot needs, relative to its root.
+pub fn required_files() -> Vec<String> {
+    vec![
+        format!("{DIT_DIR}/config.json"),
+        format!("{DIT_DIR}/model.safetensors"),
+        format!("{DIT_DIR}/silence_latent.pt"),
+        format!("{TEXT_ENCODER_DIR}/config.json"),
+        format!("{TEXT_ENCODER_DIR}/model.safetensors"),
+        format!("{TEXT_ENCODER_DIR}/tokenizer.json"),
+        format!("{VAE_DIR}/config.json"),
+        format!("{VAE_DIR}/diffusion_pytorch_model.safetensors"),
+    ]
+}
+
+/// Files a planner LM needs, relative to the snapshot root.
+pub fn planner_files(dir: &str) -> Vec<String> {
+    ["config.json", "model.safetensors", "tokenizer.json"]
+        .iter()
+        .map(|f| format!("{dir}/{f}"))
+        .collect()
+}
+
+/// True for a repo laid out like ACE-Step 1.5: the turbo DiT and the VAE as
+/// subfolders. Checked before the generic safetensors sniff, which would
+/// otherwise take the root `config.json` for a text model.
+pub fn is_acestep_layout(filenames: &[&str]) -> bool {
+    filenames.contains(&format!("{DIT_DIR}/config.json").as_str())
+        && filenames.contains(&format!("{VAE_DIR}/diffusion_pytorch_model.safetensors").as_str())
+}
+
+/// A resolved snapshot: the root directory every component sits under, and
+/// which planner (if any) it has.
+#[derive(Debug, Clone)]
+pub struct AceStepFiles {
+    pub root: PathBuf,
+    pub planner: Option<String>,
+}
+
+/// One generation's parameters, already validated against the model's limits.
+#[derive(Debug, Clone)]
+pub struct AudioGenParams {
+    pub prompt: String,
+    /// Empty for an instrumental.
+    pub lyrics: String,
+    pub duration_secs: f64,
+    pub seed: u64,
+    pub steps: Option<usize>,
+    pub guidance_scale: Option<f64>,
+    /// Run the song planner first (when the model has one).
+    pub planner: bool,
+    pub temperature: Option<f64>,
+    /// Lyric language code; `None` lets the planner (or `"unknown"`) decide.
+    pub language: Option<String>,
+}
+
+/// One generation's output: CPU-side stereo samples in `[-1, 1]`.
+#[derive(Debug, Clone)]
+pub struct AudioGenResult {
+    pub left: Vec<f32>,
+    pub right: Vec<f32>,
+    pub sample_rate: u32,
+    pub timing: cortex_core::audio::AudioTiming,
+}
+
+impl AudioGenResult {
+    pub fn duration_secs(&self) -> f64 {
+        self.left.len() as f64 / self.sample_rate.max(1) as f64
+    }
+}
+
+/// The loaded pipeline, owned by the device worker.
+pub struct AceStepEngine {
+    pipeline: AceStepPipeline,
+    has_planner: bool,
+}
+
+fn read_json<T: serde::de::DeserializeOwned>(path: &Path) -> Result<T> {
+    let f = std::fs::File::open(path).with_context(|| format!("open {}", path.display()))?;
+    serde_json::from_reader(f).with_context(|| format!("parse {}", path.display()))
+}
+
+impl AceStepEngine {
+    /// Build every component on `device`: bf16 on a GPU, f32 on the CPU.
+    pub fn load(files: &AceStepFiles, device: &Device) -> Result<Self> {
+        let dtype = if device.is_cpu() {
+            DType::F32
+        } else {
+            DType::BF16
+        };
+        let root = &files.root;
+        let dit = root.join(DIT_DIR);
+        let te = root.join(TEXT_ENCODER_DIR);
+        let vae = root.join(VAE_DIR);
+
+        let dit_config: AceStepConfig = read_json(&dit.join("config.json"))?;
+        let vae_config: VaeConfig = read_json(&vae.join("config.json"))?;
+        let text_config: qwen3::Config = read_json(&te.join("config.json"))?;
+
+        // SAFETY: the snapshot is a read-only cache the daemon owns; the
+        // files are not modified while mapped.
+        let dit_vb = unsafe {
+            VarBuilder::from_mmaped_safetensors(&[dit.join("model.safetensors")], dtype, device)?
+        };
+        let vae_vb = unsafe {
+            VarBuilder::from_mmaped_safetensors(
+                &[vae.join("diffusion_pytorch_model.safetensors")],
+                dtype,
+                device,
+            )?
+        };
+        // The embedding model's tensors carry a `model.` prefix the bare
+        // qwen3 encoder does not expect.
+        let text_vb = unsafe {
+            VarBuilder::from_mmaped_safetensors(&[te.join("model.safetensors")], dtype, device)?
+        }
+        .rename_f(|n| n.strip_prefix("model.").unwrap_or(n).to_string());
+
+        let silence_latent = super::arch::acestep::pipeline::load_silence_latent(
+            &dit.join("silence_latent.pt"),
+            dit_config.audio_acoustic_hidden_dim,
+            dtype,
+            device,
+        )?;
+        let mut pipeline = AceStepPipeline::new(
+            &dit_config,
+            &vae_config,
+            &text_config,
+            dit_vb,
+            vae_vb,
+            text_vb,
+            &te.join("tokenizer.json"),
+            silence_latent,
+            false,
+        )?;
+
+        let has_planner = match &files.planner {
+            Some(dir) => {
+                let lm = root.join(dir);
+                let lm_config: qwen3::Config = read_json(&lm.join("config.json"))?;
+                let lm_vb = unsafe {
+                    VarBuilder::from_mmaped_safetensors(
+                        &[lm.join("model.safetensors")],
+                        dtype,
+                        device,
+                    )?
+                };
+                pipeline.load_lm(&lm_config, lm_vb, &lm.join("tokenizer.json"))?;
+                true
+            }
+            None => false,
+        };
+        Ok(Self {
+            pipeline,
+            has_planner,
+        })
+    }
+
+    pub fn has_planner(&self) -> bool {
+        self.has_planner
+    }
+
+    /// Render one song.
+    pub fn generate(&mut self, p: &AudioGenParams) -> Result<AudioGenResult> {
+        let mut timing = cortex_core::audio::AudioTiming::default();
+
+        let (lm_hints, lm_metadata) = if p.planner && self.has_planner {
+            let t = Instant::now();
+            let cfg = LmConfig {
+                temperature: p.temperature.unwrap_or(0.85),
+                seed: p.seed,
+                ..Default::default()
+            };
+            let (hints, meta) =
+                self.pipeline
+                    .generate_lm_hints(&p.prompt, &p.lyrics, p.duration_secs, &cfg)?;
+            timing.plan_ms = t.elapsed().as_millis() as u64;
+            let metadata = LmMetadata {
+                bpm: meta.get("bpm").cloned(),
+                timesignature: meta.get("timesignature").cloned(),
+                keyscale: meta.get("keyscale").cloned(),
+                caption: meta.get("caption").cloned(),
+                language: meta
+                    .get("language")
+                    .or_else(|| meta.get("vocal_language"))
+                    .cloned(),
+            };
+            (Some(hints), Some(metadata))
+        } else {
+            (None, None)
+        };
+
+        let params = GenerationParams {
+            duration_secs: p.duration_secs,
+            seed: Some(p.seed),
+            num_steps: Some(p.steps.unwrap_or(DEFAULT_STEPS)),
+            guidance_scale: p
+                .guidance_scale
+                .unwrap_or(GenerationParams::default().guidance_scale),
+            shift: TURBO_SHIFT,
+            is_turbo: true,
+            lm_hints_25hz: lm_hints,
+            lm_metadata,
+            infer_method: sampling::InferMethod::Ode,
+            language: p.language.clone(),
+            ..Default::default()
+        };
+        let t = Instant::now();
+        let out = self.pipeline.text2music(&p.prompt, &p.lyrics, &params)?;
+        let audio = out.audio.to_dtype(DType::F32)?.to_device(&Device::Cpu)?;
+        timing.generate_ms = t.elapsed().as_millis() as u64;
+
+        let mut left = audio.get(0)?.to_vec1::<f32>()?;
+        let mut right = audio.get(1)?.to_vec1::<f32>()?;
+        if left.iter().chain(&right).any(|s| !s.is_finite()) {
+            anyhow::bail!("generation produced non-finite samples");
+        }
+        peak_normalize(&mut left, &mut right, PEAK_DB);
+        Ok(AudioGenResult {
+            left,
+            right,
+            sample_rate: out.sample_rate,
+            timing,
+        })
+    }
+}
+
+/// Scale both channels so the loudest sample sits at `target_db` dBFS, as
+/// upstream does before saving. Near-silence is left alone rather than
+/// amplified into noise.
+pub fn peak_normalize(left: &mut [f32], right: &mut [f32], target_db: f64) {
+    let peak = left
+        .iter()
+        .chain(right.iter())
+        .fold(0f32, |m, s| m.max(s.abs()));
+    if peak < 1e-6 {
+        return;
+    }
+    let gain = (10f64.powf(target_db / 20.0) / peak as f64) as f32;
+    for s in left.iter_mut().chain(right.iter_mut()) {
+        *s = (*s * gain).clamp(-1.0, 1.0);
+    }
+}
+
+/// Encode stereo samples as a 16-bit PCM WAV file.
+pub fn encode_wav_pcm16(left: &[f32], right: &[f32], sample_rate: u32) -> Vec<u8> {
+    let frames = left.len().min(right.len());
+    let data_len = (frames * 4) as u32;
+    let mut out = Vec::with_capacity(44 + frames * 4);
+    out.extend_from_slice(b"RIFF");
+    out.extend_from_slice(&(36 + data_len).to_le_bytes());
+    out.extend_from_slice(b"WAVE");
+    out.extend_from_slice(b"fmt ");
+    out.extend_from_slice(&16u32.to_le_bytes()); // fmt chunk size
+    out.extend_from_slice(&1u16.to_le_bytes()); // PCM
+    out.extend_from_slice(&2u16.to_le_bytes()); // channels
+    out.extend_from_slice(&sample_rate.to_le_bytes());
+    out.extend_from_slice(&(sample_rate * 4).to_le_bytes()); // byte rate
+    out.extend_from_slice(&4u16.to_le_bytes()); // block align
+    out.extend_from_slice(&16u16.to_le_bytes()); // bits per sample
+    out.extend_from_slice(b"data");
+    out.extend_from_slice(&data_len.to_le_bytes());
+    let pcm = |s: f32| -> [u8; 2] {
+        ((s.clamp(-1.0, 1.0) * i16::MAX as f32).round() as i16).to_le_bytes()
+    };
+    for i in 0..frames {
+        out.extend_from_slice(&pcm(left[i]));
+        out.extend_from_slice(&pcm(right[i]));
+    }
+    out
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn layout_detection() {
+        let mut files: Vec<String> = required_files();
+        files.push("config.json".into());
+        let refs: Vec<&str> = files.iter().map(String::as_str).collect();
+        assert!(is_acestep_layout(&refs));
+        assert!(!is_acestep_layout(&["config.json", "model.safetensors"]));
+        assert!(!is_acestep_layout(&["acestep-v15-turbo/config.json"]));
+    }
+
+    #[test]
+    fn peak_normalize_hits_target_and_spares_silence() {
+        let (mut l, mut r) = (vec![0.5, -0.25], vec![0.1, 0.0]);
+        peak_normalize(&mut l, &mut r, -1.0);
+        let target = 10f32.powf(-1.0 / 20.0);
+        assert!((l[0] - target).abs() < 1e-6);
+        assert!((l[1] + target / 2.0).abs() < 1e-6);
+        let (mut l, mut r) = (vec![1e-8, 0.0], vec![0.0, -1e-8]);
+        peak_normalize(&mut l, &mut r, -1.0);
+        assert_eq!(l, vec![1e-8, 0.0]);
+    }
+
+    #[test]
+    fn wav_header_and_samples() {
+        let wav = encode_wav_pcm16(&[1.0, -1.0], &[0.0, 0.5], 48_000);
+        assert_eq!(wav.len(), 44 + 8);
+        assert_eq!(&wav[0..4], b"RIFF");
+        assert_eq!(u32::from_le_bytes(wav[4..8].try_into().unwrap()), 36 + 8);
+        assert_eq!(&wav[8..12], b"WAVE");
+        assert_eq!(u16::from_le_bytes(wav[22..24].try_into().unwrap()), 2);
+        assert_eq!(u32::from_le_bytes(wav[24..28].try_into().unwrap()), 48_000);
+        assert_eq!(u32::from_le_bytes(wav[40..44].try_into().unwrap()), 8);
+        let s = |i: usize| i16::from_le_bytes(wav[44 + 2 * i..46 + 2 * i].try_into().unwrap());
+        assert_eq!((s(0), s(1), s(2), s(3)), (i16::MAX, 0, -i16::MAX, 16384));
+    }
+}
+
+/// End-to-end render against real weights, on the CPU. Ignored by default:
+/// it needs an ACE-Step 1.5 snapshot on disk and minutes of CPU time.
+///
+///   ACESTEP_SNAPSHOT=/path/to/Ace-Step1.5 \
+///     cargo test -p neuron --lib -- --ignored acestep_renders_on_cpu --nocapture
+#[cfg(test)]
+mod hardware {
+    use super::*;
+    use crate::harness::arch::acestep::pipeline::AceStepPipeline;
+
+    #[test]
+    #[ignore = "needs ACESTEP_SNAPSHOT and minutes of CPU"]
+    fn acestep_renders_on_cpu() {
+        let Ok(root) = std::env::var("ACESTEP_SNAPSHOT") else {
+            eprintln!("ACESTEP_SNAPSHOT not set; skipping");
+            return;
+        };
+        let files = AceStepFiles {
+            root: root.into(),
+            planner: None,
+        };
+        let t = Instant::now();
+        let mut engine = AceStepEngine::load(&files, &Device::Cpu).expect("load");
+        eprintln!("loaded in {:.1}s", t.elapsed().as_secs_f32());
+        let params = AudioGenParams {
+            prompt: "Warm lo-fi hip hop beat with mellow Rhodes chords, 82 BPM".into(),
+            lyrics: String::new(),
+            duration_secs: 10.0,
+            seed: 42,
+            steps: None,
+            guidance_scale: None,
+            planner: false,
+            temperature: None,
+            language: None,
+        };
+        let t = Instant::now();
+        let out = engine.generate(&params).expect("generate");
+        eprintln!(
+            "rendered {:.2}s of audio in {:.1}s ({:?})",
+            out.duration_secs(),
+            t.elapsed().as_secs_f32(),
+            out.timing
+        );
+        assert_eq!(out.sample_rate, 48_000);
+        assert!(
+            (out.duration_secs() - 10.0).abs() < 0.5,
+            "{}",
+            out.duration_secs()
+        );
+        let rms = (out.left.iter().map(|s| s * s).sum::<f32>() / out.left.len() as f32).sqrt();
+        assert!(rms > 1e-3, "silent render, rms {rms}");
+        if let Ok(path) = std::env::var("ACESTEP_OUT") {
+            std::fs::write(
+                &path,
+                encode_wav_pcm16(&out.left, &out.right, out.sample_rate),
+            )
+            .expect("write wav");
+            eprintln!("wrote {path}");
+        }
+        // Same seed, same song.
+        let again = engine.generate(&params).expect("generate again");
+        assert_eq!(out.left, again.left, "same seed must render identically");
+    }
+
+    /// Relative L2 distance between two tensors of the same shape.
+    fn rel_l2(a: &candle_core::Tensor, b: &candle_core::Tensor) -> f32 {
+        let a = a.to_dtype(DType::F32).unwrap().flatten_all().unwrap();
+        let b = b.to_dtype(DType::F32).unwrap().flatten_all().unwrap();
+        let diff = (&a - &b)
+            .unwrap()
+            .sqr()
+            .unwrap()
+            .sum_all()
+            .unwrap()
+            .to_scalar::<f32>()
+            .unwrap();
+        let norm = b
+            .sqr()
+            .unwrap()
+            .sum_all()
+            .unwrap()
+            .to_scalar::<f32>()
+            .unwrap();
+        (diff / norm.max(1e-20)).sqrt()
+    }
+
+    /// Parity with the upstream PyTorch reference (#381) on the CPU in f32,
+    /// from the reference's own noise: prompt token ids exactly, then the
+    /// denoised latents and the decoded waveform within float tolerance.
+    #[test]
+    #[ignore = "needs ACESTEP_SNAPSHOT and minutes of CPU"]
+    fn acestep_matches_reference_cpu_f32() {
+        let Ok(root) = std::env::var("ACESTEP_SNAPSHOT") else {
+            eprintln!("ACESTEP_SNAPSHOT not set; skipping");
+            return;
+        };
+        let cases = std::env::var("ACESTEP_CASES")
+            .unwrap_or_else(|_| "instr_10s_noplan,vocal_es_20s_noplan".into());
+        let mut engine = AceStepEngine::load(
+            &AceStepFiles {
+                root: root.into(),
+                planner: None,
+            },
+            &Device::Cpu,
+        )
+        .expect("load");
+        let dir =
+            std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("src/harness/testdata/acestep");
+        for case in cases.split(',') {
+            let meta: serde_json::Value = serde_json::from_slice(
+                &std::fs::read(dir.join(format!("{case}__cpu_f32.json"))).unwrap(),
+            )
+            .unwrap();
+            // Read tensors by name: the file also holds BOOL masks, which
+            // candle cannot load.
+            let file = unsafe {
+                candle_core::safetensors::MmapedSafetensors::new(
+                    dir.join(format!("{case}__cpu_f32.safetensors")),
+                )
+            }
+            .unwrap();
+            let want: std::collections::HashMap<&str, candle_core::Tensor> = [
+                "text_token_ids",
+                "lyric_token_ids",
+                "noise",
+                "target_latents",
+                "vae_waveform",
+            ]
+            .into_iter()
+            .map(|k| (k, file.load(k, &Device::Cpu).unwrap()))
+            .collect();
+            let c = &meta["case"];
+            let caption = c["caption"].as_str().unwrap();
+            let lyrics = c["lyrics"].as_str().unwrap();
+            let language = c["language"].as_str().unwrap();
+            let duration = c["duration"].as_f64().unwrap();
+
+            let (text, lyric) =
+                AceStepPipeline::format_prompts(caption, lyrics, duration, None, Some(language));
+            let (text_ids, lyric_ids) = engine.pipeline.prompt_token_ids(&text, &lyric).unwrap();
+            let ids = |k: &str| -> Vec<u32> {
+                want[k]
+                    .flatten_all()
+                    .unwrap()
+                    .to_vec1::<i64>()
+                    .unwrap()
+                    .into_iter()
+                    .map(|v| v as u32)
+                    .collect()
+            };
+            assert_eq!(text_ids, ids("text_token_ids"), "{case}: text token ids");
+            assert_eq!(lyric_ids, ids("lyric_token_ids"), "{case}: lyric token ids");
+
+            let params = GenerationParams {
+                duration_secs: duration,
+                seed: None,
+                num_steps: Some(DEFAULT_STEPS),
+                shift: TURBO_SHIFT,
+                is_turbo: true,
+                infer_method: sampling::InferMethod::Ode,
+                initial_noise: Some(want["noise"].clone()),
+                language: Some(language.to_string()),
+                ..Default::default()
+            };
+            let t = Instant::now();
+            let out = engine
+                .pipeline
+                .text2music(caption, lyrics, &params)
+                .unwrap();
+            let latents = rel_l2(&out.latents, &want["target_latents"]);
+            let wave_ref = &want["vae_waveform"];
+            let n = wave_ref.dim(2).unwrap();
+            let wave = out.audio.narrow(1, 0, n).unwrap().unsqueeze(0).unwrap();
+            let waveform = rel_l2(&wave, wave_ref);
+            eprintln!(
+                "{case}: latents rel L2 {latents:.2e}, waveform rel L2 {waveform:.2e} ({:.1}s)",
+                t.elapsed().as_secs_f32()
+            );
+            assert!(latents < 1e-3, "{case}: latents rel L2 {latents}");
+            assert!(waveform < 1e-2, "{case}: waveform rel L2 {waveform}");
+        }
+    }
+}
@@ -146,6 +146,8 @@ pub enum LoadedHandle {
    /// Laya decision model (#336): typed questions in, calibrated
    /// option scores out, no text generation.
    Decision(Arc<super::decision_model::LoadedDecisionModel>),
+    /// ACE-Step music generation pipeline (#383).
+    Audio(Arc<LoadedAudioModel>),
 }

 impl LoadedHandle {
@@ -155,6 +157,7 @@ impl LoadedHandle {
            #[cfg(feature = "cuda")]
            LoadedHandle::Tp(m) => &m.model_id,
            LoadedHandle::Image(m) => &m.model_id,
+            LoadedHandle::Audio(m) => &m.model_id,
            LoadedHandle::Decision(m) => &m.model_id,
        }
    }
@@ -166,6 +169,7 @@ impl LoadedHandle {
            #[cfg(feature = "cuda")]
            LoadedHandle::Tp(m) => &m.spec,
            LoadedHandle::Image(m) => &m.spec,
+            LoadedHandle::Audio(m) => &m.spec,
            LoadedHandle::Decision(m) => &m.spec,
        }
    }
@@ -176,6 +180,7 @@ impl LoadedHandle {
            #[cfg(feature = "cuda")]
            LoadedHandle::Tp(m) => m.devices.clone(),
            LoadedHandle::Image(m) => m.devices.clone(),
+            LoadedHandle::Audio(m) => m.devices.clone(),
            LoadedHandle::Decision(m) => m.devices.clone(),
        }
    }
@@ -190,6 +195,7 @@ impl LoadedHandle {
            #[cfg(feature = "cuda")]
            LoadedHandle::Tp(m) => m.poisoned.load(Ordering::Acquire),
            LoadedHandle::Image(m) => m.poisoned.load(Ordering::Acquire),
+            LoadedHandle::Audio(m) => m.poisoned.load(Ordering::Acquire),
            LoadedHandle::Decision(m) => m.poisoned.load(Ordering::Acquire),
        }
    }
@@ -202,6 +208,7 @@ impl LoadedHandle {
            #[cfg(feature = "cuda")]
            LoadedHandle::Tp(m) => (m.admission.in_flight(), m.admission.queue_depth()),
            LoadedHandle::Image(m) => (m.admission.in_flight(), m.admission.queue_depth()),
+            LoadedHandle::Audio(m) => (m.admission.in_flight(), m.admission.queue_depth()),
            LoadedHandle::Decision(m) => (m.admission.in_flight(), m.admission.queue_depth()),
        }
    }
@@ -215,6 +222,7 @@ impl LoadedHandle {
            #[cfg(feature = "cuda")]
            LoadedHandle::Tp(m) => (m.admission.max_in_flight(), m.admission.max_queue_depth()),
            LoadedHandle::Image(m) => (m.admission.max_in_flight(), m.admission.max_queue_depth()),
+            LoadedHandle::Audio(m) => (m.admission.max_in_flight(), m.admission.max_queue_depth()),
            LoadedHandle::Decision(m) => {
                (m.admission.max_in_flight(), m.admission.max_queue_depth())
            }
@@ -241,6 +249,10 @@ impl LoadedHandle {
                m.admission.anon_in_flight(),
                m.admission.anon_max_in_flight(),
            ),
+            LoadedHandle::Audio(m) => (
+                m.admission.anon_in_flight(),
+                m.admission.anon_max_in_flight(),
+            ),
            LoadedHandle::Decision(m) => (
                m.admission.anon_in_flight(),
                m.admission.anon_max_in_flight(),
@@ -255,6 +267,7 @@ impl LoadedHandle {
            #[cfg(feature = "cuda")]
            LoadedHandle::Tp(m) => m.admission.close(),
            LoadedHandle::Image(m) => m.admission.close(),
+            LoadedHandle::Audio(m) => m.admission.close(),
            LoadedHandle::Decision(m) => m.admission.close(),
        }
    }
@@ -267,6 +280,7 @@ impl LoadedHandle {
            #[cfg(feature = "cuda")]
            LoadedHandle::Tp(m) => m.admission.rejections(),
            LoadedHandle::Image(m) => m.admission.rejections(),
+            LoadedHandle::Audio(m) => m.admission.rejections(),
            LoadedHandle::Decision(m) => m.admission.rejections(),
        }
    }
@@ -279,6 +293,7 @@ impl LoadedHandle {
            #[cfg(feature = "cuda")]
            LoadedHandle::Tp(m) => &m.admission,
            LoadedHandle::Image(m) => &m.admission,
+            LoadedHandle::Audio(m) => &m.admission,
            LoadedHandle::Decision(m) => &m.admission,
        };
        (admission.kv_budget_mb(), admission.kv_available_mb())
@@ -293,7 +308,9 @@ impl LoadedHandle {
            LoadedHandle::Tp(m) => &m.prefill_rate,
            // Image models have no token throughput; steps/sec lives in
            // the per-request timing instead.
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => return (0.0, 0.0),
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => {
+                return (0.0, 0.0);
+            }
        };
        (ema.get().unwrap_or(0.0), ema.decode().unwrap_or(0.0))
    }
@@ -308,12 +325,17 @@ impl LoadedHandle {
        if matches!(self, LoadedHandle::Decision(_)) {
            return vec!["decision".to_string()];
        }
+        if matches!(self, LoadedHandle::Audio(_)) {
+            return vec![cortex_core::audio::CAPABILITY.to_string()];
+        }
        let mut caps = vec!["text".to_string()];
        let has_vision = match self {
            LoadedHandle::Single(m) => m.has_vision,
            #[cfg(feature = "cuda")]
            LoadedHandle::Tp(m) => m.has_vision,
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => unreachable!("handled above"),
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => {
+                unreachable!("handled above")
+            }
        };
        if has_vision {
            caps.push("vision".to_string());
@@ -359,7 +381,9 @@ impl LoadedHandle {
            LoadedHandle::Tp(m) => (m.last_free_mb.load(Ordering::Acquire), m.context_profile),
            // Image generation has its own VRAM path and no prefill
            // floor, so this check would say nothing useful about it.
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => return None,
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => {
+                return None;
+            }
        };
        // `0` is the CPU-load and missing-worker sentinel — the same one
        // `check_vram` skips on. Nothing to assert either way.
@@ -440,7 +464,7 @@ impl LoadedHandle {
            LoadedHandle::Single(m) => m.tool_call_tokens.is_some(),
            #[cfg(feature = "cuda")]
            LoadedHandle::Tp(m) => m.tool_call_tokens.is_some(),
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => false,
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => false,
        }
    }

@@ -451,7 +475,7 @@ impl LoadedHandle {
            LoadedHandle::Single(m) => m.reasoning_tokens.is_some(),
            #[cfg(feature = "cuda")]
            LoadedHandle::Tp(m) => m.reasoning_tokens.is_some(),
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => false,
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => false,
        }
    }

@@ -462,7 +486,7 @@ impl LoadedHandle {
            LoadedHandle::Single(m) => m.preserve_thinking,
            #[cfg(feature = "cuda")]
            LoadedHandle::Tp(m) => m.preserve_thinking,
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => None,
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => None,
        }
    }

@@ -472,7 +496,9 @@ impl LoadedHandle {
            LoadedHandle::Single(m) => m.reasoning_efforts.clone(),
            #[cfg(feature = "cuda")]
            LoadedHandle::Tp(m) => m.reasoning_efforts.clone(),
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => Default::default(),
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => {
+                Default::default()
+            }
        }
    }

@@ -498,7 +524,7 @@ impl LoadedHandle {
            #[cfg(feature = "cuda")]
            LoadedHandle::Tp(m) => m.query_vram_tightest_free_mb().await,
            // No context-limit derivation for image models; nothing to cache.
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => return,
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => return,
        };
        // Don't clobber a good cached value with a transient `0`
        // (worker gone/poisoned sentinel).
@@ -507,7 +533,7 @@ impl LoadedHandle {
                LoadedHandle::Single(m) => m.last_free_mb.store(free, Ordering::Release),
                #[cfg(feature = "cuda")]
                LoadedHandle::Tp(m) => m.last_free_mb.store(free, Ordering::Release),
-                LoadedHandle::Image(_) | LoadedHandle::Decision(_) => {}
+                LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => {}
            }
        }
    }
@@ -536,7 +562,9 @@ impl LoadedHandle {
                m.prefill_rate.get(),
            ),
            // Token limits are meaningless for image generation.
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => return None,
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => {
+                return None;
+            }
        };
        let rate = rate.unwrap_or(cfg.bootstrap_prefill_tok_per_sec);
        // The KV budget admission will actually grant. Advertising past
@@ -546,7 +574,7 @@ impl LoadedHandle {
            LoadedHandle::Single(m) => m.admission.kv_budget_mb(),
            #[cfg(feature = "cuda")]
            LoadedHandle::Tp(m) => m.admission.kv_budget_mb(),
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => {
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => {
                unreachable!("derived_limit returned above")
            }
        };
@@ -567,7 +595,7 @@ impl LoadedHandle {
                LoadedHandle::Single(m) => &m.derived_input_cap,
                #[cfg(feature = "cuda")]
                LoadedHandle::Tp(m) => &m.derived_input_cap,
-                LoadedHandle::Image(_) | LoadedHandle::Decision(_) => {
+                LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => {
                    unreachable!("derived_limit returned above")
                }
            };
@@ -633,6 +661,29 @@ impl LoadedImageModel {
    }
 }

+/// A loaded ACE-Step music pipeline (#383). Every component lives in the
+/// device worker's audio slab, addressed via [`Self::handle`]; this
+/// async-side struct holds admission bookkeeping and the model's limits.
+/// A render is one `Job::GenerateAudio` round-trip.
+pub struct LoadedAudioModel {
+    pub model_id: String,
+    pub spec: ModelSpec,
+    pub devices: Vec<u32>,
+    pub worker: Arc<super::device_worker::DeviceWorkerHandle>,
+    pub handle: super::device_worker::AudioHandle,
+    /// Whether the snapshot carried a song planner.
+    pub has_planner: bool,
+    /// What a request may ask of this model; validated before admission.
+    pub limits: cortex_core::audio::AudioLimits,
+    /// Set when a render fails at the driver level: subsequent requests
+    /// fast-reject and `/models` reports `poisoned`.
+    pub poisoned: Arc<AtomicBool>,
+    /// Admission control (#53) with `max_in_flight` forced to 1 — a render
+    /// saturates the device. The cortex job queue is the real backlog.
+    pub admission: super::admission::AdmissionController,
+    pub inference_lock: Arc<tokio::sync::Mutex<()>>,
+}
+
 /// A loaded model with its tokenizer, device placement, and architecture-
 /// specific weights. The `arch` is `Arc<Mutex<>>` so the lock guard can be
 /// moved into `spawn_blocking` for synchronous candle forward passes.
@@ -3307,7 +3358,7 @@ impl CandleHarness {
            LoadedHandle::Tp(m) => {
                return self.inference_tp_stream(m, request, principal, wire).await;
            }
-            LoadedHandle::Image(_) | LoadedHandle::Decision(_) => {
+            LoadedHandle::Image(_) | LoadedHandle::Decision(_) | LoadedHandle::Audio(_) => {
                return Err(InferenceError::WrongModality {
                    model_id: request.model.clone(),
                });
@@ -4031,6 +4082,11 @@ impl Harness for CandleHarness {
        if matches!(plan.format, super::preflight::SourceFormat::DecisionModel) {
            return self.load_decision_model(spec, &source_id).await;
        }
+        // ACE-Step music models (#383) take the audio path: one device,
+        // unquantized (preflight rejected tp>1 and quant).
+        if matches!(plan.format, super::preflight::SourceFormat::AceStep) {
+            return self.load_audio_model(spec, &source_id).await;
+        }

        let tp_size = spec.tensor_parallel.unwrap_or(1);
        if tp_size > 1 {
@@ -4402,6 +4458,17 @@ impl Harness for CandleHarness {
                // them (the device worker, for a CUDA load).
                dm.release().await;
            }
+            LoadedHandle::Audio(audio) => {
+                // Drop the pipeline on the worker thread so its CUDA
+                // tensors release on the context that allocated them.
+                if let Err(e) = audio.worker.drop_audio(audio.handle).await {
+                    tracing::warn!(
+                        model = %model_id,
+                        error = %e,
+                        "audio unload: DropAudio RPC failed (pipeline may leak in worker slab)"
+                    );
+                }
+            }
            #[cfg(feature = "cuda")]
            LoadedHandle::Tp(tp) => {
                // Try to recover the inner TpLoadedModel so we can move
@@ -4576,15 +4643,117 @@ impl CandleHarness {
        Ok(())
    }

-    /// Load a Laya decision model (#336): resolve the checkpoint's files,
-    /// build it on the device worker (or the CPU), and register it with
-    /// `capabilities: ["decision"]`.
+    /// Load an ACE-Step music pipeline (#383): resolve the snapshot's
+    /// components, build them on the device worker, and register the
+    /// model with `capabilities: ["audio"]`.
+    #[cfg_attr(not(feature = "cuda"), allow(unused_variables, unreachable_code))]
+    async fn load_audio_model(
+        &self,
+        spec: &ModelSpec,
+        source_id: &cortex_core::source::ModelSourceId,
+    ) -> Result<()> {
+        let devices = spec.devices.clone().unwrap_or_else(|| vec![0]);
+        let device = Self::pick_device(&devices)?;
+        let worker: Arc<super::device_worker::DeviceWorkerHandle> = match &device {
+            #[cfg(feature = "cuda")]
+            Device::Cuda(_) => self.ensure_device_worker(devices[0]).await?,
+            // On CPU builds this arm is the whole match: a song on the CPU
+            // takes many minutes and is not a serving target.
+            _ => anyhow::bail!(
+                "music generation requires a CUDA device; '{}' resolved to a CPU device",
+                spec.model_id
+            ),
+        };
+
+        let files = self.resolve_audio_files(source_id).await?;
+        let has_planner = files.planner.is_some();
+        let handle = worker
+            .load_audio(files, spec.model_id.clone())
+            .await
+            .map_err(|e| anyhow::anyhow!("worker load_audio: {e:#}"))?;
+
+        // Admission (#53) with the in-flight ceiling forced to 1: a render
+        // saturates the device. cortex queues jobs ahead of this.
+        let admission_cfg = crate::config::AdmissionConfig {
+            max_in_flight: 1,
+            ..self.admission_cfg.clone()
+        };
+        let loaded = LoadedAudioModel {
+            model_id: spec.model_id.clone(),
+            spec: spec.clone(),
+            devices,
+            worker,
+            handle,
+            has_planner,
+            limits: super::audio::ACESTEP_LIMITS,
+            poisoned: Arc::new(AtomicBool::new(false)),
+            admission: crate::harness::admission::AdmissionController::new(&admission_cfg),
+            inference_lock: Arc::new(tokio::sync::Mutex::new(())),
+        };
+        let mut models = self.models.write().await;
+        models.insert(spec.model_id.clone(), LoadedHandle::Audio(Arc::new(loaded)));
+        tracing::info!(model = %spec.model_id, planner = has_planner, "audio pipeline loaded");
+        Ok(())
+    }
+
+    /// Resolve an ACE-Step snapshot into a local root, downloading through
+    /// hf-hub as needed: every required component, plus the first planner
+    /// LM the repo has.
+    async fn resolve_audio_files(
+        &self,
+        source_id: &cortex_core::source::ModelSourceId,
+    ) -> Result<super::audio::AceStepFiles> {
+        let api = self.hf_api_for(&source_id.scheme)?;
+        let repo = api.model(source_id.repo_path());
+        let display_id = source_id.to_string();
+        let mut first = None;
+        for file in super::audio::required_files() {
+            let path = repo
+                .get(&file)
+                .await
+                .with_context(|| format!("fetch {file} from {display_id}"))?;
+            if first.is_none() {
+                first = Some((path, file));
+            }
+        }
+        // hf-hub materialises a snapshot as a tree mirroring the repo, so
+        // the root is the first file's path less its repo-relative part.
+        let (path, rel) = first.context("ACE-Step needs at least one file")?;
+        let mut root = path;
+        for _ in rel.split('/') {
+            root = root
+                .parent()
+                .context("snapshot path shorter than its repo path")?
+                .to_path_buf();
+        }
+        let mut planner = None;
+        for dir in super::audio::PLANNER_DIRS {
+            let mut ok = true;
+            for file in super::audio::planner_files(dir) {
+                if let Err(e) = repo.get(&file).await {
+                    tracing::info!(repo = %display_id, file, error = %format!("{e:#}"), "planner file unavailable");
+                    ok = false;
+                    break;
+                }
+            }
+            if ok {
+                planner = Some(dir.to_string());
+                break;
+            }
+        }
+        Ok(super::audio::AceStepFiles { root, planner })
+    }
+
+    /// Load a Laya decision model (#336): resolve the checkpoints' files,
+    /// build them on the device worker (or the CPU), and register the
+    /// model with `capabilities: ["decision"]`.
    ///
-    /// Only the repo-root checkpoint is loaded. A repo can carry more in
-    /// subfolders (`multilingual/`, `typed-decisions/`); which of those a
-    /// model serves, and how requests are routed between them, is a
-    /// separate decision (#339) — [`Self::resolve_decision_files`]
-    /// already takes the subfolder.
+    /// The repo root is required. A released family also carries
+    /// `multilingual/` and `typed-decisions/` beside it (#339); each one
+    /// the repo has is loaded too, so the `/v1/systemone` handler can
+    /// route non-English states to the multilingual checkpoint and honour
+    /// a request that pins one. A repo without them (a single fine-tune)
+    /// serves its root alone.
    async fn load_decision_model(
        &self,
        spec: &ModelSpec,
@@ -4607,20 +4776,49 @@ impl CandleHarness {
        };

        let files = self.resolve_decision_files(source_id, None).await?;
-        let checkpoint = super::decision_model::DecisionCheckpoint::load(
+        let root = super::decision_model::DecisionCheckpoint::load(
            super::decision_model::ROOT_CHECKPOINT.to_string(),
            files,
            &spec.model_id,
-            worker,
+            worker.clone(),
        )
        .await?;
-        let dtype = checkpoint.dtype;
+        let dtype = root.dtype;
+        let mut checkpoints = vec![root];
+        for sub in super::decision::route::FAMILY
+            .iter()
+            .filter(|c| **c != super::decision_model::ROOT_CHECKPOINT)
+        {
+            let files = match self.resolve_decision_files(source_id, Some(sub)).await {
+                Ok(files) => files,
+                Err(e) => {
+                    tracing::info!(
+                        model = %spec.model_id,
+                        checkpoint = sub,
+                        error = %format!("{e:#}"),
+                        "decision checkpoint not available; serving without it"
+                    );
+                    continue;
+                }
+            };
+            checkpoints.push(
+                super::decision_model::DecisionCheckpoint::load(
+                    sub.to_string(),
+                    files,
+                    &spec.model_id,
+                    worker.clone(),
+                )
+                .await?,
+            );
+        }
+        let names: Vec<&str> = checkpoints.iter().map(|c| c.name.as_str()).collect();
+        tracing::info!(model = %spec.model_id, checkpoints = ?names, "decision checkpoints loaded");

        let loaded = super::decision_model::LoadedDecisionModel {
            model_id: spec.model_id.clone(),
            spec: spec.clone(),
            devices,
-            checkpoints: vec![checkpoint],
+            checkpoints,
            poisoned: Arc::new(AtomicBool::new(false)),
            admission: crate::harness::admission::AdmissionController::new(&self.admission_cfg),
        };
@@ -4651,6 +4849,7 @@ impl CandleHarness {
            "encoder/config.json",
            "model.safetensors",
            "tokenizer/tokenizer.json",
+            "tokenizer/tokenizer_config.json",
        ] {
            let path = format!("{prefix}{file}");
            fetched.push(
@@ -4806,6 +5005,269 @@ impl CandleHarness {
        }
    }

+    /// Answer one `/v1/systemone` request (#338): pick the served decision
+    /// model and the checkpoint of it that answers, assemble the question
+    /// sequences, score them in one forward pass, and shape the answers.
+    ///
+    /// The served model is `served` when the gateway routed the request
+    /// (the `x-helexa-decision-model` header), else the model the body's
+    /// `model` names if one is loaded, else the one decision model this
+    /// node serves. The body's `model` is then read a second time, as the
+    /// reference server reads it: as an optional pin to a checkpoint of
+    /// the family (`multilingual`, `typed-decisions`, …).
+    ///
+    /// Questions are validated and tokenised before admission, so a
+    /// malformed request never takes a slot.
+    pub async fn systemone(
+        &self,
+        served: Option<&str>,
+        request: cortex_core::decisions::SystemOneRequest,
+        principal: Option<&str>,
+    ) -> Result<cortex_core::decisions::SystemOneResponse, SystemOneError> {
+        use super::decision::{self, route};
+
+        let model = match served {
+            Some(id) => self.decision_model(id).await?,
+            None => {
+                self.default_decision_model(request.model.as_deref())
+                    .await?
+            }
+        };
+        if model.poisoned.load(Ordering::Acquire) {
+            return Err(InferenceError::Other(anyhow::anyhow!(
+                "decision model '{}' is poisoned; unload and reload it",
+                model.model_id
+            ))
+            .into());
+        }
+
+        let pin = route::resolve_pin(request.model.as_deref());
+        let mut routing = route::route(&model.model_id, &request.state, &request.questions, pin);
+        let checkpoint = match model.checkpoint(Some(&routing.model)) {
+            Some(c) => c,
+            None if pin.is_some() => {
+                return Err(decision::DecisionRequestError::new(
+                    404,
+                    format!(
+                        "checkpoint '{}' is not served by '{}' on this node",
+                        routing.model, model.model_id
+                    ),
+                )
+                .into());
+            }
+            // Routed (not pinned) to a checkpoint this repo does not
+            // carry: answer with the default one, and say so.
+            None => {
+                let fallback = model.checkpoint(None).ok_or_else(|| {
+                    InferenceError::Other(anyhow::anyhow!(
+                        "decision model '{}' has no checkpoints",
+                        model.model_id
+                    ))
+                })?;
+                routing.reason = format!(
+                    "{}; the {} checkpoint is not loaded, so {} answered",
+                    routing.reason, routing.model, fallback.name
+                );
+                routing.model = fallback.name.clone();
+                routing.repo = route::repo_of(&model.model_id, &fallback.name);
+                fallback
+            }
+        };
+        let lang = route::detected_language(&routing).map(str::to_string);
+
+        let prepared = decision::prepare(
+            &request.state,
+            &request.questions,
+            &checkpoint.calibration,
+            &checkpoint.tokenizer,
+            decision::Budget {
+                max_len: request.max_len,
+                head_max_len: request.head_max_len,
+            },
+        )?;
+
+        let _permit = model.admission.enter(principal).await?;
+        let outputs: Vec<decision::RowOutput> = if prepared.rows.is_empty() {
+            Vec::new()
+        } else {
+            let batch = super::arch::laya::DecisionBatch {
+                seqs: prepared
+                    .rows
+                    .iter()
+                    .map(|r| super::arch::laya::DecisionSeq {
+                        input_ids: r.input_ids.clone(),
+                        markers: r.markers.clone(),
+                        qtype: r.qtype as u8,
+                    })
+                    .collect(),
+            };
+            checkpoint
+                .decide(batch)
+                .await
+                .map_err(|e| InferenceError::Other(anyhow::anyhow!("decision forward: {e:#}")))?
+                .into_iter()
+                .map(|row| decision::RowOutput {
+                    act_probability: row.act.first().copied().unwrap_or(0.0),
+                    logits: row.logits,
+                })
+                .collect()
+        };
+        let (answers, usage) = decision::shape(
+            &prepared,
+            &outputs,
+            &checkpoint.calibration,
+            lang.as_deref(),
+        );
+        Ok(decision::shape::response(answers, usage, Some(routing)))
+    }
+
+    /// Register an already-built decision model, bypassing the hub fetch
+    /// `load_model` does — for tests that serve local checkpoints.
+    #[cfg(test)]
+    pub(crate) async fn register_decision_model(
+        &self,
+        model: super::decision_model::LoadedDecisionModel,
+    ) {
+        self.models.write().await.insert(
+            model.model_id.clone(),
+            LoadedHandle::Decision(Arc::new(model)),
+        );
+    }
+
+    /// The decision model to serve a request no gateway routed: the one
+    /// `requested` names if it is loaded, else the only (first) loaded
+    /// decision model.
+    async fn default_decision_model(
+        &self,
+        requested: Option<&str>,
+    ) -> Result<Arc<super::decision_model::LoadedDecisionModel>, InferenceError> {
+        let models = self.models.read().await;
+        if let Some(id) = requested
+            && let Some(handle) = models.get(id)
+        {
+            return match handle {
+                LoadedHandle::Decision(m) => Ok(Arc::clone(m)),
+                _ => Err(InferenceError::WrongModality {
+                    model_id: id.to_string(),
+                }),
+            };
+        }
+        let mut decision: Vec<&Arc<super::decision_model::LoadedDecisionModel>> = models
+            .values()
+            .filter_map(|h| match h {
+                LoadedHandle::Decision(m) => Some(m),
+                _ => None,
+            })
+            .collect();
+        decision.sort_by(|a, b| a.model_id.cmp(&b.model_id));
+        decision
+            .first()
+            .map(|m| Arc::clone(m))
+            .ok_or_else(|| InferenceError::ModelNotLoaded("no decision model is loaded".into()))
+    }
+
+    /// Render one audio generation (#382). The cortex job manager sends one
+    /// of these per job and waits for the whole render; the reply is the
+    /// WAV bytes plus the result metadata cortex records on the job.
+    ///
+    /// Validated against the model's limits before admission; the render
+    /// itself is one device-worker job. A request without a seed gets a
+    /// random one, reported back so the song can be reproduced.
+    pub async fn audio_generation(
+        &self,
+        request: &cortex_core::audio::AudioGenerationRequest,
+        principal: Option<&str>,
+    ) -> Result<(Vec<u8>, cortex_core::audio::NeuronAudioResult), InferenceError> {
+        let handle = {
+            let models = self.models.read().await;
+            models.get(&request.model).cloned()
+        };
+        let audio = match handle {
+            Some(LoadedHandle::Audio(m)) => m,
+            Some(_) => {
+                return Err(InferenceError::WrongModality {
+                    model_id: request.model.clone(),
+                });
+            }
+            None => return Err(InferenceError::ModelNotLoaded(request.model.clone())),
+        };
+        if audio.poisoned.load(Ordering::Acquire) {
+            return Err(InferenceError::Other(anyhow::anyhow!(
+                "audio model '{}' is poisoned; unload and reload it",
+                request.model
+            )));
+        }
+
+        // Validate before admission so a malformed request never takes a
+        // queue slot.
+        request
+            .validate_limits(&audio.limits)
+            .map_err(|detail| InferenceError::InvalidAudioParams { detail })?;
+        // The planner runs only when asked for until its port is proven
+        // against the reference (#384); `auto` means off until then.
+        let planner = match request.planner.as_deref() {
+            Some("on") if !audio.has_planner => {
+                return Err(InferenceError::InvalidAudioParams {
+                    detail: "this model has no song planner; use planner=off or auto".into(),
+                });
+            }
+            Some("on") => true,
+            _ => false,
+        };
+        // A random seed small enough to type back in, so any render can be
+        // reproduced from what the job reports.
+        let seed = request.seed.unwrap_or_else(|| {
+            use std::hash::BuildHasher;
+            std::collections::hash_map::RandomState::new().hash_one(unix_subsec_nanos())
+                & 0x7fff_ffff
+        });
+        let params = super::audio::AudioGenParams {
+            prompt: request.prompt.clone(),
+            lyrics: if request.is_instrumental() {
+                String::new()
+            } else {
+                request.lyrics.clone().unwrap_or_default()
+            },
+            duration_secs: request.duration.unwrap_or(audio.limits.default_duration),
+            seed,
+            steps: request.steps.map(|s| s as usize),
+            guidance_scale: request.guidance_scale,
+            planner,
+            temperature: request.temperature,
+            language: request.language.clone(),
+        };
+
+        let _permit = audio.admission.enter(principal).await?;
+        let _guard = audio.inference_lock.lock().await;
+        let result = audio
+            .worker
+            .generate_audio(audio.handle, params)
+            .await
+            .map_err(|e| InferenceError::Other(anyhow::anyhow!("audio generation: {e:#}")))?;
+
+        let duration = result.duration_secs();
+        let sample_rate = result.sample_rate;
+        let timing = result.timing;
+        let wav = tokio::task::spawn_blocking(move || {
+            super::audio::encode_wav_pcm16(&result.left, &result.right, result.sample_rate)
+        })
+        .await
+        .map_err(|e| InferenceError::Other(anyhow::anyhow!("wav encode task: {e:#}")))?;
+        Ok((
+            wav,
+            cortex_core::audio::NeuronAudioResult {
+                duration,
+                sample_rate,
+                channels: 2,
+                seed,
+                usage: cortex_core::audio::AudioUsage {
+                    helexa_audio_units: duration,
+                    helexa_timing: Some(timing),
+                },
+            },
+        ))
+    }
+
    /// Run one text-to-image generation (#198). This is the internal
    /// path the `/v1/images/generations` handler (#200) fronts;
    /// admission, validation, tokenization, and PNG encoding happen
@@ -6409,6 +6871,10 @@ pub enum InferenceError {
    /// Maps to 400 — client error, not server fault.
    #[error("invalid image generation parameters: {detail}")]
    InvalidImageParams { detail: String },
+    /// Audio generation parameters the model cannot serve (#383): a
+    /// duration out of range, lyrics for a model that does not sing.
+    #[error("invalid audio generation parameters: {detail}")]
+    InvalidAudioParams { detail: String },
    /// The loaded model's chat template could not render the request
    /// (e.g. a message / tool-call structure it rejects). Returned only
    /// when the request carried tools — silently degrading to a
@@ -6430,6 +6896,33 @@ pub enum InferenceError {
    Other(#[from] anyhow::Error),
 }

+/// Why a `/v1/systemone` request was not answered (#338): a request the
+/// reference server would also reject (its status and `detail` text), or
+/// a serving failure shared with the other endpoints.
+#[derive(Debug)]
+pub enum SystemOneError {
+    Request(super::decision::DecisionRequestError),
+    Inference(InferenceError),
+}
+
+impl From<super::decision::DecisionRequestError> for SystemOneError {
+    fn from(e: super::decision::DecisionRequestError) -> Self {
+        SystemOneError::Request(e)
+    }
+}
+
+impl From<InferenceError> for SystemOneError {
+    fn from(e: InferenceError) -> Self {
+        SystemOneError::Inference(e)
+    }
+}
+
+impl From<super::admission::AdmissionRejection> for SystemOneError {
+    fn from(e: super::admission::AdmissionRejection) -> Self {
+        SystemOneError::Inference(e.into())
+    }
+}
+
 impl From<super::admission::AdmissionRejection> for InferenceError {
    fn from(rejection: super::admission::AdmissionRejection) -> Self {
        use super::admission::AdmissionRejection;
@@ -60,6 +60,8 @@ pub fn classify_files(
        // A decision model answers typed questions (`/v1/systemone`);
        // it generates no text, so it must not claim "text".
        SourceFormat::DecisionModel => vec!["decision".to_string()],
+        // A music model (#383) renders audio on `/v1/audio/generations`.
+        SourceFormat::AceStep => vec![cortex_core::audio::CAPABILITY.to_string()],
        // A repo with no recognised weights (tokenizer-only, or an empty
        // entry) is not servable, so claiming "text" would be a lie of
        // the same kind this module exists to remove.
@@ -171,4 +173,13 @@ mod tests {
        let dir = std::env::temp_dir().join("helexa-capability-probe-missing");
        assert!(from_cache(&dir, "does-not/exist").is_none());
    }
+
+    #[test]
+    fn acestep_snapshot_is_audio_not_text() {
+        let mut files = crate::harness::audio::required_files();
+        files.push("config.json".into());
+        files.push("acestep-5Hz-lm-1.7B/model.safetensors".into());
+        let refs: Vec<&str> = files.iter().map(String::as_str).collect();
+        assert_eq!(classify_files(&refs, || None), vec!["audio".to_string()]);
+    }
 }
@@ -0,0 +1,233 @@
+//! Per-checkpoint decision settings, read from `rl_agent_config.json`.
+//!
+//! Parsed with [`Json`], not `serde_json`: the fitted temperatures are
+//! written with 17 significant digits, and `serde_json`'s default float
+//! parser is not correctly rounded (it reads `1.7601518630981445` as
+//! `…443`).
+
+use cortex_core::decisions::Json;
+use std::collections::HashMap;
+
+/// Default token budgets when a checkpoint's config omits them.
+pub const DEFAULT_MAX_LEN: usize = 512;
+pub const DEFAULT_HEAD_MAX_LEN: usize = 192;
+
+/// Bounds on a usable calibration temperature. A fitted value below 1
+/// sharpens instead of softening; the shipped `choice:11+` bucket is 0.1006,
+/// which would publish a coin flip as a certainty, so the reference refuses
+/// to apply anything outside this range.
+pub const TEMP_MIN: f64 = 0.5;
+pub const TEMP_MAX: f64 = 5.0;
+
+/// A temperature confined to [`TEMP_MIN`, `TEMP_MAX`]; a non-finite or
+/// non-numeric value falls back to 1.0.
+pub fn clamp_temperature(t: Option<f64>) -> f64 {
+    match t {
+        Some(t) if t.is_finite() => t.clamp(TEMP_MIN, TEMP_MAX),
+        _ => 1.0,
+    }
+}
+
+/// Python's `float(v)` for a JSON value: numbers, bools, and numeric strings.
+fn py_float(v: &Json) -> Option<f64> {
+    match v {
+        Json::Int(d) => d.parse().ok(),
+        Json::Float(x) => Some(*x),
+        Json::Bool(b) => Some(if *b { 1.0 } else { 0.0 }),
+        Json::String(s) => s.trim().parse().ok(),
+        _ => None,
+    }
+}
+
+/// One set of calibration temperatures: per question type, overridden per
+/// `(type, option-count bucket)` (see [`temp_bucket`]). Already clamped.
+#[derive(Debug, Clone, PartialEq)]
+pub struct Temperatures {
+    /// Indexed by [`super::QType`] (`choice`, `score`, `noul`).
+    pub by_type: [f64; 3],
+    /// `"choice:3-5"` → temperature, and so on.
+    pub by_options: HashMap<String, f64>,
+}
+
+impl Temperatures {
+    /// The temperature for a question of type index `qtype` with `k` options.
+    pub fn for_question(&self, qtype: usize, k: usize) -> f64 {
+        self.by_options
+            .get(&temp_bucket(qtype, k))
+            .copied()
+            .unwrap_or(self.by_type[qtype])
+    }
+
+    fn from_json(
+        temperature: Option<&Json>,
+        by_options: Option<&Json>,
+        fallback: Option<[f64; 3]>,
+    ) -> Result<Self, String> {
+        let raw = match temperature {
+            None => fallback.unwrap_or([1.0; 3]).map(Some),
+            Some(Json::Array(a)) if a.len() == 3 => {
+                [py_float(&a[0]), py_float(&a[1]), py_float(&a[2])]
+            }
+            Some(other) => {
+                return Err(format!(
+                    "temperature must be a list of 3 floats, got {}",
+                    other.py_repr()
+                ));
+            }
+        };
+        let by_options = match by_options {
+            Some(Json::Object(m)) => m
+                .iter()
+                .map(|(k, v)| (k.clone(), clamp_temperature(py_float(v))))
+                .collect(),
+            _ => HashMap::new(),
+        };
+        Ok(Self {
+            by_type: raw.map(clamp_temperature),
+            by_options,
+        })
+    }
+}
+
+/// The calibration bucket of a question: its type and option count.
+pub fn temp_bucket(qtype: usize, k: usize) -> String {
+    let size = match k {
+        0..=2 => "2",
+        3..=5 => "3-5",
+        6..=10 => "6-10",
+        _ => "11+",
+    };
+    format!("{}:{size}", super::QType::NAMES[qtype])
+}
+
+/// What the sequence builder and the calibration need from a checkpoint.
+#[derive(Debug, Clone, PartialEq)]
+pub struct DecisionConfig {
+    /// Total sequence budget.
+    pub max_len: usize,
+    /// Budget for the question text and every option together.
+    pub head_max_len: usize,
+    pub temperatures: Temperatures,
+    /// Per-language overrides, keyed by the lower-cased primary subtag
+    /// (`"de"` for `de-AT`). Empty for the shipped checkpoints; an operator
+    /// may fit them per language.
+    pub lang_temperatures: HashMap<String, Temperatures>,
+}
+
+impl DecisionConfig {
+    /// Parse a checkpoint's `rl_agent_config.json`.
+    pub fn from_json(text: &str) -> Result<Self, String> {
+        let v = Json::parse(text.as_bytes()).map_err(|e| format!("rl_agent_config.json: {e}"))?;
+        let usize_or = |key: &str, default: usize| match v.get(key) {
+            Some(Json::Int(d)) => d.parse().unwrap_or(default),
+            _ => default,
+        };
+        Ok(Self {
+            max_len: usize_or("max_len", DEFAULT_MAX_LEN),
+            head_max_len: usize_or("head_max_len", DEFAULT_HEAD_MAX_LEN),
+            temperatures: Temperatures::from_json(
+                v.get("temperature"),
+                v.get("temperature_by_options"),
+                None,
+            )?,
+            lang_temperatures: HashMap::new(),
+        })
+    }
+
+    /// Add per-language temperatures, the shape the reference takes them in:
+    /// `{"de": {"temperature": [..3], "temperature_by_options": {..}}}`. A
+    /// language that omits `temperature` inherits the checkpoint's raw
+    /// per-type values.
+    pub fn with_lang_temperatures(mut self, spec: &Json) -> Result<Self, String> {
+        let Json::Object(m) = spec else {
+            return Err("lang_temperatures must be an object".into());
+        };
+        for (lang, cfg) in m {
+            let t = Temperatures::from_json(
+                cfg.get("temperature"),
+                cfg.get("temperature_by_options"),
+                Some(self.temperatures.by_type),
+            )
+            .map_err(|e| format!("language override {lang:?}: {e}"))?;
+            let primary = lang.split('-').next().unwrap_or(lang).to_lowercase();
+            self.lang_temperatures.insert(primary, t);
+        }
+        Ok(self)
+    }
+
+    /// The temperatures to apply for a request detected (or declared) as
+    /// `lang`, which may carry a region (`de-AT`).
+    pub fn temperatures_for(&self, lang: Option<&str>) -> &Temperatures {
+        lang.and_then(|l| {
+            let primary = l.split('-').next().unwrap_or(l).to_lowercase();
+            self.lang_temperatures.get(&primary)
+        })
+        .unwrap_or(&self.temperatures)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const ENGLISH: &str = r#"{"encoder": "answerdotai/ModernBERT-large", "head_layers": 2,
+        "max_len": 512, "head_max_len": 192,
+        "temperature": [1.6369030475616455, 1.2514300346374512, 1.983399510383606],
+        "temperature_by_options": {"choice:3-5": 1.7601518630981445, "choice:11+": 0.10058280825614929}}"#;
+
+    #[test]
+    fn shipped_choice_11_bucket_is_clamped() {
+        let c = DecisionConfig::from_json(ENGLISH).unwrap();
+        assert_eq!(c.temperatures.for_question(0, 12), TEMP_MIN);
+        assert_eq!(c.temperatures.for_question(0, 4), 1.7601518630981445);
+        // No bucket: the per-type value.
+        assert_eq!(c.temperatures.for_question(1, 7), 1.2514300346374512);
+    }
+
+    #[test]
+    fn buckets_split_where_the_reference_splits() {
+        let b: Vec<_> = [1, 2, 3, 5, 6, 10, 11]
+            .iter()
+            .map(|&k| temp_bucket(2, k))
+            .collect();
+        assert_eq!(
+            b,
+            [
+                "noul:2",
+                "noul:2",
+                "noul:3-5",
+                "noul:3-5",
+                "noul:6-10",
+                "noul:6-10",
+                "noul:11+"
+            ]
+        );
+    }
+
+    #[test]
+    fn language_overrides_match_on_primary_subtag() {
+        let c = DecisionConfig::from_json(ENGLISH)
+            .unwrap()
+            .with_lang_temperatures(
+                &Json::parse(br#"{"DE": {"temperature": [1.0, 1.4, 9.0]}}"#).unwrap(),
+            )
+            .unwrap();
+        assert_eq!(
+            c.temperatures_for(Some("de-AT")).by_type,
+            [1.0, 1.4, TEMP_MAX]
+        );
+        assert_eq!(c.temperatures_for(Some("fr")), &c.temperatures);
+        assert_eq!(c.temperatures_for(None), &c.temperatures);
+    }
+
+    #[test]
+    fn bad_temperatures_are_refused_or_neutralised() {
+        assert!(DecisionConfig::from_json(r#"{"temperature": [1.0]}"#).is_err());
+        let c = DecisionConfig::from_json(r#"{"temperature": [" 2.5 ", "x", null]}"#).unwrap();
+        assert_eq!(c.temperatures.by_type, [2.5, 1.0, 1.0]);
+        assert_eq!(
+            (c.max_len, c.head_max_len),
+            (DEFAULT_MAX_LEN, DEFAULT_HEAD_MAX_LEN)
+        );
+    }
+}
@@ -0,0 +1,249 @@
+//! `CandleHarness::systemone` end to end: request in, routed, validated,
+//! scored and shaped response out (#338).
+//!
+//! The tiny checkpoint (`testdata/laya/tiny/`) exercises the plumbing in
+//! CI. The released-weights test replays every reference case through
+//! the same entry point the HTTP handler calls and compares the whole
+//! response — answers, usage and routing — with what the reference
+//! server returned; it needs `LAYA_REFERENCE_DIR`.
+
+use std::path::{Path, PathBuf};
+use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
+
+use cortex_core::decisions::{Json, SystemOneRequest, limits};
+
+use crate::harness::candle::{CandleHarness, SystemOneError};
+use crate::harness::decision_model::{CheckpointFiles, DecisionCheckpoint, LoadedDecisionModel};
+
+fn testdata() -> PathBuf {
+    Path::new(env!("CARGO_MANIFEST_DIR")).join("src/harness/testdata/laya")
+}
+
+fn harness() -> Arc<CandleHarness> {
+    CandleHarness::new(
+        "http://localhost:13131".into(),
+        &crate::config::CandleHarnessConfig::default(),
+    )
+}
+
+/// A decision model made of the given `(routing name, directory)`
+/// checkpoints, loaded on the CPU in f32.
+async fn decision_model(model_id: &str, checkpoints: &[(&str, PathBuf)]) -> LoadedDecisionModel {
+    let mut loaded = Vec::new();
+    for (name, dir) in checkpoints {
+        let files = CheckpointFiles {
+            tokenizer: if dir.join("tokenizer/tokenizer.json").exists() {
+                dir.join("tokenizer/tokenizer.json")
+            } else {
+                dir.join("tokenizer.json")
+            },
+            dir: dir.clone(),
+        };
+        loaded.push(
+            DecisionCheckpoint::load(name.to_string(), files, model_id, None)
+                .await
+                .unwrap_or_else(|e| panic!("load {name}: {e:#}")),
+        );
+    }
+    LoadedDecisionModel {
+        model_id: model_id.to_string(),
+        spec: serde_json::from_value(serde_json::json!({
+            "model_id": model_id,
+            "harness": "candle",
+        }))
+        .expect("model spec"),
+        devices: vec![0],
+        checkpoints: loaded,
+        poisoned: Arc::new(AtomicBool::new(false)),
+        admission: crate::harness::admission::AdmissionController::new(
+            &crate::config::AdmissionConfig::default(),
+        ),
+    }
+}
+
+fn request(body: &str) -> SystemOneRequest {
+    SystemOneRequest::parse(body.as_bytes(), limits::DEFAULT_MAX_TOKEN_BUDGET)
+        .expect("valid request")
+}
+
+async fn tiny_harness() -> Arc<CandleHarness> {
+    let h = harness();
+    h.register_decision_model(
+        decision_model("test/tiny", &[("english", testdata().join("tiny"))]).await,
+    )
+    .await;
+    h
+}
+
+#[tokio::test]
+async fn answers_every_question_in_order_with_exact_usage() {
+    let h = tiny_harness().await;
+    let req = request(
+        r#"{"state": "a b a", "questions": {
+            "z_last_alphabetically": {"type": "noul", "instructions": "a?"},
+            "a_first": {"type": "choice", "instructions": "b?", "criteria": ["a", "b", "a b"]}
+        }}"#,
+    );
+    let resp = h.systemone(None, req, None).await.expect("answered");
+    let json = serde_json::to_value(&resp).unwrap();
+    // Answers come back in request order, not sorted.
+    let keys: Vec<&str> = resp.answers.0.iter().map(|(k, _)| k.as_str()).collect();
+    assert_eq!(keys, ["z_last_alphabetically", "a_first"]);
+    assert_eq!(json["answers"]["z_last_alphabetically"]["type"], "noul");
+    let probs = json["answers"]["a_first"]["probabilities"]
+        .as_object()
+        .unwrap();
+    let total: f64 = probs.values().map(|p| p.as_f64().unwrap()).sum();
+    assert!((total - 1.0).abs() < 1e-3, "{probs:?}");
+    // Usage is the exact number of tokens encoded, output tokens zero.
+    assert!(json["usage"]["input_tokens"].as_u64().unwrap() > 0);
+    assert_eq!(json["usage"]["output_tokens"], 0);
+    assert_eq!(json["routing"]["model"], "english");
+    assert_eq!(json["routing"]["repo"], "test/tiny");
+}
+
+#[tokio::test]
+async fn pinning_a_checkpoint_the_model_lacks_is_404() {
+    let h = tiny_harness().await;
+    let req = request(
+        r#"{"state": "a", "model": "multilingual",
+            "questions": {"q": {"type": "noul", "instructions": "a?"}}}"#,
+    );
+    match h.systemone(None, req, None).await {
+        Err(SystemOneError::Request(e)) => assert_eq!(e.status, 404, "{}", e.detail),
+        other => panic!("expected a 404, got {other:?}"),
+    }
+}
+
+#[tokio::test]
+async fn routing_to_a_missing_checkpoint_falls_back_and_says_so() {
+    // Devanagari routes to `multilingual`, which this model lacks.
+    let h = tiny_harness().await;
+    let req = request(
+        r#"{"state": "मुझे पैसे वापस चाहिए",
+            "questions": {"q": {"type": "noul", "instructions": "a?"}}}"#,
+    );
+    let resp = h.systemone(None, req, None).await.expect("answered");
+    let routing = resp.routing.expect("routing");
+    assert_eq!(routing.model, "english");
+    assert!(
+        routing
+            .reason
+            .contains("multilingual checkpoint is not loaded"),
+        "{}",
+        routing.reason
+    );
+}
+
+#[tokio::test]
+async fn invalid_questions_are_rejected_before_admission() {
+    let h = tiny_harness().await;
+    let req =
+        request(r#"{"state": "a", "questions": {"q": {"type": "maybe", "instructions": "a?"}}}"#);
+    match h.systemone(None, req, None).await {
+        Err(SystemOneError::Request(e)) => assert_eq!(e.status, 422, "{}", e.detail),
+        other => panic!("expected a 422, got {other:?}"),
+    }
+}
+
+#[tokio::test]
+async fn a_gateway_routed_model_that_is_not_a_decision_model_is_refused() {
+    let h = tiny_harness().await;
+    let req = request(r#"{"state": "a", "questions": {}}"#);
+    match h.systemone(Some("not/loaded"), req, None).await {
+        Err(SystemOneError::Inference(crate::harness::candle::InferenceError::ModelNotLoaded(
+            id,
+        ))) => {
+            assert_eq!(id, "not/loaded")
+        }
+        other => panic!("expected ModelNotLoaded, got {other:?}"),
+    }
+}
+
+/// Every reference case through `systemone`, compared with the reference
+/// server's whole response. Routing and usage must match exactly;
+/// probabilities to within the engine's measured f32 distance from the
+/// reference (a last-digit wobble after 4-decimal rounding); the chosen
+/// answer exactly.
+#[tokio::test]
+#[ignore = "needs the released weights: set LAYA_REFERENCE_DIR"]
+async fn released_family_answers_like_the_reference_server() {
+    let root = PathBuf::from(std::env::var("LAYA_REFERENCE_DIR").expect("LAYA_REFERENCE_DIR"));
+    let h = harness();
+    h.register_decision_model(
+        decision_model(
+            "convaiinnovations/laya",
+            &[
+                ("english", root.clone()),
+                ("multilingual", root.join("multilingual")),
+                ("typed-decisions", root.join("typed-decisions")),
+            ],
+        )
+        .await,
+    )
+    .await;
+
+    let text = std::fs::read_to_string(testdata().join("reference.json")).unwrap();
+    let reference = Json::parse(text.as_bytes()).unwrap();
+    let Some(Json::Array(cases)) = reference.get("cases") else {
+        panic!("no cases")
+    };
+    let mut max_dp = 0.0f64;
+    for case in cases {
+        let name = case.get("name").and_then(Json::as_str).unwrap();
+        let body = serde_json::to_vec(case.get("request").unwrap()).unwrap();
+        let req = SystemOneRequest::parse(&body, limits::DEFAULT_MAX_TOKEN_BUDGET).unwrap();
+        let resp = h
+            .systemone(Some("convaiinnovations/laya"), req, None)
+            .await
+            .unwrap_or_else(|e| panic!("{name}: {e:?}"));
+        let want_json = case.get("response").unwrap();
+        // Routing and usage compared as serialized text through `Json`:
+        // exact, key order included. serde_json would re-round the
+        // reference's floats on the way in.
+        for key in ["routing", "usage", "model"] {
+            let got = match key {
+                "routing" => serde_json::to_string(&resp.routing).unwrap(),
+                "usage" => serde_json::to_string(&resp.usage).unwrap(),
+                _ => serde_json::to_string(&resp.model).unwrap(),
+            };
+            let want = serde_json::to_string(want_json.get(key).unwrap()).unwrap();
+            assert_eq!(got, want, "{name}: {key}");
+        }
+        let got = serde_json::to_value(&resp).unwrap();
+        let want: serde_json::Value =
+            serde_json::from_str(&serde_json::to_string(want_json).unwrap()).unwrap();
+        let (ga, wa) = (
+            got["answers"].as_object().unwrap(),
+            want["answers"].as_object().unwrap(),
+        );
+        assert_eq!(
+            ga.keys().collect::<Vec<_>>(),
+            wa.keys().collect::<Vec<_>>(),
+            "{name}: answer ids"
+        );
+        for (qid, w) in wa {
+            let g = &ga[qid];
+            for key in ["choice", "type", "legend"] {
+                assert_eq!(g.get(key), w.get(key), "{name}/{qid}: {key}");
+            }
+            if let Some(wp) = w["probabilities"].as_object() {
+                for (opt, p) in wp {
+                    let d = (g["probabilities"][opt].as_f64().unwrap() - p.as_f64().unwrap()).abs();
+                    max_dp = max_dp.max(d);
+                }
+            }
+            for key in ["noul", "score", "confidence", "answer_confidence"] {
+                if let (Some(gv), Some(wv)) = (
+                    g.get(key).and_then(|v| v.as_f64()),
+                    w.get(key).and_then(|v| v.as_f64()),
+                ) {
+                    max_dp = max_dp.max((gv - wv).abs());
+                }
+            }
+        }
+    }
+    eprintln!("released family: max |Δ| over probabilities and confidences {max_dp:.1e}");
+    assert!(max_dp <= 5e-4, "max |Δ| {max_dp:e}");
+}
@@ -0,0 +1,59 @@
+//! Laya decision models: everything around the forward pass (#335).
+//!
+//! A Laya checkpoint answers typed questions over a *state* (text, a JSON
+//! document, a conversation) in one encoder forward pass: each question
+//! becomes one sequence,
+//!
+//! ```text
+//! [CLS] "<type> question: <instructions>" [SEP] [MASK] opt0 [MASK] opt1 … [SEP] state [SEP]
+//! ```
+//!
+//! and the decision head scores the hidden state at each option's `[MASK]`
+//! marker. This module is the part that touches no tensor:
+//!
+//! - [`question`] — validate and normalise the question definitions, and
+//!   render each option's text;
+//! - [`sequence`] — assemble the token sequences under the model's token
+//!   budgets ([`prepare`]);
+//! - [`shape`] — turn per-option logits into calibrated, typed answers
+//!   ([`shape()`]);
+//! - [`config`] / [`tokenizer`] — the per-checkpoint pieces those need.
+//!
+//! The reference is the `laya` Python SDK (`laya/common.py`,
+//! `laya/agent.py`, `laya/serve.py`), whose behaviour this reproduces token
+//! for token: the parity tests replay fixtures recorded from it
+//! (`script/laya-reference.py`). The calibration arithmetic follows numpy's
+//! float32 evaluation order, so rounded probabilities match to the last
+//! published digit.
+//!
+//! The wire types (`/v1/systemone`) live in [`cortex_core::decisions`].
+
+pub mod config;
+pub mod question;
+pub mod route;
+pub mod sequence;
+pub mod shape;
+pub mod tokenizer;
+
+#[cfg(test)]
+mod handler_tests;
+
+pub use config::DecisionConfig;
+pub use cortex_core::decisions::DecisionRequestError;
+pub use question::{QType, Question};
+pub use sequence::{Budget, OptionStats, Prepared, Row, prepare};
+pub use shape::{RowOutput, shape};
+pub use tokenizer::{DecisionTokenizer, SpecialTokens, TextEncoder};
+
+/// The constant `model` field of every response: the decision head's name.
+/// Which checkpoint answered is reported in `routing`.
+pub const DECISION_MODEL_NAME: &str = "laya-rl-agent";
+
+/// A request error in the reference's `ValueError` class: the server
+/// answers it with 422 and the message as `detail`.
+pub(crate) fn unprocessable(detail: impl Into<String>) -> DecisionRequestError {
+    DecisionRequestError::new(422, detail)
+}
+
+#[cfg(test)]
+mod parity_tests;
@@ -0,0 +1,209 @@
+//! Replays the reference fixtures (`testdata/laya/reference.json`, recorded
+//! by `script/laya-reference.py`) through [`prepare`] and [`shape`].
+//!
+//! Token ids and marker positions must match exactly. The answers and usage
+//! must serialize to exactly the same JSON — same keys, same order, same
+//! rounded values — when fed the reference's own raw logits.
+//!
+//! The English tokenizer is committed. The multilingual one is 34 MB, so the
+//! multilingual cases run only when `LAYA_ML_TOKENIZER_DIR` names a copy of
+//! the checkpoint's `multilingual/tokenizer/` directory.
+
+use super::*;
+use cortex_core::decisions::{Json, SystemOneRequest, limits};
+use std::path::{Path, PathBuf};
+
+fn testdata() -> PathBuf {
+    Path::new(env!("CARGO_MANIFEST_DIR")).join("src/harness/testdata/laya")
+}
+
+fn reference() -> Json {
+    Json::parse(&std::fs::read(testdata().join("reference.json")).expect("fixtures")).expect("json")
+}
+
+fn to_value(j: &Json) -> serde_json::Value {
+    serde_json::to_value(j).expect("serializable")
+}
+
+fn tokenizer_for(checkpoint: &str) -> Option<DecisionTokenizer> {
+    let dir = match checkpoint {
+        // typed-decisions ships the English tokenizer, byte for byte.
+        "english" | "typed-decisions" => testdata().join("tokenizer"),
+        "multilingual" => PathBuf::from(std::env::var_os("LAYA_ML_TOKENIZER_DIR")?),
+        other => panic!("unknown checkpoint {other}"),
+    };
+    Some(DecisionTokenizer::from_dir(&dir).expect("tokenizer"))
+}
+
+fn config_for(reference: &Json, checkpoint: &str) -> DecisionConfig {
+    let cfg = reference
+        .get("header")
+        .and_then(|h| h.get("checkpoints"))
+        .and_then(|c| c.get(checkpoint))
+        .and_then(|c| c.get("cfg"))
+        .expect("checkpoint cfg");
+    DecisionConfig::from_json(&serde_json::to_string(cfg).unwrap()).unwrap()
+}
+
+fn array(j: &Json) -> &[Json] {
+    match j {
+        Json::Array(a) => a,
+        other => panic!("expected a list, got {other:?}"),
+    }
+}
+
+fn num(j: &Json) -> f64 {
+    match j {
+        Json::Int(d) => d.parse().unwrap(),
+        Json::Float(x) => *x,
+        other => panic!("expected a number, got {other:?}"),
+    }
+}
+
+/// Replay every case whose tokenizer is available; return how many ran per
+/// checkpoint.
+fn replay(filter: impl Fn(&str) -> bool) -> Vec<(String, usize)> {
+    let reference = reference();
+    let mut ran: Vec<(String, usize)> = Vec::new();
+    for case in array(reference.get("cases").unwrap()) {
+        let name = case.get("name").and_then(Json::as_str).unwrap();
+        if !filter(name) {
+            continue;
+        }
+        let response = case.get("response").unwrap();
+        let routing = response.get("routing").unwrap();
+        let checkpoint = routing.get("model").and_then(Json::as_str).unwrap();
+        let Some(tok) = tokenizer_for(checkpoint) else {
+            continue;
+        };
+        let cfg = config_for(&reference, checkpoint);
+
+        let body = serde_json::to_vec(case.get("request").unwrap()).unwrap();
+        let req = SystemOneRequest::parse(&body, limits::DEFAULT_MAX_TOKEN_BUDGET)
+            .unwrap_or_else(|e| panic!("{name}: {e}"));
+        let budget = Budget {
+            max_len: req.max_len,
+            head_max_len: req.head_max_len,
+        };
+        let prepared = prepare(&req.state, &req.questions, &cfg, &tok, budget)
+            .unwrap_or_else(|e| panic!("{name}: {e}"));
+
+        let items = array(case.get("items").unwrap());
+        assert_eq!(prepared.rows.len(), items.len(), "{name}: row count");
+        let mut outputs = Vec::new();
+        for (i, (row, item)) in prepared.rows.iter().zip(items).enumerate() {
+            let ids: Vec<u32> = array(item.get("input_ids").unwrap())
+                .iter()
+                .map(|v| num(v) as u32)
+                .collect();
+            let markers: Vec<usize> = array(item.get("markers").unwrap())
+                .iter()
+                .map(|v| num(v) as usize)
+                .collect();
+            assert_eq!(row.input_ids, ids, "{name}: row {i} input_ids");
+            assert_eq!(row.markers, markers, "{name}: row {i} markers");
+            assert_eq!(row.qtype as usize, num(item.get("qtype").unwrap()) as usize);
+            let opts = item.get("options").unwrap();
+            let s = prepared.stats[i];
+            assert_eq!(
+                (s.options, s.options_distinct, s.tokens_per_option),
+                (
+                    num(opts.get("options").unwrap()) as usize,
+                    num(opts.get("options_distinct").unwrap()) as usize,
+                    match opts.get("tokens_per_option").unwrap() {
+                        Json::Null => None,
+                        v => Some(num(v) as usize),
+                    }
+                ),
+                "{name}: row {i} option stats"
+            );
+            outputs.push(RowOutput {
+                logits: array(item.get("logits").unwrap())
+                    .iter()
+                    .map(|v| num(v) as f32)
+                    .collect(),
+                act_probability: num(&array(item.get("act").unwrap())[0]) as f32,
+            });
+        }
+
+        let lang = routing
+            .get("detection")
+            .and_then(|d| d.get("language"))
+            .and_then(Json::as_str);
+        let (answers, usage) = shape(&prepared, &outputs, &cfg, lang);
+        // Serialized strings, not values: key order is part of the contract.
+        assert_eq!(
+            serde_json::to_string(&answers).unwrap(),
+            serde_json::to_string(response.get("answers").unwrap()).unwrap(),
+            "{name}: answers"
+        );
+        assert_eq!(
+            serde_json::to_value(&usage).unwrap(),
+            to_value(response.get("usage").unwrap()),
+            "{name}: usage"
+        );
+        match ran.iter_mut().find(|(c, _)| c == checkpoint) {
+            Some((_, n)) => *n += 1,
+            None => ran.push((checkpoint.to_string(), 1)),
+        }
+    }
+    ran
+}
+
+#[test]
+fn every_english_fixture_matches_the_reference() {
+    let ran = replay(|_| true);
+    let count = |c: &str| ran.iter().find(|(n, _)| n == c).map_or(0, |(_, n)| *n);
+    // The corpus routes 19 cases to English and 1 to typed-decisions; a
+    // drop means a case silently stopped running.
+    assert_eq!(count("english"), 19, "{ran:?}");
+    assert_eq!(count("typed-decisions"), 1, "{ran:?}");
+    if std::env::var_os("LAYA_ML_TOKENIZER_DIR").is_some() {
+        assert_eq!(count("multilingual"), 8, "{ran:?}");
+    }
+}
+
+#[test]
+fn empty_questions_need_no_forward() {
+    let reference = reference();
+    let tok = tokenizer_for("english").unwrap();
+    let cfg = config_for(&reference, "english");
+    let prepared = prepare(
+        &Json::String("x".into()),
+        &[],
+        &cfg,
+        &tok,
+        Budget::default(),
+    )
+    .unwrap();
+    assert!(prepared.rows.is_empty());
+    let (answers, usage) = shape(&prepared, &[], &cfg, None);
+    assert_eq!(serde_json::to_string(&answers).unwrap(), "{}");
+    assert_eq!(
+        serde_json::to_string(&usage).unwrap(),
+        r#"{"input_tokens":0,"output_tokens":0}"#
+    );
+}
+
+#[test]
+fn options_that_cannot_all_fit_are_refused() {
+    let reference = reference();
+    let tok = tokenizer_for("english").unwrap();
+    let cfg = config_for(&reference, "english");
+    let questions = Json::parse(
+        br#"{"q": {"type": "choice", "instructions": "i", "criteria": ["a", "b", "c", "d"]}}"#,
+    )
+    .unwrap();
+    let Json::Object(questions) = questions else {
+        unreachable!()
+    };
+    let budget = Budget {
+        max_len: Some(10),
+        head_max_len: None,
+    };
+    let err = prepare(&Json::String("x".into()), &questions, &cfg, &tok, budget).unwrap_err();
+    assert_eq!(
+        (err.status, err.detail.as_str()),
+        (422, "question 'q' options exceed head_max_len=192")
+    );
+}
@@ -0,0 +1,458 @@
+//! Question validation, normalisation and option rendering.
+//!
+//! Every rejection here is a 422 whose message names the question and what
+//! to fix — the same text the reference returns, so a client written
+//! against it sees the same errors.
+
+use super::{DecisionRequestError, unprocessable};
+use cortex_core::decisions::{HashKey, Json};
+use std::collections::HashMap;
+
+/// The three decision primitives. The discriminant is the index the model's
+/// type embedding and the calibration tables use.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum QType {
+    /// Pick one of the caller's labels.
+    Choice = 0,
+    /// An ordinal level on the caller's rubric.
+    Score = 1,
+    /// Yes/no: the probability that a statement holds.
+    Noul = 2,
+}
+
+impl QType {
+    pub const NAMES: [&'static str; 3] = ["choice", "score", "noul"];
+
+    pub fn name(self) -> &'static str {
+        Self::NAMES[self as usize]
+    }
+
+    fn parse(s: &str) -> Option<Self> {
+        match s {
+            "choice" => Some(Self::Choice),
+            "score" => Some(Self::Score),
+            "noul" => Some(Self::Noul),
+            _ => None,
+        }
+    }
+}
+
+/// A validated question, normalised the way the reference normalises it.
+#[derive(Debug, Clone, PartialEq)]
+pub struct Question {
+    pub id: String,
+    pub kind: QType,
+    /// The instruction text: the caller's string, or the JSON rendering of
+    /// a structured value.
+    pub instructions: String,
+    /// Each option's rendered text, in answer-index order. For `noul` this
+    /// is always `[false, true]`, so index 1 is P(true).
+    pub options: Vec<String>,
+    /// `choice`: each option's label as the caller wrote it (the answer key).
+    /// `score`: each level's description (the legend). `noul`: empty.
+    pub labels: Vec<Json>,
+}
+
+impl Question {
+    /// Validate one question definition and normalise it.
+    pub fn from_definition(id: &str, def: &Json) -> Result<Self, DecisionRequestError> {
+        check_question(id, def)?;
+        Ok(to_internal(id, def))
+    }
+}
+
+/// Validate and normalise every question, in request order. The first
+/// invalid question is reported; nothing is normalised until all pass.
+pub fn questions_from_request(
+    questions: &[(String, Json)],
+) -> Result<Vec<Question>, DecisionRequestError> {
+    for (id, def) in questions {
+        check_question(id, def)?;
+    }
+    Ok(questions
+        .iter()
+        .map(|(id, def)| to_internal(id, def))
+        .collect())
+}
+
+fn q(id: &str) -> String {
+    Json::String(id.to_string()).py_repr()
+}
+
+fn check_question(id: &str, def: &Json) -> Result<(), DecisionRequestError> {
+    let qid = q(id);
+    if !matches!(def, Json::Object(_)) {
+        return Err(unprocessable(format!(
+            "question {qid}: definition must be a dict, got {}",
+            def.py_type_name()
+        )));
+    }
+    let t = def.get("type").unwrap_or(&Json::Null);
+    let Some(kind) = t.as_str().and_then(QType::parse) else {
+        return Err(unprocessable(format!(
+            "question {qid}: unknown type {}; use one of ['choice', 'noul', 'score']",
+            t.py_repr()
+        )));
+    };
+    if def.get("instructions").is_none() {
+        return Err(unprocessable(format!(
+            "question {qid}: no 'instructions'; add the text the model should answer"
+        )));
+    }
+    let crit = def.get("criteria");
+    match kind {
+        QType::Choice => check_choice(&qid, crit)?,
+        QType::Score => {
+            let Some(Json::Array(levels)) = crit else {
+                return Err(unprocessable(format!(
+                    "question {qid}: a score question takes 'criteria' as a list of level \
+                     descriptions, index 0 first"
+                )));
+            };
+            if levels.is_empty() {
+                return Err(unprocessable(format!(
+                    "question {qid}: a score question needs at least one level"
+                )));
+            }
+            if let Some(i) = levels.iter().position(|l| *l == Json::Null) {
+                return Err(unprocessable(format!(
+                    "question {qid}: score level {i} is null; give every level a description, \
+                     index 0 first"
+                )));
+            }
+        }
+        QType::Noul => match crit {
+            None | Some(Json::Null) => {}
+            Some(Json::Object(m)) => {
+                let mut keys: Vec<String> = m.iter().map(|(k, _)| k.to_lowercase()).collect();
+                keys.sort();
+                keys.dedup();
+                if keys.iter().any(|k| k != "true" && k != "false") {
+                    let listed = Json::Array(keys.into_iter().map(Json::String).collect());
+                    return Err(unprocessable(format!(
+                        "question {qid}: a noul question takes 'criteria' keyed only \
+                         'true'/'false' (either or both, and omitted is fine), got {}. Those \
+                         keys are the option texts the model reads; any other key was silently \
+                         dropped and replaced with the defaults. If you want the answer worded \
+                         differently, keep 'criteria' keyed 'true'/'false' and set 'labels' \
+                         instead.",
+                        listed.py_repr()
+                    )));
+                }
+            }
+            Some(_) => {
+                return Err(unprocessable(format!(
+                    "question {qid}: a noul question takes 'criteria' as a dict with optional \
+                     'true'/'false' descriptions, or omits it"
+                )));
+            }
+        },
+    }
+    if let Some(labels) = def.get("labels") {
+        if kind != QType::Noul {
+            return Err(unprocessable(format!(
+                "question {qid}: 'labels' is only supported for noul questions"
+            )));
+        }
+        resolve_noul_labels(labels).map_err(|e| unprocessable(format!("question {qid}: {e}")))?;
+    }
+    Ok(())
+}
+
+fn check_choice(qid: &str, crit: Option<&Json>) -> Result<(), DecisionRequestError> {
+    let labels: Vec<Json> = match crit {
+        Some(Json::Object(m)) => m.iter().map(|(k, _)| Json::String(k.clone())).collect(),
+        Some(Json::Array(a)) => a.clone(),
+        _ => {
+            return Err(unprocessable(format!(
+                "question {qid}: a choice question takes 'criteria' as a dict of label -> \
+                 description, or a list of labels"
+            )));
+        }
+    };
+    if labels.is_empty() {
+        return Err(unprocessable(format!(
+            "question {qid}: a choice question needs at least one criterion"
+        )));
+    }
+    for (i, label) in labels.iter().enumerate() {
+        match label {
+            Json::Array(_) | Json::Object(_) => {
+                return Err(unprocessable(format!(
+                    "question {qid}: choice label {i} is a {}; a label is rendered as option \
+                     text and used as the answer key, so it must be a scalar (a string, number \
+                     or bool), got {}",
+                    label.py_type_name(),
+                    label.py_repr()
+                )));
+            }
+            Json::Null => {
+                return Err(unprocessable(format!(
+                    "question {qid}: choice label {i} is null; a label is rendered as option \
+                     text and used as the answer key, so it must be a string, number or bool \
+                     -- a null label renders as the text \"None\" while its answer key is \
+                     \"null\""
+                )));
+            }
+            _ => {}
+        }
+    }
+    // Only the list form can repeat a key: an object's keys are unique.
+    if matches!(crit, Some(Json::Array(_))) {
+        let mut seen: HashMap<HashKey, usize> = HashMap::new();
+        for (i, label) in labels.iter().enumerate() {
+            let key = label.hash_key().expect("scalars are hashable");
+            if key == HashKey::Unique {
+                continue;
+            }
+            if let Some(first) = seen.get(&key) {
+                return Err(unprocessable(format!(
+                    "question {qid}: choice label {i} ({}) repeats label {first}; the labels are \
+                     the answer keys, so every option needs its own (1, 1.0 and True are one key)",
+                    label.py_repr()
+                )));
+            }
+            seen.insert(key, i);
+        }
+    }
+    Ok(())
+}
+
+const DEFAULT_NOUL_FALSE: &str = "no, the statement does not hold";
+const DEFAULT_NOUL_TRUE: &str = "yes, the statement holds";
+
+/// The words a noul question shows the model for its two outcomes.
+fn resolve_noul_labels(labels: &Json) -> Result<(String, String), &'static str> {
+    const BAD: &str =
+        "noul labels must map exactly 'false' and 'true' to distinct non-empty strings";
+    let Json::Object(m) = labels else {
+        return if *labels == Json::Null {
+            Ok(("false".into(), "true".into()))
+        } else {
+            Err(BAD)
+        };
+    };
+    let mut keys: Vec<&str> = m.iter().map(|(k, _)| k.as_str()).collect();
+    keys.sort_unstable();
+    if keys != ["false", "true"] {
+        return Err(BAD);
+    }
+    let (Some(Json::String(f)), Some(Json::String(t))) = (labels.get("false"), labels.get("true"))
+    else {
+        return Err(BAD);
+    };
+    let (f, t) = (py_strip(f), py_strip(t));
+    if f.is_empty() || t.is_empty() || f == t {
+        return Err(BAD);
+    }
+    Ok((f.to_string(), t.to_string()))
+}
+
+/// Python's `str.strip()` with no argument: Unicode whitespace.
+fn py_strip(s: &str) -> &str {
+    s.trim_matches(|c: char| c.is_whitespace() || ('\u{1c}'..='\u{1f}').contains(&c))
+}
+
+/// `render_criterion`: a string as itself, anything else as JSON.
+fn render_criterion(v: &Json) -> String {
+    match v {
+        Json::String(s) => s.clone(),
+        other => other.dumps(),
+    }
+}
+
+/// Only `None` and `""` mean "no description"; `0` and `false` are values.
+fn is_blank(v: &Json) -> bool {
+    matches!(v, Json::Null) || matches!(v, Json::String(s) if s.is_empty())
+}
+
+fn to_internal(id: &str, def: &Json) -> Question {
+    let kind = def
+        .get("type")
+        .and_then(Json::as_str)
+        .and_then(QType::parse)
+        .expect("validated");
+    let instructions = match def.get("instructions").expect("validated") {
+        Json::String(s) => s.clone(),
+        other => other.dumps(),
+    };
+    let crit = def.get("criteria");
+    let (options, labels) = match kind {
+        QType::Choice => {
+            let pairs: Vec<(Json, Json)> = match crit {
+                Some(Json::Object(m)) => m
+                    .iter()
+                    .map(|(k, v)| (Json::String(k.clone()), v.clone()))
+                    .collect(),
+                Some(Json::Array(a)) => a.iter().map(|l| (l.clone(), Json::Null)).collect(),
+                _ => unreachable!("validated"),
+            };
+            let options = pairs
+                .iter()
+                .map(|(k, v)| {
+                    if is_blank(v) {
+                        k.py_str()
+                    } else {
+                        format!("{}: {}", k.py_str(), render_criterion(v))
+                    }
+                })
+                .collect();
+            (options, pairs.into_iter().map(|(k, _)| k).collect())
+        }
+        QType::Score => {
+            let Some(Json::Array(levels)) = crit else {
+                unreachable!("validated")
+            };
+            let options = levels
+                .iter()
+                .enumerate()
+                .map(|(i, c)| format!("level {i}: {}", render_criterion(c)))
+                .collect();
+            (options, levels.clone())
+        }
+        QType::Noul => {
+            // Keys are matched case-insensitively; a repeat after
+            // lower-casing keeps the later value, as a dict would.
+            let mut desc: HashMap<String, Json> = HashMap::new();
+            if let Some(Json::Object(m)) = crit {
+                for (k, v) in m {
+                    desc.insert(k.to_lowercase(), v.clone());
+                }
+            }
+            let (fl, tl) =
+                resolve_noul_labels(def.get("labels").unwrap_or(&Json::Null)).expect("validated");
+            let render = |key: &str, default: &str| match desc.get(key) {
+                Some(v) if !is_blank(v) => render_criterion(v),
+                _ => default.to_string(),
+            };
+            (
+                vec![
+                    format!("{fl}: {}", render("false", DEFAULT_NOUL_FALSE)),
+                    format!("{tl}: {}", render("true", DEFAULT_NOUL_TRUE)),
+                ],
+                Vec::new(),
+            )
+        }
+    };
+    Question {
+        id: id.to_string(),
+        kind,
+        instructions,
+        options,
+        labels,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn def(s: &str) -> Json {
+        Json::parse(s.as_bytes()).unwrap()
+    }
+
+    fn err(s: &str) -> String {
+        Question::from_definition("q", &def(s)).unwrap_err().detail
+    }
+
+    #[test]
+    fn options_render_as_the_reference_renders_them() {
+        let c = Question::from_definition(
+            "q",
+            &def(r#"{"type": "choice", "instructions": {"a": 1},
+                     "criteria": {"x": "", "y": null, "z": 0, "w": {"d": [1, 2.0]}, "v": "text"}}"#),
+        )
+        .unwrap();
+        assert_eq!(c.instructions, r#"{"a": 1}"#);
+        assert_eq!(
+            c.options,
+            ["x", "y", "z: 0", r#"w: {"d": [1, 2.0]}"#, "v: text"]
+        );
+
+        let l = Question::from_definition(
+            "q",
+            &def(r#"{"type": "choice", "instructions": "i", "criteria": [1, 2.5, false, "s"]}"#),
+        )
+        .unwrap();
+        assert_eq!(l.options, ["1", "2.5", "False", "s"]);
+
+        let s = Question::from_definition(
+            "q",
+            &def(r#"{"type": "score", "instructions": "i", "criteria": ["calm", {"k": false}]}"#),
+        )
+        .unwrap();
+        assert_eq!(s.options, ["level 0: calm", r#"level 1: {"k": false}"#]);
+
+        let n = Question::from_definition(
+            "q",
+            &def(
+                r#"{"type": "noul", "instructions": "i", "criteria": {"TRUE": "held", "false": ""},
+                     "labels": {"false": " no ", "true": "yes"}}"#,
+            ),
+        )
+        .unwrap();
+        assert_eq!(
+            n.options,
+            ["no: no, the statement does not hold", "yes: held"]
+        );
+    }
+
+    #[test]
+    fn invalid_questions_name_the_question_and_the_fix() {
+        assert_eq!(
+            err("[]"),
+            "question 'q': definition must be a dict, got list"
+        );
+        assert_eq!(
+            err(r#"{"type": "pick"}"#),
+            "question 'q': unknown type 'pick'; use one of ['choice', 'noul', 'score']"
+        );
+        assert_eq!(
+            err(r#"{"type": "noul"}"#),
+            "question 'q': no 'instructions'; add the text the model should answer"
+        );
+        assert!(
+            err(r#"{"type": "choice", "instructions": "i", "criteria": []}"#)
+                .ends_with("needs at least one criterion")
+        );
+        assert!(
+            err(r#"{"type": "choice", "instructions": "i", "criteria": [[1]]}"#)
+                .starts_with("question 'q': choice label 0 is a list;")
+        );
+        assert!(
+            err(r#"{"type": "choice", "instructions": "i", "criteria": ["a", null]}"#)
+                .starts_with("question 'q': choice label 1 is null;")
+        );
+        assert!(
+            err(r#"{"type": "choice", "instructions": "i", "criteria": [1, "x", true]}"#)
+                .starts_with("question 'q': choice label 2 (True) repeats label 0;")
+        );
+        assert_eq!(
+            err(r#"{"type": "score", "instructions": "i", "criteria": ["a", null]}"#),
+            "question 'q': score level 1 is null; give every level a description, index 0 first"
+        );
+        assert!(
+            err(r#"{"type": "noul", "instructions": "i", "criteria": {"yes": "x"}}"#)
+                .contains("got ['yes'].")
+        );
+        assert_eq!(
+            err(r#"{"type": "score", "instructions": "i", "criteria": ["a"], "labels": {}}"#),
+            "question 'q': 'labels' is only supported for noul questions"
+        );
+        assert_eq!(
+            err(r#"{"type": "noul", "instructions": "i", "labels": {"false": "x", "true": " x"}}"#),
+            "question 'q': noul labels must map exactly 'false' and 'true' to distinct non-empty strings"
+        );
+    }
+
+    #[test]
+    fn nan_labels_never_collide() {
+        assert!(
+            Question::from_definition(
+                "q",
+                &def(r#"{"type": "choice", "instructions": "i", "criteria": [NaN, NaN]}"#),
+            )
+            .is_ok()
+        );
+    }
+}
@@ -0,0 +1,301 @@
+//! Choosing which checkpoint of a Laya family answers a request (#339).
+//!
+//! A released Laya repo carries three checkpoints: the English one at
+//! its root, `multilingual/` and `typed-decisions/`. The English
+//! checkpoint cannot read other scripts — upstream measured Khmer at
+//! 0.000 accuracy with 0.952 confidence — so an unpinned request is
+//! routed by the language of its state. This follows the reference's
+//! `Router._route` (`laya/router.py`) and its HTTP server's
+//! `_resolve_model` (`laya/serve.py`), including the `reason` strings.
+//!
+//! Precedence, as upstream: a checkpoint the caller names in `model`,
+//! then the detected script and language, then the default (English).
+//! Upstream's opt-in typed-decisions workflow routing and caller
+//! language hints are not exposed over HTTP by the reference, so they
+//! are not here either; the matched workflow is still reported.
+
+use cortex_core::decisions::{DecisionRouting, Json};
+
+use crate::harness::lang_detect;
+
+/// The root checkpoint's routing name.
+pub const ENGLISH: &str = "english";
+pub const MULTILINGUAL: &str = "multilingual";
+pub const TYPED_DECISIONS: &str = "typed-decisions";
+
+/// Where an unpinned request with no usable language signal goes.
+pub const DEFAULT_CHECKPOINT: &str = ENGLISH;
+
+/// Every checkpoint a released family carries, root first.
+pub const FAMILY: [&str; 3] = [ENGLISH, MULTILINGUAL, TYPED_DECISIONS];
+
+/// The reference's `_ALIASES`: names people are likely to type.
+const ALIASES: [(&str, &str); 10] = [
+    ("en", ENGLISH),
+    ("laya", ENGLISH),
+    ("default", ENGLISH),
+    ("multi", MULTILINGUAL),
+    ("ml", MULTILINGUAL),
+    ("laya-multilingual", MULTILINGUAL),
+    ("typed", TYPED_DECISIONS),
+    ("typed_decisions", TYPED_DECISIONS),
+    ("laya-typed-decisions", TYPED_DECISIONS),
+    ("decisions", TYPED_DECISIONS),
+];
+
+/// The reference server's `_PUBLISHED_MODEL_IDS`: the Hugging Face ids
+/// the non-root checkpoints are published under.
+const PUBLISHED: [(&str, &str); 2] = [
+    ("convaiinnovations/laya-multilingual", MULTILINGUAL),
+    ("convaiinnovations/laya-typed-decisions", TYPED_DECISIONS),
+];
+
+/// Question-id signatures of the four typed-decisions workflows.
+const WORKFLOWS: [(&str, [&str; 5]); 4] = [
+    (
+        "agent_trace_observability",
+        ["action", "needs_review", "outcome", "risk", "urgency"],
+    ),
+    (
+        "customer_service",
+        ["action", "category", "churn_risk", "needs_human", "urgency"],
+    ),
+    (
+        "invoice_processing",
+        [
+            "discrepancy_severity",
+            "disposition",
+            "duplicate",
+            "matches_order",
+            "urgency",
+        ],
+    ),
+    (
+        "security_incidents",
+        [
+            "credential_compromise",
+            "disposition",
+            "severity",
+            "true_positive",
+            "urgency",
+        ],
+    ),
+];
+
+/// The checkpoint a request's `model` field pins, if it names one.
+///
+/// Anything else — a Jev service name like `jev-1`, the family's own
+/// repo id, an empty string — pins nothing and the request is routed.
+pub fn resolve_pin(model: Option<&str>) -> Option<&'static str> {
+    let model = model?;
+    if model.is_empty() {
+        return None;
+    }
+    let key = py_strip(model).to_lowercase();
+    if let Some((_, ckpt)) = PUBLISHED.iter().find(|(id, _)| *id == key) {
+        return Some(ckpt);
+    }
+    let key = ALIASES
+        .iter()
+        .find(|(alias, _)| *alias == key)
+        .map(|(_, ckpt)| *ckpt)
+        .unwrap_or(key.as_str());
+    FAMILY.iter().copied().find(|c| *c == key)
+}
+
+/// The typed-decisions workflow whose question ids these are exactly.
+pub fn match_workflow(questions: &[(String, Json)]) -> Option<&'static str> {
+    let ids: std::collections::BTreeSet<&str> = questions.iter().map(|(q, _)| q.as_str()).collect();
+    WORKFLOWS.iter().find_map(|(name, sig)| {
+        let sig: std::collections::BTreeSet<&str> = sig.iter().copied().collect();
+        (ids == sig).then_some(*name)
+    })
+}
+
+/// Decide which checkpoint answers, and say why.
+///
+/// `model_id` is the served family's repo id, used for the `repo` field.
+/// The result names a checkpoint whether or not this node has it
+/// loaded; the caller decides what to do about one it lacks.
+pub fn route(
+    model_id: &str,
+    state: &Json,
+    questions: &[(String, Json)],
+    pin: Option<&'static str>,
+) -> DecisionRouting {
+    if let Some(key) = pin {
+        return DecisionRouting {
+            model: key.to_string(),
+            repo: repo_of(model_id, key),
+            reason: format!("explicit model={}", py_repr(key)),
+            detection: None,
+            workflow: None,
+        };
+    }
+    let workflow = match_workflow(questions).map(str::to_string);
+    let det = lang_detect::analyse(state);
+    let (key, reason) = if det.script == "unknown" {
+        (
+            DEFAULT_CHECKPOINT,
+            format!("no letters detected in state; using default ({DEFAULT_CHECKPOINT})"),
+        )
+    } else if det.script != "latin" {
+        (
+            MULTILINGUAL,
+            format!(
+                "non-Latin script ({}, {:.0}% of letters); the English checkpoint cannot read it",
+                det.script,
+                100.0 * det.non_latin_fraction
+            ),
+        )
+    } else if !det.is_english {
+        let reason = if let Some(segment) = &det.mixed_segment {
+            format!(
+                "Latin script, mostly English, but a line or field reads as {} ({}); \
+                 the English checkpoint cannot read it",
+                det.language.map(py_repr).unwrap_or_else(|| "None".into()),
+                py_repr(&segment.chars().take(60).collect::<String>()),
+            )
+        } else if let Some(lang) = det.language {
+            format!(
+                "Latin script but language looks like {}, not English",
+                py_repr(lang)
+            )
+        } else {
+            format!(
+                "Latin script, language not identified but {:.0}% non-English letters; \
+                 not safe for the English checkpoint",
+                100.0 * det.diacritic_rate
+            )
+        };
+        (MULTILINGUAL, reason)
+    } else if det.language_undecided {
+        (
+            DEFAULT_CHECKPOINT,
+            format!(
+                "Latin script, language not identified and no non-English letters; \
+                 using default ({DEFAULT_CHECKPOINT})"
+            ),
+        )
+    } else {
+        (ENGLISH, "English Latin text".to_string())
+    };
+    DecisionRouting {
+        model: key.to_string(),
+        repo: repo_of(model_id, key),
+        reason,
+        detection: Some(analysis_json(&det)),
+        workflow,
+    }
+}
+
+/// The language per-language calibration keys on: the detected one.
+pub fn detected_language(routing: &DecisionRouting) -> Option<&str> {
+    routing
+        .detection
+        .as_ref()
+        .and_then(|d| d.get("language"))
+        .and_then(Json::as_str)
+}
+
+/// `<repo>` for the root checkpoint, `<repo>/<subfolder>` otherwise.
+pub fn repo_of(model_id: &str, checkpoint: &str) -> String {
+    if checkpoint == ENGLISH {
+        model_id.to_string()
+    } else {
+        format!("{model_id}/{checkpoint}")
+    }
+}
+
+/// The analysis as an order-preserving [`Json`], so the response keeps
+/// the reference's field order.
+fn analysis_json(det: &lang_detect::Analysis) -> Json {
+    let text = serde_json::to_string(det).unwrap_or_else(|_| "null".into());
+    Json::parse(text.as_bytes()).unwrap_or(Json::Null)
+}
+
+fn py_repr(s: &str) -> String {
+    Json::String(s.to_string()).py_repr()
+}
+
+/// Python's `str.strip()`: Unicode whitespace, which includes the
+/// U+001C–U+001F separators Rust's `trim` leaves alone.
+fn py_strip(s: &str) -> &str {
+    s.trim_matches(|c: char| c.is_whitespace() || ('\u{1c}'..='\u{1f}').contains(&c))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn pins_follow_the_reference_server() {
+        assert_eq!(resolve_pin(Some("multilingual")), Some(MULTILINGUAL));
+        assert_eq!(resolve_pin(Some(" ML ")), Some(MULTILINGUAL));
+        assert_eq!(
+            resolve_pin(Some("convaiinnovations/laya-typed-decisions")),
+            Some(TYPED_DECISIONS)
+        );
+        assert_eq!(resolve_pin(Some("laya")), Some(ENGLISH));
+        // A Jev service name, the family's own id, and nothing: no pin.
+        assert_eq!(resolve_pin(Some("jev-1")), None);
+        assert_eq!(resolve_pin(Some("convaiinnovations/laya")), None);
+        assert_eq!(resolve_pin(Some("")), None);
+        assert_eq!(resolve_pin(None), None);
+    }
+
+    #[test]
+    fn workflow_needs_the_exact_id_set() {
+        let qs = |ids: &[&str]| -> Vec<(String, Json)> {
+            ids.iter().map(|i| (i.to_string(), Json::Null)).collect()
+        };
+        assert_eq!(
+            match_workflow(&qs(&[
+                "urgency",
+                "action",
+                "category",
+                "churn_risk",
+                "needs_human"
+            ])),
+            Some("customer_service")
+        );
+        assert_eq!(match_workflow(&qs(&["urgency"])), None);
+    }
+
+    /// Every case in the reference fixtures routes to the checkpoint, with
+    /// the reason and detection, the reference reported.
+    #[test]
+    fn routes_every_reference_case_like_the_reference() {
+        let text = std::fs::read_to_string(
+            std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
+                .join("src/harness/testdata/laya/reference.json"),
+        )
+        .expect("read reference.json");
+        let reference = Json::parse(text.as_bytes()).expect("parse reference.json");
+        let Some(Json::Array(cases)) = reference.get("cases") else {
+            panic!("no cases");
+        };
+        for case in cases {
+            let name = case.get("name").and_then(Json::as_str).unwrap();
+            let request = case.get("request").unwrap();
+            let body = serde_json::to_vec(request).unwrap();
+            let req = cortex_core::decisions::SystemOneRequest::parse(
+                &body,
+                cortex_core::decisions::limits::DEFAULT_MAX_TOKEN_BUDGET,
+            )
+            .unwrap();
+            let routing = route(
+                "convaiinnovations/laya",
+                &req.state,
+                &req.questions,
+                resolve_pin(req.model.as_deref()),
+            );
+            let want = case.get("response").unwrap().get("routing").unwrap();
+            assert_eq!(
+                serde_json::to_string(&routing).unwrap(),
+                serde_json::to_string(want).unwrap(),
+                "{name}"
+            );
+        }
+    }
+}
@@ -0,0 +1,316 @@
+//! Sequence assembly: one token sequence per question, under the model's
+//! token budgets.
+//!
+//! ```text
+//! [CLS] "<type> question: <instructions>" [SEP] [MASK] opt0 [MASK] opt1 … [SEP] state [SEP]
+//! ```
+//!
+//! The budget rules, in the order they apply:
+//!
+//! 1. Each option is `[MASK]` plus at most 48 tokens of its text.
+//! 2. `head_max_len` bounds the instruction text and all options together.
+//!    If fewer than 16 tokens would remain for the instructions, every
+//!    option is cut to the same `max(4, (head_max_len - 16) / options)`
+//!    tokens (marker included).
+//! 3. The instructions get what is left of `head_max_len`, but never fewer
+//!    than 8 tokens.
+//! 4. The state fills the rest of `max_len`. A conversation (a JSON list)
+//!    keeps its newest turns, so it is cut from the left; anything else is
+//!    cut from the right.
+//! 5. The whole sequence is cut to `max_len`, dropping any marker that no
+//!    longer fits — a question whose markers do not all fit is refused.
+
+use super::question::{QType, Question, questions_from_request};
+use super::tokenizer::TextEncoder;
+use super::{DecisionConfig, DecisionRequestError, unprocessable};
+use cortex_core::decisions::Json;
+
+/// Option text tokens kept before the head budget is considered.
+pub const MAX_OPTION_TOKENS: usize = 48;
+/// Head tokens the instructions are guaranteed before options are shrunk.
+const MIN_INSTRUCTION_ROOM: usize = 16;
+/// Per-option floor (marker included) when options are shrunk.
+const MIN_TOKENS_PER_OPTION: usize = 4;
+/// Instruction tokens kept however many the options take.
+const MIN_INSTRUCTION_TOKENS: usize = 8;
+
+/// Per-request overrides of the checkpoint's token budgets.
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
+pub struct Budget {
+    pub max_len: Option<usize>,
+    pub head_max_len: Option<usize>,
+}
+
+/// One model input row: one question over the state.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct Row {
+    pub input_ids: Vec<u32>,
+    /// Position of each option's `[MASK]` marker, in option order.
+    pub markers: Vec<usize>,
+    pub qtype: QType,
+}
+
+/// What the head budget did to a question's options.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct OptionStats {
+    pub options: usize,
+    /// Options whose token span is still unique after the cut. Fewer than
+    /// `options` means the model can no longer tell some of them apart.
+    pub options_distinct: usize,
+    /// The per-option cap applied, when the options had to be shrunk.
+    pub tokens_per_option: Option<usize>,
+}
+
+/// A request turned into model input: one [`Row`] per question, in request
+/// order. With no questions there are no rows and nothing to run.
+#[derive(Debug, Clone, PartialEq)]
+pub struct Prepared {
+    pub questions: Vec<Question>,
+    pub rows: Vec<Row>,
+    pub stats: Vec<OptionStats>,
+    /// Every token the forward pass encodes: the metered input.
+    pub input_tokens: usize,
+    /// The id batched rows are padded with.
+    pub pad_id: u32,
+}
+
+/// The text a state is tokenized from: a string as-is, anything else as
+/// `json.dumps(state, ensure_ascii=False)`.
+pub fn serialize_state(state: &Json) -> String {
+    match state {
+        Json::String(s) => s.clone(),
+        other => other.dumps(),
+    }
+}
+
+/// Validate a request's questions and build their model input.
+///
+/// `questions` are the request's definitions in request order; `cfg` and
+/// `tok` belong to the checkpoint the request was routed to.
+pub fn prepare(
+    state: &Json,
+    questions: &[(String, Json)],
+    cfg: &DecisionConfig,
+    tok: &dyn TextEncoder,
+    budget: Budget,
+) -> Result<Prepared, DecisionRequestError> {
+    let special = tok.special();
+    if questions.is_empty() {
+        return Ok(Prepared {
+            questions: Vec::new(),
+            rows: Vec::new(),
+            stats: Vec::new(),
+            input_tokens: 0,
+            pad_id: special.pad_id,
+        });
+    }
+    let questions = questions_from_request(questions)?;
+    let max_len = budget.max_len.unwrap_or(cfg.max_len);
+    let head_max_len = budget.head_max_len.unwrap_or(cfg.head_max_len);
+    let encode = |text: &str| {
+        tok.encode(text)
+            .map_err(|e| DecisionRequestError::new(500, format!("tokenizer failed: {e}")))
+    };
+    // Tokenized once and sliced per question.
+    let state_ids = encode(&serialize_state(state).replace(&special.mask_token, " "))?;
+    let truncate_left = matches!(state, Json::Array(_));
+    let frame = Frame {
+        cls: special.cls_id,
+        sep: special.sep_id,
+        max_len,
+        head_max_len,
+    };
+
+    let mut rows = Vec::with_capacity(questions.len());
+    let mut stats = Vec::with_capacity(questions.len());
+    for question in &questions {
+        let instructions = question.instructions.replace(&special.mask_token, " ");
+        let head_ids = encode(&format!(
+            "{} question: {instructions}",
+            question.kind.name()
+        ))?;
+        let mut option_ids = Vec::with_capacity(question.options.len());
+        for option in &question.options {
+            let mut ids = vec![special.mask_id];
+            let text = encode(&format!(" {}", option.replace(&special.mask_token, " ")))?;
+            ids.extend(text.into_iter().take(MAX_OPTION_TOKENS));
+            option_ids.push(ids);
+        }
+        let (row, s) = frame.assemble(head_ids, option_ids, &state_ids, truncate_left);
+        if row.markers.len() != question.options.len() {
+            return Err(unprocessable(format!(
+                "question {} options exceed head_max_len={head_max_len}",
+                Json::String(question.id.clone()).py_repr()
+            )));
+        }
+        rows.push(Row {
+            qtype: question.kind,
+            ..row
+        });
+        stats.push(s);
+    }
+    let input_tokens = rows.iter().map(|r| r.input_ids.len()).sum();
+    Ok(Prepared {
+        questions,
+        rows,
+        stats,
+        input_tokens,
+        pad_id: special.pad_id,
+    })
+}
+
+/// The fixed parts of every sequence for one request.
+struct Frame {
+    cls: u32,
+    sep: u32,
+    max_len: usize,
+    head_max_len: usize,
+}
+
+impl Frame {
+    /// The pure budget arithmetic of one sequence, over already-tokenized parts.
+    fn assemble(
+        &self,
+        mut head_ids: Vec<u32>,
+        mut option_ids: Vec<Vec<u32>>,
+        state_ids: &[u32],
+        truncate_left: bool,
+    ) -> (Row, OptionStats) {
+        let Frame {
+            cls,
+            sep,
+            max_len,
+            head_max_len,
+        } = *self;
+        let used = |o: &[Vec<u32>]| o.iter().map(Vec::len).sum::<usize>();
+        // Signed: options can overrun the head budget.
+        let mut opt_budget = head_max_len as isize - used(&option_ids) as isize;
+        let mut tokens_per_option = None;
+        if opt_budget < MIN_INSTRUCTION_ROOM as isize {
+            let per = MIN_TOKENS_PER_OPTION
+                .max(head_max_len.saturating_sub(MIN_INSTRUCTION_ROOM) / option_ids.len().max(1));
+            tokens_per_option = Some(per);
+            for o in &mut option_ids {
+                o.truncate(per);
+            }
+            opt_budget = head_max_len as isize - used(&option_ids) as isize;
+        }
+        head_ids.truncate(opt_budget.max(MIN_INSTRUCTION_TOKENS as isize) as usize);
+
+        let mut ids = Vec::with_capacity(max_len);
+        ids.push(cls);
+        ids.extend(head_ids);
+        ids.push(sep);
+        let mut markers = Vec::with_capacity(option_ids.len());
+        for o in &option_ids {
+            markers.push(ids.len());
+            ids.extend_from_slice(o);
+        }
+        ids.push(sep);
+        let room = max_len.saturating_sub(ids.len() + 1);
+        let state = if truncate_left {
+            &state_ids[state_ids.len().saturating_sub(room)..]
+        } else {
+            &state_ids[..room.min(state_ids.len())]
+        };
+        ids.extend_from_slice(state);
+        ids.push(sep);
+        ids.truncate(max_len);
+        markers.retain(|&m| m < max_len);
+
+        let mut distinct = option_ids.clone();
+        distinct.sort();
+        distinct.dedup();
+        (
+            Row {
+                input_ids: ids,
+                markers,
+                qtype: QType::Choice,
+            },
+            OptionStats {
+                options: option_ids.len(),
+                options_distinct: distinct.len(),
+                tokens_per_option,
+            },
+        )
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const CLS: u32 = 1;
+    const SEP: u32 = 2;
+    const M: u32 = 9;
+
+    fn frame(max_len: usize, head_max_len: usize) -> Frame {
+        Frame {
+            cls: CLS,
+            sep: SEP,
+            max_len,
+            head_max_len,
+        }
+    }
+
+    fn opts(lens: &[usize]) -> Vec<Vec<u32>> {
+        lens.iter()
+            .enumerate()
+            .map(|(i, &n)| {
+                let mut o = vec![M];
+                o.extend((0..n - 1).map(|j| 100 + (i * 50 + j) as u32));
+                o
+            })
+            .collect()
+    }
+
+    #[test]
+    fn roomy_budget_keeps_everything() {
+        let (row, s) = frame(64, 32).assemble(vec![10, 11], opts(&[2, 3]), &[50, 51], false);
+        assert_eq!(
+            row.input_ids,
+            [CLS, 10, 11, SEP, M, 100, M, 150, 151, SEP, 50, 51, SEP]
+        );
+        assert_eq!(row.markers, [4, 6]);
+        assert_eq!(
+            (s.options, s.options_distinct, s.tokens_per_option),
+            (2, 2, None)
+        );
+    }
+
+    #[test]
+    fn crowded_head_shrinks_options_evenly_and_keeps_eight_instruction_tokens() {
+        // 3 options of 10 = 30 tokens against a 40-token head: 10 left < 16.
+        let head: Vec<u32> = (10..40).collect();
+        let (row, s) = frame(128, 40).assemble(head, opts(&[10, 10, 10]), &[], false);
+        // per = max(4, (40 - 16) / 3) = 8; 40 - 24 = 16 instruction tokens.
+        assert_eq!(s.tokens_per_option, Some(8));
+        assert_eq!(row.markers, [18, 26, 34]);
+        let (row, s) = frame(256, 40).assemble((10..40).collect(), opts(&[10; 12]), &[], false);
+        // per = max(4, 24 / 12 = 2) = 4; the budget goes negative, 8 tokens survive.
+        assert_eq!(s.tokens_per_option, Some(4));
+        assert_eq!(row.markers[0], 1 + 8 + 1);
+    }
+
+    #[test]
+    fn state_is_cut_from_the_side_the_state_shape_asks_for() {
+        let state: Vec<u32> = (50..60).collect();
+        let (right, _) = frame(8, 32).assemble(vec![], opts(&[1]), &state, false);
+        let (left, _) = frame(8, 32).assemble(vec![], opts(&[1]), &state, true);
+        // [CLS] [SEP] [MASK] [SEP] = 4 tokens, 1 reserved for the last [SEP]: room 3.
+        assert_eq!(right.input_ids, [CLS, SEP, M, SEP, 50, 51, 52, SEP]);
+        assert_eq!(left.input_ids, [CLS, SEP, M, SEP, 57, 58, 59, SEP]);
+        // No room at all: none of the state, not all of it.
+        let (none, _) = frame(5, 32).assemble(vec![], opts(&[1]), &state, true);
+        assert_eq!(none.input_ids, [CLS, SEP, M, SEP, SEP]);
+    }
+
+    #[test]
+    fn markers_past_max_len_are_dropped_and_collisions_counted() {
+        let same = vec![vec![M, 7, 7], vec![M, 7, 7]];
+        let (row, s) = frame(5, 64).assemble(vec![], same, &[], false);
+        // [CLS] [SEP] then markers at 2 and 5; only the first fits in 5.
+        assert_eq!(row.markers, [2]);
+        assert_eq!((s.options, s.options_distinct), (2, 1));
+    }
+}
@@ -0,0 +1,318 @@
+//! Per-option logits → calibrated, typed answers.
+//!
+//! The reference computes this in numpy on the float32 logits the model
+//! returns, and publishes each probability rounded to 4 decimal places.
+//! Reproducing the rounded values exactly means reproducing the float32
+//! arithmetic, in numpy's order: division by a float32 temperature, a
+//! max-subtracted `exp`, and numpy's *pairwise* summation (not a running
+//! sum) for every reduction. The expected score is the exception: numpy
+//! promotes `arange(k) * p` to float64.
+
+use super::question::QType;
+use super::sequence::Prepared;
+use super::{DECISION_MODEL_NAME, DecisionConfig};
+use cortex_core::decisions::{
+    CollapsedOptions, DecisionAction, DecisionAnswer, DecisionUsage, OrderedMap, SystemOneResponse,
+};
+
+/// The model's output for one [`super::Row`].
+#[derive(Debug, Clone, PartialEq)]
+pub struct RowOutput {
+    /// The scorer's raw logit at each option marker, before calibration.
+    /// Entries past the row's marker count (batch padding) are ignored.
+    pub logits: Vec<f32>,
+    /// The act head's probability of acting: `softmax(act_logits)[0]`.
+    pub act_probability: f32,
+}
+
+/// Build the answers and usage for a prepared request from the model's
+/// per-row outputs (`outputs[i]` belongs to `prepared.rows[i]`).
+///
+/// `lang` selects per-language calibration where the checkpoint has it:
+/// the language the request was declared or detected as.
+pub fn shape(
+    prepared: &Prepared,
+    outputs: &[RowOutput],
+    cfg: &DecisionConfig,
+    lang: Option<&str>,
+) -> (OrderedMap<DecisionAnswer>, DecisionUsage) {
+    assert_eq!(
+        outputs.len(),
+        prepared.rows.len(),
+        "one model output per prepared row"
+    );
+    let temps = cfg.temperatures_for(lang);
+    let mut answers = Vec::with_capacity(outputs.len());
+    let mut collapsed = Vec::new();
+    for ((question, row), (out, stats)) in prepared
+        .questions
+        .iter()
+        .zip(&prepared.rows)
+        .zip(outputs.iter().zip(&prepared.stats))
+    {
+        let k = row.markers.len();
+        let p = calibrate(
+            &out.logits[..k],
+            temps.for_question(question.kind as usize, k),
+        );
+        let answer_confidence = round4(
+            p.iter()
+                .copied()
+                .fold(f32::NEG_INFINITY, f32::max)
+                .clamp(0.0, 1.0) as f64,
+        );
+        let action = DecisionAction {
+            act_probability: round4(out.act_probability as f64),
+        };
+        let probabilities = |key: &dyn Fn(usize) -> String| {
+            OrderedMap(
+                p.iter()
+                    .enumerate()
+                    .map(|(i, &v)| (key(i), round4(v as f64)))
+                    .collect(),
+            )
+        };
+        let answer = match question.kind {
+            QType::Choice => DecisionAnswer::Choice {
+                choice: question.labels[argmax(&p)].clone(),
+                probabilities: probabilities(&|i| {
+                    question.labels[i]
+                        .json_key()
+                        .expect("choice labels are scalars")
+                }),
+                confidence: round4(entropy_confidence(&p) as f64),
+                answer_confidence,
+                action,
+            },
+            QType::Score => {
+                // int64 * float32 promotes to float64 in numpy.
+                let expected: Vec<f64> = p
+                    .iter()
+                    .enumerate()
+                    .map(|(i, &v)| i as f64 * v as f64)
+                    .collect();
+                DecisionAnswer::Score {
+                    score: round4(pairwise_sum(&expected)),
+                    legend: OrderedMap(
+                        question
+                            .labels
+                            .iter()
+                            .enumerate()
+                            .map(|(i, c)| (i.to_string(), c.clone()))
+                            .collect(),
+                    ),
+                    probabilities: probabilities(&|i| i.to_string()),
+                    confidence: round4(entropy_confidence(&p) as f64),
+                    answer_confidence,
+                    action,
+                }
+            }
+            QType::Noul => {
+                let yes = p[1] as f64;
+                DecisionAnswer::Noul {
+                    noul: round4(yes),
+                    confidence: round4(yes.max(1.0 - yes)),
+                    answer_confidence,
+                    action,
+                }
+            }
+        };
+        answers.push((question.id.clone(), answer));
+        if stats.options_distinct < stats.options {
+            collapsed.push((
+                question.id.clone(),
+                CollapsedOptions {
+                    total: stats.options,
+                    distinct: stats.options_distinct,
+                    tokens_per_option: stats.tokens_per_option,
+                },
+            ));
+        }
+    }
+    let usage = DecisionUsage {
+        input_tokens: prepared.input_tokens,
+        output_tokens: 0,
+        options: (!collapsed.is_empty()).then_some(OrderedMap(collapsed)),
+    };
+    (OrderedMap(answers), usage)
+}
+
+/// Assemble the full response once routing is known.
+pub fn response(
+    answers: OrderedMap<DecisionAnswer>,
+    usage: DecisionUsage,
+    routing: Option<cortex_core::decisions::DecisionRouting>,
+) -> SystemOneResponse {
+    SystemOneResponse {
+        model: DECISION_MODEL_NAME.to_string(),
+        answers,
+        usage,
+        routing,
+    }
+}
+
+/// Temperature-scaled softmax, as the reference evaluates it in numpy:
+/// `z = logits / t`, `e = exp(z - z.max())`, `p = e / e.sum()`.
+///
+/// numpy's SIMD `expf` and libm's disagree by one ulp on about a third of
+/// inputs. That is invisible after the 4-decimal rounding on every
+/// reference fixture, and far below the engine's own distance from the
+/// reference's logits, so it is not reproduced; the steps either side of
+/// the `exp` are.
+fn calibrate(logits: &[f32], t: f64) -> Vec<f32> {
+    let z = scale(logits, t);
+    let max = z.iter().copied().fold(f32::NEG_INFINITY, f32::max);
+    let e: Vec<f32> = z.iter().map(|&v| (v - max).exp()).collect();
+    normalize(&e)
+}
+
+/// `logits / t` in float32: numpy casts the Python-float temperature to the
+/// array's dtype rather than promoting the array.
+fn scale(logits: &[f32], t: f64) -> Vec<f32> {
+    let t = t as f32;
+    logits.iter().map(|&l| l / t).collect()
+}
+
+/// `e / e.sum()` with numpy's pairwise float32 sum.
+fn normalize(e: &[f32]) -> Vec<f32> {
+    let sum = pairwise_sum(e);
+    e.iter().map(|&v| v / sum).collect()
+}
+
+/// `1 - H(p) / log(k)`, clipped to [0, 1]: how concentrated the whole
+/// distribution is. Not calibrated; see `answer_confidence` for that.
+fn entropy_confidence(p: &[f32]) -> f32 {
+    let k = p.len();
+    if k < 2 {
+        return 1.0;
+    }
+    let terms: Vec<f32> = p.iter().map(|&v| v * v.clamp(1e-12, 1.0).ln()).collect();
+    let ent = -pairwise_sum(&terms);
+    (1.0 - ent / (k as f64).ln() as f32).clamp(0.0, 1.0)
+}
+
+/// The first index of the largest value, as `argmax` picks it.
+fn argmax(p: &[f32]) -> usize {
+    let mut best = 0;
+    for (i, &v) in p.iter().enumerate() {
+        if v > p[best] {
+            best = i;
+        }
+    }
+    best
+}
+
+/// Python's `round(x, 4)`: correctly rounded, ties to even on the exact
+/// binary value. Rust's fixed-precision formatting rounds the same way.
+fn round4(x: f64) -> f64 {
+    format!("{x:.4}").parse().expect("formatted float parses")
+}
+
+/// numpy's pairwise summation over a contiguous array, which the rounded
+/// outputs depend on: a plain loop below 8 elements, eight interleaved
+/// accumulators up to 128, and recursive halving (at a multiple of 8)
+/// beyond.
+fn pairwise_sum<T>(a: &[T]) -> T
+where
+    T: Copy + std::ops::Add<Output = T> + Default,
+{
+    const BLOCK: usize = 128;
+    let n = a.len();
+    if n < 8 {
+        let mut res = T::default();
+        for &v in a {
+            res = res + v;
+        }
+        res
+    } else if n <= BLOCK {
+        let mut r = [a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7]];
+        let mut i = 8;
+        while i < n - n % 8 {
+            for j in 0..8 {
+                r[j] = r[j] + a[i + j];
+            }
+            i += 8;
+        }
+        let mut res = ((r[0] + r[1]) + (r[2] + r[3])) + ((r[4] + r[5]) + (r[6] + r[7]));
+        while i < n {
+            res = res + a[i];
+            i += 1;
+        }
+        res
+    } else {
+        let mut n2 = n / 2;
+        n2 -= n2 % 8;
+        pairwise_sum(&a[..n2]) + pairwise_sum(&a[n2..])
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn pairwise_sum_differs_from_a_running_sum_where_numpy_does() {
+        // 1e8 swamps every 1.0 added to it one at a time in f32; pairwise
+        // grouping keeps them. `int(np.float32 array.sum())` of this array is 100000008.
+        let mut a = vec![1.0f32; 16];
+        a[0] = 1e8;
+        let running = a.iter().fold(0.0f32, |s, &v| s + v);
+        assert_eq!(running, 1e8);
+        assert_eq!(pairwise_sum(&a), 100_000_008.0);
+    }
+
+    /// Rows of `testdata/laya/numpy_softmax.txt`: the reference's
+    /// calibration run in numpy on random float32 logits, with the bits of
+    /// every intermediate. Bit equality pins the evaluation order — a
+    /// running sum instead of numpy's pairwise one, or dividing by the
+    /// temperature in float64, changes low bits that the 4-decimal rounding
+    /// usually hides.
+    fn numpy_rows() -> Vec<(f64, [Vec<f32>; 4])> {
+        include_str!("../testdata/laya/numpy_softmax.txt")
+            .lines()
+            .filter(|l| !l.starts_with('#') && !l.trim().is_empty())
+            .map(|l| {
+                let mut f = l.split('|');
+                let t: f64 = f.next().unwrap().parse().unwrap();
+                let bits = |s: &str| -> Vec<f32> {
+                    s.split(',')
+                        .map(|h| f32::from_bits(u32::from_str_radix(&h[2..], 16).unwrap()))
+                        .collect()
+                };
+                let v: Vec<Vec<f32>> = f.map(bits).collect();
+                (t, [v[0].clone(), v[1].clone(), v[2].clone(), v[3].clone()])
+            })
+            .collect()
+    }
+
+    fn bits(v: &[f32]) -> Vec<u32> {
+        v.iter().map(|x| x.to_bits()).collect()
+    }
+
+    #[test]
+    fn calibration_matches_numpy_either_side_of_exp() {
+        let rows = numpy_rows();
+        assert_eq!(rows.len(), 4);
+        for (t, [logits, z, e, p]) in rows {
+            assert_eq!(
+                bits(&scale(&logits, t)),
+                bits(&z),
+                "k={} t={t}: z",
+                logits.len()
+            );
+            assert_eq!(
+                bits(&normalize(&e)),
+                bits(&p),
+                "k={} t={t}: p",
+                logits.len()
+            );
+        }
+    }
+
+    #[test]
+    fn round4_ties_to_even_on_the_binary_value() {
+        assert_eq!(round4(0.03125), 0.0312);
+        assert_eq!(round4(1.00005), 1.0001); // 1.00005 is just above the tie in binary
+        assert_eq!(round4(0.99995), 1.0);
+    }
+}
@@ -0,0 +1,81 @@
+//! The tokenizer a decision checkpoint ships, and its special tokens.
+
+use std::path::Path;
+
+/// The special tokens the sequence layout uses, resolved to ids.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct SpecialTokens {
+    /// The option marker's text (`[MASK]`, `<mask>` for mmBERT). Literal
+    /// occurrences in caller text are replaced with a space before
+    /// tokenizing, so no caller can forge a marker.
+    pub mask_token: String,
+    pub cls_id: u32,
+    pub sep_id: u32,
+    pub mask_id: u32,
+    /// Padding for batched rows; not part of any sequence.
+    pub pad_id: u32,
+}
+
+/// Text → token ids, with no special tokens added.
+pub trait TextEncoder {
+    fn encode(&self, text: &str) -> Result<Vec<u32>, String>;
+    fn special(&self) -> &SpecialTokens;
+}
+
+/// A checkpoint's `tokenizer/` directory: `tokenizer.json` plus the special
+/// token names in `tokenizer_config.json`.
+pub struct DecisionTokenizer {
+    tok: tokenizers::Tokenizer,
+    special: SpecialTokens,
+}
+
+impl DecisionTokenizer {
+    pub fn from_dir(dir: &Path) -> anyhow::Result<Self> {
+        let mut tok = tokenizers::Tokenizer::from_file(dir.join("tokenizer.json"))
+            .map_err(|e| anyhow::anyhow!("{}: {e}", dir.join("tokenizer.json").display()))?;
+        // The Python tokenizer call the reference makes pads and truncates
+        // nothing unless asked; a tokenizer.json that carries either setting
+        // would otherwise apply it here.
+        tok.with_truncation(None)
+            .map_err(|e| anyhow::anyhow!("disabling truncation: {e}"))?;
+        tok.with_padding(None);
+        let cfg_path = dir.join("tokenizer_config.json");
+        let cfg: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(&cfg_path)?)
+            .map_err(|e| anyhow::anyhow!("{}: {e}", cfg_path.display()))?;
+        let name = |key: &str| -> anyhow::Result<String> {
+            cfg.get(key)
+                .and_then(|v| {
+                    v.as_str()
+                        .or_else(|| v.get("content").and_then(|c| c.as_str()))
+                })
+                .map(str::to_owned)
+                .ok_or_else(|| anyhow::anyhow!("{}: no {key}", cfg_path.display()))
+        };
+        let id = |token: &str| -> anyhow::Result<u32> {
+            tok.token_to_id(token)
+                .ok_or_else(|| anyhow::anyhow!("special token {token:?} is not in the vocabulary"))
+        };
+        let mask_token = name("mask_token")?;
+        let special = SpecialTokens {
+            cls_id: id(&name("cls_token")?)?,
+            sep_id: id(&name("sep_token")?)?,
+            mask_id: id(&mask_token)?,
+            pad_id: id(&name("pad_token")?)?,
+            mask_token,
+        };
+        Ok(Self { tok, special })
+    }
+}
+
+impl TextEncoder for DecisionTokenizer {
+    fn encode(&self, text: &str) -> Result<Vec<u32>, String> {
+        self.tok
+            .encode(text, false)
+            .map(|e| e.get_ids().to_vec())
+            .map_err(|e| e.to_string())
+    }
+
+    fn special(&self) -> &SpecialTokens {
+        &self.special
+    }
+}
@@ -19,14 +19,16 @@ use std::sync::atomic::AtomicBool;

 use anyhow::{Context, Result};
 use candle_core::DType;
-use tokenizers::Tokenizer;

 use super::arch::laya::{CheckpointConfig, DecisionBatch, DecisionRow, LayaModel};
+use super::decision::{DecisionConfig, DecisionTokenizer};
 use super::device_worker::{DecisionHandle, DeviceWorkerHandle};
 use cortex_core::harness::ModelSpec;

-/// Name of the checkpoint at the repo root.
-pub const ROOT_CHECKPOINT: &str = "root";
+/// Routing name of the checkpoint at the repo root. A released Laya
+/// family's root is its English checkpoint, and upstream routes to it
+/// under that name.
+pub const ROOT_CHECKPOINT: &str = super::decision::route::ENGLISH;

 /// Where one checkpoint's weights live and how to reach them.
 enum Backend {
@@ -50,10 +52,14 @@ pub struct CheckpointFiles {
 }

 pub struct DecisionCheckpoint {
-    /// [`ROOT_CHECKPOINT`] or the subfolder name.
+    /// The checkpoint's routing name: [`ROOT_CHECKPOINT`] for the repo
+    /// root, else its subfolder (`multilingual`, `typed-decisions`).
    pub name: String,
    pub config: CheckpointConfig,
-    pub tokenizer: Tokenizer,
+    /// Budgets and fitted temperatures from `rl_agent_config.json`, as
+    /// the prompt builder and calibration read them.
+    pub calibration: DecisionConfig,
+    pub tokenizer: DecisionTokenizer,
    /// Precision the weights were loaded in.
    pub dtype: DType,
    backend: Backend,
@@ -69,8 +75,17 @@ impl DecisionCheckpoint {
        worker: Option<Arc<DeviceWorkerHandle>>,
    ) -> Result<Self> {
        let config = CheckpointConfig::from_dir(&files.dir)?;
-        let tokenizer = Tokenizer::from_file(&files.tokenizer)
-            .map_err(|e| anyhow::anyhow!("load tokenizer {}: {e}", files.tokenizer.display()))?;
+        let calibration_path = files.dir.join("rl_agent_config.json");
+        let calibration = DecisionConfig::from_json(
+            &std::fs::read_to_string(&calibration_path)
+                .with_context(|| format!("read {}", calibration_path.display()))?,
+        )
+        .map_err(|e| anyhow::anyhow!("{e}"))?;
+        let tokenizer_dir = files
+            .tokenizer
+            .parent()
+            .context("tokenizer.json has no parent directory")?;
+        let tokenizer = DecisionTokenizer::from_dir(tokenizer_dir)?;
        let (dtype, backend) = match worker {
            Some(worker) => {
                let wanted = device_dtype(config.laya.amp_dtype.as_deref());
@@ -93,6 +108,7 @@ impl DecisionCheckpoint {
        Ok(Self {
            name,
            config,
+            calibration,
            tokenizer,
            dtype,
            backend,
@@ -19,7 +19,7 @@ use crate::harness::candle::ModelArch;
 #[cfg(feature = "cuda")]
 use crate::harness::device_worker::jobs::TpHandle;
 use crate::harness::device_worker::jobs::{
-    ArchHandle, DecisionHandle, DenseLoad, ImageHandle, ImageInput, Job, KvSnapshotId,
+    ArchHandle, AudioHandle, DecisionHandle, DenseLoad, ImageHandle, ImageInput, Job, KvSnapshotId,
 };
 use crate::harness::image::ZImagePipeline;
 #[cfg(feature = "cuda")]
@@ -71,6 +71,11 @@ struct DeviceWorkerState {
    decision_models: HashMap<DecisionHandle, Box<LayaModel>>,
    /// Counter for minting fresh `DecisionHandle`s.
    next_decision_handle: u64,
+    /// ACE-Step music pipeline slab (#383). Same lifecycle discipline as
+    /// `image_models`, in its own namespace.
+    audio_models: HashMap<AudioHandle, Box<crate::harness::audio::AceStepEngine>>,
+    /// Counter for minting fresh `AudioHandle`s.
+    next_audio_handle: u64,
    /// Leader's NCCL state. Populated by `Job::NcclInit`; the
    /// underlying `Comm`'s libnccl handle lives bound to this thread
    /// for its entire lifetime. Subprocess workers maintain their own
@@ -813,6 +818,57 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
                };
                let _ = reply.send(result);
            }
+            Job::LoadAudio {
+                files,
+                model_id,
+                reply,
+            } => {
+                let result = crate::harness::audio::AceStepEngine::load(&files, &state.device)
+                    .with_context(|| format!("load ace-step pipeline for {model_id}"))
+                    .map(|engine| {
+                        let handle = AudioHandle(state.next_audio_handle);
+                        state.next_audio_handle = state.next_audio_handle.wrapping_add(1);
+                        tracing::info!(
+                            device_index,
+                            model_id,
+                            planner = engine.has_planner(),
+                            handle = handle.0,
+                            "device worker: audio pipeline loaded"
+                        );
+                        state.audio_models.insert(handle, Box::new(engine));
+                        handle
+                    });
+                trim_device_pool(&state);
+                let _ = reply.send(result);
+            }
+            Job::DropAudio { handle, reply } => {
+                let removed = state.audio_models.remove(&handle);
+                let was_present = removed.is_some();
+                // Explicit drop on this thread, as for DropImage.
+                drop(removed);
+                trim_device_pool(&state);
+                tracing::debug!(
+                    device_index,
+                    handle = handle.0,
+                    was_present,
+                    "device worker: audio pipeline dropped"
+                );
+                let _ = reply.send(());
+            }
+            Job::GenerateAudio {
+                handle,
+                params,
+                reply,
+            } => {
+                let result = match state.audio_models.get_mut(&handle) {
+                    Some(engine) => engine.generate(&params),
+                    None => Err(anyhow::anyhow!("no audio pipeline for handle {}", handle.0)),
+                };
+                // Song lengths vary per request, so release this render's
+                // pooled transients before the next one, as images do.
+                trim_device_pool(&state);
+                let _ = reply.send(result);
+            }
            Job::Shutdown => unreachable!("Shutdown should break above"),
        }
    }
@@ -887,6 +943,8 @@ fn init_state(device_index: u32) -> DeviceWorkerState {
            next_image_handle: 1,
            decision_models: HashMap::new(),
            next_decision_handle: 1,
+            audio_models: HashMap::new(),
+            next_audio_handle: 1,
            nccl: NcclState::new(),
            tp_models: HashMap::new(),
            next_tp_handle: 1,
@@ -907,6 +965,8 @@ fn init_state(device_index: u32) -> DeviceWorkerState {
            next_image_handle: 1,
            decision_models: HashMap::new(),
            next_decision_handle: 1,
+            audio_models: HashMap::new(),
+            next_audio_handle: 1,
            nccl: NcclState::new(),
        }
    }
@@ -2044,6 +2104,15 @@ fn drain_poisoned(job: Job, device_index: u32) {
        Job::Decide { reply, .. } => {
            let _ = reply.send(Err(err()));
        }
+        Job::LoadAudio { reply, .. } => {
+            let _ = reply.send(Err(err()));
+        }
+        Job::DropAudio { reply, .. } => {
+            let _ = reply.send(());
+        }
+        Job::GenerateAudio { reply, .. } => {
+            let _ = reply.send(Err(err()));
+        }
        Job::ForwardLogitsWithImages { reply, .. } => {
            let _ = reply.send(Err(err()));
        }
@@ -68,6 +68,15 @@ pub struct ImageHandle(pub u64);
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct DecisionHandle(pub u64);

+/// Opaque handle to an ACE-Step music pipeline ([`AceStepEngine`]) stored
+/// in the worker thread's audio slab (#383). Same discipline as
+/// [`ImageHandle`]: construction, every generation and `Drop` happen on the
+/// worker thread, and `Job::DropAudio` is the only way to release it.
+///
+/// [`AceStepEngine`]: crate::harness::audio::AceStepEngine
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct AudioHandle(pub u64);
+
 /// Opaque handle to a prefix-cache snapshot (#11) stored worker-side
 /// next to the model slab. Scoped to the `ArchHandle` it was captured
 /// from — `Job::DropArch` drops every snapshot under its handle. The
@@ -585,6 +594,27 @@ pub enum Job {
        batch: crate::harness::arch::laya::DecisionBatch,
        reply: oneshot::Sender<Result<Vec<crate::harness::arch::laya::DecisionRow>>>,
    },
+    /// Load an ACE-Step snapshot (#383) — already on local disk — into
+    /// the audio slab. File resolution (hf-hub) happens on the async side.
+    LoadAudio {
+        files: crate::harness::audio::AceStepFiles,
+        model_id: String,
+        reply: oneshot::Sender<Result<AudioHandle>>,
+    },
+    /// Remove a music pipeline from the slab and drop it on the worker
+    /// thread. Mirrors `DropImage`.
+    DropAudio {
+        handle: AudioHandle,
+        reply: oneshot::Sender<()>,
+    },
+    /// Render one song (#383): planner, text encoding, denoise and VAE
+    /// decode in one job, since a render saturates the device. The reply
+    /// is CPU-side samples — no device tensor crosses back.
+    GenerateAudio {
+        handle: AudioHandle,
+        params: crate::harness::audio::AudioGenParams,
+        reply: oneshot::Sender<Result<crate::harness::audio::AudioGenResult>>,
+    },
    /// Tell the worker to break its dispatch loop and exit. Any jobs
    /// queued after this in the channel reply `Err` to their oneshot
    /// senders (the senders are dropped on the worker's exit, which
@@ -51,7 +51,9 @@ use tokio::sync::oneshot;

 #[cfg(feature = "cuda")]
 pub use jobs::TpHandle;
-pub use jobs::{ArchHandle, DecisionHandle, DenseLoad, ImageHandle, Job, KvSnapshotId};
+pub use jobs::{
+    ArchHandle, AudioHandle, DecisionHandle, DenseLoad, ImageHandle, Job, KvSnapshotId,
+};

 /// Errors returned by `DeviceWorkerHandle` submit methods.
 #[derive(Debug, thiserror::Error)]
@@ -337,6 +339,82 @@ impl DeviceWorkerHandle {
        }
    }

+    /// Load an ACE-Step music pipeline (#383) on the worker thread.
+    pub async fn load_audio(
+        &self,
+        files: crate::harness::audio::AceStepFiles,
+        model_id: String,
+    ) -> Result<AudioHandle, WorkerError> {
+        if self.poisoned.load(Ordering::Acquire) {
+            return Err(WorkerError::Poisoned {
+                device_index: self.device_index,
+            });
+        }
+        let (reply_tx, reply_rx) = oneshot::channel();
+        self.tx
+            .send(Job::LoadAudio {
+                files,
+                model_id,
+                reply: reply_tx,
+            })
+            .map_err(|_| WorkerError::Gone {
+                device_index: self.device_index,
+            })?;
+        match reply_rx.await {
+            Ok(result) => result.map_err(WorkerError::from),
+            Err(_) => Err(WorkerError::Gone {
+                device_index: self.device_index,
+            }),
+        }
+    }
+
+    /// Drop the music pipeline for `handle` on the worker thread. Same
+    /// poison-tolerant contract as [`Self::drop_arch`].
+    pub async fn drop_audio(&self, handle: AudioHandle) -> Result<(), WorkerError> {
+        let (reply_tx, reply_rx) = oneshot::channel();
+        self.tx
+            .send(Job::DropAudio {
+                handle,
+                reply: reply_tx,
+            })
+            .map_err(|_| WorkerError::Gone {
+                device_index: self.device_index,
+            })?;
+        reply_rx.await.map_err(|_| WorkerError::Gone {
+            device_index: self.device_index,
+        })
+    }
+
+    /// Render one song (#383) on the worker thread. The reply is CPU-side
+    /// samples; no device tensor crosses back.
+    pub async fn generate_audio(
+        &self,
+        handle: AudioHandle,
+        params: crate::harness::audio::AudioGenParams,
+    ) -> Result<crate::harness::audio::AudioGenResult, WorkerError> {
+        if self.poisoned.load(Ordering::Acquire) {
+            return Err(WorkerError::Poisoned {
+                device_index: self.device_index,
+            });
+        }
+        let (reply_tx, reply_rx) = oneshot::channel();
+        self.tx
+            .send(Job::GenerateAudio {
+                handle,
+                params,
+                reply: reply_tx,
+            })
+            .map_err(|_| WorkerError::Gone {
+                device_index: self.device_index,
+            })?;
+        match reply_rx.await {
+            Ok(result) => result.map_err(WorkerError::from),
+            Err(_) => Err(WorkerError::Gone {
+                device_index: self.device_index,
+            }),
+        }
+    }
+
    /// Load one Laya checkpoint directory (#336) on the worker thread.
    /// Returns the handle and the dtype the weights were loaded in.
    pub async fn load_decision(
@@ -2,15 +2,18 @@

 pub mod admission;
 pub mod arch;
+pub mod audio;
 pub mod candle;
 pub mod capability;
 pub mod chat_template;
 pub mod context_limit;
+pub mod decision;
 pub mod decision_model;
 pub mod device_worker;
 pub mod engine;
 pub mod image;
 pub mod isq_cache;
+pub mod lang_detect;
 pub mod prefix_cache;
 pub mod preflight;
 pub mod preprocess;
@@ -59,6 +59,12 @@ pub enum SourceFormat {
    /// `config.json`. Served by the decision path — single device,
    /// unquantized. Takes precedence over the safetensors sniff.
    DecisionModel,
+    /// An ACE-Step 1.5 music repo (#383): the turbo DiT, text encoder,
+    /// VAE and planner as subfolders of one snapshot. Served by the audio
+    /// path — single device, unquantized for now. Takes precedence over
+    /// the safetensors sniff, which would read the root `config.json` as a
+    /// text model.
+    AceStep,
 }

 /// Output of `preflight` for a load that can proceed. Carries the
@@ -258,6 +264,22 @@ pub async fn preflight(
            tp_size,
            picked_quant_file: None,
        }),
+        (SourceFormat::AceStep, tp, _) if tp > 1 => Err(PreflightError::Unsupported {
+            model_id: source_id.to_string(),
+            reason: format!(
+                "music models load on a single device; tensor_parallel={tp} is not supported"
+            ),
+        }),
+        (SourceFormat::AceStep, _, Some(q)) => Err(PreflightError::Unsupported {
+            model_id: source_id.to_string(),
+            reason: format!("music models load unquantized for now; remove quant={q:?}"),
+        }),
+        (SourceFormat::AceStep, _, None) => Ok(PlacementPlan {
+            model_id: source_id.to_string(),
+            format: format.clone(),
+            tp_size,
+            picked_quant_file: None,
+        }),

        (SourceFormat::DenseSafetensors { .. } | SourceFormat::Mixed { .. }, _, _) => {
            Ok(PlacementPlan {
@@ -327,6 +349,9 @@ pub fn classify(filenames: &[&str]) -> SourceFormat {
    if is_decision_layout(filenames) {
        return SourceFormat::DecisionModel;
    }
+    if crate::harness::audio::is_acestep_layout(filenames) {
+        return SourceFormat::AceStep;
+    }
    let mut gguf_quants: Vec<String> = filenames
        .iter()
        .filter(|f| f.to_lowercase().ends_with(".gguf"))
@@ -707,6 +732,20 @@ mod tests {
                tp_size,
                picked_quant_file: None,
            }),
+            (SourceFormat::AceStep, tp, _) if tp > 1 => Err(PreflightError::Unsupported {
+                model_id: source_id.to_string(),
+                reason: format!("music models load on a single device; tensor_parallel={tp}"),
+            }),
+            (SourceFormat::AceStep, _, Some(q)) => Err(PreflightError::Unsupported {
+                model_id: source_id.to_string(),
+                reason: format!("music models load unquantized for now; remove quant={q:?}"),
+            }),
+            (SourceFormat::AceStep, _, None) => Ok(PlacementPlan {
+                model_id: source_id.to_string(),
+                format: format.clone(),
+                tp_size,
+                picked_quant_file: None,
+            }),
            (SourceFormat::DenseSafetensors { .. } | SourceFormat::Mixed { .. }, _, _) => {
                Ok(PlacementPlan {
                    model_id: source_id.to_string(),
@@ -751,6 +790,10 @@ mod tests {
    #[test]
    fn classify_laya_layout() {
        assert_eq!(classify(&LAYA_FILES), SourceFormat::DecisionModel);
+        let mut ace = crate::harness::audio::required_files();
+        ace.push("config.json".into());
+        let ace: Vec<&str> = ace.iter().map(String::as_str).collect();
+        assert_eq!(classify(&ace), SourceFormat::AceStep);
        // A causal LM that also ships an rl config is still a text model.
        let lm = [
            "config.json",
@@ -0,0 +1,140 @@
+# ACE-Step 1.5 reference traces
+
+Parity fixtures for neuron's ACE-Step port (#383, #384, #385), recorded by
+`script/acestep-reference.py` from the upstream PyTorch pipeline:
+
+- code: `ace-step/ACE-Step-1.5` @ `ca1e85fe9430179831e6bc6be790c332190a3866`
+- weights: `ACE-Step/Ace-Step1.5` @ `19671f406d603126926c1b7e2adc169acbcade22`
+  (`acestep-v15-turbo`, `acestep-5Hz-lm-1.7B`, `Qwen3-Embedding-0.6B`, `vae`)
+- torch 2.10.0+cu128, transformers 4.57.6, pytorch_wavelets 1.3.0
+
+Each `<case>__<trace>.safetensors` holds the tensors, and the matching
+`.json` holds the metadata: the case inputs, decoded prompts, the sampler
+settings `generate_audio` actually received, the timestep schedule, planner
+prompts and outputs, per-tensor stats (shape, dtype, NaN/Inf counts, absmax,
+mean, std), the first tensor to contain a NaN, drift against `cpu_f32`, wall
+time and peak VRAM. `index.json` summarises every trace.
+
+## Cases
+
+| case | content | duration | planner |
+|---|---|---|---|
+| `instr_10s_noplan` | instrumental | 10 s | off |
+| `vocal_en_30s_noplan` | English vocals | 30 s | off |
+| `vocal_en_30s_plan` | English vocals | 30 s | on (1.7B, greedy) |
+| `vocal_es_20s_noplan` | Spanish vocals | 20 s | off |
+| `vocal_en_120s_noplan` | English vocals | 120 s | off |
+
+"Planner off" turns off `thinking` and all three `use_cot_*` flags, so the
+LM is never called. "Planner on" runs the LM greedily (`lm_temperature=0`),
+with upstream's default CFG (2.0 for the code phase) and constrained
+decoding.
+
+## Traces
+
+| trace | device | dtype | what it is |
+|---|---|---|---|
+| `cpu_f32` | CPU | float32 | the parity reference: deterministic and complete |
+| `cuda_bf16` | RTX 4090 | bfloat16 | what upstream runs on Ampere and later |
+| `cuda_f32` | RTX 4090 | float32 | true f32 on the GPU, TF32 off: the Pascal compute recipe |
+| `cuda_f16` | RTX 4090 | float16 | upstream's own pre-Ampere path (fp16 + eager attention), forced |
+
+All traces of a case start from the **same noise**: `cpu_f32` records it,
+and the GPU traces load it instead of drawing their own. CUDA and CPU
+`torch.Generator`s differ at the same seed. A port should likewise take
+`noise` as an input when checking parity.
+
+`cuda_f32` is loaded on the CPU and then moved to the GPU. Loading on CUDA
+and casting to f32 afterwards is not equivalent. A CUDA load rounds every
+RoPE `inv_freq` buffer and the float32 `silence_latent` to bf16, and a later
+cast cannot restore them. That alone moved a single DiT forward by 11%
+(rel L2) from the CPU result.
+
+## Results
+
+| case | `cuda_f32` vs CPU | `cuda_bf16` vs CPU | `cuda_f16` |
+|---|---|---|---|
+| `instr_10s_noplan` | 6e-6 | 0.08 | clean (0.018) |
+| `vocal_es_20s_noplan` | 1e-5 | 0.17 | **NaN** |
+| `vocal_en_30s_noplan` | 8e-4 | 0.41 | **NaN** |
+| `vocal_en_120s_noplan` | 2e-4 | 0.63 | **NaN** |
+| `vocal_en_30s_plan` | — | 1.01 (different plan) | — |
+
+Values are the relative L2 error of `target_latents` against `cpu_f32`.
+
+- **fp16 vocals NaN, instrumentals don't.** This is upstream issue #1055,
+  reproduced on Ada. The first non-finite value is in the lyric encoder's
+  last layer: `layers.7.mlp.down_proj` overflows float16 (the residual
+  stream is already ~3.6k at layer 6). Keeping RoPE in f32 does not help.
+  "[Instrumental]" is a short lyric that stays in range. Upstream raises
+  rather than decoding; the trace records everything up to that point.
+- **True f32 on the GPU matches the CPU** to about 1e-5 per forward, and
+  1e-5 to 8e-4 end to end. The f32 recipe is numerically sound.
+- **Upstream's bf16 path is far from f32**: 8% to 63% end to end, growing
+  with length. Part of that is the bf16-rounded `inv_freq` and silence
+  latent described above, not only bf16 arithmetic. A candle port that
+  computes RoPE in f32 should be judged against `cpu_f32`/`cuda_f32`.
+  `cuda_bf16` is upstream's behaviour, not a tolerance target.
+- **The planner's code phase is fragile under precision change.** The CoT
+  phase (164 tokens) is identical between `cpu_f32` and `cuda_bf16`, but
+  the very first audio code differs, so the whole plan differs.
+
+Peak VRAM on the 4090 (DiT, text encoder and VAE; plus the 1.7B LM where
+loaded): bf16 with the LM loaded, 10.2 GiB; f32 without the LM, 12.6 GiB;
+fp16 without the LM, 6.2–6.7 GiB.
+
+## Tensors
+
+`cpu_f32` and `cuda_f32` store float32; `cuda_bf16` and `cuda_f16` store
+their own dtype, which is exact and half the size. Shapes are for batch 1,
+and `T` is the number of 25 Hz latent frames.
+
+| name | shape | what |
+|---|---|---|
+| `text_token_ids` | `[1, Lt]` | DiT text prompt, tokenized by the Qwen3-Embedding tokenizer (prompt text in the JSON) |
+| `text_hidden_states` | `[1, Lt, 1024]` | Qwen3-Embedding-0.6B `last_hidden_state` |
+| `lyric_token_ids` | `[1, Ll]` | lyric prompt (`# Languages\n<lang>\n\n# Lyric\n<lyrics><\|endoftext\|>`) |
+| `lyric_embeds` | `[1, Ll, 1024]` | `embed_tokens` lookup only, with no transformer pass |
+| `encoder_hidden_states` | `[1, Ll+1+Lt, 2048]` | packed condition: lyric-encoder output (8 layers), then timbre (1 token), then projected text |
+| `encoder_attention_mask` | `[1, Ll+1+Lt]` | packed mask |
+| `gen_chunk_masks`, `gen_is_covers` | | `generate_audio` inputs |
+| `gen_precomputed_lm_hints_25Hz` | `[1, T, 64]` | planner cases: LM audio codes, detokenized to 25 Hz |
+| `context_latents` | `[1, T, 128]` | src latents (silence, or LM hints) concatenated with chunk masks |
+| `decoder_attention_mask` | `[1, T]` | |
+| `noise` | `[1, T, 64]` | initial `xt` (CPU traces only; GPU traces reuse it) |
+| `step{i}_xt`, `step{i}_vt` | `[1, T, 64]` | DiT input and predicted velocity at sampler step `i` |
+| `target_latents` | `[1, T, 64]` | sampler output, after the final step's `x0` and DCW |
+| `vae_input` | `[1, 64, T]` | what the VAE decoder received |
+| `vae_waveform` | `[1, 2, 24000]` | VAE output, first 0.5 s (48 kHz stereo), before normalisation. The pipeline output is this × `normalisation_gain` (JSON) |
+| `planner{n}_input_ids`, `planner{n}_output_ids` | | LM phase `n` (0 = CoT, 1 = codes): prompt ids (with CFG, conditional and unconditional rows, left-padded) and generated ids |
+| `planner_audio_codes` | `[N]` | the 5 Hz audio codes the planner produced (N = 5 × seconds) |
+
+The lyric encoder's own output is not stored. With an all-true mask it is
+exactly the first `Ll` rows of `encoder_hidden_states`, and the JSON
+records that this was checked (`lyric_encoder_out_is_encoder_prefix`).
+
+Steps saved (`steps_saved` in the JSON):
+
+- **CPU trace:** every step for the 10 s case; steps 0, 4 and 7 for the
+  20–30 s cases; step 0 for the 120 s case.
+- **GPU traces:** `step0_vt`, `step7_vt` and `target_latents`. The 120 s
+  case keeps only `target_latents`. A step where a NaN first appears is
+  always kept.
+- `step0_xt` is never stored, because it is `noise`.
+- The 120 s case omits `context_latents` (silence plus the chunk mask) and
+  `vae_input` (`target_latents` transposed).
+
+Every tensor that was dropped still has its stats in the JSON.
+
+## Regenerating
+
+In the ACE-Step venv, with the checkpoints already on disk (fetched through
+your registry; the script refuses to download). Record `cpu_f32` first,
+because the GPU traces load its noise:
+
+```sh
+python -I script/acestep-reference.py --src <ACE-Step-1.5 checkout> --trace cpu_f32 --out <this dir>
+python -I script/acestep-reference.py --src ... --trace cuda_bf16 --out ...
+python -I script/acestep-reference.py --src ... --trace cuda_f32 --out ...
+python -I script/acestep-reference.py --src ... --trace cuda_f16 --out ...
+```
@@ -0,0 +1,238 @@
+{
+ "traces": {
+  "instr_10s_noplan__cpu_f32": {
+   "case": "instr_10s_noplan",
+   "trace": "cpu_f32",
+   "dtype": "float32",
+   "latent_frames": 250,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": null,
+   "error": null,
+   "wall_s": 17.4,
+   "vram_peak_gib": null,
+   "bytes": 2569385
+  },
+  "instr_10s_noplan__cuda_bf16": {
+   "case": "instr_10s_noplan",
+   "trace": "cuda_bf16",
+   "dtype": "bfloat16",
+   "latent_frames": 250,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": 0.0816783532500267,
+   "error": null,
+   "wall_s": 5.07,
+   "vram_peak_gib": 10.15,
+   "bytes": 96321
+  },
+  "instr_10s_noplan__cuda_f16": {
+   "case": "instr_10s_noplan",
+   "trace": "cuda_f16",
+   "dtype": "float16",
+   "latent_frames": 250,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": 0.018279578536748886,
+   "error": null,
+   "wall_s": 7.37,
+   "vram_peak_gib": 6.69,
+   "bytes": 96313
+  },
+  "instr_10s_noplan__cuda_f32": {
+   "case": "instr_10s_noplan",
+   "trace": "cuda_f32",
+   "dtype": "float32",
+   "latent_frames": 250,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": 6.1245264078024775e-06,
+   "error": null,
+   "wall_s": 10.95,
+   "vram_peak_gib": 12.63,
+   "bytes": 384401
+  },
+  "vocal_en_120s_noplan__cpu_f32": {
+   "case": "vocal_en_120s_noplan",
+   "trace": "cpu_f32",
+   "dtype": "float32",
+   "latent_frames": 3000,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": null,
+   "error": null,
+   "wall_s": 144.06,
+   "vram_peak_gib": null,
+   "bytes": 4603028
+  },
+  "vocal_en_120s_noplan__cuda_bf16": {
+   "case": "vocal_en_120s_noplan",
+   "trace": "cuda_bf16",
+   "dtype": "bfloat16",
+   "latent_frames": 3000,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": 0.6251527070999146,
+   "error": null,
+   "wall_s": 40.29,
+   "vram_peak_gib": 10.18,
+   "bytes": 384169
+  },
+  "vocal_en_120s_noplan__cuda_f16": {
+   "case": "vocal_en_120s_noplan",
+   "trace": "cuda_f16",
+   "dtype": "float16",
+   "latent_frames": 3000,
+   "nan_anywhere": true,
+   "first_nan": "lyric_encoder_out",
+   "target_latents_rel_l2_vs_cpu_f32": Infinity,
+   "error": "Generation produced NaN or Inf latents (shape=[1, 3000, 64], dtype=torch.float16, device=cuda:0, nan=192000, inf=0).\nCommon causes and fixes:\n  1. LoRA/adapter trained on an older model version \u2014 retrain or update the adapter.\n  2. Checkpoint/config mismatch \u2014 verify model checkpoints match this release.\n  3. Unsupported quantization/backend \u2014 try running with --backend pt.\n  4. CPU offload left parameters on wrong device \u2014 restart and regenerate.\n  5. Float16 overflow on pre-Ampere GPU \u2014 set ACESTEP_DTYPE=float32.",
+   "wall_s": 18.43,
+   "vram_peak_gib": 6.42,
+   "bytes": 384169
+  },
+  "vocal_en_120s_noplan__cuda_f32": {
+   "case": "vocal_en_120s_noplan",
+   "trace": "cuda_f32",
+   "dtype": "float32",
+   "latent_frames": 3000,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": 0.00023412087466567755,
+   "error": null,
+   "wall_s": 80.42,
+   "vram_peak_gib": 12.64,
+   "bytes": 960249
+  },
+  "vocal_en_30s_noplan__cpu_f32": {
+   "case": "vocal_en_30s_noplan",
+   "trace": "cpu_f32",
+   "dtype": "float32",
+   "latent_frames": 750,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": null,
+   "error": null,
+   "wall_s": 39.68,
+   "vram_peak_gib": null,
+   "bytes": 3636105
+  },
+  "vocal_en_30s_noplan__cuda_bf16": {
+   "case": "vocal_en_30s_noplan",
+   "trace": "cuda_bf16",
+   "dtype": "bfloat16",
+   "latent_frames": 750,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": 0.4098935127258301,
+   "error": null,
+   "wall_s": 11.14,
+   "vram_peak_gib": 10.16,
+   "bytes": 288321
+  },
+  "vocal_en_30s_noplan__cuda_f16": {
+   "case": "vocal_en_30s_noplan",
+   "trace": "cuda_f16",
+   "dtype": "float16",
+   "latent_frames": 750,
+   "nan_anywhere": true,
+   "first_nan": "lyric_encoder_out",
+   "target_latents_rel_l2_vs_cpu_f32": Infinity,
+   "error": "Generation produced NaN or Inf latents (shape=[1, 750, 64], dtype=torch.float16, device=cuda:0, nan=48000, inf=0).\nCommon causes and fixes:\n  1. LoRA/adapter trained on an older model version \u2014 retrain or update the adapter.\n  2. Checkpoint/config mismatch \u2014 verify model checkpoints match this release.\n  3. Unsupported quantization/backend \u2014 try running with --backend pt.\n  4. CPU offload left parameters on wrong device \u2014 restart and regenerate.\n  5. Float16 overflow on pre-Ampere GPU \u2014 set ACESTEP_DTYPE=float32.",
+   "wall_s": 3.77,
+   "vram_peak_gib": 6.16,
+   "bytes": 288321
+  },
+  "vocal_en_30s_noplan__cuda_f32": {
+   "case": "vocal_en_30s_noplan",
+   "trace": "cuda_f32",
+   "dtype": "float32",
+   "latent_frames": 750,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": 0.0008115487289614975,
+   "error": null,
+   "wall_s": 23.05,
+   "vram_peak_gib": 12.63,
+   "bytes": 768401
+  },
+  "vocal_en_30s_plan__cpu_f32": {
+   "case": "vocal_en_30s_plan",
+   "trace": "cpu_f32",
+   "dtype": "float32",
+   "latent_frames": 750,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": null,
+   "error": null,
+   "wall_s": 97.87,
+   "vram_peak_gib": null,
+   "bytes": 5055700
+  },
+  "vocal_en_30s_plan__cuda_bf16": {
+   "case": "vocal_en_30s_plan",
+   "trace": "cuda_bf16",
+   "dtype": "bfloat16",
+   "latent_frames": 750,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": 1.0133154392242432,
+   "error": null,
+   "wall_s": 72.41,
+   "vram_peak_gib": 10.18,
+   "bytes": 394505
+  },
+  "vocal_es_20s_noplan__cpu_f32": {
+   "case": "vocal_es_20s_noplan",
+   "trace": "cpu_f32",
+   "dtype": "float32",
+   "latent_frames": 500,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": null,
+   "error": null,
+   "wall_s": 25.68,
+   "vram_peak_gib": null,
+   "bytes": 2843838
+  },
+  "vocal_es_20s_noplan__cuda_bf16": {
+   "case": "vocal_es_20s_noplan",
+   "trace": "cuda_bf16",
+   "dtype": "bfloat16",
+   "latent_frames": 500,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": 0.16632689535617828,
+   "error": null,
+   "wall_s": 7.86,
+   "vram_peak_gib": 10.17,
+   "bytes": 192321
+  },
+  "vocal_es_20s_noplan__cuda_f16": {
+   "case": "vocal_es_20s_noplan",
+   "trace": "cuda_f16",
+   "dtype": "float16",
+   "latent_frames": 500,
+   "nan_anywhere": true,
+   "first_nan": "lyric_encoder_out",
+   "target_latents_rel_l2_vs_cpu_f32": Infinity,
+   "error": "Generation produced NaN or Inf latents (shape=[1, 500, 64], dtype=torch.float16, device=cuda:0, nan=32000, inf=0).\nCommon causes and fixes:\n  1. LoRA/adapter trained on an older model version \u2014 retrain or update the adapter.\n  2. Checkpoint/config mismatch \u2014 verify model checkpoints match this release.\n  3. Unsupported quantization/backend \u2014 try running with --backend pt.\n  4. CPU offload left parameters on wrong device \u2014 restart and regenerate.\n  5. Float16 overflow on pre-Ampere GPU \u2014 set ACESTEP_DTYPE=float32.",
+   "wall_s": 6.31,
+   "vram_peak_gib": 6.16,
+   "bytes": 192321
+  },
+  "vocal_es_20s_noplan__cuda_f32": {
+   "case": "vocal_es_20s_noplan",
+   "trace": "cuda_f32",
+   "dtype": "float32",
+   "latent_frames": 500,
+   "nan_anywhere": false,
+   "first_nan": null,
+   "target_latents_rel_l2_vs_cpu_f32": 1.0157832548429724e-05,
+   "error": null,
+   "wall_s": 15.1,
+   "vram_peak_gib": 12.63,
+   "bytes": 576401
+  }
+ }
+}
@@ -0,0 +1,594 @@
+{
+ "stats": {
+  "text_token_ids": {
+   "shape": [
+    1,
+    72
+   ],
+   "dtype": "int64",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 151643.0,
+   "mean": 14568.125,
+   "std": 30829.681640625
+  },
+  "text_hidden_states": {
+   "shape": [
+    1,
+    72,
+    1024
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 32.86247253417969,
+   "mean": 0.021975461393594742,
+   "std": 3.2490692138671875
+  },
+  "lyric_token_ids": {
+   "shape": [
+    1,
+    15
+   ],
+   "dtype": "int64",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 151643.0,
+   "mean": 29993.19921875,
+   "std": 52921.25
+  },
+  "lyric_embeds": {
+   "shape": [
+    1,
+    15,
+    1024
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.1767578125,
+   "mean": 2.041088191617746e-05,
+   "std": 0.03075374662876129
+  },
+  "lyric_encoder_out": {
+   "shape": [
+    1,
+    15,
+    2048
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.524937629699707,
+   "mean": -9.119734022533521e-05,
+   "std": 0.11492618173360825
+  },
+  "encoder_hidden_states": {
+   "shape": [
+    1,
+    88,
+    2048
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 58.46120834350586,
+   "mean": 0.006924757268279791,
+   "std": 3.4267163276672363
+  },
+  "encoder_attention_mask": {
+   "shape": [
+    1,
+    88
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_src_latents": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.726634979248047,
+   "mean": -0.2051306664943695,
+   "std": 0.9637550711631775
+  },
+  "gen_chunk_masks": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_is_covers": {
+   "shape": [
+    1
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.0,
+   "mean": 0.0,
+   "std": 0.0
+  },
+  "gen_text_attention_mask": {
+   "shape": [
+    1,
+    72
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_lyric_attention_mask": {
+   "shape": [
+    1,
+    15
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_refer_audio_acoustic_hidden_states_packed": {
+   "shape": [
+    1,
+    750,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.726634979248047,
+   "mean": -0.20505407452583313,
+   "std": 0.963009774684906
+  },
+  "gen_refer_audio_order_mask": {
+   "shape": [
+    1
+   ],
+   "dtype": "int64",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.0,
+   "mean": 0.0,
+   "std": 0.0
+  },
+  "noise": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.068782329559326,
+   "mean": -0.001519594225101173,
+   "std": 1.0108991861343384
+  },
+  "decoder_attention_mask": {
+   "shape": [
+    1,
+    250
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "context_latents": {
+   "shape": [
+    1,
+    250,
+    128
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.726634979248047,
+   "mean": 0.39743465185165405,
+   "std": 0.9096668362617493
+  },
+  "target_latents": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.789435386657715,
+   "mean": -0.037685856223106384,
+   "std": 1.038394808769226
+  },
+  "vae_input": {
+   "shape": [
+    1,
+    64,
+    250
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.789435386657715,
+   "mean": -0.037685852497816086,
+   "std": 1.038394808769226
+  },
+  "vae_output": {
+   "shape": [
+    1,
+    2,
+    480000
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0992342233657837,
+   "mean": 0.0012270656879991293,
+   "std": 0.16220736503601074
+  },
+  "step0_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.068782329559326,
+   "mean": -0.001519594225101173,
+   "std": 1.0108991861343384
+  },
+  "step0_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.779356002807617,
+   "mean": 0.029778892174363136,
+   "std": 1.0888593196868896
+  },
+  "step1_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.017130374908447,
+   "mean": -0.0014519158285111189,
+   "std": 0.9958045482635498
+  },
+  "step1_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.739595413208008,
+   "mean": 0.0698830708861351,
+   "std": 1.1297178268432617
+  },
+  "step2_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.936816453933716,
+   "mean": -0.00226192525587976,
+   "std": 0.9712164402008057
+  },
+  "step2_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 7.189512729644775,
+   "mean": 0.048355162143707275,
+   "std": 1.1502469778060913
+  },
+  "step3_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.813173294067383,
+   "mean": -0.0036722831428050995,
+   "std": 0.9389426112174988
+  },
+  "step3_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 7.870788097381592,
+   "mean": 0.04673989862203598,
+   "std": 1.2650023698806763
+  },
+  "step4_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.5883963108062744,
+   "mean": -0.006106651853770018,
+   "std": 0.8947911858558655
+  },
+  "step4_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 8.094154357910156,
+   "mean": 0.03310832753777504,
+   "std": 1.3123292922973633
+  },
+  "step5_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.4403023719787598,
+   "mean": -0.008855825290083885,
+   "std": 0.8410162329673767
+  },
+  "step5_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 8.550355911254883,
+   "mean": 0.04433209449052811,
+   "std": 1.3689327239990234
+  },
+  "step6_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.3948163986206055,
+   "mean": -0.014476503245532513,
+   "std": 0.7926256656646729
+  },
+  "step6_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 8.73010540008545,
+   "mean": 0.05107048898935318,
+   "std": 1.4008506536483765
+  },
+  "step7_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.267222881317139,
+   "mean": -0.024307560175657272,
+   "std": 0.8039591312408447
+  },
+  "step7_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 9.258280754089355,
+   "mean": 0.04459429532289505,
+   "std": 1.5695395469665527
+  },
+  "final_audio_full": {
+   "shape": [
+    2,
+    480000
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.8912509083747864,
+   "mean": 0.0009948957012966275,
+   "std": 0.13151653110980988
+  },
+  "vae_waveform": {
+   "shape": [
+    1,
+    2,
+    24000
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.9028035402297974,
+   "mean": 0.001543997903354466,
+   "std": 0.23970195651054382
+  }
+ },
+ "planner": [],
+ "decoder_calls": 8,
+ "text_prompt": "# Instruction\nFill the audio semantic mask based on the given conditions:\n\n# Caption\nwarm lo-fi hip hop beat, dusty vinyl drums, mellow rhodes chords, relaxed\n\n# Metas\n- bpm: N/A\n- timesignature: N/A\n- keyscale: N/A\n- duration: 10 seconds\n<|endoftext|>\n<|endoftext|>",
+ "lyric_prompt": "# Languages\nunknown\n\n# Lyric\n[Instrumental]<|endoftext|><|endoftext|>",
+ "generate_audio_kwargs": {
+  "seed": [
+   1234
+  ],
+  "non_cover_text_hidden_states": null,
+  "non_cover_text_attention_mask": null,
+  "precomputed_lm_hints_25Hz": null,
+  "audio_cover_strength": 1.0,
+  "cover_noise_strength": 0.0,
+  "infer_method": "ode",
+  "infer_steps": 8,
+  "diffusion_guidance_scale": 1.0,
+  "use_adg": false,
+  "cfg_interval_start": 0.0,
+  "cfg_interval_end": 1.0,
+  "shift": 3.0,
+  "repaint_mask": null,
+  "clean_src_latents": null,
+  "repaint_crossfade_frames": 12,
+  "repaint_injection_ratio": 0.5,
+  "sampler_mode": "euler",
+  "velocity_norm_threshold": 0.0,
+  "velocity_ema_factor": 0.0,
+  "dcw_enabled": true,
+  "dcw_mode": "double",
+  "dcw_scaler": 0.05,
+  "dcw_high_scaler": 0.02,
+  "dcw_wavelet": "haar",
+  "retake_seed": null,
+  "retake_variance": 0.0
+ },
+ "noise_seed": [
+  1234
+ ],
+ "vae_path": "tiled_decode",
+ "first_nan": null,
+ "timesteps": [
+  1.0,
+  0.9545454382896423,
+  0.8999999761581421,
+  0.8333333134651184,
+  0.75,
+  0.6428571343421936,
+  0.5,
+  0.30000001192092896
+ ],
+ "steps_saved": {
+  "xt": [
+   1,
+   2,
+   3,
+   4,
+   5,
+   6,
+   7
+  ],
+  "vt": [
+   0,
+   1,
+   2,
+   3,
+   4,
+   5,
+   6,
+   7
+  ]
+ },
+ "lyric_encoder_out_is_encoder_prefix": true,
+ "normalisation_gain": 0.8107925653457642,
+ "case": {
+  "name": "instr_10s_noplan",
+  "caption": "warm lo-fi hip hop beat, dusty vinyl drums, mellow rhodes chords, relaxed",
+  "lyrics": "[Instrumental]",
+  "instrumental": true,
+  "language": "unknown",
+  "duration": 10.0,
+  "planner": false,
+  "seed": 1234
+ },
+ "trace": "cpu_f32",
+ "sample_rate": 48000,
+ "latent_frames": 250,
+ "wall_s": 17.4,
+ "planner_metadata": {
+  "cot_bpm": null,
+  "cot_keyscale": "",
+  "cot_timesignature": "",
+  "cot_duration": null,
+  "cot_vocal_language": "unknown",
+  "cot_caption": "",
+  "cot_lyrics": ""
+ },
+ "nan_anywhere": false,
+ "dtype": "float32",
+ "time_costs": {
+  "dit_encoder_time_cost": 0.5645701885223389,
+  "dit_diffusion_time_cost": 5.696354866027832,
+  "dit_diffusion_per_step_time_cost": 0.712044358253479,
+  "dit_total_time_cost": 10.253296375274658,
+  "dit_offload_time_cost": 0.0,
+  "dit_vae_decode_time_cost": 3.9923713207244873,
+  "pipeline_total_time": 10.253296375274658
+ }
+}
@@ -0,0 +1,628 @@
+{
+ "stats": {
+  "text_token_ids": {
+   "shape": [
+    1,
+    72
+   ],
+   "dtype": "int64",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 151643.0,
+   "mean": 14568.125,
+   "std": 30829.6796875
+  },
+  "text_hidden_states": {
+   "shape": [
+    1,
+    72,
+    1024
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 33.0,
+   "mean": 0.021900618448853493,
+   "std": 3.2504899501800537
+  },
+  "lyric_token_ids": {
+   "shape": [
+    1,
+    15
+   ],
+   "dtype": "int64",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 151643.0,
+   "mean": 29993.201171875,
+   "std": 52921.24609375
+  },
+  "lyric_embeds": {
+   "shape": [
+    1,
+    15,
+    1024
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.1767578125,
+   "mean": 2.0410874640219845e-05,
+   "std": 0.03075374662876129
+  },
+  "lyric_encoder_out": {
+   "shape": [
+    1,
+    15,
+    2048
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.53125,
+   "mean": -8.808567508822307e-05,
+   "std": 0.11486581712961197
+  },
+  "encoder_hidden_states": {
+   "shape": [
+    1,
+    88,
+    2048
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 58.75,
+   "mean": 0.007172065321356058,
+   "std": 3.4257657527923584
+  },
+  "encoder_attention_mask": {
+   "shape": [
+    1,
+    88
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_src_latents": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.71875,
+   "mean": -0.20533928275108337,
+   "std": 0.9645420908927917
+  },
+  "gen_chunk_masks": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_is_covers": {
+   "shape": [
+    1
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.0,
+   "mean": 0.0,
+   "std": 0.0
+  },
+  "gen_text_attention_mask": {
+   "shape": [
+    1,
+    72
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_lyric_attention_mask": {
+   "shape": [
+    1,
+    15
+   ],
+   "dtype": "bool",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "gen_refer_audio_acoustic_hidden_states_packed": {
+   "shape": [
+    1,
+    750,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.71875,
+   "mean": -0.20527024567127228,
+   "std": 0.9638134837150574
+  },
+  "gen_refer_audio_order_mask": {
+   "shape": [
+    1
+   ],
+   "dtype": "int64",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.0,
+   "mean": 0.0,
+   "std": 0.0
+  },
+  "noise": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.0625,
+   "mean": -0.0015215945895761251,
+   "std": 1.0108590126037598
+  },
+  "decoder_attention_mask": {
+   "shape": [
+    1,
+    250
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.0,
+   "mean": 1.0,
+   "std": 0.0
+  },
+  "context_latents": {
+   "shape": [
+    1,
+    250,
+    128
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.71875,
+   "mean": 0.3973304033279419,
+   "std": 0.9101529121398926
+  },
+  "target_latents": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.8125,
+   "mean": -0.03683716058731079,
+   "std": 1.0391204357147217
+  },
+  "vae_input": {
+   "shape": [
+    1,
+    64,
+    250
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.8125,
+   "mean": -0.03683716803789139,
+   "std": 1.0391203165054321
+  },
+  "vae_output": {
+   "shape": [
+    1,
+    2,
+    480000
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 1.09375,
+   "mean": 0.0019740236457437277,
+   "std": 0.16050520539283752
+  },
+  "step0_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.0625,
+   "mean": -0.0015215944731608033,
+   "std": 1.0108590126037598
+  },
+  "step0_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.75,
+   "mean": 0.029024723917245865,
+   "std": 1.087060570716858
+  },
+  "step1_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.0,
+   "mean": -0.0014959202380850911,
+   "std": 0.9948191046714783
+  },
+  "step1_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 6.78125,
+   "mean": 0.0699673667550087,
+   "std": 1.131856918334961
+  },
+  "step2_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.90625,
+   "mean": -0.002318344544619322,
+   "std": 0.9700552821159363
+  },
+  "step2_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 7.15625,
+   "mean": 0.04690621793270111,
+   "std": 1.1479475498199463
+  },
+  "step3_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.796875,
+   "mean": -0.0036648816894739866,
+   "std": 0.9381669163703918
+  },
+  "step3_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 7.875,
+   "mean": 0.0431983508169651,
+   "std": 1.2662577629089355
+  },
+  "step4_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.5625,
+   "mean": -0.0058504026383161545,
+   "std": 0.8947983980178833
+  },
+  "step4_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 8.125,
+   "mean": 0.0328899547457695,
+   "std": 1.3014971017837524
+  },
+  "step5_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.421875,
+   "mean": -0.008506354875862598,
+   "std": 0.8428772687911987
+  },
+  "step5_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 8.5625,
+   "mean": 0.044462598860263824,
+   "std": 1.3666576147079468
+  },
+  "step6_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 3.359375,
+   "mean": -0.014202592894434929,
+   "std": 0.7932853102684021
+  },
+  "step6_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 8.75,
+   "mean": 0.048620857298374176,
+   "std": 1.3998051881790161
+  },
+  "step7_xt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 4.25,
+   "mean": -0.023524314165115356,
+   "std": 0.8046457171440125
+  },
+  "step7_vt": {
+   "shape": [
+    1,
+    250,
+    64
+   ],
+   "dtype": "bfloat16",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 9.3125,
+   "mean": 0.044245440512895584,
+   "std": 1.5749262571334839
+  },
+  "final_audio_full": {
+   "shape": [
+    2,
+    480000
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.8912509083747864,
+   "mean": 0.0016085486859083176,
+   "std": 0.13078895211219788
+  },
+  "vae_waveform": {
+   "shape": [
+    1,
+    2,
+    24000
+   ],
+   "dtype": "float32",
+   "nan": 0,
+   "inf": 0,
+   "absmax": 0.90625,
+   "mean": 0.0024766994174569845,
+   "std": 0.2409236878156662
+  }
+ },
+ "planner": [],
+ "decoder_calls": 8,
+ "text_prompt": "# Instruction\nFill the audio semantic mask based on the given conditions:\n\n# Caption\nwarm lo-fi hip hop beat, dusty vinyl drums, mellow rhodes chords, relaxed\n\n# Metas\n- bpm: N/A\n- timesignature: N/A\n- keyscale: N/A\n- duration: 10 seconds\n<|endoftext|>\n<|endoftext|>",
+ "lyric_prompt": "# Languages\nunknown\n\n# Lyric\n[Instrumental]<|endoftext|><|endoftext|>",
+ "generate_audio_kwargs": {
+  "seed": [
+   1234
+  ],
+  "non_cover_text_hidden_states": null,
+  "non_cover_text_attention_mask": null,
+  "precomputed_lm_hints_25Hz": null,
+  "audio_cover_strength": 1.0,
+  "cover_noise_strength": 0.0,
+  "infer_method": "ode",
+  "infer_steps": 8,
+  "diffusion_guidance_scale": 1.0,
+  "use_adg": false,
+  "cfg_interval_start": 0.0,
+  "cfg_interval_end": 1.0,
+  "shift": 3.0,
+  "repaint_mask": null,
+  "clean_src_latents": null,
+  "repaint_crossfade_frames": 12,
+  "repaint_injection_ratio": 0.5,
+  "sampler_mode": "euler",
+  "velocity_norm_threshold": 0.0,
+  "velocity_ema_factor": 0.0,
+  "dcw_enabled": true,
+  "dcw_mode": "double",
+  "dcw_scaler": 0.05,
+  "dcw_high_scaler": 0.02,
+  "dcw_wavelet": "haar",
+  "retake_seed": null,
+  "retake_variance": 0.0
+ },
+ "noise_source": "cpu_f32",
+ "noise_seed": [
+  1234
+ ],
+ "vae_path": "tiled_decode",
+ "first_nan": null,
+ "drift_vs_cpu_f32": {
+  "target_latents": {
+   "max_abs": 1.1789467334747314,
+   "rel_l2": 0.0816783532500267,
+   "cos": 0.9966683387756348
+  },
+  "step0_vt": {
+   "max_abs": 0.9140055179595947,
+   "rel_l2": 0.0775710716843605,
+   "cos": 0.9969895482063293
+  },
+  "step1_vt": {
+   "max_abs": 1.0068531036376953,
+   "rel_l2": 0.04007680341601372,
+   "cos": 0.9992018342018127
+  },
+  "step2_vt": {
+   "max_abs": 2.0684454441070557,
+   "rel_l2": 0.09504576772451401,
+   "cos": 0.9954776763916016
+  },
+  "step3_vt": {
+   "max_abs": 1.8307780027389526,
+   "rel_l2": 0.1291956901550293,
+   "cos": 0.9916636943817139
+  },
+  "step4_vt": {
+   "max_abs": 4.880147457122803,
+   "rel_l2": 0.14142011106014252,
+   "cos": 0.989952802658081
+  },
+  "step5_vt": {
+   "max_abs": 1.726083755493164,
+   "rel_l2": 0.10134051740169525,
+   "cos": 0.9948592185974121
+  },
+  "step6_vt": {
+   "max_abs": 2.1955137252807617,
+   "rel_l2": 0.10752364993095398,
+   "cos": 0.9942163825035095
+  },
+  "step7_vt": {
+   "max_abs": 1.8701152801513672,
+   "rel_l2": 0.09535709768533707,
+   "cos": 0.9954763650894165
+  }
+ },
+ "timesteps": [
+  1.0,
+  0.953125,
+  0.8984375,
+  0.83203125,
+  0.75,
+  0.64453125,
+  0.5,
+  0.30078125
+ ],
+ "steps_saved": {
+  "xt": [],
+  "vt": [
+   0,
+   7
+  ]
+ },
+ "normalisation_gain": 0.8148579597473145,
+ "case": {
+  "name": "instr_10s_noplan",
+  "caption": "warm lo-fi hip hop beat, dusty vinyl drums, mellow rhodes chords, relaxed",
+  "lyrics": "[Instrumental]",
+  "instrumental": true,
+  "language": "unknown",
+  "duration": 10.0,
+  "planner": false,
+  "seed": 1234
+ },
+ "trace": "cuda_bf16",
+ "sample_rate": 48000,
+ "latent_frames": 250,
+ "wall_s": 5.07,
+ "planner_metadata": {
+  "cot_bpm": null,
+  "cot_keyscale": "",
+  "cot_timesignature": "",
+  "cot_duration": null,
+  "cot_vocal_language": "unknown",
+  "cot_caption": "",
+  "cot_lyrics": ""
+ },
+ "nan_anywhere": false,
+ "dtype": "bfloat16",
+ "time_costs": {
+  "dit_encoder_time_cost": 0.12479805946350098,
+  "dit_diffusion_time_cost": 1.8494272232055664,
+  "dit_diffusion_per_step_time_cost": 0.2311784029006958,
+  "dit_total_time_cost": 3.6470632553100586,
+  "dit_offload_time_cost": 0.0,
+  "dit_vae_decode_time_cost": 1.6728379726409912,
+  "pipeline_total_time": 3.6470632553100586
+ },
+ "vram_peak_gib": 10.15
+}
--- a/Show More
+++ b/Show More