Skip to content

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

Open
gqf2008 wants to merge 18 commits into
0xShug0:mainfrom
gqf2008:feat/audio8-tts-falcon-h1-01b
Open

gqf2008 wants to merge 18 commits into
0xShug0:mainfrom
gqf2008:feat/audio8-tts-falcon-h1-01b

Conversation

@gqf2008

@gqf2008 gqf2008 commented Sep 4, 2026

Copy link
Copy Markdown
Contributor

Summary

Companion to #333 (Audio8-TTS 0.6B): native support for the Audio8-TTS-Preview-0.1B checkpoint, whose AR stage is a Falcon-H1 hybrid — stateful Mamba2 + grouped-query attention per token — instead of the 0.6B's pure transformer.

Port

  • falcon_forward_step: RMSNorm → (Mamba2 ‖ GQA attention) → residual → gated FFN, stateful per token: conv/SSM states + KV cache. Mamba2 path: in_proj[z,xBC,dt] split → ssm_conv (kernel flipped to match HF causal_conv1d) → ssm_scan → D → silu(z) gate → out_proj. Attention: q/k/v → RoPE NEOX (base 1e11) → flash attention with KV cache.
  • Multiplier semantics fixed against the reference: lm_head_multiplier does not apply to the compact semantic head; embedding_multiplier applies to (text_emb + codebook_sum) jointly.
  • Decisive correctness bug: the host KV cache used a seq-scaled head stride while appending only the new token, so from the second step on, every head > 0 read corrupted K/V. Fixed in falcon_kv_cache.h + regression test.

Performance (M4, Metal, q8_0)

  • 0.1B session RTF ~1.3 → ~0.34 (19s audio in ~5.9s excluding load), byte-identical output across runs (seed-pinned).
  • AR: per-token graph moved to a dedicated CPU backend (600-node graphs were dispatch-bound on Metal), then zero-copy state/constants, padded host KV buffers with in-graph slot writes, and per-capacity-bucket graph reuse (~781ms/generation compute, ~2ms overhead).
  • codec: stride-1 conv1d Metal fast path on time-fast layouts, decoder blocks channel-fast (saves 46 transposes + 26 bias repeats per pass), new GGML_OP_MUL_MAT_ACC (accumulate-in-place per-tap GEMMs), and a fused GGML_OP_SNAKE_1D op replacing the 5-kernel/11-pass activation chain (codec compute −21.7%). 0.6B codec decode 4.8s → ~2.25s at the first step of this series; all gated work paths kept byte-identical output.
  • ggml: ssm_scan reduction read garbage when sgptg < NW (Metal); ggml-metal.metal CRLF renormalized first for a clean diff.

Verification

  • Logits argmax parity vs transformers fp32 trust_remote_code reference (first-frame argmatch; bf16-GGUF noise stays below top-2 gap except known ties).
  • ASR round-trip (Qwen3-ASR): synthesized "你好" → "你好。"; 605-position long run: SSM state peak ~1.1e3, zero NaN, on both CPU and Metal backends.
  • Every perf step verified output byte-identical (or env-gated fallback identical) to the prior step; regression test audio8_tts_falcon_kv_cache_test covers the KV stride bug.
  • Status doc: docs/community_models/audio8_tts_falcon_h1_status.md (root causes for the three mid-port misdiagnoses included).

…d port

Implement the previously-stubbed Falcon-H1 slow-AR path:
- stateful per-token forward (falcon_forward_step): RMSNorm -> (Mamba2 || GQA
  attention) -> residual -> gated FFN, with conv/SSM states and KV cache
- Mamba2: in_proj [z,xBC,dt] split -> ssm_conv (kernel flipped to match HF
  causal_conv1d) -> ssm_scan -> D -> silu(z) gate -> out_proj
- GQA attention: q/k/v proj -> RoPE NEOX (base 1e11) -> flash_attn with KV cache
- fix: lm_head_multiplier does NOT apply to ArkttsModel compact semantic head
- fix: embedding_multiplier applies to (text_emb + codebook_sum) together

Verified against transformers reference: logits scale/argmax match; CPU runs
627 tokens stable.

(The ggml-metal.metal change from the original af7bcd4 was split into the two
following commits: a mechanical LF renormalize, then the ssm_scan fix.)
- ggml_set_output(scan): keep the SSM state tail alive for host read-back
  (gallocr was reusing the scan result buffer, corrupting the read state)
- conv1d flip: GGUF layout is [d_conv,1,conv_dim], not HF [conv_dim,1,d_conv]
- embedding_multiplier applies to (text_emb + codebook_sum) together
- lm_head_multiplier does NOT apply to ArkttsModel compact semantic head
- ggml-metal ssm_scan: fix reduction garbage for d_state>32, n_t<sgptg

Result: 0.1b logits match transformers reference scale/argmax; CPU runs
185 tokens before a dt-clamping-related zero-out. Still TODO: clamp dt_softplus
to [time_step_min, time_step_max] to prevent state death on long sequences.
Document the two remaining issues (logits argmax mismatch vs reference,
recurrent SSM state blow-up) plus the debugging approach and cleanup notes,
for hand-off to a follow-up agent.
The fast (audio-codebook) AR runs one graph submission per generated
codebook token, making it submit+sync latency bound on GPU backends:
~2.9ms/step on Metal versus ~0.6ms/step for the same graph on CPU.
When the main backend is a GPU, give the fast AR a dedicated CPU
backend: byte-copy the fast-layer projections and fast_output onto it
(same ggml type, lossless for q8_0/f16/f32), build the fast graph,
state buffers, and constants cache there, and route run() at it. The
slow path (Qwen or Falcon-H1) and codec stay on the main backend; the
only cross-backend traffic is small host vectors.

Measured on M4/Metal, q8_0, 62-char prompt (278 frames):
- 0.1B: fast AR 7.28s -> 1.68s, session 16.1s -> 10.2s (RTF ~1.3 -> ~0.8)
- 0.6B: fast AR 8.6s  -> 3.39s, session 17.4s -> 12.7s (RTF 1.26 -> ~1.0)
- --backend cpu unchanged; ASR round-trip verbatim on both backends
Split audio8_tts.codec_decode_ms into graph_build_ms (graph
construction, weight upload, gallocr reserve) and graph_compute_ms
(submit + GPU execution), and log the node count. On M4/Metal q8_0
the 1105-node decode graph builds in ~200ms and computes in ~4.8s,
so per-phase attribution now points at GPU execution directly.
The audio codecs store activations as [frames, channels] (time-fast),
the transpose of the LLM layout ggml kernels are tuned for. On that
layout ggml_conv_1d's im2col is a strided gather (kernel taps sit
C*4 bytes apart, ~16x read amplification), costing ~200ms per conv at
[569k, 96] on M4 -- ~4.5s of the audio8_tts codec decode.

Add a Metal-only fast path in Conv1dModule for padding=0, stride=1,
batch=1, contiguous F32 input: transpose the input to channel-fast
once, run one contiguous GEMM per kernel tap over shifted views, and
transpose the accumulator back. Weights are regrouped to per-tap rows
with a single cont(permute).

audio8_tts long-text A/B/A (normalized by fast_graph_ms load
indicator, quiet machine): codec graph compute 4970-4981ms ->
2248-2279ms (~2.2x), session wall 9.9s -> 6.6-6.9s. Same-seed output
byte-identical across runs; in-graph parity vs ggml_conv_1d
(sum_abs_diff ~ 0); ASR round-trips unchanged.
The falcon_forward_step graph is ~600 tiny nodes per token; on Metal it
is dispatch-latency bound (~5.9 ms/step) while the same graph computes
in ~2.0 ms on CPU (measured both ways via new falcon_step_* profile
timers). Mirror the fast-AR treatment: retarget the Falcon-H1 layer
weights (incl. ssm_A/ssm_D, so the per-step A/D reads become plain CPU
memcpys instead of GPU->host syncs) and the semantic head onto the
existing dedicated CPU backend, and call falcon_forward_step with it.
--backend cpu behavior is unchanged (no retarget, same backend).

Long-text session: ar_generate 6284 -> 3040 ms, session wall ~6.8s ->
5.3s (RTF ~0.48 -> ~0.38). Audio duration byte-identical, ASR
round-trips verbatim on both backends.
…slots in-graph

The zero-copy path no longer uploads the KV cache per step nor reads the fresh
k/v back for a host-side append: per-layer caches live in geometrically grown
padded buffers ([head_dim, cap, n_kv]) bound as external leaves; the new
token's k/v are copied into slot seq by in-graph ggml_cpy (expanded before the
attention nodes so the write precedes the read on sequential backends), and
flash_attn_ext reads a strided prefix view over slots [0, seq+1) (CPU flash
only requires contiguous rows). This removes the per-step concat of the full
cache prefix as well as the host re-striding append.

Measured (M4, quiet, interleaved A/B): falcon download 92->0.8 ms,
compute -20 ms, ar_generate -100 ms; same-seed output wavs bit-identical.
The exact-size cache + concat path is kept for non-host backends.
The per-token step graph no longer depends on the sequence length: flash
attention reads the full padded KV cache under an -inf mask (masked slots
contribute exactly zero to the softmax, so the reduction is bitwise identical
to the exact-prefix one), and the fresh k/v land in their slot via
ggml_set_rows with the slot index read from host memory. The graph, its
context, and its allocator are baked once per 128-slot bucket
(FalconStepPlan) and reused for every step in the bucket; the per-step feed
is a 2 KB embedding memcpy plus three scalar writes (position, slot, mask).

Bucket sizes stay below the CPU flash kernel's split-KV threshold (512) as
long as possible so the masked padded reduction stays in the same code path
as the exact-prefix one.

Measured (M4, interleaved A/B vs the per-step-build version): falcon
build 90->0.7 ms, gallocr 60->0.6 ms, init 0.7->0.005 ms per generation
(~3 bucket rebuilds total); same-seed output wavs bit-identical on the
380-step benchmark, metal/cpu gates and ASR round-trip pass. The per-call
upload/download fallback for non-host backends is unchanged.
…ul_mat_acc

The channel-fast per-tap conv ran one ggml_mul_mat per kernel tap and folded
the partials in with ggml_add, so every non-first tap paid a temporary write
plus the add's three-way traffic (read acc + read partial + write acc). Add a
new ggml op, GGML_OP_MUL_MAT_ACC, computing acc += a * b with the result a view
of acc (the ggml_cpy in-place idiom), and switch taps 1..K-1 to it.

The new Metal kernel_mul_mm_acc mirrors kernel_mul_mm through the MMA (both the
simdgroup and tensor variants), so the partial products are bit-identical; only
the epilogue differs, folding the result tile into the destination with one F32
add per element (matching ggml_add). The simdgroup variant always stages the
tile through threadgroup memory so partial tiles clip identically; the tensor
variant loads the destination tile into a second cooperative tensor and adds
element-wise. Add-only per open-closed: new op, kernels, encoder, pipeline
getter, and supports-op cases; no existing function modified. The CPU backend
gets a naive single-threaded reference forward (the op is exercised on Metal
only; the per-tap path is Metal-gated).

Measured (M4, interleaved A/B vs the mul_mat+add chain under residual background
load): codec graph_compute_ms -~400 ms median (-~570 ms min), consistent with
the ~365 ms predicted by the skip-taps probe; output wavs bit-identical across
13/13 same-seed runs, CPU backend path untouched.
…odec

Add a dedicated ggml elementwise op computing x + sin^2(alpha*x)/alpha in a
single pass (Metal kernel + scalar CPU reference) and route the codec's
snake1d through it, replacing a 5-kernel / 11-pass elementwise chain over
the largest decoder tensors.

Audio8-TTS 0.6B on M4 (342-char zh, seed 1234, interleaved A/B n=8):
codec.graph_compute_ms 2624 -> 2056 min-to-min (-21.7%, 8/8 pairwise),
wall -669 ms; falcon fast_graph unchanged. Numerics: scalar-reference
max_abs 1.9e-6 (metal sin ulp); full-utterance wav vs the chain max
int16 delta 12. AUDIO8_TTS_CODEC_SNAKE_FUSED=0 falls back to the chain.
…ckend

Bind loop-invariant weights (norms, conv kernel/bias, A=-exp(A_log), expanded
D) and all recurrent state (conv, ssm, KV cache) to the per-token step graph
as external views of their host vectors; gallocr skips tensors whose data is
set externally, so the per-step upload collapses from ~300 ms/generation to
~0.15 ms. The conv/ssm next-state write-backs now run in-graph as ggml_cpy
into the host state vectors (each cpy depends on the nodes that consumed the
old state, so reads strictly precede writes), eliminating the per-step state
read-back. A/D are resolved once per generation instead of via four
tensor_get calls per layer per step (the build-time pair was dead code and is
removed). Embedding/ids/pos are bound directly as well.

Measured (M4, interleaved A/B under background load, fast_graph as the load
indicator): falcon upload 300->0.15 ms, download 182->107 ms, compute +~170 ms
(the in-graph write-backs), ar_generate -150..-350 ms; same-seed output wavs
are bit-identical, metal/cpu gates and ASR round-trip pass. Non-CPU backends
keep the previous explicit upload/download path.
The per-tap conv fast path still paid two transposes per conv plus a
materialized bias repeat: 52 transposes + 26 repeats per decode. The
ops between convs (snake, residual add) are all elementwise and hence
layout-agnostic, so the whole block region -- snake, convT upsample,
three residual units -- can run channel-fast [channels, frames] with
transposes only at region edges (6 instead of 52).

- conv_modules: extract the per-tap GEMM core; expose raw channel-fast
  helpers conv1d_pertap_channel_fast (bias via broadcast add, no
  repeat) and conv_transpose1d_col2im_channel_fast (skips the col2im
  path's internal transpose). Guard the module fast path on F32 weights
  (the per-tap weight views assume 4-byte rows).
- codec: chain each decoder block through the raw helpers behind an
  env kill-switch (AUDIO8_TTS_CODEC_CHANNEL_FAST=0), causal pad via
  scaled-to-zero prefix columns + concat; falls back to the module path
  elsewhere.

Same-seed output is byte-identical to the previous commit across
repeated runs (pure data-movement change). Interleaved A/B under
background load, fast_graph_ms-normalized: codec graph compute ratio
1.73 -> 1.45 (~550ms saved).
Each token's output is the sum over all simdgroups of that token's partial
sums (shared_sums[t*NW + g] for g in 0..sgptg-1). The previous
simd_sum(shared_sums[sgitg*NW + tiisg]) read garbage columns whenever
sgptg < NW (e.g. d_state=64 -> sgptg=2) with few tokens, corrupting the
SSM state. Compute the token sum redundantly on every thread instead.

Split out of af7bcd4 so the fix is reviewable on its own.
Three root causes made the Falcon-H1 slow-AR path diverge from the
transformers reference (first-frame argmax 3620 vs 2732, ASR round-trip
said the wrong word) and blew the SSM state up on long sequences:

1. conv1d kernel was flipped at load time, but ggml ssm_conv and the HF
   FalconH1 decode (nn.Conv1d prefill and the cached torch.sum path
   alike) use the same cross-correlation orientation with window[0] the
   oldest frame. Feed the GGUF kernel unflipped.
2. sx (conv window) and k_r/v (fresh K/V) were read back on the host
   without ggml_set_output, so gallocr reused their buffers and the
   conv/SSM/KV states were fed garbage every step. Pin exactly those
   three tensors per layer.
3. The host KV cache used the current sequence length as the per-head
   stride while appending only the new token, so from the second token
   on the append overwrote the previous head blocks and every head past
   the first read corrupted context. The resulting residual-stream
   garbage - not the recurrent scan itself - drove the SSM state to
   1e15..1e18 around token 180. append_falcon_kv_token now re-lays the
   cached tokens into the new stride before appending.

Verified against an f32 transformers 4.57.6 reference forced onto the
recurrent path: per-layer conv/SSM/KV states match within bf16 rounding
over the whole prompt, first-frame argmax = 2732, synthesized speech
round-trips through ASR verbatim, and a 605-position generation stays
bounded (max state ~1.1e3, no NaN) on both CPU and Metal backends.

Adds a regression test for the KV cache re-stride that fails against
the previous append logic at the second token.
@gqf2008
gqf2008 force-pushed the feat/audio8-tts-falcon-h1-01b branch from 9c852cb to fcba3a4 Compare September 4, 2026 06:26
@0xShug0

0xShug0 commented Sep 11, 2026

Copy link
Copy Markdown
Owner

@gqf2008 I will ping you when I'm ready to review this PR. There will be some GGML changes and I'm still testing them, which may cause conflicts with this PR.

@gqf2008

gqf2008 commented Sep 15, 2026

Copy link
Copy Markdown
Contributor Author

Updated the branch to drop a stray line-ending renormalization that was inflating this diff.

  • The PR showed +24934/-22329; the real change is +2733/-128. The extra ~22k lines came from one chore: renormalize CRLF -> LF commit plus the merge that followed it, which together rewrote external/ggml/include/ggml.h, external/ggml/src/ggml.c and external/ggml/src/ggml-metal/ggml-metal.metal end to end.
  • That commit is gone, and those three files now keep the same CRLF/LF layout as main, so the diff contains only the actual work (Falcon-H1 forward port, GGML_OP_SNAKE_1D, conv1d per-tap GEMM accumulation, Metal stride-1 conv1d fast path, ssm_scan reduction fix).
  • Content is unchanged: git diff --ignore-cr-at-eol <new head> <old head> is empty — the only difference is line endings. Messages, authors and dates are preserved (17 commits, was 18).

On your 2026-09-11 note about GGML changes conflicting here: the only external/ggml change on main since then is #508 (Vulkan fill dispatch — ggml-vulkan.cpp + fill.comp), which does not overlap this PR's ggml-metal/*, ggml.h or ggml.c. git merge-tree against current main (024297b) merges cleanly as of today.

One thing worth knowing, since it bit me here: .gitattributes declares * text=auto eol=lf, but these vendored files are CRLF/LF-mixed in every blob on main — the most recent commit touching ggml.h (62a9441) still stores 2811 CRLF lines, so git add does not renormalize them. Editing one of these files therefore rewrites the whole file unless the endings are preserved explicitly. That's why I restored the mixed layout instead of normalizing; if you'd rather have them LF everywhere, that's a one-time git add --renormalize on main and I'll keep this PR out of it.

Still based on a10738a (behind main, merges clean). Say the word if you'd prefer a rebase onto current main instead.

@0xShug0

0xShug0 commented Sep 15, 2026

Copy link
Copy Markdown
Owner

@gqf2008 #532 and #533 haven’t been merged yet...I've decided to merge them after the next stable release that includes YuE2, possibly tomorrow, to avoid introducing regressions right before the release. I will ping you.

Don’t worry about LF normalizing. I've given up on that :)

# Conflicts:
#	CMakeLists.txt
#	external/ggml/include/ggml.h
#	external/ggml/src/ggml.c
@gqf2008
gqf2008 force-pushed the feat/audio8-tts-falcon-h1-01b branch from af2e01d to b6c8458 Compare September 15, 2026 04:00
@0xShug0 0xShug0 added the ggml label Sep 15, 2026
…n-h1-01b

Resolve the GGML op-enum tail conflict: main added
CONV_3D_CONCAT_PAD_SPATIAL_GEMM, RMS_NORM_CHANNELS(_SILU),
RMS_NORM_CHANNELS_ADD_BIAS_SILU and ROPE_INTERLEAVED_PAIRS (GGML_OP_COUNT
112); this branch appends MUL_MAT_ACC and SNAKE_1D. Keep main's op order
stable and append the audio8_tts ops at the tail, with GGML_OP_COUNT and
both op-name/symbol tables in sync at 114.
@gqf2008

gqf2008 commented Sep 18, 2026

Copy link
Copy Markdown
Contributor Author

Merged current main (82b4dc3) into the branch to clear the conflict.

The conflict was the GGML op-enum tail: main added CONV_3D_CONCAT_PAD_SPATIAL_GEMM, RMS_NORM_CHANNELS(_SILU), RMS_NORM_CHANNELS_ADD_BIAS_SILU and ROPE_INTERLEAVED_PAIRS (GGML_OP_COUNT == 112) while this branch appends MUL_MAT_ACC and SNAKE_1D. Kept main's op order stable and appended the two audio8_tts ops at the tail, with GGML_OP_COUNT and both name/symbol tables in sync at 114.

Verification on this merge (macOS 26.5.2, Apple M4, Metal):

  • CMake RelWithDebInfo build with ENGINE_ENABLE_METAL=ON + ENGINE_BUILD_TESTS=ON: 958/958 targets built.
  • ctest: 35/35 passed.
  • End-to-end with the merged tree: audio8_tts 0.1B synthesizes "你好,这是一个测试。" and audio8_asr 0.1B transcribes it back as "你好,这是一个。" The Metal pipeline log shows kernel_mul_mm_acc_f32_f32(_bci) and kernel_snake_1d_f32 being compiled and dispatched, so the new ops are live on Metal.

The PR diff against main is now clean again: 18 files, +2733/-128, no line-ending noise.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants