feat(audio8_tts): Falcon-H1 0.1B (Mamba2 + attention) support and codec/AR performance

docs/community_models/audio8_tts.md

@@ -4,7 +4,7 @@ Audio8 TTS Preview 0.6B (Qwen backbone) and 0.1B (Falcon-H1 hybrid Mamba2+attent
 S2 Pro: a slow semantic transformer generates speech semantics, a fast codebook transformer expands each semantic step into a full codec frame, and a neural codec renders 44.1 kHz audio. The native path executes all three
 stages directly on ggml with no Python dependency.
 
-> **Status 2026-08-29:** `0.6B` Qwen is fully native, CPU-validated via SenseVoice ASR round-trip (`The quick brown fox…`, `你好，欢迎使用audio8。`, `Artificial intelligence…`). `0.1B` Falcon-H1 is weight-complete and builds natively (GGUF `slow.embed_tokens` + `24× mamba/attention` + `semantic_output`), but the slow AR forward is a documented stub pending the Mamba2 port — see `docs/FALCON_H1_0.1B_PORT_PLAN.md` and `src/community_models/audio8_tts/ar.cpp:861` `TODO(Falcon-H1)`.
+> **Status 2026-09-01:** `0.6B` Qwen is fully native, CPU-validated via SenseVoice ASR round-trip (`The quick brown fox…`, `你好，欢迎使用audio8。`, `Artificial intelligence…`). `0.1B` Falcon-H1 now has a **stateful native slow-AR forward** (Mamba2 + hybrid GQA attention, branch `feat/audio8-tts-falcon-h1-mamba2`), but it is **not yet correct for synthesis** — two open issues (logits argmax mismatch vs transformers reference, and recurrent SSM state blow-up on long sequences). See [audio8_tts_falcon_h1_status.md](audio8_tts_falcon_h1_status.md) for details and next steps.
 
 | Field | Value |
 |---|---|

docs/community_models/audio8_tts_falcon_h1_status.md

@@ -0,0 +1,116 @@
+# Audio8 TTS 0.1B (Falcon-H1) — Port Status
+
+> **Status 2026-09-01 (resolved):** The Falcon-H1 slow-AR path is a **stateful
+> native implementation** (Mamba2 + hybrid GQA attention) and now matches the
+> transformers reference: first-frame semantic argmax = 2732, per-step argmax
+> parity over the whole prompt, ASR round-trip of synthesized "你好" returns
+> "你好。", and long generation (~600 positions) is numerically stable. The
+> 0.6B Qwen path is unaffected and fully functional.
+
+## What has been done
+
+Branch `feat/audio8-tts-falcon-h1-mamba2`.
+
+- `src/community_models/audio8_tts/ar.cpp`
+  - `FalconH1StepState` + `init_falcon_step_state`: per-layer conv/SSM states
+    and attention KV cache.
+  - `falcon_forward_step`: stateful single-token forward —
+    `RMSNorm -> (Mamba2 || GQA attention) -> residual -> gated FFN -> RMSNorm
+    -> semantic_output`, matching `transformers.models.falcon_h1`.
+  - `build_falcon_embedding_step`: `(text_emb + codebook_sum) *
+    embedding_multiplier` (multiplier applies to the whole sum).
+  - `generate()` falcon branch: token-by-token prefill + generation.
+- `src/community_models/audio8_tts/falcon_kv_cache.h`
+  - `append_falcon_kv_token`: host KV-cache append with per-head re-stride
+    (see "Root causes" below). Covered by `audio8_tts_falcon_kv_cache_test`.
+- `external/ggml/src/ggml-metal/ggml-metal.metal`
+  - Fixed `kernel_ssm_scan_f32` reduction: the old
+    `simd_sum(shared_sums[sgitg*NW + tiisg])` read garbage columns when
+    `sgptg < NW` (happens for `d_state=64` with `n_t=1`). Replaced with an
+    explicit loop summing `shared_sums[(i2+sgitg)*NW + g]` over `g < sgptg`.
+
+## Resolved Issue 1 — logits argmax mismatch vs transformers
+
+**Symptom (before):** first generated semantic code was wrong (argmax 3620
+instead of 2732; ASR round-trip said "三星" instead of "你好").
+
+**Root causes (three, all fixed):**
+
+1. **conv1d kernel flip was wrong.** `ggml_ssm_conv` computes
+   `y[c] = sum_k w[k,c]*window[k,c]` with `window[0]` the OLDEST frame —
+   the same orientation as HF (`nn.Conv1d` prefill and the cached
+   `torch.sum(conv_states * w, dim=-1)` decode are both cross-correlation).
+   The GGUF tensor `[d_conv,1,conv_dim]` is the HF `[conv_dim,1,d_conv]`
+   weight with unchanged flat bytes, i.e. already in the layout ssm_conv
+   wants. An earlier "fix" that flipped the kernel taps corrupted the x/B/C
+   split every step. Fix: feed the kernel unflipped (`load_falcon_layer`,
+   `conv1d_kernel`).
+2. **Unprotected host read-back of intermediate tensors.** `sx` (conv
+   window) and `k_r`/`v` (fresh K/V) are graph intermediates whose buffers
+   gallocr reuses; reading them back without `ggml_set_output` returned
+   garbage and corrupted conv state / KV cache every step. Fix:
+   `ggml_set_output` on exactly those three tensors per layer (pinning ~300
+   tensors corrupts the whole graph — pin only what is read back).
+3. **KV cache head-stride bug (the decisive one).** The host cache used the
+   *current* sequence length as the per-head stride while appending only the
+   new token: at step 1 the new head-0 token was written over token 0's
+   head-1 block, so every head past the first read corrupted context from
+   the second token on (head 0 was always correct, which masked the bug).
+   Fix: `append_falcon_kv_token` re-lays existing entries into the new
+   stride before appending. Regression test:
+   `tests/unittests/test_audio8_tts_falcon_kv_cache.cpp` (fails with the old
+   algorithm at the second append, passes after).
+
+**Verification:** bf16 GGUF vs f32 HF reference (`transformers==4.57.6`,
+recurrent path forced for every token): per-layer conv/SSM/K/V states match
+within bf16 rounding over the full 23-token prompt; argmax matches at every
+prompt position except one knife-edge tie (ref top-2 margin 0.03 vs bf16
+logit noise 0.19). First frame: argmax 2732 (logit 26.524 vs ref 26.557).
+End-to-end: synthesized "你好" transcribes back as "你好。" (Qwen3-ASR), on
+both CPU and Metal backends.
+
+## Resolved Issue 2 — recurrent SSM state blow-up on long sequences
+
+**Symptom (before):** ~180 tokens in, per-layer states reached 1e15..1e18,
+then logits went to zero / NaN.
+
+**Root cause:** not the recurrent scan itself — the corrupted KV cache
+(Issue 1, cause 3) fed garbage attention output into the residual stream,
+which drove `x`/`dt` of the Mamba2 branch into regime where the state
+exploded. The f32 HF reference running the same recurrent math stays bounded
+(states ~270 over the prompt), which ruled out the "inherent weak-decay"
+theory previously recorded here.
+
+**Verification:** 140-character text → 525 generated frames (position 605):
+max per-layer SSM state ≈ 1.1e3, zero NaN, clean EOS, and the audio
+transcribes back to the input text verbatim. No chunked-scan or dt-clamp
+mitigation was needed; HF's recurrent fallback does not apply the
+`time_step_min/max` clamp either (`time_step_limit` is hardcoded
+`(0.0, inf)` in `modeling_falcon_h1.py`).
+
+## Notes for the next agent
+
+- The parity harness used for the fix (an HF golden-dump script forcing the
+  recurrent path token-by-token, plus a differ) was session tooling and is not
+  committed. To rebuild it: run `modeling_arktts` under `transformers==4.57.6`
+  with each `layer.mamba.forward` replaced by the `use_precomputed_states`
+  recurrent branch, dump `cache.conv_states/ssm_states/key_cache` per step,
+  and diff against `state.ssm_states/conv_states/k_cache/v_cache` read back in
+  `falcon_forward_step` (same flat layouts: ssm `[s + 64d + 2048h]`, conv rows
+  oldest→newest, kv `d + 64*(t + T*h)`).
+- `A_log` / `D` / `dt_bias` / `conv1d.weight` must load with
+  `assets::TensorStorageType::F32` (GGUF stores them quantized; `Native`
+  keeps the quantized type and `ggml_backend_tensor_get` then reads out of
+  bounds).
+- The GGUF layout for `conv1d.weight` is `[d_conv, 1, conv_dim]`, which is
+  the HF `[conv_dim, 1, d_conv]` weight with unchanged flat bytes. Feed it to
+  `ggml_ssm_conv` **unflipped** (see Resolved Issue 1, cause 1).
+- `ggml_set_output` on a graph intermediate pins its buffer so host
+  read-back is safe — but mass-pinning hundreds of tensors corrupts the
+  whole graph (all-zero logits). Pin only the tensors actually read back.
+- Reference environment: `uv venv` + `uv pip install torch "transformers>=4.57,<5"`;
+  `transformers>=5` renames `FalconHybridMambaAttentionDynamicCache` and
+  breaks the 0.1B remote code.
+- Known remaining gap: the fast-AR codebooks during generation mean the C++
+  rollout cannot be compared token-by-token against a reference that feeds
+  zero codebook rows; parity was established over the prompt + first frame.

external/ggml/include/ggml.h

@@ -604,6 +604,9 @@ extern "C" {
         GGML_OP_MUL_MAT_ADD,
         GGML_OP_MUL_MAT_ADD_RELU,
         GGML_OP_IM2COL_ASYM,
+        // audio8_tts codec per-tap accumulation and fused snake (audio8 PR).
+        GGML_OP_MUL_MAT_ACC,
+        GGML_OP_SNAKE_1D,
 
         GGML_OP_COUNT,
     };
@@ -1447,6 +1450,21 @@ extern "C" {
             struct ggml_context * ctx,
             struct ggml_tensor  * a,
             struct ggml_tensor  * b);
+
+    // accumulate matrix multiplication in-place: acc += a * b
+    // result is a view of acc (which must have the shape of a * b), so the
+    // accumulation lands directly in acc's memory without a separate add pass
+    GGML_API struct ggml_tensor * ggml_mul_mat_acc(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * b,
+            struct ggml_tensor  * acc);
+
+    // fused snake activation: dst = a + sin(a * alpha)^2 / alpha, alpha broadcast per channel
+    GGML_API struct ggml_tensor * ggml_snake_1d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            struct ggml_tensor  * alpha);
 
     GGML_API struct ggml_tensor * ggml_mul_mat_pack4(
             struct ggml_context * ctx,

external/ggml/src/ggml-cpu/ggml-cpu.c

@@ -1708,6 +1708,93 @@ static void ggml_compute_forward_mul_mat_id(
     }
 }
 
+// reference implementation of the accumulate-in-place matmul (dst aliases src[2]):
+// dst += src0 * src1. The op is only exercised on Metal; this plain single-threaded
+// loop exists so the CPU backend stays correct if a graph containing it is ever run.
+static void ggml_compute_forward_mul_mat_acc(
+        const struct ggml_compute_params * params,
+        struct ggml_tensor * dst) {
+    const struct ggml_tensor * src0 = dst->src[0]; // a [K, M]
+    const struct ggml_tensor * src1 = dst->src[1]; // b [K, N]
+
+    GGML_ASSERT(dst->type  == GGML_TYPE_F32);
+    GGML_ASSERT(src0->type == GGML_TYPE_F32);
+    GGML_ASSERT(src1->type == GGML_TYPE_F32);
+
+    if (params->ith != 0) {
+        return;
+    }
+
+    const int64_t K = src0->ne[0];
+    const int64_t M = src0->ne[1];
+    const int64_t N = src1->ne[1];
+
+    GGML_ASSERT(src1->ne[0] == K);
+    GGML_ASSERT(dst->ne[0] == M && dst->ne[1] == N);
+    GGML_ASSERT(src0->ne[2] == 1 && src0->ne[3] == 1);
+    GGML_ASSERT(src1->ne[2] == 1 && src1->ne[3] == 1);
+    GGML_ASSERT(dst->ne[2]  == 1 && dst->ne[3]  == 1);
+
+    const char * A = (const char *) src0->data;
+    const char * B = (const char *) src1->data;
+    char       * C = (char *)       dst->data;
+
+    for (int64_t n = 0; n < N; ++n) {
+        for (int64_t m = 0; m < M; ++m) {
+            float sum = 0.0f;
+            for (int64_t k = 0; k < K; ++k) {
+                const float av = *(const float *) (A + k*src0->nb[0] + m*src0->nb[1]);
+                const float bv = *(const float *) (B + k*src1->nb[0] + n*src1->nb[1]);
+                sum += av * bv;
+            }
+            float * cv = (float *) (C + m*dst->nb[0] + n*dst->nb[1]);
+            *cv += sum;
+        }
+    }
+}
+
+// ggml_compute_forward_snake_1d
+//
+// fused snake activation: y = x + sin(x * alpha)^2 / alpha, alpha broadcast per channel.
+// naive single-threaded reference so the CPU backend stays correct if a graph containing
+// this op is ever run there.
+static void ggml_compute_forward_snake_1d(
+        const struct ggml_compute_params * params,
+        struct ggml_tensor * dst) {
+    const struct ggml_tensor * src0 = dst->src[0]; // x [C, T], channels on the fast axis
+    const struct ggml_tensor * src1 = dst->src[1]; // alpha [C, 1]
+
+    if (params->ith != 0) {
+        return;
+    }
+
+    GGML_ASSERT(dst->type  == GGML_TYPE_F32);
+    GGML_ASSERT(src0->type == GGML_TYPE_F32);
+    GGML_ASSERT(src1->type == GGML_TYPE_F32);
+    GGML_ASSERT(ggml_is_contiguous(src0));
+    GGML_ASSERT(ggml_is_contiguous(src1));
+    GGML_ASSERT(ggml_is_contiguous(dst));
+    GGML_ASSERT(src0->ne[2] == 1 && src0->ne[3] == 1);
+    GGML_ASSERT(src1->ne[0] == src0->ne[0] && src1->ne[1] == 1);
+
+    const int64_t nc = src0->ne[0];
+    const int64_t nt = src0->ne[1];
+
+    const float * x = (const float *) src0->data;
+    const float * a = (const float *) src1->data;
+    float       * y = (float       *) dst->data;
+
+    for (int64_t t = 0; t < nt; t++) {
+        for (int64_t c = 0; c < nc; c++) {
+            const float av = a[c];
+            const float xv = x[t*nc + c];
+            const float ax = xv * av;
+            const float s  = sinf(ax);
+            y[t*nc + c] = xv + (s*s)/av;
+        }
+    }
+}
+
 /////////////////////////////////
 
 static void ggml_compute_forward(struct ggml_compute_params * params, struct ggml_tensor * tensor) {
@@ -1840,6 +1927,14 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_mul_mat(params, tensor);
             } break;
+        case GGML_OP_MUL_MAT_ACC:
+            {
+                ggml_compute_forward_mul_mat_acc(params, tensor);
+            } break;
+        case GGML_OP_SNAKE_1D:
+            {
+                ggml_compute_forward_snake_1d(params, tensor);
+            } break;
         case GGML_OP_MUL_MAT_ID:
             {
                 ggml_compute_forward_mul_mat_id(params, tensor);
@@ -2339,6 +2434,16 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
             {
                 n_tasks = n_threads;
             } break;
+        case GGML_OP_MUL_MAT_ACC:
+            {
+                // reference implementation is single-threaded
+                n_tasks = 1;
+            } break;
+        case GGML_OP_SNAKE_1D:
+            {
+                // reference implementation is single-threaded
+                n_tasks = 1;
+            } break;
         case GGML_OP_GET_ROWS:
         case GGML_OP_SET_ROWS:
             {

external/ggml/src/ggml-metal/ggml-metal-device.cpp

@@ -352,6 +352,26 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_unary(ggml_metal
     return res;
 }
 
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_snake_1d(ggml_metal_library_t lib, const ggml_tensor * op) {
+    GGML_ASSERT(op->op == GGML_OP_SNAKE_1D);
+    GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32);
+    GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32);
+    GGML_ASSERT(op->type == GGML_TYPE_F32);
+
+    char base[256];
+    char name[256];
+
+    snprintf(base, 256, "kernel_snake_1d_f32");
+    snprintf(name, 256, "%s", base);
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
+    }
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glu(ggml_metal_library_t lib, const ggml_tensor * op) {
     GGML_ASSERT(ggml_is_contiguous_1(op->src[0]));
 
@@ -814,6 +834,72 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm(ggml_meta
     return res;
 }
 
+// accumulate-in-place variant of kernel_mul_mm (tensor-core path only): identical tiling and
+// threadgroup usage to ggml_metal_library_get_pipeline_mul_mm, just a different kernel that adds
+// the result tile into the destination instead of overwriting it.
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_acc(ggml_metal_library_t lib, const ggml_tensor * op) {
+    char base[256];
+    char name[256];
+
+    const ggml_type tsrc0 = op->src[0]->type;
+    const ggml_type tsrc1 = op->src[1]->type;
+
+    const bool bc_inp = op->src[0]->ne[0] % 32 != 0;
+
+    constexpr int NRA = SZ_SIMDGROUP * N_MM_BLOCK_Y * N_MM_SIMD_GROUP_Y;
+    constexpr int NRB = SZ_SIMDGROUP * N_MM_BLOCK_X * N_MM_SIMD_GROUP_X;
+
+    const bool bc_out = (op->ne[0] % NRA != 0 || op->ne[1] % NRB != 0);
+
+    GGML_ASSERT(op->src[1]->ne[2] <= INT16_MAX && op->src[1]->ne[3] <= INT16_MAX);
+    const int16_t ne12 = (int16_t) op->src[1]->ne[2];
+    const int16_t ne13 = (int16_t) op->src[1]->ne[3];
+    const int16_t r2   = (int16_t) (ne12 / op->src[0]->ne[2]);
+    const int16_t r3   = (int16_t) (ne13 / op->src[0]->ne[3]);
+
+    snprintf(base, 256, "kernel_mul_mm_acc_%s_%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1));
+    snprintf(name, 256, "%s_bci=%d_bco=%d_ne12=%d_ne13=%d_r2=%d_r3=%d",
+             base, bc_inp, bc_out, ne12, ne13, r2, r3);
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+    if (!res.pipeline) {
+        ggml_metal_cv_t cv = ggml_metal_cv_init();
+
+        ggml_metal_cv_set_bool(cv, bc_inp, FC_MUL_MM + 0);
+        ggml_metal_cv_set_bool(cv, bc_out, FC_MUL_MM + 1);
+        ggml_metal_cv_set_int16(cv, ne12,  FC_MUL_MM + 2);
+        ggml_metal_cv_set_int16(cv, ne13,  FC_MUL_MM + 3);
+        ggml_metal_cv_set_int16(cv, r2,    FC_MUL_MM + 4);
+        ggml_metal_cv_set_int16(cv, r3,    FC_MUL_MM + 5);
+
+        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
+
+        ggml_metal_cv_free(cv);
+    }
+
+    const bool has_tensor = ggml_metal_device_get_props(ggml_metal_library_get_device(lib))->has_tensor;
+
+    if (has_tensor) {
+        res.nr0 = NRA;
+        res.nr1 = NRB;
+
+        // threadgroup memory holds the dequantized A tile only (the epilogue accumulates
+        // through per-thread registers, no extra shared memory)
+        res.smem = NRA * N_MM_NK_TOTAL * sizeof(ggml_fp16_t);
+    } else {
+        res.nr0 = 64;
+        res.nr1 = 32;
+
+        // the accumulate epilogue always stages the result tile through threadgroup memory
+        // (NR0 * NR1 floats), which subsumes the sa/sb region
+        res.smem = 8192;
+    }
+
+    res.nsg = N_MM_SIMD_GROUP_X * N_MM_SIMD_GROUP_Y;
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_metal_library_t lib, const ggml_tensor * op) {
     GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
     GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne);