feat(audio8_tts): Falcon-H1 0.1B (Mamba2 + attention) support and codec/AR performance - #444
Open
gqf2008 wants to merge 17 commits into
Open
feat(audio8_tts): Falcon-H1 0.1B (Mamba2 + attention) support and codec/AR performance#444gqf2008 wants to merge 17 commits into
gqf2008 wants to merge 17 commits into
Conversation
…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.)
Mechanical line-ending normalization per .gitattributes (text=auto eol=lf); no semantic change (diff -w against the parent is empty). Split out of af7bcd4 so the ssm_scan fix that follows is visible as its own hunk.
- 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
force-pushed
the
feat/audio8-tts-falcon-h1-01b
branch
from
September 4, 2026 06:26
9c852cb to
fcba3a4
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.lm_head_multiplierdoes not apply to the compact semantic head;embedding_multiplierapplies to (text_emb + codebook_sum) jointly.falcon_kv_cache.h+ regression test.Performance (M4, Metal, q8_0)
GGML_OP_MUL_MAT_ACC(accumulate-in-place per-tap GEMMs), and a fusedGGML_OP_SNAKE_1Dop 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.ssm_scanreduction read garbage whensgptg < NW(Metal);ggml-metal.metalCRLF renormalized first for a clean diff.Verification
trust_remote_codereference (first-frame argmatch; bf16-GGUF noise stays below top-2 gap except known ties).audio8_tts_falcon_kv_cache_testcovers the KV stride bug.docs/community_models/audio8_tts_falcon_h1_status.md(root causes for the three mid-port misdiagnoses included).