lfm3_omni: LFM3-Omni-2B on Metal, speech in, text and speech out (part of #1627)

docs/recipes/README.md

@@ -24,6 +24,7 @@ For agent-assisted contributions, read the [TensorFold contributor skill](SKILL.
 | Kolibri-1 | [Native CUDA, block FP8, research checkpoint](kolibri-1.md) |
 | LFM2.5-24B-A2B | [Native Metal, MLX 4-bit, shared rounds](lfm2.5-24b-a2b.md) |
 | Spark-X2.6-4B | [Native CUDA](spark-x2.6-4b-cuda.md), [native Metal, the bf16 release, contexts to 1M](spark-x2.6-4b-metal.md) |
+| LFM3-Omni-2B | [Native Metal, speech in, text and speech out](lfm3-omni-metal.md) |
 
 ## Capability floor
 

docs/recipes/lfm3-omni-metal.md

@@ -0,0 +1,112 @@
+# LFM3-Omni-2B speech to speech, native Metal
+
+`lfm3_omni` runs LiquidAI's [`LiquidAI/LFM3-Omni-2B`](https://huggingface.co/LiquidAI/LFM3-Omni-2B) (revision
+`5797d07c125503bd4f4c0e7f915ae2d3d703e8de`) on Apple silicon: a spoken turn in, interleaved text and 24 kHz speech out,
+streamed. The weights are the bf16 release as stored, the bundled Liquid codec (`audio_tokenizer/`) included. The
+engine is `zig/src/families/lfm3_omni/metal/` over one metallib (`zig/kernels/metal/lfm3_omni/`). Checked on an M5
+Ultra; chips without tensor units (M1 to M4) are untested. The vision tower is not loaded.
+
+```bash
+hf download LiquidAI/LFM3-Omni-2B --revision 5797d07c125503bd4f4c0e7f915ae2d3d703e8de --local-dir "$HOME/models/lfm3-omni-2b"
+# the Metal compiler: xcodebuild -downloadComponent MetalToolchain (macOS 26 and later), with Xcode selected
+zig build lfm3-omni-generate -Doptimize=ReleaseFast -Dcpu=apple_m1
+# one turn: a WAV in (any rate, resampled to 16 kHz), the reply as a WAV and its text
+zig-out/bin/lfm3-omni-generate --model "$HOME/models/lfm3-omni-2b" --wav question.wav --out reply.wav \
+  --system "Respond with interleaved text and audio. You are a helpful voice assistant."
+# a server: POST /v1/audio/chat, GET /health
+zig-out/bin/lfm3-omni-generate --model "$HOME/models/lfm3-omni-2b" --serve 127.0.0.1:8441 --system-file system.txt
+```
+
+About 4.4 GB of unified memory.
+
+## The endpoint
+
+`POST /v1/audio/chat` takes JSON: `audio` (a base64 WAV, required), `system`, `stream` (default true),
+`audio_temperature` (0.4), `audio_repetition_penalty` (1.2), `seed`, `max_new` (generated slots, a text token or a
+speech frame each). With `stream` the answer is server-sent events as the reply is made:
+
+```text
+data: {"type":"text","delta":"I'm on it,"}
+data: {"type":"audio","format":"pcm16","sample_rate":24000,"data":"<base64 s16le>"}
+...
+data: {"type":"done","text":"...","frames":87,"speech_s":6.96,"first_text_ms":36.0,"first_audio_ms":61.2,"total_ms":640.4}
+data: [DONE]
+```
+
+Without it, one JSON object with the same fields and `audio`, the whole reply as a base64 WAV. The server takes one
+request at a time; it is the tool's own (not a route of `tensorfold-native serve`, whose speech route is text in).
+Text is greedy, as the reference's `do_sample=False`; speech frames are sampled with the card's settings.
+
+## What runs
+
+- Log-mel (16 kHz, 400-point periodic Hann, 128 Slaney bins) as fp32 products on the device; the speech encoder's
+  three stride-2 convolutions, its 17 LFM2 layers and the projector.
+- The 24-layer LFM3 backbone (gated short conv, full and sliding attention), the prompt in chunks of 512 rows.
+- The depth decoder's three mask passes over a frame's 12 codebook slots, each codebook's pick on the GPU (repetition
+  penalty, greedy or a Gumbel-max draw, exactly a softmax sample).
+- The Liquid codec's decoder and its spectrum head, the inverse STFT as an fp32 product; overlap-add on the host.
+- Products of many rows (prompt chunks, the encoder, codec chunks) on the tensor unit through the prefill GEMM's
+  blocks (`zig/kernels/metal/prefill/gemm_nax.metal`), fp32 sums; a GPU without tensor units takes a
+  simdgroup-matrix product (`LFM3_OMNI_NAX=0` forces it).
+- Up to 16 rows (a backbone row, the depth decoder's 12 slots): a weight-streaming product, the heads' norm, rotary
+  and cache write in one dispatch, one-row attention with LFM3's output gate inside, SwiGLU as the product's
+  epilogue; the greedy text pick in two stages.
+- One serial command buffer from one copy back to the next: a backbone row runs while the host does the overlap-add;
+  the first frame is decoded alone, then every 8.
+
+Every op rounds to bf16 where the reference's op rounds. `LFM3_OMNI_PROFILE=1` logs each phase's GPU time per reply.
+
+## Checks
+
+`tools/lfm3_omni/` holds the checks' clients (Python dev tools, not part of the engine): `make_utterances.sh` makes
+the 16 synthetic spoken turns of `utterances.txt` with macOS voices; `reference_dump.py` and `reference_depth.py` dump
+the reference's stages (transformers with the checkpoint's own code); `parity.sh` runs `lfm3-omni-generate --dump`,
+`--force` and `--depth-check` against them and `summarize_parity.py` prints the table below;
+`speech_chat_client.py` streams turns through the server and times them; `reference_speech.py` does the same with the
+reference runtime; `score_replies.py` scores replies (UTMOS22 strong through SpeechMOS, and the checkpoint's own ASR
+against each reply's text).
+
+## Receipt
+
+M5 Ultra (256 GB), macOS 27.0.1, Xcode 27.1; the reference is `transformers` 5.21.0 with the checkpoint's own code on
+PyTorch 2.14.1 (MPS, bf16), on the same Mac. The 16 turns of `tools/lfm3_omni/utterances.txt` (macOS voices), a
+voice-assistant system text, identity codebook order, greedy text and codes:
+
+| stage | check | Metal vs MPS bf16 |
+| --- | --- | --- |
+| log-mel | cosine (min), max abs | 1.000000, 3.1e-4 |
+| speech encoder | cosine (min of 16) | 0.99999 |
+| projector | cosine (min) | 0.99990 |
+| prompt | token ids | 16/16 identical |
+| backbone prefill | row cosine (median of medians, min of medians) | 0.99982, 0.99977 |
+| text | teacher-forced greedy top-1 | 520/524 (99.2%) |
+| depth decoder | on the reference's own states: greedy agreement, logits cosine median / min | 95.5%, 0.999994 / 0.99989 |
+| codes | teacher-forced through the whole loop | 89.2% |
+| codec | the reference's codes to audio: SNR median (min); spectrum rows cosine (min) | 36.6 dB (25.1 dB); 0.99994 |
+
+Text and code misses are near ties (within 2 bf16 ulps of the reference's pick); 15% of the reference's own greedy
+codebook picks are exact bf16 ties, so free-running codes part after a frame or two. Against MPS fp32 the port sits
+where the reference's own bf16 run sits (encoder 0.999995, prefill rows 0.99988 median, for both).
+
+Replies with the card's audio settings (temperature 0.4, penalty 1.2, seeds 1000 + i, the server and
+`speech_chat_client.py`; the reference through `reference_speech.py`, its first audio being the first sampled frame
+plus that frame's codec decode). Two rounds, the servers and the reference alternated; other services were resident
+on the Mac but idle (runs paused while its interactive voice service was in use):
+
+| | TensorFold Metal | reference runtime (MPS bf16) |
+| --- | --- | --- |
+| first text, median | 0.034 s | |
+| first audio, median | 0.063 s | 0.172 to 0.191 s |
+| whole reply, median (6.5 s of speech) | 0.67 to 0.68 s | 2.41 to 2.54 s |
+| UTMOS22 strong, median (min) | 4.43 (4.35) | 4.44 (4.15) |
+| speech vs its own text, WER | 1.5% | 1.2% |
+
+WER is the checkpoint's own ASR on each reply's audio against the reply's streamed text.
+
+
+## Licence
+
+The weights are under Liquid AI's LFM Open License v1.0 (`license: other`, `lfm1.0` on the model card; the text is the
+`LICENSE` file in the checkpoint). Use, changes and redistribution are allowed; commercial use is licensed only to an
+organisation under US$10 million in annual revenue, and a qualified non-profit may use it for research or other
+non-commercial purposes regardless. TensorFold's code stays Apache-2.0.

tools/lfm3_omni/compare_stages.py

@@ -0,0 +1,78 @@
+"""Stage parity of lfm3-omni-generate's dumps against reference_dump.py's: python compare_stages.py REF_DIR TF_DIR."""
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+
+MIXED_START = 125015
+
+
+def stat(a, b):
+    """Cosine, relative L2 and largest absolute difference in float64."""
+    a, b = a.astype(np.float64).ravel(), b.astype(np.float64).ravel()
+    cos = float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-30))
+    rel = float(np.linalg.norm(a - b) / (np.linalg.norm(a) + 1e-30))
+    return {"cos": round(cos, 6), "rel_l2": round(rel, 5), "max_abs": round(float(np.abs(a - b).max()), 5)}
+
+
+def first_split(a, b):
+    """The length of the common prefix of two sequences."""
+    m = min(len(a), len(b))
+    return next((i for i in range(m) if np.any(a[i] != b[i])), m)
+
+
+def compare(r, t):
+    """Every stage the two dumps share."""
+    def raw(name, dtype=np.float32):
+        p = t / name
+        return np.fromfile(p, dtype=dtype) if p.exists() else None
+
+    out = {}
+    mel = np.load(r / "mel.npy")
+    tm = raw("mel.f32").reshape(128, -1)
+    out["mel"] = stat(mel, tm[:, : mel.shape[1]]) | {"frames": [mel.shape[1], tm.shape[1]]}
+    enc = np.load(r / "enc.npy")
+    te = raw("enc.f32").reshape(-1, enc.shape[1])
+    out["encoder"] = stat(enc, te) | {"tokens": [enc.shape[0], te.shape[0]]}
+    emb = np.load(r / "audio_emb.npy")
+    out["projector"] = stat(emb, raw("audio_emb.f32").reshape(emb.shape))
+    ids, ti = np.load(r / "input_ids.npy"), raw("input_ids.i32", np.int32)
+    out["prompt_ids_equal"] = bool(len(ti) == len(ids) + 1 and (ti[: len(ids)] == ids).all() and ti[-1] == MIXED_START)
+    ph = np.load(r / "prefill_hidden.npy")
+    tph = raw("prefill_hidden.f32").reshape(-1, ph.shape[1])
+    n = ph.shape[0]
+    out["prefill_hidden"] = stat(ph, tph[:n]) | {"last_row": stat(ph[-1], tph[n - 1])}
+    out["prefill_hidden"]["min_row_cos"] = min(stat(ph[i], tph[i])["cos"] for i in range(n))
+    dl = np.load(r / "depth_logits0.npy")
+    tdl = raw("depth_logits0.f32").reshape(dl.shape)
+    out["depth_logits_frame0"] = stat(dl, tdl) | {"argmax_equal": int((dl.argmax(-1) == tdl.argmax(-1)).sum())}
+    seq = np.load(r / "gen_seq.npy")
+    ts = np.concatenate([[MIXED_START], raw("gen_seq.i32", np.int32)])
+    div = first_split(seq, ts)
+    out["free_run_slots"] = {"ref": int(len(seq)), "tf": int(len(ts)), "identical_prefix": int(div)}
+    codes, tc = np.load(r / "gen_codes.npy").T, raw("gen_codes.i32", np.int32).reshape(-1, 12)
+    same = int(first_split(codes, tc))
+    out["free_run_frames"] = {"ref": int(len(codes)), "tf": int(len(tc)), "identical_prefix": same}
+    sh, tsh = np.load(r / "step_hidden.npy"), raw("step_hidden.f32").reshape(-1, ph.shape[1])
+    k = min(len(sh), len(tsh), div)
+    min_cos = round(min(stat(sh[i], tsh[i])["cos"] for i in range(k)), 6) if k else None
+    out["step_hidden_before_divergence"] = {"steps": int(k), "min_cos": min_cos, "first": stat(sh[0], tsh[0])}
+    cl = raw("codec_lin.f32")
+    if cl is not None and (r / "codec_lin.npy").exists():
+        rl = np.load(r / "codec_lin.npy")
+        cl = cl.reshape(-1, rl.shape[1])
+        k = min(len(rl), len(cl))
+        out["codec_lin"] = stat(rl[:k], cl[:k]) | {"rows": [len(rl), len(cl)]}
+        rw, tw = np.load(r / "reply.npy"), raw("reply.f32")
+        k = min(len(rw), len(tw))
+        snr = 10 * np.log10(np.sum(rw[:k] ** 2) / (np.sum((rw[:k] - tw[:k]) ** 2) + 1e-30))
+        out["codec_wave"] = {"samples": [len(rw), len(tw)], "snr_db": round(float(snr), 2)} | stat(rw[:k], tw[:k])
+    return out
+
+
+if __name__ == "__main__":
+    R, T = Path(sys.argv[1]), Path(sys.argv[2])
+    res = compare(R, T)
+    print(json.dumps(res, indent=1))
+    json.dump(res, open(T / "parity.json", "w"), indent=1)

tools/lfm3_omni/make_utterances.sh

@@ -0,0 +1,17 @@
+#!/bin/sh
+# The 16 synthetic spoken turns (macOS `say`, eight system voices in turn) as 16 kHz WAVs: make_utterances.sh OUT_DIR
+set -eu
+here=$(cd "$(dirname "$0")" && pwd)
+out=$1
+mkdir -p "$out"
+: > "$out/refs.tsv"
+set -- Samantha Daniel Karen Moira Rishi Tessa Fred Albert
+i=0
+while IFS= read -r line; do
+  [ -z "$line" ] && continue
+  v=$(eval echo "\${$((i % 8 + 1))}")
+  f=$(printf "u%02d.wav" "$i")
+  say -v "$v" -o "$out/$f" --data-format=LEI16@16000 "$line"
+  printf "%s\t%s\t%s\n" "$f" "$v" "$line" >> "$out/refs.tsv"
+  i=$((i + 1))
+done < "$here/utterances.txt"

tools/lfm3_omni/parity.sh

@@ -0,0 +1,20 @@
+#!/bin/sh
+# Stage parity over WAV_DIR/u*.wav: parity.sh BIN MODEL REF_DIR OUT_DIR SYSTEM_FILE WAV_DIR [DEV] [DTYPE] ($PYTHON)
+set -u
+here=$(cd "$(dirname "$0")" && pwd)
+B=$1; M=$2; REF=$3; OUT=$4; SYS=$5; WAVS=$6; DEV=${7:-cuda}; DT=${8:-bf16}; P=${PYTHON:-python3}
+busy() { [ -n "${LFM3_OMNI_BUSY:-}" ] && sh -c "$LFM3_OMNI_BUSY"; }
+for w in "$WAVS"/u*.wav; do
+  u=$(basename "$w" .wav); mkdir -p "$OUT/$u"
+  while busy; do sleep 5; done
+  "$B" --model "$M" --wav "$w" --system-file "$SYS" --audio-temperature 0 --identity-order 1 --dump "$OUT/$u" \
+    --codec-ref "$REF/$u" > "$OUT/$u/run.log" 2>&1
+  "$P" "$here/compare_stages.py" "$REF/$u" "$OUT/$u" > /dev/null
+  "$B" --model "$M" --wav "$w" --system-file "$SYS" --audio-temperature 0 --identity-order 1 --force "$REF/$u" \
+    > "$OUT/$u/forced.log" 2>&1
+  [ -f "$REF/$u/ref_depth.npy" ] || "$P" "$here/reference_depth.py" "$M" "$REF/$u" prep "$DEV" "$DT" > /dev/null
+  "$B" --model "$M" --depth-check "$REF/$u" > "$OUT/$u/depth.log" 2>&1
+  "$P" "$here/reference_depth.py" "$M" "$REF/$u" score > "$OUT/$u/depth_parity.json"
+  echo "$u $(grep FORCED "$OUT/$u/forced.log") $(cat "$OUT/$u/depth_parity.json")"
+done
+"$P" "$here/summarize_parity.py" "$REF" "$OUT"

tools/lfm3_omni/refcommon.py

@@ -0,0 +1,50 @@
+"""The reference runtime for LFM3-Omni-2B's checks: transformers with the checkpoint's own code, on cuda, mps or cpu."""
+import os
+import subprocess
+import time
+from pathlib import Path
+
+import torch
+
+
+def dtype_of(name):
+    """bf16 (the release) or f32 (the reference's own rounding noise)."""
+    return {"bf16": torch.bfloat16, "f32": torch.float32}[name]
+
+
+def sync(dev):
+    """Waits for the device so host timers cover its work."""
+    if dev == "cuda":
+        torch.cuda.synchronize()
+    elif dev == "mps":
+        torch.mps.synchronize()
+
+
+def load(model_dir, dev, dtype):
+    """The model and processor, loaded on the CPU then moved (an MPS fp32 load straight to the device can fail)."""
+    from transformers import AutoModelForMultimodalLM, AutoProcessor
+    model = AutoModelForMultimodalLM.from_pretrained(Path(model_dir), dtype=dtype, device_map="cpu",
+                                                     trust_remote_code=True)
+    model = model.eval().to(dev)
+    proc = AutoProcessor.from_pretrained(Path(model_dir), trust_remote_code=True)
+    return model, proc
+
+
+def chat(proc, system, wav, dev):
+    """The processor's inputs for a system text and one spoken turn."""
+    conv = [{"role": "system", "content": system}, {"role": "user", "content": [{"type": "audio", "audio": str(wav)}]}]
+    out = proc.apply_chat_template(conv, add_generation_prompt=True, tokenize=True, return_dict=True,
+                                   return_tensors="pt")
+    return out.to(dev)
+
+
+def wait_idle():
+    """Waits while $LFM3_OMNI_BUSY (a command) exits 0, so a shared machine's interactive work goes first."""
+    cmd = os.environ.get("LFM3_OMNI_BUSY")
+    while cmd and subprocess.run(cmd, shell=True).returncode == 0:
+        time.sleep(5)
+
+
+def identity_order():
+    """Greedy codes in identity codebook order: the reveal order is the only randomness left under greedy audio."""
+    torch.randperm = lambda n, *a, device=None, **k: torch.arange(n, device=device)

tools/lfm3_omni/reference_depth.py

@@ -0,0 +1,68 @@
+"""The depth decoder on the reference reply's states: reference_depth.py MODEL DUMP prep|score [DEV] [DTYPE]."""
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+
+import refcommon as rc
+
+AUDIO_OUT = 125032
+
+
+def inputs(d):
+    """The state before each reference frame and its codes, as lfm3-omni-generate --depth-check reads them."""
+    seq, sh, codes = np.load(d / "gen_seq.npy"), np.load(d / "step_hidden.npy"), np.load(d / "gen_codes.npy").T
+    slots = [j for j in range(1, len(seq)) if seq[j] == AUDIO_OUT][: len(codes)]
+    return np.stack([sh[j - 1] for j in slots]).astype(np.float32), codes[: len(slots)].astype(np.int32)
+
+
+def prep(model_dir, d, dev, dtype):
+    """Writes depth_h.f32 and depth_codes.i32, and the reference's three passes of logits per frame."""
+    h_all, c_all = inputs(d)
+    h_all.tofile(d / "depth_h.f32")
+    c_all.tofile(d / "depth_codes.i32")
+    model, _ = rc.load(model_dir, dev, dtype)
+    dd, out = model.model.depth_decoder, []
+    with torch.no_grad():
+        for h, c in zip(h_all, c_all):
+            ht = torch.from_numpy(h)[None].to(dev, dtype)
+            ct = torch.zeros((1, 12), dtype=torch.long, device=dev)
+            masked = torch.ones((1, 12), dtype=torch.bool, device=dev)
+            rows = []
+            for p in range(3):
+                rows.append(dd(ht, ct, masked)[0].float().cpu().numpy())
+                for q in range(4 * p, 4 * p + 4):
+                    ct[0, q] = int(c[q])
+                    masked[0, q] = False
+            out.append(np.stack(rows))
+    np.save(d / "ref_depth.npy", np.stack(out))
+
+
+def score(d):
+    """Greedy agreement and logits cosine of depth_check.f32 against the reference's passes."""
+    ref = np.load(d / "ref_depth.npy")
+    tf = np.fromfile(d / "depth_check.f32", dtype=np.float32).reshape(ref.shape)
+    agree = total = 0
+    cos = []
+    for p in range(3):
+        for q in range(4 * p, 4 * p + 4):
+            a, b = ref[:, p, q], tf[:, p, q]
+            agree += int((a.argmax(-1) == b.argmax(-1)).sum())
+            total += len(a)
+            for x, y in zip(a.astype(np.float64), b.astype(np.float64)):
+                cos.append(float(x @ y / (np.linalg.norm(x) * np.linalg.norm(y))))
+    r = {"frames": int(ref.shape[0]), "argmax_agree": agree, "scored": total, "rate": round(agree / total, 4),
+         "logit_cos_min": round(min(cos), 6), "logit_cos_median": round(float(np.median(cos)), 6)}
+    print(json.dumps(r))
+    json.dump(r, open(d / "depth_parity.json", "w"))
+
+
+if __name__ == "__main__":
+    D = Path(sys.argv[2])
+    if sys.argv[3] == "prep":
+        dev = sys.argv[4] if len(sys.argv) > 4 else "cuda"
+        prep(sys.argv[1], D, dev, rc.dtype_of(sys.argv[5] if len(sys.argv) > 5 else "bf16"))
+    else:
+        score(D)