diff --git a/CMakeLists.txt b/CMakeLists.txt index 70a643c99..534b13a3a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1691,6 +1691,18 @@ audiocpp_add_model(breeze_tts engine::models::breeze_tts::make_breeze_tts_loader ) +audiocpp_add_model(vibeasr + SOURCES + src/community_models/vibeasr/assets.cpp + src/community_models/vibeasr/vae_encoder.cpp + src/community_models/vibeasr/lm_decoder.cpp + src/community_models/vibeasr/session.cpp + INCLUDES + engine/community_models/vibeasr/session.h + LOADERS + engine::community_models::vibeasr::make_vibeasr_loader +) + set(AUDIOCPP_ENABLED_MODELS "") if (AUDIOCPP_MODEL_SET STREQUAL "full") set(AUDIOCPP_ENABLED_MODELS ${AUDIOCPP_MODEL_TARGETS}) @@ -2500,6 +2512,56 @@ if (ENGINE_BUILD_TESTS) target_link_libraries(test_granite5asr_golden_transcription PRIVATE OpenMP::OpenMP_CXX) endif() + if (vibeasr IN_LIST AUDIOCPP_LINKED_MODELS) + add_executable(test_vibeasr_vae_encoder + tests/vibeasr/test_vibeasr_vae_encoder.cpp + ) + target_compile_definitions(test_vibeasr_vae_encoder PRIVATE + ENGINE_REPO_ROOT="${CMAKE_CURRENT_SOURCE_DIR}" + ) + target_link_libraries(test_vibeasr_vae_encoder PRIVATE engine_runtime ggml) + target_include_directories(test_vibeasr_vae_encoder PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) + if (ENGINE_ENABLE_OPENMP) + target_link_libraries(test_vibeasr_vae_encoder PRIVATE OpenMP::OpenMP_CXX) + endif() + add_test( + NAME test_vibeasr_vae_encoder + COMMAND test_vibeasr_vae_encoder + --model ${CMAKE_CURRENT_SOURCE_DIR}/models/vibeasr/vibeasr-vae-encoder-i8_s.gguf + --audio ${CMAKE_CURRENT_SOURCE_DIR}/assets/asr_validation/librispeech/librispeech_test_clean_6930-75918-0000.wav + ) + # Needs the converted 703 MB encoder package, which a normal checkout + # does not have; the probe exits 125 (skip) instead of failing. Pass + # --reference-acoustic / --reference-semantic by hand to also check + # parity against a VibeASR.cpp dump. + set_tests_properties(test_vibeasr_vae_encoder PROPERTIES + SKIP_RETURN_CODE 125 + TIMEOUT 300 + ) + + add_executable(test_vibeasr_asr + tests/vibeasr/test_vibeasr_asr.cpp + ) + target_compile_definitions(test_vibeasr_asr PRIVATE + ENGINE_REPO_ROOT="${CMAKE_CURRENT_SOURCE_DIR}" + ) + target_link_libraries(test_vibeasr_asr PRIVATE engine_runtime ggml) + target_include_directories(test_vibeasr_asr PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) + if (ENGINE_ENABLE_OPENMP) + target_link_libraries(test_vibeasr_asr PRIVATE OpenMP::OpenMP_CXX) + endif() + add_test( + NAME test_vibeasr_asr + COMMAND test_vibeasr_asr --threads 8 + ) + # Same story as the encoder probe, plus the 993 MB decoder: exits 125 + # (skip) unless both converted GGUFs sit in models/vibeasr/. + set_tests_properties(test_vibeasr_asr PROPERTIES + SKIP_RETURN_CODE 125 + TIMEOUT 600 + ) + endif() + if (audio8_asr IN_LIST AUDIOCPP_LINKED_MODELS) add_executable(test_audio8_asr_units tests/audio8_asr/test_audio8_asr_units.cpp diff --git a/README.md b/README.md index 7907c4647..aef40aa1a 100644 --- a/README.md +++ b/README.md @@ -166,6 +166,7 @@ Community model ports live under `community_models` to make the ownership bounda | **sense_asr** | ASR | auto, zh, en, yue, ja, ko, pt, ru, es, it, fr, de, nl, pl, tr, ar, hi, vi, th, id, ms, fa, nospeech | GGUF Q8, Stream | Jason Chen [@jasonchen31](https://github.com/jasonchen31), [@LauraGPT](https://github.com/LauraGPT) / FunASR | [SenseVoice-Small](docs/community_models/sense_asr.md) offline/streaming SAN-M + CTC transcription with event/emotion/language tags and ITN | | **soprano_tts** | TTS | en | GGUF Q8, Stream | [@drzsdrtfg](https://github.com/drzsdrtfg) | [Soprano-1.1-80M](https://huggingface.co/WalkingCat/Soprano-1.1-80M-GGUF) ultra-lightweight TTS with Qwen3 LM + Vocos decoder | | **vietneu_tts** | TTS, Clone | vi, en | GGUF | Phuoc [@phuocnguyen90](https://github.com/phuocnguyen90) | [VieNeu-TTS-v3-Turbo](docs/community_models/vietneu_tts.md) TTS and voice cloning support | +| **vibeasr** | ASR | en | GGUF I8_S + I2_S | [@XsquirrelC](https://github.com/XsquirrelC) | [VibeASR](docs/community_models/vibeasr.md) fully quantized port of [VibeASR.cpp](https://github.com/microsoft/VibeASR.cpp): VibeVoice acoustic/semantic tokenizers on INT8 weights and INT8 activations, feeding a ternary BitNet Qwen2 decoder. Offline, CPU only | | **voxcpm1** | TTS, Clone | zh, en, ja, ko | GGUF Q8, Stream | [@jasonchen31](https://github.com/jasonchen31) | [VoxCPM1](docs/community_models/voxcpm1.md) tokenizer-free 0.5B TTS with 16 kHz output, streaming, and continuation-mode voice cloning | ## Docker diff --git a/docs/asr.md b/docs/asr.md index ce9c92817..ffafa7b08 100644 --- a/docs/asr.md +++ b/docs/asr.md @@ -329,6 +329,11 @@ chunking, server usage, and validation notes. VibeVoice ASR is an offline ASR model with greedy, sampling, and beam-search decode paths. It can return transcription text and structured segment/speaker-turn output when the model produces timestamps. +A fully quantized port of the same model — INT8 activations through the encoder, +ternary BitNet weights in the decoder — lives under community models as +`vibeasr`: see [VibeASR](community_models/vibeasr.md). It is not a separate +model, only a CPU-only alternative numeric pipeline for the same weights. + | Field | Value | |---|---| | Family | `vibevoice_asr` | diff --git a/docs/community_models/models.md b/docs/community_models/models.md index d88164887..4c49c49a2 100644 --- a/docs/community_models/models.md +++ b/docs/community_models/models.md @@ -35,3 +35,4 @@ Practical expectations: | **sense_asr** | ASR | auto, zh, en, yue, ja, ko, pt, ru, es, it, fr, de, nl, pl, tr, ar, hi, vi, th, id, ms, fa, nospeech | Jason Chen [@jasonchen31](https://github.com/jasonchen31), [@LauraGPT](https://github.com/LauraGPT) / FunASR | [SenseVoice-Small](sense_asr.md) offline/streaming SAN-M + CTC transcription with event/emotion/language tags and ITN | | **vietneu_tts** | TTS, voice cloning | vi, en | Phuoc [@phuocnguyen90](https://github.com/phuocnguyen90) | [VieNeu-TTS-v3-Turbo](vietneu_tts.md) TTS and voice cloning support | | **moss_voicegen** | Voice design | en, zh | Joost [@jrohde](https://github.com/jrohde) | [MOSS-VoiceGenerator](moss_voicegen.md) voice design from a written instruction, on the MOSS delay architecture | +| **vibeasr** | ASR | en | [@XsquirrelC](https://github.com/XsquirrelC) | [VibeASR](vibeasr.md) fully quantized port of [VibeASR.cpp](https://github.com/microsoft/VibeASR.cpp): the VibeVoice acoustic/semantic tokenizers on INT8 weights *and* INT8 activations through the fused `GGML_TYPE_I8_S` ops, feeding a ternary `GGML_TYPE_I2_S` BitNet Qwen2 decoder. Offline, CPU only | diff --git a/docs/community_models/vibeasr.md b/docs/community_models/vibeasr.md new file mode 100644 index 000000000..c97a142a2 --- /dev/null +++ b/docs/community_models/vibeasr.md @@ -0,0 +1,364 @@ +# VibeASR in audio.cpp + +[VibeASR.cpp](https://github.com/microsoft/VibeASR.cpp) is Microsoft's CPU-first +port of the VibeVoice ASR stack, quantized end to end for edge inference: the +audio VAE encoder runs on INT8 weights *and* INT8 activations, and the Qwen2 +decoder runs on BitNet-style ternary weights. This entry ports both halves, so +`--family vibeasr` transcribes end to end on CPU. + +## Relation to the existing `vibevoice_asr` family + +audio.cpp already ships [VibeVoice ASR](../asr.md#vibevoice-asr) in the core model +tree, and it is the same model: the same acoustic/semantic causal ConvNeXt +tokenizers, the same connectors, the same Qwen2 decoder. That family runs F32 / +Q8_0 weights through the generic ggml ops. + +What VibeASR.cpp adds is a different *numeric pipeline* for that architecture, +not a different architecture: + +| | `vibevoice_asr` (core) | this entry | +|---|---|---| +| Encoder weights | F32 / Q8_0 | `GGML_TYPE_I8_S`, one F32 scale per tensor | +| Encoder activations | F32 | INT8 throughout; every stage requantizes | +| Ops | generic ggml | the five fused I8_S ops (`ggml_mul_mat_add`, `ggml_mul_mat_add_relu`, `ggml_add_scaled`, `ggml_rms_norm_scaled`, `ggml_im2col_asym`) | +| Decoder weights | Q8_0 Qwen2 | ternary `GGML_TYPE_I2_S`, 993 MB for a 1.5B decoder | +| Backends | CPU, CUDA, Metal | CPU only — the I8_S and I2_S kernels have no GPU variants | +| Decode | greedy, sampling, beam search | greedy | +| Output | text, segments, speaker turns | text | + +Both are offline-only. + +So this is an alternative execution path for weights that were quantized +upstream, useful where the INT8/ternary package is the point: no F32 activations +anywhere, integer dot products, and a decoder that fits in under 1 GB. + +It stays a separate community entry rather than becoming a weight path inside +`vibevoice_asr`, because the two share no encoder graph code: every activation +there is I8_S and every node is one of the fused CPU-only ops, so folding it in +would put a second, mutually exclusive graph builder and a second backend policy +behind one family's loader. The reuse that is worth having — tokenizer +vocabulary, prompt layout, feature-injection order, audio normalization — is data +and conventions, and this entry follows `vibevoice_asr` on all of it. The decoder +half needs no new graph code at all: it is +`modules::QwenCausalDecoderModule` unchanged, because every projection goes +through `LinearModule`'s plain `ggml_mul_mat`, which dispatches on the weight +type. + +## Architecture + +### Encoder + +Both branches are identical in shape and differ only in latent width: + +- **Input**: mono 24 kHz waveform in `[-1, 1]`, quantized to a single I8_S + tensor (one scale for the whole waveform, `amax` floored at 1e-5 to match + upstream). +- **7 stages**, strides `{1, 2, 2, 4, 5, 5, 8}` (upstream `encoder_ratios` + `[8, 5, 5, 4, 2, 2]` reversed, with a stride-1 stem), so **3200 samples per + frame** — 7.5 frames per second at 24 kHz. Channels `32 → 64 → 128 → 256 → + 512 → 1024 → 2048`, depths `3-3-3-3-3-3-8`. +- Each stage starts with a **strided causal conv** (left pad `K - stride`, right + pad 0) and then runs its ConvNeXt-style blocks: RMSNorm → depthwise conv → + layer scale → residual → RMSNorm → FC1 → ReLU → FC2 → layer scale → residual. +- **Latent head**: causal conv to `vae_dim` — 64 acoustic, 128 semantic. +- **Connector**: `FC1 → RMSNorm → FC2`, both 1536 wide, i.e. the decoder hidden + size. Output is `[frames][1536]` for each branch. + +Two details the port copies rather than corrects: + +- RMSNorm epsilon is **1e-5 everywhere**, including the norms the checkpoint + metadata labels 1e-6. Upstream hardcodes it and the published weights were + validated that way. +- The converter left-pads the 7-tap depthwise kernels with leading zeros up to a + SIMD-friendly width. Convolving with the padded width and a matching causal + left pad is bit-exact with convolving the unpadded kernel, so the geometry is + read back from the weight shapes rather than from metadata. + +The encoder geometry is derived from the tensor table (which block tensors +exist, what shape each weight has), not from GGUF KV metadata — the same +approach upstream takes, and it keeps the loader working for any checkpoint with +this topology. + +### Decoder + +A stock Qwen2 causal decoder, geometry read from the LM GGUF's KV block: 28 +layers, hidden 1536, intermediate 8960, 12 heads over 2 KV heads, head_dim 128, +RMSNorm eps 1e-6, RoPE theta 1e6, context 65536. The checkpoint has no +`qwen2.attention.key_length`, so `head_dim` comes from +`qwen2.rope.dimension_count`, which for this model equals +`embedding_length / head_count`; the loader cross-checks +`head_dim * head_count == embedding_length` and validates the declared geometry +against `token_embd.weight`'s shape. + +Weight types are mixed on purpose, exactly as published: + +| Tensors | Type | +|---|---| +| `blk.N.{attn_q,attn_k,attn_v,attn_output,ffn_gate,ffn_up,ffn_down}.weight` | `I2_S` (ternary) | +| `token_embd.weight` | Q6_K | +| `output.weight` | F16 | +| norms and `blk.N.attn_{q,k,v}.bias` | F32 | + +`I2_S` packs `{-1, 0, +1}` as codes `{0, 1, 2}`, 128 values per 32-byte group, +over the whole flat tensor, with one F32 absmax scale after the payload. The +kernel asserts `ne00 % 128 == 0`; hidden 1536 and intermediate 8960 both satisfy +it, and the weight is always 2-D by the time `LinearModule` calls +`ggml_mul_mat`. + +### Prompt + +Qwen2.5 ChatML, assembled to match `VibeASR.cpp/utils/prompt_builder.h` token for +token: + +``` +<|im_start|>system\nYou are a helpful assistant that transcribes audio input into text output in JSON format.<|im_end|>\n +<|im_start|>user\n<|speech_start|><|speech_pad|>×N<|speech_end|>\nThis is a 3.50 seconds audio, please transcribe it.<|im_end|>\n +``` + +- The special tokens are inserted by numeric id (151643–151648), not through the + tokenizer, because the GGUF vocabulary still carries Qwen2.5's original text + for those slots while the embedding rows are the ones VibeVoice trained. Every + text segment is tokenized with `parse_special = false`. +- There is deliberately **no generation prompt**: the model emits its own + `<|im_start|>assistant\n` header, and the session strips that leading triple + before decoding, as upstream does. +- `N` is the encoder frame count. Upstream builds `ceil(samples / 3200)` pads but + prefills only `min(pads, frames)` of them, so emitting exactly `frames` pads + produces the same sequence. +- The `<|speech_pad|>` rows are replaced in-graph by a `ggml_set_rows` over the + embedding lookup, with the speech features being the **element-wise sum** of + the acoustic and semantic connector outputs — both are 1536 wide, which is what + makes the sum well-defined. +- `output_format=json` swaps the instruction for `please transcribe it with these + keys: Start, End, Speaker, Content`; `context=...` switches to the + `with extra info:` suffix variant. + +Decoding is greedy, stopping at `<|im_end|>` or `<|endoftext|>`. Upstream's +default is temperature 0.7 / top-p 0.9 sampling with `--greedy` as an opt-in; +this port only implements the deterministic path, which is what parity is +measured against. + +### Audio front end + +Mixdown to mono, resample to 24 kHz, RMS-normalize to −25 dBFS with `eps = 1e-6`, +then divide by `max_abs` if it exceeded 1.0. This is audio.cpp's own +`vibevoice_asr` front end, not upstream's: VibeASR.cpp resamples with a naive +linear kernel and omits the clamp. For a clip already at 24 kHz the two agree; +for anything else the resampler differs and so do the encoder features (see +[Parity](#parity)). + +## Usage + +VibeASR.cpp already ships both halves quantized, so there is nothing to +re-quantize. The two forks only disagree on the numeric type *ids* — the VibeASR +fork put I2_S/I8_S at 36/37, which upstream ggml had already spent on the retired +`IQ4_NL_4_4` / `IQ4_NL_4_8` slots, so audio.cpp registers them at 43/42. The +converter rewrites the 4-byte type field in each tensor info and copies +everything else through byte for byte: + +```bash +# inspect first +python3 tools/community_models/convert_vibeasr_gguf.py \ + --input vibeasr-vae-encoder-i8_s.gguf --list + +# fix both GGUFs in place (703 MB encoder, 993 MB decoder) +python3 tools/community_models/convert_vibeasr_gguf.py \ + --input models/vibeasr/vibeasr-vae-encoder-i8_s.gguf --in-place +python3 tools/community_models/convert_vibeasr_gguf.py \ + --input models/vibeasr/vibeasr-lm-i2_s-embed-q6_k.gguf --in-place + +# confirm an already-converted package needs no further remapping +python3 tools/community_models/convert_vibeasr_gguf.py \ + --input models/vibeasr/vibeasr-lm-i2_s-embed-q6_k.gguf --check +``` + +Use `--output ` instead of `--in-place` to keep the original. + +The package is two GGUFs plus the tokenizer, so `--model` points at the LM GGUF +and the spec is resolved from the repo — the same invocation shape as +[`minimax_h3`](minimax_h3.md): + +``` +models/vibeasr/ +├── vibeasr-vae-encoder-i8_s.gguf +├── vibeasr-lm-i2_s-embed-q6_k.gguf +├── tokenizer.json +└── tokenizer_config.json +``` + +```bash +cmake -B build -DCMAKE_BUILD_TYPE=Release +cmake --build build -j --target audiocpp_cli + +./build/bin/audiocpp_cli \ + --task asr \ + --family vibeasr \ + --model models/vibeasr/vibeasr-lm-i2_s-embed-q6_k.gguf \ + --model-spec-override model_specs \ + --backend cpu \ + --threads 8 \ + --audio assets/asr_validation/librispeech/librispeech_test_clean_6930-75918-0000.wav \ + --metrics +``` + +``` +text_output=Concord returned to its place amidst the tents. +metrics.wall_ms=1284.65 +metrics.rtf=0.36652 +``` + +Request options: `output_format` (`text` | `json`), `context` (a string folded +into the prompt to bias recognition), `max_new_tokens` (default 1024). Session +options: `vibeasr.encoder_graph_arena_mb` (64), +`vibeasr.prefill_graph_arena_mb` (256), `vibeasr.decode_graph_arena_mb` (256). + +Note that `output_format=json` returns an empty transcript on short +single-speaker clips — the model emits an immediate end-of-turn. VibeASR.cpp +behaves identically on the same input; this port does not paper over it. + +## Tests + +```bash +cmake -B build -DCMAKE_BUILD_TYPE=Release -DENGINE_BUILD_MODEL_TESTS=ON +cmake --build build -j --target test_vibeasr_asr test_vibeasr_vae_encoder + +# end to end: loader, session, prompt, both graphs, greedy decode +./build/bin/test_vibeasr_asr --threads 8 + +# encoder only: shape, finiteness, frame count, and optional upstream parity +./build/bin/test_vibeasr_vae_encoder \ + --model models/vibeasr/vibeasr-vae-encoder-i8_s.gguf \ + --audio assets/asr_validation/librispeech/librispeech_test_clean_6930-75918-0000.wav \ + --threads 8 +``` + +Both exit 125 (SKIP) when the checkpoint is missing, so they are safe in ctest. +`i2_s_mul_mat_test` and `i8_s_fused_ops_test` cover the kernels themselves +against plain-loop references and need no checkpoint. + +## Parity + +### End to end + +Four LibriSpeech clips, greedy on both sides, against VibeASR.cpp's own +`asr_infer --greedy` on the same two GGUFs: + +| Clip | VibeASR.cpp | this port | +|---|---|---| +| test-clean 6930-75918-0000 | `Concord returned to its place amidst the tents.` | identical | +| test-clean 6930-75918-0001 | `The english forwarded to the french baskets of flowers, of which they had made a plentiful provision to greet the arrival of the young princess. The french, in return, invited the english to a supper, which was to be given the next day.` | identical | +| test-other 7902-96591-0001 | `Don't cry, he said. I was obliged to come.` | identical | +| test-other 7902-96591-0000 | `I'm from the cut or lying off the coast.` | `I'm from the cutter lying off the coast.` | + +Three of four match token for token. The fourth diverges because these clips are +16 kHz and the two resamplers differ — this port uses soxr, upstream uses naive +linear interpolation — which perturbs the encoder features enough to flip one +greedy argmax. (Reference text: `I AM FROM THE CUTTER LYING OFF THE COAST`.) A +clip already at 24 kHz skips resampling entirely and does not have this failure +mode. + +### Encoder + +The reference dump is raw F32, `frames * dim`, row-major, produced by calling +`vae_encode_acoustic` / `vae_encode_semantic` from VibeASR.cpp's own `vae.h` on +the same WAV: + +```bash +./build/bin/test_vibeasr_vae_encoder \ + --model models/vibeasr/vibeasr-vae-encoder-i8_s.gguf \ + --audio assets/asr_validation/librispeech/librispeech_test_clean_6930-75918-0000.wav \ + --reference-acoustic ref_acoustic.f32 \ + --reference-semantic ref_semantic.f32 \ + --threads 8 +``` + +3.505 s LibriSpeech clip fed at its native 16 kHz, 17 frames × 1536 per branch: + +| Branch | max abs | mean abs | cosine | +|---|---|---|---| +| acoustic | 1.478 (12.1% of range) | 0.0930 (0.76% of range) | 0.99238739 | +| semantic | 2.526 (9.5% of range) | 0.1804 (0.68% of range) | 0.98475210 | + +**Layer by layer, stage 0 is bit-exact** — every int8 byte and every scale +matches, which is what pins the layouts, the causal padding, the kernel padding, +and the weight mapping. The first divergence is 5 of 1,794,560 elements one int8 +step apart at an identical scale, entering stage 1, and it grows from there +because each of the remaining stages requantizes. + +Bit-exactness is not reachable and the tolerances say so. audio.cpp stores each +per-tensor scale as a multiplier (`amax/127`, dequantize by multiplying) while +VibeASR.cpp stores its reciprocal (`127/amax`, dequantize by dividing) — the +same number to within the last float bit, which is enough to flip a value that +sits on a rounding boundary. Upstream also rounds ties to even in its vector +body but away from zero in its scalar tail, so no single convention reproduces it +exactly. + +To calibrate what that is worth, nudging **one** input sample by one int8 step +and re-running VibeASR.cpp against *itself* moves its own output by cosine +0.99592 (acoustic) / 0.98700 (semantic) — the graph amplifies a single LSB about +as far as the two implementations differ from each other. The probe therefore +gates on mean-abs-relative ≤ 2% and cosine ≥ 0.98; anything tighter would be +testing rounding luck. + +`i8_s_fused_ops_test` covers the op arithmetic itself against plain-loop +references, including the in-band scale surviving `ggml_cont(ggml_permute(...))` +— the encoder flips activations between channel-major and length-major +constantly, and a byte copy that drops the scale leaves the values right and +everything downstream off by an arbitrary factor. + +## Measured performance + +Release build, CPU backend, 24 vCPU AMD EPYC 7V13, 3.505 s clip resampled to +24 kHz (26 speech frames, 72-token prompt, 13 generated tokens): + +| Threads | encoder (both branches) | prefill | decode | wall | RTF | +|---|---|---|---|---|---| +| 8 | 758 ms | 214 ms | 239 ms | 1285 ms | 0.367 | +| 1 | 4652 ms | 1415 ms | 941 ms | 7081 ms | 2.020 | + +The encoder dominates: it is run twice, once per branch, and it processes raw +samples rather than tokens. Decode is about 18 ms/token at 8 threads. + +The encoder-only probe reports 546 ms for both branches at 8 threads because it +feeds the clip at its native 16 kHz (17 frames); the session resamples to 24 kHz +first (26 frames). + +Peak RSS is 2.20 GB against 1.70 GB of weights: `BackendWeightStore` stages each +tensor before upload, so weight loading briefly holds roughly two copies of the +tensor being uploaded. Graph arenas are 64 MB (encoder) + 256 MB (prefill) + +256 MB (decode) by default. + +## Status + +Ported: + +- I8_S VAE encoder graph, both branches, CPU backend. +- Ternary I2_S matmul kernel and the Qwen2 decoder graph on top of it, with + prefill + static-cache single-step decode. +- Prompt assembly, speech-feature injection, greedy decode, tokenizer, loader, + session, and `--family vibeasr`. +- GGUF type remapping tool and geometry-from-tensors asset loader. +- End-to-end and encoder parity probes against upstream, plus op-level unit + tests. + +Known limitations: + +- **CPU only.** The fused I8_S ops and the I2_S matmul have no CUDA or Metal + kernels; the session pins the backend to CPU. +- **Offline only**, like `vibevoice_asr` itself. Upstream's encoder is causal, so + streaming is implementable, but the state machine is not ported. +- **Greedy only.** Upstream's sampling path (temperature, top-p) is not ported, + and neither is `vibevoice_asr`'s beam search. +- **Text only.** No `--segments-out` / `--turns-out` equivalent; `output_format=json` + is a prompt variant, not structured decoding. +- The package is two GGUFs, so it needs `--model ` plus + `--model-spec-override model_specs` rather than a directory path. +- Bit-exact parity with upstream is out of reach by design; see + [Parity](#parity). + +## Upstream + +- Model port: (`src/vae.cpp`, + `src/lm.cpp`, `src/asr_server.cpp`, `utils/prompt_builder.h`) +- Base model: VibeVoice ASR, also in tree as [`vibevoice_asr`](../asr.md#vibevoice-asr) +- Weights: diff --git a/include/engine/community_models/vibeasr/assets.h b/include/engine/community_models/vibeasr/assets.h new file mode 100644 index 000000000..d7b77d580 --- /dev/null +++ b/include/engine/community_models/vibeasr/assets.h @@ -0,0 +1,109 @@ +#pragma once + +// VibeASR assets: the I8_S audio VAE encoder and the ternary I2_S Qwen2 decoder. +// +// Ported from https://github.com/microsoft/VibeASR.cpp (src/vae.cpp, src/lm.cpp). + +#include "engine/framework/assets/resource_bundle.h" +#include "engine/framework/assets/tensor_source.h" + +#include +#include +#include +#include +#include + +namespace engine::community_models::vibeasr { + +// One ConvNeXt-style block inside a stage. +struct VaeBlockConfig { + int64_t channels = 0; + // Padded depthwise kernel width. The converter left-pads the real kernel + // (7 taps) up to a SIMD-friendly width with leading zeros, so convolving + // with the padded width and a matching causal left pad is bit-exact with + // convolving the unpadded kernel. + int64_t kernel_size = 0; + int64_t ffn_hidden = 0; +}; + +struct VaeStageConfig { + // Strided causal conv that enters the stage. + int64_t downsample_kernel_size = 0; + int64_t downsample_stride = 0; + int64_t in_channels = 0; + int64_t out_channels = 0; + std::vector blocks; +}; + +// One of the two encoder branches (acoustic / semantic). Both share the layout +// and differ only in latent width and stage depths. +struct VaeBranchConfig { + std::string prefix; // "acoustic" or "semantic" + std::vector stages; + int64_t head_kernel_size = 0; // padded causal kernel of the latent head + int64_t latent_dim = 0; // head output width + int64_t connector_hidden = 0; // connector output width, i.e. LM hidden size + int64_t total_stride = 0; // product of the stage strides + + // Downsampling factor from waveform samples to encoder frames. + [[nodiscard]] int64_t frames_for_samples(int64_t num_samples) const; +}; + +struct VibeASRVaeConfig { + VaeBranchConfig acoustic; + VaeBranchConfig semantic; + // VibeASR's graph hardcodes 1e-5 for every RMS norm, including the ones the + // checkpoint metadata labels 1e-6. The published weights were validated + // against the hardcoded value, so the port keeps it. + float rms_norm_eps = 1e-5f; +}; + +struct VibeASRVaeAssets { + std::shared_ptr source; + VibeASRVaeConfig config; +}; + +// Derives the encoder geometry from the tensor table instead of GGUF metadata: +// stage depths come from which block tensors are present, channel counts and +// kernel widths from the weight shapes. That keeps the loader working for any +// VibeASR VAE checkpoint with this topology, and avoids trusting metadata the +// reference implementation itself ignores. +VibeASRVaeConfig derive_vae_config(const assets::TensorSource & source); + +std::shared_ptr load_vibeasr_vae_assets(const std::filesystem::path & model_path); + +// Same, for a tensor source already opened from a resource bundle. +std::shared_ptr make_vibeasr_vae_assets( + std::shared_ptr source); + +// Decoder geometry. Unlike the encoder, none of this is recoverable from the +// tensor shapes alone -- head_dim, rope_theta and the RMS norm epsilon are not +// implied by any weight -- so it comes from the GGUF qwen2.* metadata block. +struct VibeASRLmConfig { + int64_t vocab_size = 0; + int64_t hidden_size = 0; + int64_t intermediate_size = 0; + int64_t num_hidden_layers = 0; + int64_t num_attention_heads = 0; + int64_t num_key_value_heads = 0; + int64_t head_dim = 0; + int64_t max_position_embeddings = 0; + // 1e-6 for the published checkpoint. Note this is *not* the encoder's + // epsilon: the VAE graph hardcodes 1e-5 (see VibeASRVaeConfig). + float rms_norm_eps = 1e-6f; + float rope_theta = 1e6f; +}; + +// The two GGUF halves plus the tokenizer files, as named by model_specs/vibeasr.json. +struct VibeASRAssets { + assets::ResourceBundle resources; + std::shared_ptr vae; + std::shared_ptr lm_weights; + VibeASRLmConfig lm; +}; + +VibeASRLmConfig derive_lm_config(const assets::TensorSource & source); + +std::shared_ptr load_vibeasr_assets(const std::filesystem::path & model_path); + +} // namespace engine::community_models::vibeasr diff --git a/include/engine/community_models/vibeasr/lm_decoder.h b/include/engine/community_models/vibeasr/lm_decoder.h new file mode 100644 index 000000000..7080f4b1d --- /dev/null +++ b/include/engine/community_models/vibeasr/lm_decoder.h @@ -0,0 +1,62 @@ +#pragma once + +// VibeASR language model: the Qwen2 causal decoder whose projections are stored +// as ternary GGML_TYPE_I2_S. Speech features from the VAE encoder replace the +// prompt's <|speech_pad|> placeholders before prefill. +// +// Ported from https://github.com/microsoft/VibeASR.cpp (src/lm.cpp). + +#include "engine/community_models/vibeasr/assets.h" +#include "engine/framework/core/execution_context.h" + +#include +#include +#include +#include + +namespace engine::community_models::vibeasr { + +struct VibeASRLmPrompt { + std::vector input_ids; + // Positions in input_ids occupied by <|speech_pad|>, in order. + std::vector speech_positions; +}; + +// Summed acoustic + semantic connector output, row-major [tokens][hidden_size]. +struct VibeASRSpeechEmbeddings { + int64_t tokens = 0; + int64_t hidden_size = 0; + std::vector values; +}; + +struct VibeASRGenerationOptions { + int64_t max_new_tokens = 1024; + std::vector eos_token_ids; +}; + +class VibeASRLmRuntime { +public: + VibeASRLmRuntime( + std::shared_ptr weights_source, + const VibeASRLmConfig & config, + core::ExecutionContext & execution, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes, + size_t weight_context_bytes); + ~VibeASRLmRuntime(); + + VibeASRLmRuntime(const VibeASRLmRuntime &) = delete; + VibeASRLmRuntime & operator=(const VibeASRLmRuntime &) = delete; + + // Greedy decode. Stops at any eos id or after max_new_tokens. + std::vector generate( + const VibeASRLmPrompt & prompt, + const VibeASRSpeechEmbeddings & speech, + const VibeASRGenerationOptions & options); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::community_models::vibeasr diff --git a/include/engine/community_models/vibeasr/session.h b/include/engine/community_models/vibeasr/session.h new file mode 100644 index 000000000..e35ef8f8a --- /dev/null +++ b/include/engine/community_models/vibeasr/session.h @@ -0,0 +1,60 @@ +#pragma once + +// Offline ASR session for the VibeASR package: I8_S VAE encoder -> ternary I2_S +// Qwen2 decoder, with VibeASR.cpp's ChatML prompt around the speech features. +// +// Ported from https://github.com/microsoft/VibeASR.cpp (src/asr_server.cpp, +// utils/prompt_builder.h). + +#include "engine/community_models/vibeasr/assets.h" +#include "engine/community_models/vibeasr/lm_decoder.h" +#include "engine/community_models/vibeasr/vae_encoder.h" +#include "engine/framework/model_spec/metadata.h" +#include "engine/framework/runtime/model.h" +#include "engine/framework/runtime/session_base.h" +#include "engine/framework/tokenizers/llama_bpe.h" + +#include +#include + +namespace engine::community_models::vibeasr { + +std::shared_ptr make_vibeasr_loader(); + +class VibeASRSession final : public runtime::RuntimeSessionBase, public runtime::IOfflineVoiceTaskSession { +public: + VibeASRSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets, + std::shared_ptr contract); + ~VibeASRSession() override; + + std::string family() const override; + runtime::VoiceTaskKind task_kind() const override; + runtime::RunMode run_mode() const override; + void prepare(const runtime::SessionPreparationRequest & request) override; + runtime::TaskResult run(const runtime::TaskRequest & request) override; + +private: + struct RequestOptions { + std::string output_format = "text"; + std::string context; + int64_t max_new_tokens = 1024; + }; + + RequestOptions parse_request_options(const runtime::TaskRequest & request) const; + runtime::AudioBuffer normalize(const runtime::AudioBuffer & audio) const; + VibeASRSpeechEmbeddings encode_speech(const std::vector & samples); + VibeASRLmPrompt build_prompt(int64_t speech_tokens, float duration_seconds, const RequestOptions & options) const; + std::string decode_tokens(const std::vector & token_ids) const; + + runtime::TaskSpec task_; + std::shared_ptr assets_; + std::shared_ptr contract_; + std::shared_ptr tokenizer_; + VibeASRVaeEncoderRuntime encoder_; + VibeASRLmRuntime lm_; +}; + +} // namespace engine::community_models::vibeasr diff --git a/include/engine/community_models/vibeasr/vae_encoder.h b/include/engine/community_models/vibeasr/vae_encoder.h new file mode 100644 index 000000000..dd4021dfe --- /dev/null +++ b/include/engine/community_models/vibeasr/vae_encoder.h @@ -0,0 +1,88 @@ +#pragma once + +// VibeASR audio VAE encoder: a ConvNeXt-style causal encoder that turns a mono +// waveform into LM-width features, running end to end in GGML_TYPE_I8_S. +// +// Ported from https://github.com/microsoft/VibeASR.cpp (src/vae.cpp). + +#include "engine/community_models/vibeasr/assets.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/core/module.h" + +#include +#include +#include + +namespace engine::community_models::vibeasr { + +struct VaeBlockWeights { + core::TensorValue mixer_norm; // [channels] + core::TensorValue mixer_conv_weight; // [channels, 1, kernel_size], I8_S + core::TensorValue mixer_conv_bias; // [channels] + core::TensorValue mixer_gamma; // [channels] + core::TensorValue ffn_norm; // [channels] + core::TensorValue ffn_fc1_weight; // [ffn_hidden, channels], I8_S + core::TensorValue ffn_fc1_bias; // [ffn_hidden] + core::TensorValue ffn_fc2_weight; // [channels, ffn_hidden], I8_S + core::TensorValue ffn_fc2_bias; // [channels] + core::TensorValue ffn_gamma; // [channels] +}; + +struct VaeStageWeights { + core::TensorValue downsample_weight; // [out_channels, in_channels, kernel_size], I8_S + core::TensorValue downsample_bias; // [out_channels] + std::vector blocks; +}; + +struct VaeBranchWeights { + std::vector stages; + core::TensorValue head_weight; // [latent_dim, channels, kernel_size], I8_S + core::TensorValue head_bias; // [latent_dim] + core::TensorValue connector_fc1_weight; // [connector_hidden, latent_dim], I8_S + core::TensorValue connector_fc1_bias; // [connector_hidden] + core::TensorValue connector_norm; // [connector_hidden] + core::TensorValue connector_fc2_weight; // [connector_hidden, connector_hidden], I8_S + core::TensorValue connector_fc2_bias; // [connector_hidden] +}; + +struct VibeASRVaeEncoderWeights { + VaeBranchWeights acoustic; + VaeBranchWeights semantic; +}; + +// Encoder output, row-major [frames][dim]. +struct VaeEncoderFeatures { + int64_t frames = 0; + int64_t dim = 0; + std::vector values; +}; + +class VibeASRVaeEncoderRuntime { +public: + VibeASRVaeEncoderRuntime( + std::shared_ptr assets, + engine::core::ExecutionContext & execution_context, + size_t graph_arena_bytes = 64ull * 1024ull * 1024ull); + + // Both branches consume the same waveform, sampled at 24 kHz and scaled to + // [-1, 1], and produce connector_hidden-wide features. + VaeEncoderFeatures encode_acoustic(const std::vector & samples); + VaeEncoderFeatures encode_semantic(const std::vector & samples); + + const VibeASRVaeAssets & assets() const noexcept { return *assets_; } + +private: + VaeEncoderFeatures encode( + const VaeBranchConfig & config, + const VaeBranchWeights & weights, + const std::vector & samples); + + std::shared_ptr assets_; + engine::core::ExecutionContext * execution_context_ = nullptr; + engine::core::BackendWeightStore weight_store_; + VibeASRVaeEncoderWeights weights_; + size_t graph_arena_bytes_; +}; + +} // namespace engine::community_models::vibeasr diff --git a/model_specs/vibeasr.json b/model_specs/vibeasr.json new file mode 100644 index 000000000..4c25db9c8 --- /dev/null +++ b/model_specs/vibeasr.json @@ -0,0 +1,139 @@ +{ + "schema_version": 1, + "family": "vibeasr", + "display_name": "VibeVoice-ASR-BitNet", + "description": "VibeASR.cpp's CPU-first VibeVoice ASR port: an INT8 (I8_S) audio VAE encoder feeding a ternary (I2_S) Qwen2 decoder, ported to audio.cpp.", + "category": "asr", + "status": "community", + "tasks": [ + "asr" + ], + "modes": [ + "offline" + ], + "languages": [ + "en", + "zh", + "fr", + "it", + "ko", + "pt", + "vi" + ], + "capabilities": {}, + "options": { + "request": [ + { + "name": "output_format", + "type": "enum", + "description": "Prompt suffix asked of the decoder: plain transcription text, or JSON rows with Start/End/Speaker/Content.", + "values": [ + "text", + "json" + ], + "required": false, + "default": "text" + }, + { + "name": "context", + "type": "string", + "description": "Extra context injected into the prompt (names, jargon) to bias the transcription.", + "required": false, + "default": "" + }, + { + "name": "max_new_tokens", + "type": "int", + "description": "Cap on decoded tokens for one request.", + "required": false, + "min": 1, + "default": 1024 + } + ], + "session": [ + { + "name": "encoder_graph_arena_mb", + "type": "int", + "description": "VAE encoder graph arena size in MB.", + "required": false, + "min": 16, + "default": 64 + }, + { + "name": "prefill_graph_arena_mb", + "type": "int", + "description": "Decoder prefill graph arena size in MB.", + "required": false, + "min": 16, + "default": 256 + }, + { + "name": "decode_graph_arena_mb", + "type": "int", + "description": "Decoder single-step graph arena size in MB.", + "required": false, + "min": 16, + "default": 256 + } + ], + "load": [] + }, + "runtime": { + "tags": [ + "gguf", + "cpu" + ] + }, + "packages": [ + { + "id": "vibeasr_bitnet_i2_s", + "display_name": "VibeVoice-ASR-BitNet I8_S encoder + I2_S decoder", + "description": "Upstream VibeASR.cpp GGUF package. The two GGUFs carry the VibeASR ggml fork's type ids and need one pass of tools/community_models/convert_vibeasr_gguf.py --in-place before audio.cpp can load them.", + "default": true, + "format": "gguf", + "precision": "native", + "target_directory": "VibeVoice-ASR-BitNet", + "files": [ + "vibeasr-vae-encoder-i8_s.gguf", + "vibeasr-lm-i2_s-embed-q6_k.gguf", + "tokenizer.json", + "tokenizer_config.json" + ], + "download": { + "kind": "huggingface_snapshot", + "repo": "microsoft/VibeVoice-ASR-BitNet", + "revision": "main", + "gated": false + } + } + ], + "dependencies": [], + "ui": { + "recommended_package": "vibeasr_bitnet_i2_s", + "tags": [ + "ASR", + "GGUF" + ], + "docs": [ + "docs/community_models/vibeasr.md" + ], + "summary": "INT8 encoder plus ternary Qwen2 decoder transcription on CPU." + }, + "sources": [ + { + "format": "gguf", + "roots": { + "model": "." + }, + "files": { + "tokenizer_json": "model:tokenizer.json", + "tokenizer_config": "model:tokenizer_config.json" + }, + "optional_files": {}, + "tensors": { + "vae_weights": "model:vibeasr-vae-encoder-i8_s.gguf", + "lm_weights": "model:vibeasr-lm-i2_s-embed-q6_k.gguf" + } + } + ] +} diff --git a/src/community_models/vibeasr/assets.cpp b/src/community_models/vibeasr/assets.cpp new file mode 100644 index 000000000..fcccdc04e --- /dev/null +++ b/src/community_models/vibeasr/assets.cpp @@ -0,0 +1,273 @@ +#include "engine/community_models/vibeasr/assets.h" + +#include "engine/framework/model_spec/package.h" + +#include + +#include +#include +#include +#include +#include + +namespace engine::community_models::vibeasr { +namespace { + +// VibeASR's AudioVAEEncoder fixes the stride schedule in code (it is not part of +// the checkpoint), giving a total downsampling factor of 3200 samples per frame. +constexpr int64_t kDownsampleStrides[] = {1, 2, 2, 4, 5, 5, 8}; +constexpr size_t kNumStages = sizeof(kDownsampleStrides) / sizeof(kDownsampleStrides[0]); + +std::vector require_shape( + const assets::TensorSource & source, + const std::string & name, + size_t expected_rank) { + auto shape = source.require_metadata(name).shape; + if (shape.size() != expected_rank) { + throw std::runtime_error( + "VibeASR VAE tensor " + name + " has rank " + std::to_string(shape.size()) + + ", expected " + std::to_string(expected_rank)); + } + return shape; +} + +std::string block_prefix(const std::string & branch, size_t stage, size_t block) { + return branch + ".stages." + std::to_string(stage) + "." + std::to_string(block); +} + +VaeBlockConfig derive_block(const assets::TensorSource & source, const std::string & prefix) { + VaeBlockConfig block; + // Depthwise kernel is stored as [channels, 1, kernel_size]. + const auto mixer = require_shape(source, prefix + ".mixer.conv.conv.conv.weight", 3); + block.channels = mixer[0]; + block.kernel_size = mixer[2]; + if (mixer[1] != 1) { + throw std::runtime_error("VibeASR VAE mixer conv at " + prefix + " is not depthwise"); + } + // Linear weights are stored as [out_features, in_features]. + const auto fc1 = require_shape(source, prefix + ".ffn.linear1.weight", 2); + const auto fc2 = require_shape(source, prefix + ".ffn.linear2.weight", 2); + block.ffn_hidden = fc1[0]; + if (fc1[1] != block.channels || fc2[0] != block.channels || fc2[1] != block.ffn_hidden) { + throw std::runtime_error("VibeASR VAE FFN shapes at " + prefix + " are inconsistent"); + } + return block; +} + +VaeBranchConfig derive_branch(const assets::TensorSource & source, const std::string & prefix) { + VaeBranchConfig branch; + branch.prefix = prefix; + branch.total_stride = 1; + + int64_t expected_in_channels = 1; // raw mono waveform + for (size_t stage = 0; stage < kNumStages; ++stage) { + const std::string downsample = + prefix + ".downsample_layers." + std::to_string(stage) + ".0.conv.conv.weight"; + if (!source.has_tensor(downsample)) { + throw std::runtime_error("VibeASR VAE checkpoint is missing " + downsample); + } + // Conv weight is stored as [out_channels, in_channels, kernel_size]. + const auto shape = require_shape(source, downsample, 3); + + VaeStageConfig config; + config.out_channels = shape[0]; + config.in_channels = shape[1]; + config.downsample_kernel_size = shape[2]; + config.downsample_stride = kDownsampleStrides[stage]; + if (config.in_channels != expected_in_channels) { + throw std::runtime_error("VibeASR VAE stage " + std::to_string(stage) + " channel count does not chain"); + } + if (config.downsample_kernel_size < config.downsample_stride) { + throw std::runtime_error("VibeASR VAE stage " + std::to_string(stage) + " kernel is shorter than its stride"); + } + + for (size_t block = 0;; ++block) { + const std::string block_name = block_prefix(prefix, stage, block); + if (!source.has_tensor(block_name + ".norm.weight")) { + break; + } + auto derived = derive_block(source, block_name); + if (derived.channels != config.out_channels) { + throw std::runtime_error("VibeASR VAE block " + block_name + " width does not match its stage"); + } + config.blocks.push_back(derived); + } + if (config.blocks.empty()) { + throw std::runtime_error("VibeASR VAE stage " + std::to_string(stage) + " has no blocks"); + } + + branch.total_stride *= config.downsample_stride; + expected_in_channels = config.out_channels; + branch.stages.push_back(std::move(config)); + } + + const auto head = require_shape(source, prefix + ".head.conv.conv.weight", 3); + branch.latent_dim = head[0]; + branch.head_kernel_size = head[2]; + if (head[1] != expected_in_channels) { + throw std::runtime_error("VibeASR VAE head input width does not match the last stage"); + } + + const auto fc1 = require_shape(source, prefix + "_connector.fc1.weight", 2); + const auto fc2 = require_shape(source, prefix + "_connector.fc2.weight", 2); + branch.connector_hidden = fc1[0]; + if (fc1[1] != branch.latent_dim || fc2[0] != branch.connector_hidden || + fc2[1] != branch.connector_hidden) { + throw std::runtime_error("VibeASR VAE " + prefix + " connector shapes are inconsistent"); + } + return branch; +} + +// assets::TensorSource exposes tensors, not the GGUF KV block, and the decoder +// geometry lives entirely in the KV block. Reading it directly is what the other +// community entries do (see sense_asr/assets.cpp). +class GgufMetadataReader { +public: + explicit GgufMetadataReader(const std::filesystem::path & path) { + gguf_init_params params{}; + params.no_alloc = true; + params.ctx = nullptr; + gguf_context * gguf = gguf_init_from_file(path.string().c_str(), params); + if (gguf == nullptr) { + throw std::runtime_error("Failed to read VibeASR GGUF metadata from " + path.string()); + } + ctx_.reset(gguf); + } + + int64_t require_u32(const char * key) const { + const int64_t id = gguf_find_key(ctx_.get(), key); + if (id < 0) { + throw std::runtime_error(std::string("VibeASR LM GGUF is missing ") + key); + } + return static_cast(gguf_get_val_u32(ctx_.get(), id)); + } + + float require_f32(const char * key) const { + const int64_t id = gguf_find_key(ctx_.get(), key); + if (id < 0) { + throw std::runtime_error(std::string("VibeASR LM GGUF is missing ") + key); + } + return gguf_get_val_f32(ctx_.get(), id); + } + + std::string kv_str(const char * key, std::string fallback) const { + const int64_t id = gguf_find_key(ctx_.get(), key); + return id < 0 ? std::move(fallback) : std::string(gguf_get_val_str(ctx_.get(), id)); + } + +private: + struct GgufDeleter { + void operator()(gguf_context * ctx) const noexcept { + if (ctx != nullptr) { + gguf_free(ctx); + } + } + }; + + std::unique_ptr ctx_; +}; + +} // namespace + +int64_t VaeBranchConfig::frames_for_samples(int64_t num_samples) const { + // Every stage is a causal conv with left padding kernel_size - stride, so + // its output length is ggml_calc_conv_output_size() with that padding. + int64_t length = num_samples; + for (const auto & stage : stages) { + const int64_t padding = stage.downsample_kernel_size - stage.downsample_stride; + length = (length + padding - stage.downsample_kernel_size) / stage.downsample_stride + 1; + if (length <= 0) { + return 0; + } + } + return length; +} + +VibeASRVaeConfig derive_vae_config(const assets::TensorSource & source) { + VibeASRVaeConfig config; + config.acoustic = derive_branch(source, "acoustic"); + config.semantic = derive_branch(source, "semantic"); + if (config.acoustic.connector_hidden != config.semantic.connector_hidden) { + throw std::runtime_error("VibeASR VAE branches disagree on the connector width"); + } + return config; +} + +std::shared_ptr load_vibeasr_vae_assets(const std::filesystem::path & model_path) { + return make_vibeasr_vae_assets(engine::assets::open_tensor_source(model_path)); +} + +std::shared_ptr make_vibeasr_vae_assets( + std::shared_ptr source) { + auto assets = std::make_shared(); + assets->config = derive_vae_config(*source); + assets->source = std::move(source); + return assets; +} + +VibeASRLmConfig derive_lm_config(const assets::TensorSource & source) { + const GgufMetadataReader reader(source.source_path()); + + const std::string architecture = reader.kv_str("general.architecture", ""); + if (architecture != "qwen2") { + throw std::runtime_error( + "VibeASR LM GGUF declares architecture '" + architecture + "', expected qwen2"); + } + + VibeASRLmConfig config; + config.vocab_size = reader.require_u32("qwen2.vocab_size"); + config.hidden_size = reader.require_u32("qwen2.embedding_length"); + config.intermediate_size = reader.require_u32("qwen2.feed_forward_length"); + config.num_hidden_layers = reader.require_u32("qwen2.block_count"); + config.num_attention_heads = reader.require_u32("qwen2.attention.head_count"); + config.num_key_value_heads = reader.require_u32("qwen2.attention.head_count_kv"); + config.max_position_embeddings = reader.require_u32("qwen2.context_length"); + // The checkpoint has no attention.key_length: Qwen2 stores the per-head width + // only as the RoPE dimension count, which for this model equals + // embedding_length / head_count. + config.head_dim = reader.require_u32("qwen2.rope.dimension_count"); + config.rms_norm_eps = reader.require_f32("qwen2.attention.layer_norm_rms_epsilon"); + config.rope_theta = reader.require_f32("qwen2.rope.freq_base"); + + if (config.head_dim * config.num_attention_heads != config.hidden_size) { + throw std::runtime_error("VibeASR LM head_dim * head_count does not match embedding_length"); + } + if (config.num_key_value_heads <= 0 || config.num_attention_heads % config.num_key_value_heads != 0) { + throw std::runtime_error("VibeASR LM head_count is not a multiple of head_count_kv"); + } + if (config.num_hidden_layers <= 0) { + throw std::runtime_error("VibeASR LM declares no layers"); + } + + // Cross-check the metadata against the one tensor whose shape pins both dims. + const auto embedding = source.require_metadata("token_embd.weight").shape; + if (embedding.size() != 2 || embedding[0] != config.vocab_size || embedding[1] != config.hidden_size) { + throw std::runtime_error("VibeASR LM token_embd.weight does not match the declared geometry"); + } + return config; +} + +std::shared_ptr load_vibeasr_assets(const std::filesystem::path & model_path) { + auto assets = std::make_shared(); + assets->resources = engine::model_spec::load_resource_bundle_for_family(model_path, "vibeasr"); + + // A GGUF still carrying the VibeASR fork's type ids (36/37) fails deep inside + // the reader with an unhelpful message, so name the fix here. + auto open = [&assets](const char * id) { + try { + return assets->resources.open_tensor_source(id); + } catch (const std::exception & error) { + throw std::runtime_error( + std::string("VibeASR could not open the '") + id + "' GGUF (" + error.what() + + "). If it came straight from huggingface.co/microsoft/VibeVoice-ASR-BitNet, run " + "tools/community_models/convert_vibeasr_gguf.py --in-place on it first."); + } + }; + + assets->vae = make_vibeasr_vae_assets(open("vae_weights")); + assets->lm_weights = open("lm_weights"); + assets->lm = derive_lm_config(*assets->lm_weights); + return assets; +} + +} // namespace engine::community_models::vibeasr diff --git a/src/community_models/vibeasr/lm_decoder.cpp b/src/community_models/vibeasr/lm_decoder.cpp new file mode 100644 index 000000000..35bde5ef5 --- /dev/null +++ b/src/community_models/vibeasr/lm_decoder.cpp @@ -0,0 +1,679 @@ +#include "engine/community_models/vibeasr/lm_decoder.h" + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/positional_modules.h" +#include "engine/framework/modules/transformers/qwen_causal_decoder.h" +#include "engine/framework/runtime/errors.h" +#include "engine/framework/runtime/kv_cache.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::community_models::vibeasr { +namespace { + +namespace modules = engine::modules; +using Clock = std::chrono::steady_clock; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct GgmlGallocrDeleter { + void operator()(ggml_gallocr_t alloc) const noexcept { + if (alloc != nullptr) { + ggml_gallocr_free(alloc); + } + } +}; + +struct LmLayerWeights { + core::TensorValue input_norm; + core::TensorValue q_proj; + core::TensorValue q_bias; + core::TensorValue k_proj; + core::TensorValue k_bias; + core::TensorValue v_proj; + core::TensorValue v_bias; + core::TensorValue o_proj; + core::TensorValue post_norm; + core::TensorValue gate_proj; + core::TensorValue up_proj; + core::TensorValue down_proj; +}; + +struct LmWeights { + std::shared_ptr store; + core::TensorValue token_embedding; + std::vector layers; + core::TensorValue norm; + core::TensorValue lm_head; +}; + +struct PrefillOutput { + std::vector logits; + runtime::TransformerKVState kv_state; +}; + +// I2_S is a whole-tensor quantization whose in-band F32 scale sits after the +// packed codes, which is exactly what ggml_nbytes() accounts for, so the GGUF +// payload goes to the backend byte for byte. Same contract as the encoder's +// load_i8_s_tensor(). +core::TensorValue load_i2_s_tensor( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & name, + const std::vector & expected_shape) { + const auto metadata = source.require_metadata(name); + if (metadata.dtype != "i2_s") { + throw std::runtime_error("VibeASR LM tensor " + name + " is " + metadata.dtype + ", expected i2_s"); + } + if (metadata.shape != expected_shape) { + throw std::runtime_error("VibeASR LM tensor " + name + " has an unexpected shape"); + } + + core::TensorShape shape; + shape.rank = expected_shape.size(); + for (size_t i = 0; i < shape.rank; ++i) { + shape.dims[i] = expected_shape[i]; + } + + const auto raw = source.require_tensor_data(name); + return store.make_tensor(shape, GGML_TYPE_I2_S, raw.bytes.data(), raw.bytes.size()); +} + +modules::QwenDecoderLayerWeights to_qwen_layer_weights(const LmLayerWeights & weights) { + modules::QwenDecoderLayerWeights out; + out.input_norm = {weights.input_norm, std::nullopt}; + out.self_attention.q_weight = weights.q_proj; + out.self_attention.q_bias = weights.q_bias; + out.self_attention.k_weight = weights.k_proj; + out.self_attention.k_bias = weights.k_bias; + out.self_attention.v_weight = weights.v_proj; + out.self_attention.v_bias = weights.v_bias; + out.self_attention.out_weight = weights.o_proj; + out.post_norm = {weights.post_norm, std::nullopt}; + out.mlp.gate_proj = {weights.gate_proj, std::nullopt}; + out.mlp.up_proj = {weights.up_proj, std::nullopt}; + out.mlp.down_proj = {weights.down_proj, std::nullopt}; + return out; +} + +// Plain Qwen2: attention biases, no per-head Q/K norms. Nothing here depends on +// the weight type, which is why the framework's decoder runs unmodified on I2_S +// projections -- ggml_mul_mat dispatches on the tensor type. +modules::QwenCausalDecoderConfig make_qwen_decoder_config(const VibeASRLmConfig & config) { + modules::QwenCausalDecoderConfig out; + out.stack.hidden_size = config.hidden_size; + out.stack.num_attention_heads = config.num_attention_heads; + out.stack.num_key_value_heads = config.num_key_value_heads; + out.stack.head_dim = config.head_dim; + out.stack.intermediate_size = config.intermediate_size; + out.stack.layers = config.num_hidden_layers; + out.stack.rms_norm_eps = config.rms_norm_eps; + out.stack.rope_theta = config.rope_theta; + out.stack.use_qk_norm = false; + out.stack.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + out.logits_size = config.vocab_size; + out.logits_mode = modules::QwenCausalDecoderLogitsMode::LastStep; + return out; +} + +modules::QwenCausalDecoderWeights make_qwen_decoder_weights(const LmWeights & weights) { + modules::QwenCausalDecoderWeights out; + out.stack.layers.reserve(weights.layers.size()); + for (const auto & layer : weights.layers) { + out.stack.layers.push_back(to_qwen_layer_weights(layer)); + } + out.final_norm = {weights.norm, std::nullopt}; + out.lm_head = {weights.lm_head, std::nullopt}; + return out; +} + +// Token embeddings with the encoder's speech features written over the +// <|speech_pad|> slots. Doing the overwrite in-graph with ggml_set_rows keeps +// the prompt a single I32 upload instead of a host-side embedding matrix. +core::TensorValue prompt_embeddings( + core::ModuleBuildContext & ctx, + const LmWeights & weights, + const VibeASRLmConfig & config, + ggml_tensor * token_ids, + ggml_tensor * speech_embeddings, + ggml_tensor * speech_positions, + int64_t prompt_steps, + int64_t speech_tokens) { + auto ids = core::wrap_tensor(token_ids, core::TensorShape::from_dims({prompt_steps}), GGML_TYPE_I32); + auto x = modules::EmbeddingModule({config.vocab_size, config.hidden_size}) + .build(ctx, ids, weights.token_embedding); + if (speech_tokens > 0) { + auto speech = core::wrap_tensor( + speech_embeddings, + core::TensorShape::from_dims({speech_tokens, config.hidden_size}), + GGML_TYPE_F32); + auto positions = core::wrap_tensor( + speech_positions, + core::TensorShape::from_dims({speech_tokens}), + GGML_TYPE_I64); + x = core::wrap_tensor( + ggml_set_rows(ctx.ggml, x.tensor, speech.tensor, positions.tensor), + x.shape, + GGML_TYPE_F32); + } + return core::reshape_tensor(ctx, x, core::TensorShape::from_dims({1, prompt_steps, config.hidden_size})); +} + +LmWeights load_weights( + const assets::TensorSource & source, + const VibeASRLmConfig & config, + ggml_backend_t backend, + core::BackendType backend_type, + size_t weight_context_bytes) { + LmWeights weights; + weights.store = std::make_shared( + backend, + backend_type, + "vibeasr.lm.weights", + weight_context_bytes); + + // The embedding table and the output projection are the two tensors VibeASR + // leaves unternarized -- Q6_K and F16 in the published checkpoint -- so they + // load through the framework's normal path. + weights.token_embedding = weights.store->load_tensor( + source, + "token_embd.weight", + assets::TensorStorageType::Native, + {config.vocab_size, config.hidden_size}); + + const int64_t dim = config.head_dim; + const int64_t q_dim = config.num_attention_heads * dim; + const int64_t kv_dim = config.num_key_value_heads * dim; + weights.layers.reserve(static_cast(config.num_hidden_layers)); + for (int64_t layer = 0; layer < config.num_hidden_layers; ++layer) { + const std::string prefix = "blk." + std::to_string(layer) + "."; + LmLayerWeights w; + w.input_norm = weights.store->load_f32_tensor(source, prefix + "attn_norm.weight", {config.hidden_size}); + w.q_proj = load_i2_s_tensor(*weights.store, source, prefix + "attn_q.weight", {q_dim, config.hidden_size}); + w.q_bias = weights.store->load_f32_tensor(source, prefix + "attn_q.bias", {q_dim}); + w.k_proj = load_i2_s_tensor(*weights.store, source, prefix + "attn_k.weight", {kv_dim, config.hidden_size}); + w.k_bias = weights.store->load_f32_tensor(source, prefix + "attn_k.bias", {kv_dim}); + w.v_proj = load_i2_s_tensor(*weights.store, source, prefix + "attn_v.weight", {kv_dim, config.hidden_size}); + w.v_bias = weights.store->load_f32_tensor(source, prefix + "attn_v.bias", {kv_dim}); + w.o_proj = load_i2_s_tensor(*weights.store, source, prefix + "attn_output.weight", {config.hidden_size, q_dim}); + w.post_norm = weights.store->load_f32_tensor(source, prefix + "ffn_norm.weight", {config.hidden_size}); + w.gate_proj = load_i2_s_tensor( + *weights.store, source, prefix + "ffn_gate.weight", {config.intermediate_size, config.hidden_size}); + w.up_proj = load_i2_s_tensor( + *weights.store, source, prefix + "ffn_up.weight", {config.intermediate_size, config.hidden_size}); + w.down_proj = load_i2_s_tensor( + *weights.store, source, prefix + "ffn_down.weight", {config.hidden_size, config.intermediate_size}); + weights.layers.push_back(std::move(w)); + } + + weights.norm = weights.store->load_f32_tensor(source, "output_norm.weight", {config.hidden_size}); + weights.lm_head = weights.store->load_tensor( + source, + "output.weight", + assets::TensorStorageType::Native, + {config.vocab_size, config.hidden_size}); + weights.store->upload(); + return weights; +} + +int32_t argmax_index(const std::vector & values) { + if (values.empty()) { + throw std::runtime_error("VibeASR LM cannot select from empty logits"); + } + size_t best = 0; + for (size_t i = 1; i < values.size(); ++i) { + if (values[i] > values[best]) { + best = i; + } + } + return static_cast(best); +} + +class LmWeightsRuntime { +public: + LmWeightsRuntime( + std::shared_ptr source, + VibeASRLmConfig config, + core::ExecutionContext & execution, + size_t weight_context_bytes) + : source_(std::move(source)), + config_(std::make_shared(config)), + backend_(execution.backend()), + backend_type_(execution.backend_type()), + threads_(std::max(1, execution.config().threads)), + weights_(std::make_shared(load_weights( + *source_, + *config_, + backend_, + backend_type_, + weight_context_bytes))) {} + + const VibeASRLmConfig & config() const noexcept { return *config_; } + const LmWeights & weights() const noexcept { return *weights_; } + ggml_backend_t backend() const noexcept { return backend_; } + core::BackendType backend_type() const noexcept { return backend_type_; } + int threads() const noexcept { return threads_; } + +private: + std::shared_ptr source_; + std::shared_ptr config_; + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int threads_ = 1; + std::shared_ptr weights_; +}; + +class PrefillGraph { +public: + PrefillGraph( + std::shared_ptr runtime, + int64_t prompt_steps, + int64_t speech_tokens, + size_t graph_arena_bytes) + : runtime_(std::move(runtime)), + prompt_steps_(prompt_steps), + speech_tokens_(speech_tokens) { + if (prompt_steps_ <= 0) { + throw std::runtime_error("VibeASR LM prefill requires positive prompt length"); + } + if (speech_tokens_ < 0 || speech_tokens_ > prompt_steps_) { + throw std::runtime_error("VibeASR LM prefill speech token count is invalid"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize VibeASR LM prefill graph context"); + } + const auto & config = runtime_->config(); + const auto & weights = runtime_->weights(); + core::ModuleBuildContext ctx{ctx_.get(), "vibeasr.lm.prefill", runtime_->backend_type()}; + token_ids_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, prompt_steps_); + speech_embeddings_ = ggml_new_tensor_2d( + ctx_.get(), GGML_TYPE_F32, config.hidden_size, std::max(speech_tokens_, 1)); + speech_positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I64, std::max(speech_tokens_, 1)); + auto x = prompt_embeddings( + ctx, + weights, + config, + token_ids_, + speech_embeddings_, + speech_positions_, + prompt_steps_, + speech_tokens_); + positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, prompt_steps_); + auto positions = core::wrap_tensor(positions_, core::TensorShape::from_dims({prompt_steps_}), GGML_TYPE_I32); + + auto decoder_out = modules::QwenCausalDecoderModule(make_qwen_decoder_config(config)) + .build(ctx, x, positions, make_qwen_decoder_weights(weights)); + for (const auto & layer : decoder_out.state.layers) { + if (!layer.key.has_value() || !layer.value.has_value()) { + throw std::runtime_error("VibeASR LM prefill decoder did not return K/V state"); + } + // Copy K/V out of the graph-allocated intermediates and mark them as + // outputs so the allocator cannot recycle them before run() reads + // them back. + auto * key = ggml_cpy(ctx_.get(), layer.key->tensor, ggml_dup_tensor(ctx_.get(), layer.key->tensor)); + auto * value = ggml_cpy(ctx_.get(), layer.value->tensor, ggml_dup_tensor(ctx_.get(), layer.value->tensor)); + ggml_set_output(key); + ggml_set_output(value); + keys_.push_back(key); + values_.push_back(value); + } + logits_ = decoder_out.logits.tensor; + ggml_set_output(logits_); + graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); + ggml_build_forward_expand(graph_, logits_); + for (auto * key : keys_) { + ggml_build_forward_expand(graph_, key); + } + for (auto * value : values_) { + ggml_build_forward_expand(graph_, value); + } + const auto try_alloc = [&]() { + gallocr_.reset(ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_->backend()))); + return gallocr_ != nullptr && + ggml_gallocr_reserve(gallocr_.get(), graph_) && + ggml_gallocr_alloc_graph(gallocr_.get(), graph_); + }; + if (!try_alloc() && (engine::core::trim_backend_pools(runtime_->backend()), !try_alloc())) { + throw engine::runtime::CapacityError( + "VibeASR LM prefill graph does not fit in device memory at this size (" + + std::to_string(prompt_steps_) + " prompt steps, of which " + + std::to_string(speech_tokens_) + " are speech tokens)"); + } + position_ids_ = modules::qwen_position_ids(prompt_steps_); + debug::timing_log_scalar("vibeasr.lm.prefill.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("vibeasr.lm.prefill_prompt_steps", prompt_steps_); + } + + ~PrefillGraph() { + engine::core::release_backend_graph_resources(runtime_->backend(), graph_, true); + } + + bool matches(const LmWeightsRuntime & runtime, int64_t prompt_steps, int64_t speech_tokens) const { + return runtime_.get() == &runtime && prompt_steps_ == prompt_steps && speech_tokens_ == speech_tokens; + } + + PrefillOutput run( + const std::vector & token_ids, + const std::vector & speech_embeddings, + const std::vector & speech_positions) { + const auto & config = runtime_->config(); + if (static_cast(token_ids.size()) != prompt_steps_) { + throw std::runtime_error("VibeASR LM prefill token id count mismatch"); + } + if (static_cast(speech_embeddings.size()) != speech_tokens_ * config.hidden_size) { + throw std::runtime_error("VibeASR LM prefill speech embedding size mismatch"); + } + if (static_cast(speech_positions.size()) != speech_tokens_) { + throw std::runtime_error("VibeASR LM prefill speech position count mismatch"); + } + // Re-uploaded on every run: leaves are not pinned by the graph allocator. + ggml_backend_tensor_set(positions_, position_ids_.data(), 0, position_ids_.size() * sizeof(int32_t)); + ggml_backend_tensor_set(token_ids_, token_ids.data(), 0, token_ids.size() * sizeof(int32_t)); + if (speech_tokens_ > 0) { + const std::vector positions(speech_positions.begin(), speech_positions.end()); + ggml_backend_tensor_set( + speech_embeddings_, speech_embeddings.data(), 0, speech_embeddings.size() * sizeof(float)); + ggml_backend_tensor_set(speech_positions_, positions.data(), 0, positions.size() * sizeof(int64_t)); + } + core::set_backend_threads(runtime_->backend(), runtime_->threads()); + const auto compute_start = Clock::now(); + const ggml_status status = engine::core::compute_backend_graph(runtime_->backend(), graph_); + ggml_backend_synchronize(runtime_->backend()); + debug::timing_log_scalar("vibeasr.lm.prefill.graph.compute_ms", engine::debug::elapsed_ms(compute_start, Clock::now())); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("VibeASR LM prefill graph compute failed"); + } + PrefillOutput out; + out.logits.resize(static_cast(config.vocab_size)); + ggml_backend_tensor_get(logits_, out.logits.data(), 0, out.logits.size() * sizeof(float)); + out.kv_state.current_end = prompt_steps_; + out.kv_state.layers.resize(keys_.size()); + const size_t layer_values = + static_cast(prompt_steps_ * config.num_key_value_heads * config.head_dim); + for (size_t layer = 0; layer < keys_.size(); ++layer) { + auto & state = out.kv_state.layers[layer]; + state.valid_steps = prompt_steps_; + state.key.resize(layer_values); + state.value.resize(layer_values); + ggml_backend_tensor_get(keys_[layer], state.key.data(), 0, state.key.size() * sizeof(float)); + ggml_backend_tensor_get(values_[layer], state.value.data(), 0, state.value.size() * sizeof(float)); + } + return out; + } + +private: + std::shared_ptr runtime_; + int64_t prompt_steps_ = 0; + int64_t speech_tokens_ = 0; + std::unique_ptr ctx_; + ggml_tensor * token_ids_ = nullptr; + ggml_tensor * speech_embeddings_ = nullptr; + ggml_tensor * speech_positions_ = nullptr; + ggml_tensor * positions_ = nullptr; + ggml_tensor * logits_ = nullptr; + std::vector keys_; + std::vector values_; + std::vector position_ids_; + ggml_cgraph * graph_ = nullptr; + std::unique_ptr, GgmlGallocrDeleter> gallocr_; +}; + +class DecodeGraph { +public: + DecodeGraph(std::shared_ptr runtime, int64_t cache_steps, size_t graph_arena_bytes) + : runtime_(std::move(runtime)), + cache_steps_(cache_steps) { + if (cache_steps_ <= 0) { + throw std::runtime_error("VibeASR LM decode requires positive cache length"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize VibeASR LM decode graph context"); + } + const auto & config = runtime_->config(); + const auto & weights = runtime_->weights(); + core::ModuleBuildContext ctx{ctx_.get(), "vibeasr.lm.decode", runtime_->backend_type()}; + token_id_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); + auto token_id = core::wrap_tensor(token_id_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + auto x = modules::EmbeddingModule({config.vocab_size, config.hidden_size}) + .build(ctx, token_id, weights.token_embedding); + x = core::reshape_tensor(ctx, x, core::TensorShape::from_dims({1, 1, config.hidden_size})); + positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); + auto positions = core::wrap_tensor(positions_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + cache_slot_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); + auto cache_slot = core::wrap_tensor(cache_slot_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + attention_mask_ = ggml_new_tensor_4d(ctx_.get(), GGML_TYPE_F16, cache_steps_, 1, 1, 1); + auto attention_mask = core::wrap_tensor( + attention_mask_, core::TensorShape::from_dims({1, 1, 1, cache_steps_}), GGML_TYPE_F16); + graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); + auto decoder_out = modules::QwenCausalDecoderModule(make_qwen_decoder_config(config)) + .build_static_cache_tail( + ctx, + graph_, + x, + positions, + make_qwen_decoder_weights(weights), + cache_steps_, + attention_mask, + cache_slot); + step_cache_ = std::move(decoder_out.cache); + logits_ = decoder_out.logits.tensor; + ggml_set_output(logits_); + ggml_build_forward_expand(graph_, logits_); + buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend()); + if (buffer_ == nullptr) { + engine::core::trim_backend_pools(runtime_->backend()); + buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend()); + } + if (buffer_ == nullptr) { + throw engine::runtime::CapacityError( + "VibeASR LM decode graph does not fit in device memory at " + + std::to_string(cache_steps_) + " cache steps"); + } + attention_mask_values_.assign(static_cast(cache_steps_), ggml_fp32_to_fp16(-INFINITY)); + debug::timing_log_scalar("vibeasr.lm.decode.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("vibeasr.lm.decode_cache_steps", cache_steps_); + } + + ~DecodeGraph() { + engine::core::release_backend_graph_resources(runtime_->backend(), graph_, true); + if (buffer_ != nullptr) { + ggml_backend_buffer_free(buffer_); + } + } + + bool can_run(const LmWeightsRuntime & runtime, int64_t required_steps) const { + return runtime_.get() == &runtime && cache_steps_ >= required_steps; + } + + void import_state(const runtime::TransformerKVState & state) { + step_cache_.import_state(state); + } + + std::vector run_step(int32_t token) { + const auto & config = runtime_->config(); + if (step_cache_.valid_steps() >= cache_steps_) { + throw std::runtime_error("VibeASR LM decode cache exhausted"); + } + ggml_backend_tensor_set(token_id_, &token, 0, sizeof(int32_t)); + const int32_t position = static_cast(step_cache_.current_end()); + ggml_backend_tensor_set(positions_, &position, 0, sizeof(int32_t)); + const int32_t cache_slot = static_cast(step_cache_.valid_steps()); + ggml_backend_tensor_set(cache_slot_, &cache_slot, 0, sizeof(int32_t)); + modules::write_qwen_cached_step_mask( + attention_mask_, + attention_mask_values_, + cache_steps_, + step_cache_.valid_steps(), + step_cache_.valid_steps()); + core::set_backend_threads(runtime_->backend(), runtime_->threads()); + const ggml_status status = engine::core::compute_backend_graph(runtime_->backend(), graph_); + ggml_backend_synchronize(runtime_->backend()); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("VibeASR LM decode graph compute failed"); + } + logits_buffer_.resize(static_cast(config.vocab_size)); + ggml_backend_tensor_get(logits_, logits_buffer_.data(), 0, logits_buffer_.size() * sizeof(float)); + step_cache_.advance_after_direct_append(1); + // The caller moves out of this buffer before the next step. + return std::move(logits_buffer_); + } + +private: + std::shared_ptr runtime_; + int64_t cache_steps_ = 0; + std::unique_ptr ctx_; + ggml_tensor * token_id_ = nullptr; + ggml_tensor * positions_ = nullptr; + ggml_tensor * cache_slot_ = nullptr; + ggml_tensor * attention_mask_ = nullptr; + ggml_tensor * logits_ = nullptr; + std::vector attention_mask_values_; + std::vector logits_buffer_; + runtime::TransformerKVCache step_cache_; + ggml_cgraph * graph_ = nullptr; + ggml_backend_buffer_t buffer_ = nullptr; +}; + +} // namespace + +struct VibeASRLmRuntime::Impl { + Impl( + std::shared_ptr weights_source, + const VibeASRLmConfig & config, + core::ExecutionContext & execution, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes, + size_t weight_context_bytes) + : weights(std::make_shared( + std::move(weights_source), + config, + execution, + weight_context_bytes)), + prefill_graph_arena_bytes(prefill_graph_arena_bytes), + decode_graph_arena_bytes(decode_graph_arena_bytes) {} + + void validate_speech(const VibeASRLmPrompt & prompt, const VibeASRSpeechEmbeddings & speech) const { + const auto & config = weights->config(); + if (speech.tokens > 0 && speech.hidden_size != config.hidden_size) { + throw std::runtime_error("VibeASR speech embedding hidden size mismatch"); + } + if (speech.tokens != static_cast(prompt.speech_positions.size())) { + throw std::runtime_error("VibeASR speech embedding count does not match the prompt's speech pads"); + } + if (static_cast(speech.values.size()) != speech.tokens * speech.hidden_size) { + throw std::runtime_error("VibeASR speech embedding value count mismatch"); + } + for (const int32_t position : prompt.speech_positions) { + if (position < 0 || position >= static_cast(prompt.input_ids.size())) { + throw std::runtime_error("VibeASR speech pad position out of range"); + } + } + } + + std::shared_ptr weights; + size_t prefill_graph_arena_bytes = 0; + size_t decode_graph_arena_bytes = 0; + std::unique_ptr prefill_graph; + std::unique_ptr decode_graph; +}; + +VibeASRLmRuntime::VibeASRLmRuntime( + std::shared_ptr weights_source, + const VibeASRLmConfig & config, + core::ExecutionContext & execution, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes, + size_t weight_context_bytes) + : impl_(std::make_unique( + std::move(weights_source), + config, + execution, + prefill_graph_arena_bytes, + decode_graph_arena_bytes, + weight_context_bytes)) {} + +VibeASRLmRuntime::~VibeASRLmRuntime() = default; + +std::vector VibeASRLmRuntime::generate( + const VibeASRLmPrompt & prompt, + const VibeASRSpeechEmbeddings & speech, + const VibeASRGenerationOptions & options) { + const auto & config = impl_->weights->config(); + if (prompt.input_ids.empty()) { + throw std::runtime_error("VibeASR LM prompt is empty"); + } + if (options.max_new_tokens <= 0) { + throw std::runtime_error("VibeASR max_new_tokens must be positive"); + } + const int64_t prompt_steps = static_cast(prompt.input_ids.size()); + if (prompt_steps + options.max_new_tokens > config.max_position_embeddings) { + throw std::runtime_error("VibeASR request exceeds the decoder context length"); + } + impl_->validate_speech(prompt, speech); + + if (impl_->prefill_graph == nullptr || + !impl_->prefill_graph->matches(*impl_->weights, prompt_steps, speech.tokens)) { + impl_->prefill_graph.reset(); + impl_->prefill_graph = std::make_unique( + impl_->weights, prompt_steps, speech.tokens, impl_->prefill_graph_arena_bytes); + } + auto prefill = impl_->prefill_graph->run(prompt.input_ids, speech.values, prompt.speech_positions); + + const int64_t required_cache_steps = prompt_steps + options.max_new_tokens; + if (impl_->decode_graph == nullptr || !impl_->decode_graph->can_run(*impl_->weights, required_cache_steps)) { + impl_->decode_graph.reset(); + impl_->decode_graph = + std::make_unique(impl_->weights, required_cache_steps, impl_->decode_graph_arena_bytes); + } + impl_->decode_graph->import_state(prefill.kv_state); + + const auto is_eos = [&options](int32_t token) { + return std::find(options.eos_token_ids.begin(), options.eos_token_ids.end(), token) != + options.eos_token_ids.end(); + }; + + std::vector out; + std::vector logits = std::move(prefill.logits); + const auto decode_start = Clock::now(); + for (int64_t step = 0; step < options.max_new_tokens; ++step) { + const int32_t token = argmax_index(logits); + if (is_eos(token)) { + break; + } + out.push_back(token); + logits = impl_->decode_graph->run_step(token); + } + debug::timing_log_scalar("vibeasr.lm.decode_total_ms", engine::debug::elapsed_ms(decode_start, Clock::now())); + return out; +} + +} // namespace engine::community_models::vibeasr diff --git a/src/community_models/vibeasr/session.cpp b/src/community_models/vibeasr/session.cpp new file mode 100644 index 000000000..b4d4b7989 --- /dev/null +++ b/src/community_models/vibeasr/session.cpp @@ -0,0 +1,368 @@ +#include "engine/community_models/vibeasr/session.h" + +#include "engine/framework/audio/conversion.h" +#include "engine/framework/audio/resampling.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/debug/trace.h" +#include "engine/framework/io/text.h" +#include "engine/framework/runtime/options.h" +#include "engine/framework/runtime/spec_backed_model.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::community_models::vibeasr { +namespace { + +using Clock = std::chrono::steady_clock; + +constexpr size_t kWeightContextBytes = 64ull * 1024ull * 1024ull; + +// VibeASR resamples to 24 kHz and RMS-normalizes to -25 dBFS before the encoder; +// both numbers are fixed in the reference implementation, not in the checkpoint. +constexpr int kSampleRate = 24000; +constexpr float kTargetDbFs = -25.0F; +constexpr float kNormalizeEps = 1.0e-6F; + +// Canonical HuggingFace ids for the VibeVoice special tokens. VibeASR inserts +// them numerically rather than through the tokenizer, because the GGUF vocab's +// text for these slots is Qwen2.5's original <|object_ref_start|> family while +// the embedding rows are the ones VibeVoice trained. +constexpr int32_t kEndOfText = 151643; +constexpr int32_t kImStart = 151644; +constexpr int32_t kImEnd = 151645; +constexpr int32_t kSpeechStart = 151646; +constexpr int32_t kSpeechEnd = 151647; +constexpr int32_t kSpeechPad = 151648; + +constexpr const char * kSystemPrompt = + "You are a helpful assistant that transcribes audio input into text output in JSON format."; + +std::shared_ptr require_assets(std::shared_ptr assets) { + if (assets == nullptr) { + throw std::runtime_error("VibeASR session requires assets"); + } + return assets; +} + +const engine::model_spec::ModelContract & require_contract( + const std::shared_ptr & contract) { + if (contract == nullptr) { + throw std::runtime_error("VibeASR session requires a model contract"); + } + return *contract; +} + +runtime::SessionOptions validate_session_setup( + const runtime::TaskSpec & task, + runtime::SessionOptions options, + const engine::model_spec::ModelContract & contract) { + if (task.task != runtime::VoiceTaskKind::Asr) { + throw std::runtime_error("VibeASR only supports VoiceTaskKind::Asr"); + } + if (task.mode != runtime::RunMode::Offline) { + throw std::runtime_error("VibeASR only supports offline sessions"); + } + runtime::validate_spec_backed_session_options(options, contract, "vibeasr", "VibeASR"); + return options; +} + +size_t encoder_graph_arena_bytes(const runtime::SessionOptions & options) { + return runtime::parse_size_mb_option( + options.options, {"vibeasr.encoder_graph_arena_mb"}, 64ull * 1024ull * 1024ull); +} + +size_t prefill_graph_arena_bytes(const runtime::SessionOptions & options) { + return runtime::parse_size_mb_option( + options.options, {"vibeasr.prefill_graph_arena_mb"}, 256ull * 1024ull * 1024ull); +} + +size_t decode_graph_arena_bytes(const runtime::SessionOptions & options) { + return runtime::parse_size_mb_option( + options.options, {"vibeasr.decode_graph_arena_mb"}, 256ull * 1024ull * 1024ull); +} + +std::shared_ptr load_tokenizer(const VibeASRAssets & assets) { + // No merges.txt in the published package, so the tokenizer comes from + // tokenizer.json alone. + return engine::tokenizers::load_llama_bpe_tokenizer(engine::tokenizers::LlamaBpeTokenizerSpec{ + {}, + {}, + assets.resources.require_file("tokenizer_config"), + assets.resources.require_file("tokenizer_json"), + engine::tokenizers::LlamaBpePreTokenizer::Qwen2, + }); +} + +std::string format_duration(float seconds) { + char buffer[64]; + std::snprintf(buffer, sizeof(buffer), "%.2f", static_cast(seconds)); + return std::string(buffer); +} + +} // namespace + +VibeASRSession::VibeASRSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets, + std::shared_ptr contract) + : RuntimeSessionBase(validate_session_setup(task, std::move(options), require_contract(contract))), + task_(std::move(task)), + assets_(require_assets(std::move(assets))), + contract_(std::move(contract)), + tokenizer_(load_tokenizer(*assets_)), + encoder_(assets_->vae, execution_context(), encoder_graph_arena_bytes(RuntimeSessionBase::options())), + lm_(assets_->lm_weights, + assets_->lm, + execution_context(), + prefill_graph_arena_bytes(RuntimeSessionBase::options()), + decode_graph_arena_bytes(RuntimeSessionBase::options()), + kWeightContextBytes) { + // Both weight stores have uploaded by now; drop the resident file blobs. + assets_->vae->source->release_storage(); + assets_->lm_weights->release_storage(); +} + +VibeASRSession::~VibeASRSession() = default; + +std::string VibeASRSession::family() const { + return "vibeasr"; +} + +runtime::VoiceTaskKind VibeASRSession::task_kind() const { + return task_.task; +} + +runtime::RunMode VibeASRSession::run_mode() const { + return task_.mode; +} + +void VibeASRSession::prepare(const runtime::SessionPreparationRequest & request) { + (void)request; + mark_prepared(); +} + +VibeASRSession::RequestOptions VibeASRSession::parse_request_options(const runtime::TaskRequest & request) const { + runtime::validate_spec_backed_request_options(request.options, require_contract(contract_), "VibeASR"); + RequestOptions out; + if (const auto value = runtime::find_option(request.options, {"output_format"}); value.has_value()) { + if (*value != "text" && *value != "json") { + throw std::runtime_error("VibeASR output_format must be text or json"); + } + out.output_format = *value; + } + if (const auto value = runtime::find_option(request.options, {"context"}); value.has_value()) { + out.context = *value; + } + out.max_new_tokens = runtime::parse_positive_i64_option(request.options, {"max_new_tokens"}, out.max_new_tokens); + return out; +} + +runtime::AudioBuffer VibeASRSession::normalize(const runtime::AudioBuffer & audio) const { + if (audio.samples.empty()) { + throw std::runtime_error("VibeASR requires non-empty audio"); + } + auto mono = engine::audio::mixdown_interleaved_to_mono_average(audio.samples, audio.channels); + if (audio.sample_rate != kSampleRate) { + // VibeASR.cpp resamples with a naive linear kernel; audio.cpp's soxr path + // is the better filter, so a non-24 kHz input will not match the + // reference sample for sample. + engine::audio::SoxrResampleOptions options; + options.profile = engine::audio::SoxrResampleProfile::QualityOnly; + options.output_length_policy = engine::audio::SoxrOutputLengthPolicy::ExactExpected; + options.output_padding = 256; + options.reject_empty_output = true; + options.warning_context = "VibeASR audio"; + options.fallback_description = "linear resampling"; + mono = engine::audio::resample_mono_soxr_or_linear(mono, audio.sample_rate, kSampleRate, options); + } + double sum = 0.0; + for (const float sample : mono) { + sum += static_cast(sample) * static_cast(sample); + } + const float rms = std::sqrt(static_cast(sum / std::max(mono.size(), 1))); + if (rms >= kNormalizeEps) { + const float target = std::pow(10.0F, kTargetDbFs / 20.0F); + const float gain = target / (rms + kNormalizeEps); + float max_abs = 0.0F; + for (float & sample : mono) { + sample *= gain; + max_abs = std::max(max_abs, std::abs(sample)); + } + // Not in VibeASR.cpp, which can clip on a loud clip; audio.cpp's own + // vibevoice_asr frontend clamps here and this port follows it. + if (max_abs > 1.0F) { + const float scale = max_abs + kNormalizeEps; + for (float & sample : mono) { + sample /= scale; + } + } + } + return runtime::AudioBuffer{kSampleRate, 1, std::move(mono)}; +} + +VibeASRSpeechEmbeddings VibeASRSession::encode_speech(const std::vector & samples) { + const auto encode_start = Clock::now(); + const auto acoustic = encoder_.encode_acoustic(samples); + const auto semantic = encoder_.encode_semantic(samples); + debug::timing_log_scalar("vibeasr.session.encoder_ms", engine::debug::elapsed_ms(encode_start)); + if (acoustic.frames != semantic.frames || acoustic.dim != semantic.dim) { + throw std::runtime_error("VibeASR encoder branches disagree on the feature shape"); + } + if (acoustic.dim != assets_->lm.hidden_size) { + throw std::runtime_error("VibeASR connector width does not match the decoder hidden size"); + } + + // Both connectors are LM-width, so the reference sums them element-wise. + VibeASRSpeechEmbeddings out; + out.tokens = acoustic.frames; + out.hidden_size = acoustic.dim; + out.values.resize(acoustic.values.size()); + for (size_t i = 0; i < out.values.size(); ++i) { + out.values[i] = acoustic.values[i] + semantic.values[i]; + } + return out; +} + +VibeASRLmPrompt VibeASRSession::build_prompt( + int64_t speech_tokens, + float duration_seconds, + const RequestOptions & options) const { + // Qwen2.5 ChatML, assembled exactly as VibeASR.cpp does it: + // <|im_start|>system\n{SYSTEM}<|im_end|>\n + // <|im_start|>user\n<|speech_start|><|speech_pad|>xN<|speech_end|>{suffix}<|im_end|>\n + // There is deliberately no generation prompt -- the model emits the + // <|im_start|>assistant\n header itself. + const auto encode = [this](const std::string & text) { + return tokenizer_->encode(text, false); + }; + + const std::string instruction = options.output_format == "json" + ? "please transcribe it with these keys: Start, End, Speaker, Content" + : "please transcribe it."; + std::string suffix; + if (options.context.empty()) { + suffix = "\nThis is a " + format_duration(duration_seconds) + " seconds audio, " + instruction; + } else { + suffix = "\nThis is a " + format_duration(duration_seconds) + " seconds audio, with extra info: " + + options.context + "\n\n" + + (options.output_format == "json" + ? "Please transcribe it with these keys: Start, End, Speaker, Content" + : "Please transcribe it."); + } + + const auto system_content = encode(std::string("system\n") + kSystemPrompt); + const auto newline = encode("\n"); + const auto user_prefix = encode("user\n"); + const auto user_suffix = encode(suffix); + + VibeASRLmPrompt prompt; + const auto append = [&prompt](const std::vector & ids) { + prompt.input_ids.insert(prompt.input_ids.end(), ids.begin(), ids.end()); + }; + prompt.input_ids.push_back(kImStart); + append(system_content); + prompt.input_ids.push_back(kImEnd); + append(newline); + prompt.input_ids.push_back(kImStart); + append(user_prefix); + prompt.input_ids.push_back(kSpeechStart); + // The reference builds ceil(samples / 3200) pads but only prefills + // min(pads, frames) of them, so emitting exactly `frames` pads produces the + // same sequence. + for (int64_t i = 0; i < speech_tokens; ++i) { + prompt.speech_positions.push_back(static_cast(prompt.input_ids.size())); + prompt.input_ids.push_back(kSpeechPad); + } + prompt.input_ids.push_back(kSpeechEnd); + append(user_suffix); + prompt.input_ids.push_back(kImEnd); + append(newline); + return prompt; +} + +std::string VibeASRSession::decode_tokens(const std::vector & token_ids) const { + // The prompt carries no generation prompt, so the model emits its own + // "<|im_start|>assistant\n" header; drop it exactly as the reference does. + size_t begin = 0; + const auto piece = [this](int32_t id) { return tokenizer_->decode({id}, true); }; + if (!token_ids.empty() && token_ids[0] == kImStart) { + begin = 1; + if (begin < token_ids.size() && piece(token_ids[begin]) == "assistant") { + ++begin; + if (begin < token_ids.size() && piece(token_ids[begin]) == "\n") { + ++begin; + } + } + } + + std::vector filtered; + filtered.reserve(token_ids.size() - begin); + for (size_t i = begin; i < token_ids.size(); ++i) { + const int32_t id = token_ids[i]; + if (id == kSpeechPad || id == kSpeechStart || id == kSpeechEnd || id == kEndOfText || + tokenizer_->is_control_token_id(id)) { + continue; + } + filtered.push_back(id); + } + if (filtered.empty()) { + return ""; + } + return engine::io::trim_ascii_whitespace(tokenizer_->decode(filtered, true)); +} + +runtime::TaskResult VibeASRSession::run(const runtime::TaskRequest & request) { + require_prepared("VibeASR run()"); + if (!request.audio_input.has_value()) { + throw std::runtime_error("VibeASR run() requires audio_input"); + } + const auto wall_start = Clock::now(); + const auto options = parse_request_options(request); + const auto audio = normalize(*request.audio_input); + const float duration_seconds = + static_cast(audio.samples.size()) / static_cast(kSampleRate); + + auto speech = encode_speech(audio.samples); + if (speech.tokens <= 0) { + throw std::runtime_error("VibeASR audio is too short to produce a single encoder frame"); + } + const auto prompt = build_prompt(speech.tokens, duration_seconds, options); + + VibeASRGenerationOptions generation; + generation.max_new_tokens = options.max_new_tokens; + generation.eos_token_ids = {kImEnd, kEndOfText}; + const auto generated = lm_.generate(prompt, speech, generation); + + runtime::TaskResult result; + result.text_output = runtime::Transcript{decode_tokens(generated), ""}; + debug::trace_log_scalar("vibeasr.session.speech_tokens", speech.tokens); + debug::trace_log_scalar("vibeasr.session.generated_tokens", static_cast(generated.size())); + debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(wall_start)); + return result; +} + +std::shared_ptr make_vibeasr_loader() { + runtime::SpecBackedVoiceModelConfig config; + config.family = "vibeasr"; + config.load_assets = [](const std::filesystem::path & model_path) { + return load_vibeasr_assets(model_path); + }; + config.create_session = []( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options, + std::shared_ptr assets, + std::shared_ptr contract) { + return std::make_unique(task, options, std::move(assets), std::move(contract)); + }; + return runtime::make_spec_backed_voice_loader(std::move(config)); +} + +} // namespace engine::community_models::vibeasr diff --git a/src/community_models/vibeasr/vae_encoder.cpp b/src/community_models/vibeasr/vae_encoder.cpp new file mode 100644 index 000000000..2894fd119 --- /dev/null +++ b/src/community_models/vibeasr/vae_encoder.cpp @@ -0,0 +1,379 @@ +#include "engine/community_models/vibeasr/vae_encoder.h" + +#include +#include +#include + +#include +#include +#include +#include +#include + +// The in-band tensor scale of GGML_TYPE_I8_S is internal to ggml, so the +// waveform quantizer and the feature dequantizer reach for the same declarations +// the implementation uses rather than re-deriving the layout here. Buffer sizes +// still come from the public ggml_nbytes(), which already accounts for the +// trailing scale. +extern "C" { +void ggml_i8_s_to_float (const void * x, float * y, int64_t n); +size_t ggml_i8_s_from_float(const float * x, void * y, int64_t n); +} + +namespace engine::community_models::vibeasr { +namespace { + +// The graph is ~530 nodes for the published 7-stage encoder; leave headroom for +// deeper stage stacks without making the arena reservation depend on the config. +constexpr size_t kGraphNodes = 8192; + +// Activation layout inside this file is described in ggml `ne` order, which is +// the reverse of core::TensorShape. The encoder alternates between two layouts: +// +// channel-major ne [C, L] -- what a matmul produces, what the norms and the +// FFN want, since they reduce over ne[0] +// length-major ne [L, C] -- what im2col wants, since it slides over ne[0] +// +// VibeASR's graph flips between them with permute + cont in exactly the places +// reproduced below. + +core::TensorValue load_i8_s_tensor( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & name, + const std::vector & expected_shape) { + const auto metadata = source.require_metadata(name); + if (metadata.dtype != "i8_s") { + throw std::runtime_error("VibeASR VAE tensor " + name + " is " + metadata.dtype + ", expected i8_s"); + } + if (metadata.shape != expected_shape) { + throw std::runtime_error("VibeASR VAE tensor " + name + " has an unexpected shape"); + } + + core::TensorShape shape; + shape.rank = expected_shape.size(); + for (size_t i = 0; i < shape.rank; ++i) { + shape.dims[i] = expected_shape[i]; + } + + // I8_S is a whole-tensor quantization: the GGUF payload is the int8 values + // followed by one padded F32 scale, which is exactly what ggml_nbytes() + // expects, so the bytes go to the backend untouched. + const auto raw = source.require_tensor_data(name); + return store.make_tensor(shape, GGML_TYPE_I8_S, raw.bytes.data(), raw.bytes.size()); +} + +VaeBlockWeights load_block_weights( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + const VaeBlockConfig & config) { + const int64_t channels = config.channels; + const int64_t hidden = config.ffn_hidden; + + VaeBlockWeights weights; + weights.mixer_norm = store.load_f32_tensor(source, prefix + ".norm.weight", {channels}); + weights.mixer_conv_weight = load_i8_s_tensor( + store, source, prefix + ".mixer.conv.conv.conv.weight", {channels, 1, config.kernel_size}); + weights.mixer_conv_bias = store.load_f32_tensor(source, prefix + ".mixer.conv.conv.conv.bias", {channels}); + weights.mixer_gamma = store.load_f32_tensor(source, prefix + ".gamma", {channels}); + weights.ffn_norm = store.load_f32_tensor(source, prefix + ".ffn_norm.weight", {channels}); + weights.ffn_fc1_weight = load_i8_s_tensor(store, source, prefix + ".ffn.linear1.weight", {hidden, channels}); + weights.ffn_fc1_bias = store.load_f32_tensor(source, prefix + ".ffn.linear1.bias", {hidden}); + weights.ffn_fc2_weight = load_i8_s_tensor(store, source, prefix + ".ffn.linear2.weight", {channels, hidden}); + weights.ffn_fc2_bias = store.load_f32_tensor(source, prefix + ".ffn.linear2.bias", {channels}); + weights.ffn_gamma = store.load_f32_tensor(source, prefix + ".ffn_gamma", {channels}); + return weights; +} + +VaeBranchWeights load_branch_weights( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const VaeBranchConfig & config) { + VaeBranchWeights weights; + weights.stages.reserve(config.stages.size()); + + for (size_t stage = 0; stage < config.stages.size(); ++stage) { + const auto & stage_config = config.stages[stage]; + const std::string stage_prefix = config.prefix + ".stages." + std::to_string(stage); + const std::string downsample_prefix = + config.prefix + ".downsample_layers." + std::to_string(stage) + ".0.conv.conv"; + + VaeStageWeights stage_weights; + stage_weights.downsample_weight = load_i8_s_tensor( + store, + source, + downsample_prefix + ".weight", + {stage_config.out_channels, stage_config.in_channels, stage_config.downsample_kernel_size}); + stage_weights.downsample_bias = + store.load_f32_tensor(source, downsample_prefix + ".bias", {stage_config.out_channels}); + stage_weights.blocks.reserve(stage_config.blocks.size()); + for (size_t block = 0; block < stage_config.blocks.size(); ++block) { + stage_weights.blocks.push_back(load_block_weights( + store, source, stage_prefix + "." + std::to_string(block), stage_config.blocks[block])); + } + weights.stages.push_back(std::move(stage_weights)); + } + + const int64_t last_channels = config.stages.back().out_channels; + weights.head_weight = load_i8_s_tensor( + store, source, config.prefix + ".head.conv.conv.weight", + {config.latent_dim, last_channels, config.head_kernel_size}); + weights.head_bias = store.load_f32_tensor(source, config.prefix + ".head.conv.conv.bias", {config.latent_dim}); + + const std::string connector = config.prefix + "_connector"; + weights.connector_fc1_weight = load_i8_s_tensor( + store, source, connector + ".fc1.weight", {config.connector_hidden, config.latent_dim}); + weights.connector_fc1_bias = store.load_f32_tensor(source, connector + ".fc1.bias", {config.connector_hidden}); + weights.connector_norm = store.load_f32_tensor(source, connector + ".norm.weight", {config.connector_hidden}); + weights.connector_fc2_weight = load_i8_s_tensor( + store, source, connector + ".fc2.weight", {config.connector_hidden, config.connector_hidden}); + weights.connector_fc2_bias = store.load_f32_tensor(source, connector + ".fc2.bias", {config.connector_hidden}); + return weights; +} + +// channel-major [C, L] -> length-major [L, C], and back. +ggml_tensor * transpose_layout(core::ModuleBuildContext & ctx, ggml_tensor * x) { + return ggml_cont(ctx.ggml, ggml_permute(ctx.ggml, x, 1, 0, 2, 3)); +} + +// x [C, L], gamma [C] -> [C, L]. Reduces over the channel axis, matching the +// channels-last RMSNorm of the reference implementation. +ggml_tensor * rms_norm(core::ModuleBuildContext & ctx, ggml_tensor * x, ggml_tensor * gamma, float eps) { + return ggml_rms_norm_scaled(ctx.ggml, x, gamma, eps); +} + +// x [IC, L], w [IC, OC], bias [OC] -> [OC, L]. +// +// Everything is flattened to 2D for the matmul, so the trailing ne of x carry no +// information beyond the total number of positions. +ggml_tensor * linear( + core::ModuleBuildContext & ctx, + ggml_tensor * x, + ggml_tensor * w, + ggml_tensor * bias, + bool fuse_relu) { + GGML_ASSERT(x->ne[3] == 1); + const int64_t in_features = x->ne[0]; + const int64_t out_features = w->ne[1]; + const int64_t positions = x->ne[1] * x->ne[2]; + + ggml_tensor * flat = ggml_reshape_2d(ctx.ggml, x, in_features, positions); + ggml_tensor * out = fuse_relu ? ggml_mul_mat_add_relu(ctx.ggml, w, flat, bias) + : ggml_mul_mat_add(ctx.ggml, w, flat, bias); + return ggml_reshape_2d(ctx.ggml, out, out_features, positions); +} + +// Causal Conv1d. x [L, IC, 1], w [K, IC, OC], bias [OC] -> [OC, OW]. +// +// The left pad is K - stride and the right pad is zero, which is what makes the +// stack causal; the converter left-pads short kernels with zeros so the padded +// K stays exact. +ggml_tensor * conv1d_causal( + core::ModuleBuildContext & ctx, + ggml_tensor * x, + ggml_tensor * w, + ggml_tensor * bias, + int stride) { + const int64_t kernel_size = w->ne[0]; + const int64_t in_channels = w->ne[1]; + const int64_t out_channels = w->ne[2]; + const int left_pad = static_cast(kernel_size) - stride; + GGML_ASSERT(left_pad >= 0); + + // im2col gives [IC*K, OW, N]. + ggml_tensor * cols = ggml_im2col_asym( + ctx.ggml, w, x, stride, 0, /*lp0=*/left_pad, /*rp0=*/0, /*p1=*/0, /*d0=*/1, /*d1=*/0, + /*is_2D=*/false, GGML_TYPE_I8_S); + + ggml_tensor * w2d = ggml_reshape_2d(ctx.ggml, w, kernel_size * in_channels, out_channels); + ggml_tensor * cols2d = ggml_reshape_2d(ctx.ggml, cols, cols->ne[0], cols->ne[1] * cols->ne[2]); + return ggml_mul_mat_add(ctx.ggml, w2d, cols2d, bias); +} + +// Causal depthwise Conv1d. x [L, C], w [K, 1, C], bias [C] -> [C, L]. +// +// ggml_mul_mat_add takes its depthwise contraction path when the weight is +// [K, 1, C], producing [1, L, C]; the trailing reshape and permute fold that +// back to channel-major. +ggml_tensor * conv1d_dw_causal( + core::ModuleBuildContext & ctx, + ggml_tensor * x, + ggml_tensor * w, + ggml_tensor * bias) { + const int64_t kernel_size = w->ne[0]; + + ggml_tensor * x4d = ggml_reshape_4d(ctx.ggml, x, x->ne[0], 1, x->ne[1], 1); + ggml_tensor * cols = ggml_im2col_asym( + ctx.ggml, w, x4d, /*s0=*/1, 0, /*lp0=*/static_cast(kernel_size) - 1, /*rp0=*/0, /*p1=*/0, + /*d0=*/1, /*d1=*/0, /*is_2D=*/false, GGML_TYPE_I8_S); + + ggml_tensor * out = ggml_mul_mat_add(ctx.ggml, w, cols, bias); + out = ggml_reshape_3d(ctx.ggml, out, out->ne[1], out->ne[2], 1); + return transpose_layout(ctx, out); +} + +// One ConvNeXt block. x [C, L] -> [C, L]. +ggml_tensor * build_block( + core::ModuleBuildContext & ctx, + ggml_tensor * x, + const VaeBlockWeights & weights, + float eps) { + ggml_tensor * residual = x; + ggml_tensor * h = rms_norm(ctx, x, weights.mixer_norm.tensor, eps); + h = transpose_layout(ctx, h); + h = conv1d_dw_causal(ctx, h, weights.mixer_conv_weight.tensor, weights.mixer_conv_bias.tensor); + // LayerScale folded into the residual add: h * gamma + residual. + x = ggml_add_scaled(ctx.ggml, h, residual, weights.mixer_gamma.tensor); + + residual = x; + h = rms_norm(ctx, x, weights.ffn_norm.tensor, eps); + // The I8_S FFN uses ReLU, fused into the first matmul. VibeASR's F32 + // fallback uses GELU instead; only the quantized path has published weights, + // so only ReLU is ported. + h = linear(ctx, h, weights.ffn_fc1_weight.tensor, weights.ffn_fc1_bias.tensor, /*fuse_relu=*/true); + h = linear(ctx, h, weights.ffn_fc2_weight.tensor, weights.ffn_fc2_bias.tensor, /*fuse_relu=*/false); + return ggml_add_scaled(ctx.ggml, h, residual, weights.ffn_gamma.tensor); +} + +// waveform [n_samples, 1, 1] -> features [connector_hidden, frames]. +ggml_tensor * build_branch( + core::ModuleBuildContext & ctx, + ggml_tensor * waveform, + const VaeBranchConfig & config, + const VaeBranchWeights & weights, + float eps) { + ggml_tensor * x = waveform; + + for (size_t stage = 0; stage < config.stages.size(); ++stage) { + const auto & stage_weights = weights.stages[stage]; + x = conv1d_causal( + ctx, + x, + stage_weights.downsample_weight.tensor, + stage_weights.downsample_bias.tensor, + static_cast(config.stages[stage].downsample_stride)); + for (const auto & block : stage_weights.blocks) { + x = build_block(ctx, x, block, eps); + } + // Back to length-major for the next stage's im2col (and for the head). + x = transpose_layout(ctx, x); + } + + x = conv1d_causal(ctx, x, weights.head_weight.tensor, weights.head_bias.tensor, /*stride=*/1); + + x = linear(ctx, x, weights.connector_fc1_weight.tensor, weights.connector_fc1_bias.tensor, false); + x = rms_norm(ctx, x, weights.connector_norm.tensor, eps); + return linear(ctx, x, weights.connector_fc2_weight.tensor, weights.connector_fc2_bias.tensor, false); +} + +} // namespace + +VibeASRVaeEncoderRuntime::VibeASRVaeEncoderRuntime( + std::shared_ptr assets, + engine::core::ExecutionContext & execution_context, + size_t graph_arena_bytes) + : assets_(std::move(assets)), + execution_context_(&execution_context), + weight_store_( + execution_context.backend(), + execution_context.backend_type(), + "VibeASR VAE encoder weights", + 256ull * 1024ull * 1024ull), + graph_arena_bytes_(graph_arena_bytes) { + if (assets_ == nullptr) { + throw std::runtime_error("VibeASR VAE encoder runtime requires assets"); + } + weights_.acoustic = load_branch_weights(weight_store_, *assets_->source, assets_->config.acoustic); + weights_.semantic = load_branch_weights(weight_store_, *assets_->source, assets_->config.semantic); + weight_store_.upload(); +} + +VaeEncoderFeatures VibeASRVaeEncoderRuntime::encode_acoustic(const std::vector & samples) { + return encode(assets_->config.acoustic, weights_.acoustic, samples); +} + +VaeEncoderFeatures VibeASRVaeEncoderRuntime::encode_semantic(const std::vector & samples) { + return encode(assets_->config.semantic, weights_.semantic, samples); +} + +VaeEncoderFeatures VibeASRVaeEncoderRuntime::encode( + const VaeBranchConfig & config, + const VaeBranchWeights & weights, + const std::vector & samples) { + const int64_t num_samples = static_cast(samples.size()); + const int64_t expected_frames = config.frames_for_samples(num_samples); + if (expected_frames <= 0) { + // Shorter than one encoder frame: the first stage's im2col would have no + // output column at all. + return {}; + } + + ggml_init_params params{}; + params.mem_size = graph_arena_bytes_; + params.mem_buffer = nullptr; + params.no_alloc = true; + + ggml_context * ggml_ctx = ggml_init(params); + if (ggml_ctx == nullptr) { + throw std::runtime_error("Failed to initialize GGML context for the VibeASR VAE encoder"); + } + + ggml_gallocr * galloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(execution_context_->backend())); + if (galloc == nullptr) { + ggml_free(ggml_ctx); + throw std::runtime_error("Failed to initialize GGML allocator for the VibeASR VAE encoder"); + } + + VaeEncoderFeatures features; + + try { + core::ModuleBuildContext ctx{ggml_ctx, "vibeasr_vae_encoder", execution_context_->backend_type()}; + + // The waveform enters the graph already quantized: the encoder never + // touches F32 activations, so there is no leading quantize node. + ggml_tensor * waveform = ggml_new_tensor_3d(ggml_ctx, GGML_TYPE_I8_S, num_samples, 1, 1); + ggml_set_input(waveform); + + ggml_tensor * out = build_branch(ctx, waveform, config, weights, assets_->config.rms_norm_eps); + ggml_set_output(out); + + ggml_cgraph * gf = ggml_new_graph_custom(ggml_ctx, kGraphNodes, false); + ggml_build_forward_expand(gf, out); + + if (!ggml_gallocr_alloc_graph(galloc, gf)) { + throw std::runtime_error("Failed to allocate the GGML graph for the VibeASR VAE encoder"); + } + + std::vector quantized(ggml_nbytes(waveform)); + ggml_i8_s_from_float(samples.data(), quantized.data(), num_samples); + ggml_backend_tensor_set(waveform, quantized.data(), 0, quantized.size()); + + if (ggml_backend_graph_compute(execution_context_->backend(), gf) != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Failed to compute the GGML graph for the VibeASR VAE encoder"); + } + + features.dim = out->ne[0]; + features.frames = out->ne[1]; + if (features.frames != expected_frames) { + throw std::runtime_error("VibeASR VAE encoder produced an unexpected frame count"); + } + + // The result is still I8_S, one scale for the whole feature block. + std::vector raw(ggml_nbytes(out)); + ggml_backend_tensor_get(out, raw.data(), 0, raw.size()); + features.values.resize(static_cast(features.dim * features.frames)); + ggml_i8_s_to_float(raw.data(), features.values.data(), features.dim * features.frames); + } catch (...) { + ggml_gallocr_free(galloc); + ggml_free(ggml_ctx); + throw; + } + + ggml_gallocr_free(galloc); + ggml_free(ggml_ctx); + return features; +} + +} // namespace engine::community_models::vibeasr diff --git a/tests/vibeasr/test_vibeasr_asr.cpp b/tests/vibeasr/test_vibeasr_asr.cpp new file mode 100644 index 000000000..dea782212 --- /dev/null +++ b/tests/vibeasr/test_vibeasr_asr.cpp @@ -0,0 +1,145 @@ +// End-to-end probe for the ported VibeASR pipeline: I8_S VAE encoder -> ternary +// I2_S Qwen2 decoder -> transcript. +// +// The package ships two GGUFs, so --model points at the LM GGUF and the spec is +// resolved from the repo (same convention as minimax_h3). Skips with 125 when +// the checkpoint is not installed. +// +// Upstream: https://github.com/microsoft/VibeASR.cpp + +#include "engine/framework/audio/wav_reader.h" +#include "engine/framework/io/filesystem.h" +#include "engine/framework/io/text.h" +#include "engine/framework/runtime/model.h" +#include "engine/framework/runtime/registry.h" +#include "engine/framework/runtime/session.h" + +#include +#include +#include +#include +#include +#include +#include + +#ifndef ENGINE_REPO_ROOT +#define ENGINE_REPO_ROOT "." +#endif + +namespace { + +constexpr int kExitPass = 0; +constexpr int kExitFail = 1; +constexpr int kExitSkip = 125; + +// LibriSpeech test-clean 6930-75918-0000, transcribed by VibeASR.cpp's own +// asr_infer --greedy on the same two GGUFs. +const char * kExpectedText = "Concord returned to its place amidst the tents."; + +std::filesystem::path repo_path(const std::string & relative) { + return std::filesystem::path(ENGINE_REPO_ROOT) / relative; +} + +std::string arg_value(int argc, char ** argv, const std::string & name, const std::string & fallback) { + for (int i = 1; i + 1 < argc; ++i) { + if (argv[i] == name) { + return argv[i + 1]; + } + } + return fallback; +} + +std::string normalize_text(const std::string & text) { + std::string out; + out.reserve(text.size()); + for (char ch : text) { + if (std::isalnum(static_cast(ch)) || std::isspace(static_cast(ch))) { + out.push_back(static_cast(std::tolower(static_cast(ch)))); + } + } + return engine::io::trim_ascii_whitespace(std::move(out)); +} + +} // namespace + +int main(int argc, char ** argv) { + const std::filesystem::path model_path = arg_value( + argc, argv, "--model", repo_path("models/vibeasr/vibeasr-lm-i2_s-embed-q6_k.gguf").string()); + const std::filesystem::path spec_override = arg_value( + argc, argv, "--model-spec-override", repo_path("model_specs").string()); + const std::filesystem::path audio_path = arg_value( + argc, argv, "--audio", + repo_path("assets/asr_validation/librispeech/librispeech_test_clean_6930-75918-0000.wav").string()); + const int threads = std::atoi(arg_value(argc, argv, "--threads", "4").c_str()); + + if (!engine::io::is_existing_file(model_path) || !engine::io::is_existing_file(audio_path)) { + std::fprintf( + stderr, + "SKIP: test_vibeasr_asr needs the LM GGUF at '%s' and audio at '%s'.\n" + " Fetch huggingface.co/microsoft/VibeVoice-ASR-BitNet and run\n" + " tools/community_models/convert_vibeasr_gguf.py --in-place on both GGUFs.\n", + model_path.string().c_str(), + audio_path.string().c_str()); + return kExitSkip; + } + + try { + auto registry = engine::runtime::make_default_registry(); + engine::runtime::ModelLoadRequest load_request; + load_request.model_path = model_path; + load_request.model_spec_override = spec_override; + load_request.family_hint = "vibeasr"; + auto model = registry.load(load_request); + + const engine::runtime::TaskSpec task{ + engine::runtime::VoiceTaskKind::Asr, + engine::runtime::RunMode::Offline, + }; + engine::runtime::SessionOptions session_options; + session_options.backend.threads = threads > 0 ? threads : 1; + + auto session = model->create_task_session(task, session_options); + auto * offline = dynamic_cast(session.get()); + if (offline == nullptr) { + std::cerr << "FAIL: VibeASR session is not an IOfflineVoiceTaskSession\n"; + return kExitFail; + } + + const auto wav = engine::audio::read_wav_f32(audio_path); + engine::runtime::AudioBuffer audio; + audio.sample_rate = wav.sample_rate; + audio.channels = wav.channels; + audio.samples = wav.samples; + + offline->prepare(engine::runtime::build_preparation_request(audio)); + + engine::runtime::TaskRequest request; + request.audio_input = audio; + const auto result = offline->run(request); + + if (!result.text_output.has_value()) { + std::cerr << "FAIL: VibeASR produced no text output\n"; + return kExitFail; + } + const std::string actual = result.text_output->text; + std::cout << "transcript: " << actual << "\n"; + std::cout << "expected: " << kExpectedText << "\n"; + + // Raw equality pins punctuation and casing against the reference decode; + // the normalized compare is only there to localize a failure. + if (actual != kExpectedText) { + if (normalize_text(actual) == normalize_text(kExpectedText)) { + std::cerr << "FAIL: transcript differs only in punctuation or casing\n"; + } else { + std::cerr << "FAIL: transcript mismatch\n"; + } + return kExitFail; + } + + std::cout << "PASS: VibeASR end-to-end transcription matches VibeASR.cpp\n"; + return kExitPass; + } catch (const std::exception & error) { + std::cerr << "FAIL: " << error.what() << "\n"; + return kExitFail; + } +} diff --git a/tests/vibeasr/test_vibeasr_vae_encoder.cpp b/tests/vibeasr/test_vibeasr_vae_encoder.cpp new file mode 100644 index 000000000..9d10404de --- /dev/null +++ b/tests/vibeasr/test_vibeasr_vae_encoder.cpp @@ -0,0 +1,261 @@ +// Parity probe for the ported VibeASR VAE encoder. +// +// Without --reference-* the probe only checks that the graph runs and produces a +// sane feature block. With a reference dump from VibeASR.cpp's own vae_server +// (raw float32, frames * dim, row-major) it reports max abs error, mean abs +// error, and cosine similarity, and fails outside the tolerances below. +// +// Upstream: https://github.com/microsoft/VibeASR.cpp + +#include "engine/community_models/vibeasr/vae_encoder.h" +#include "engine/framework/audio/wav_reader.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/io/filesystem.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifndef ENGINE_REPO_ROOT +#define ENGINE_REPO_ROOT "." +#endif + +namespace { + +constexpr int kExitPass = 0; +constexpr int kExitFail = 1; +constexpr int kExitSkip = 125; + +// Both encoders end in an I8_S matmul, so the whole feature block shares one +// scale: agreement is judged relative to that block's dynamic range rather than +// with an absolute epsilon. +// +// Bit-exactness is not reachable here and the tolerances reflect a measured +// noise floor rather than a guess. Every stage requantizes to int8, and the two +// implementations disagree in the last float bit of the per-tensor scale (this +// port stores amax/127 and multiplies, VibeASR.cpp stores 127/amax and divides), +// which flips a handful of values by one int8 step early on. Nudging a single +// input sample by one int8 step and re-running VibeASR.cpp against itself moves +// its own output by cosine 0.9959 (acoustic) / 0.9870 (semantic) -- i.e. the +// graph amplifies one LSB to about the same distance we see between the two +// implementations, so anything tighter would be testing rounding luck. +constexpr double kMaxMeanRelativeError = 0.02; +constexpr double kMinCosineSimilarity = 0.98; + +std::string arg_value(int argc, char ** argv, const std::string & name, const std::string & fallback) { + for (int i = 1; i + 1 < argc; ++i) { + if (argv[i] == name) { + return argv[i + 1]; + } + } + return fallback; +} + +std::vector read_f32_dump(const std::filesystem::path & path) { + std::ifstream file(path, std::ios::binary | std::ios::ate); + if (!file) { + throw std::runtime_error("cannot open reference dump: " + path.string()); + } + const auto bytes = static_cast(file.tellg()); + if (bytes % sizeof(float) != 0) { + throw std::runtime_error("reference dump is not a whole number of floats: " + path.string()); + } + std::vector values(bytes / sizeof(float)); + file.seekg(0); + file.read(reinterpret_cast(values.data()), static_cast(bytes)); + if (!file) { + throw std::runtime_error("short read on reference dump: " + path.string()); + } + return values; +} + +bool check_features( + const char * branch, + const engine::community_models::vibeasr::VaeEncoderFeatures & features, + int64_t expected_dim, + int64_t expected_frames) { + if (features.frames != expected_frames || features.dim != expected_dim) { + std::fprintf( + stderr, + "FAIL: %s features are [%lld frames, %lld dim], expected [%lld, %lld]\n", + branch, + static_cast(features.frames), + static_cast(features.dim), + static_cast(expected_frames), + static_cast(expected_dim)); + return false; + } + + double amax = 0.0; + for (float value : features.values) { + if (!std::isfinite(value)) { + std::fprintf(stderr, "FAIL: %s features contain a non-finite value\n", branch); + return false; + } + amax = std::max(amax, static_cast(std::fabs(value))); + } + if (amax == 0.0) { + std::fprintf(stderr, "FAIL: %s features are all zero\n", branch); + return false; + } + std::printf("%s: %lld frames x %lld dim, amax %.6f\n", + branch, + static_cast(features.frames), + static_cast(features.dim), + amax); + return true; +} + +bool compare_reference( + const char * branch, + const engine::community_models::vibeasr::VaeEncoderFeatures & features, + const std::filesystem::path & reference_path) { + const auto reference = read_f32_dump(reference_path); + if (reference.size() != features.values.size()) { + std::fprintf( + stderr, + "FAIL: %s reference has %zu values, encoder produced %zu\n", + branch, + reference.size(), + features.values.size()); + return false; + } + + double max_abs = 0.0; + double sum_abs = 0.0; + double reference_amax = 0.0; + double dot = 0.0; + double norm_a = 0.0; + double norm_b = 0.0; + for (size_t i = 0; i < reference.size(); ++i) { + const double a = features.values[i]; + const double b = reference[i]; + const double diff = std::fabs(a - b); + max_abs = std::max(max_abs, diff); + sum_abs += diff; + reference_amax = std::max(reference_amax, std::fabs(b)); + dot += a * b; + norm_a += a * a; + norm_b += b * b; + } + const double mean_abs = sum_abs / static_cast(reference.size()); + const double cosine = (norm_a > 0.0 && norm_b > 0.0) ? dot / std::sqrt(norm_a * norm_b) : 0.0; + const double max_relative = reference_amax > 0.0 ? max_abs / reference_amax : max_abs; + const double mean_relative = reference_amax > 0.0 ? mean_abs / reference_amax : mean_abs; + + std::printf( + "%s vs reference: max abs %.6g (%.3g of range), mean abs %.6g (%.3g of range), cosine %.8f\n", + branch, max_abs, max_relative, mean_abs, mean_relative, cosine); + + bool ok = true; + if (mean_relative > kMaxMeanRelativeError) { + std::fprintf(stderr, "FAIL: %s mean relative error %.6g exceeds %.6g\n", + branch, mean_relative, kMaxMeanRelativeError); + ok = false; + } + if (cosine < kMinCosineSimilarity) { + std::fprintf(stderr, "FAIL: %s cosine %.8f is below %.8f\n", branch, cosine, kMinCosineSimilarity); + ok = false; + } + return ok; +} + +} // namespace + +int main(int argc, char ** argv) { + const std::filesystem::path model_path = arg_value(argc, argv, "--model", ""); + const std::filesystem::path audio_path = arg_value(argc, argv, "--audio", ""); + const std::filesystem::path acoustic_reference = arg_value(argc, argv, "--reference-acoustic", ""); + const std::filesystem::path semantic_reference = arg_value(argc, argv, "--reference-semantic", ""); + const int threads = std::atoi(arg_value(argc, argv, "--threads", "4").c_str()); + + if (model_path.empty() || !engine::io::is_existing_file(model_path) || + audio_path.empty() || !engine::io::is_existing_file(audio_path)) { + std::fprintf( + stderr, + "SKIP: test_vibeasr_vae_encoder needs --model and --audio .\n" + " Convert a VibeASR.cpp checkpoint with tools/community_models/convert_vibeasr_gguf.py first.\n"); + return kExitSkip; + } + + try { + const auto wav = engine::audio::read_wav_f32(audio_path); + if (wav.channels != 1) { + std::fprintf(stderr, "SKIP: %s has %d channels, the encoder takes mono\n", + audio_path.string().c_str(), wav.channels); + return kExitSkip; + } + // The encoder is a raw-waveform stack: it accepts whatever rate the clip + // carries, and only the frame count and the reported RTF depend on it. + if (wav.sample_rate <= 0) { + std::fprintf(stderr, "SKIP: %s reports sample rate %d\n", + audio_path.string().c_str(), wav.sample_rate); + return kExitSkip; + } + + auto assets = engine::community_models::vibeasr::load_vibeasr_vae_assets(model_path); + const auto & config = assets->config; + std::printf( + "acoustic: %zu stages, total stride %lld, latent %lld, connector %lld\n", + config.acoustic.stages.size(), + static_cast(config.acoustic.total_stride), + static_cast(config.acoustic.latent_dim), + static_cast(config.acoustic.connector_hidden)); + + engine::core::BackendConfig backend_config; + backend_config.type = engine::core::BackendType::Cpu; + backend_config.threads = threads > 0 ? threads : 1; + engine::core::ExecutionContext execution_context(backend_config); + + engine::community_models::vibeasr::VibeASRVaeEncoderRuntime runtime(assets, execution_context); + + const auto num_samples = static_cast(wav.samples.size()); + const double audio_seconds = static_cast(num_samples) / static_cast(wav.sample_rate); + + const auto acoustic_start = std::chrono::steady_clock::now(); + const auto acoustic = runtime.encode_acoustic(wav.samples); + const auto semantic_start = std::chrono::steady_clock::now(); + const auto semantic = runtime.encode_semantic(wav.samples); + const auto encode_end = std::chrono::steady_clock::now(); + + const auto ms = [](auto from, auto to) { + return std::chrono::duration(to - from).count(); + }; + const double acoustic_ms = ms(acoustic_start, semantic_start); + const double semantic_ms = ms(semantic_start, encode_end); + std::printf( + "encode wall: acoustic %.1f ms, semantic %.1f ms, both %.1f ms for %.2f s of audio (RTF %.4f)\n", + acoustic_ms, semantic_ms, acoustic_ms + semantic_ms, audio_seconds, + (acoustic_ms + semantic_ms) / 1000.0 / audio_seconds); + + bool ok = true; + ok &= check_features( + "acoustic", acoustic, config.acoustic.connector_hidden, + config.acoustic.frames_for_samples(num_samples)); + ok &= check_features( + "semantic", semantic, config.semantic.connector_hidden, + config.semantic.frames_for_samples(num_samples)); + + if (!acoustic_reference.empty()) { + ok &= compare_reference("acoustic", acoustic, acoustic_reference); + } + if (!semantic_reference.empty()) { + ok &= compare_reference("semantic", semantic, semantic_reference); + } + if (acoustic_reference.empty() && semantic_reference.empty()) { + std::printf("no reference dump given: shape and sanity checks only\n"); + } + + return ok ? kExitPass : kExitFail; + } catch (const std::exception & error) { + std::fprintf(stderr, "FAIL: %s\n", error.what()); + return kExitFail; + } +} diff --git a/tools/community_models/convert_vibeasr_gguf.py b/tools/community_models/convert_vibeasr_gguf.py new file mode 100644 index 000000000..2bfb9704a --- /dev/null +++ b/tools/community_models/convert_vibeasr_gguf.py @@ -0,0 +1,248 @@ +#!/usr/bin/env python3 +"""Convert a VibeASR.cpp GGUF into an audio.cpp GGUF package. + +Upstream: https://github.com/microsoft/VibeASR.cpp +Weights: https://huggingface.co/microsoft/VibeVoice-ASR-BitNet + +Handles both halves of the published package -- the I8_S VAE encoder and the +ternary I2_S language model -- because they need exactly the same fix and +nothing else. VibeASR.cpp ships both already quantized by its own ggml fork, so +there is nothing to re-quantize here. The only thing that differs is the numeric +type id: the VibeASR fork picked 36 (I2_S) and 37 (I8_S), which upstream ggml had +already used for the retired IQ4_NL_4_4 / IQ4_NL_4_8 slots. audio.cpp therefore +registers the same two types at 42 (I8_S) and 43 (I2_S). Tensors of any other +type in the file -- the LM's Q6_K token embedding, its F16 output projection, and +every F32 norm and bias -- are already portable and pass through untouched. + +The on-disk layout is identical either way -- an I8_S tensor is `nelements` int8 +bytes followed by a single padded F32 tensor scale, an I2_S tensor is the same +with 128 ternary codes packed per 32 bytes, and ggml's GGUF writer sizes every +tensor with ggml_nbytes() -- so this tool rewrites the 4-byte type field of each +tensor info and copies everything else through byte for byte. Data offsets, the +data section, and the KV block are untouched. + +Examples: + # inspect a VibeASR GGUF without writing anything + python3 tools/community_models/convert_vibeasr_gguf.py \ + --input vibeasr-vae-encoder-i8_s.gguf --list + + # convert a downloaded package where it sits (both halves) + python3 tools/community_models/convert_vibeasr_gguf.py \ + --input VibeVoice-ASR-BitNet/vibeasr-vae-encoder-i8_s.gguf --in-place + python3 tools/community_models/convert_vibeasr_gguf.py \ + --input VibeVoice-ASR-BitNet/vibeasr-lm-i2_s-embed-q6_k.gguf --in-place + + # or write the converted copy somewhere else + python3 tools/community_models/convert_vibeasr_gguf.py \ + --input vibeasr-vae-encoder-i8_s.gguf \ + --output models/vibeasr/vae_encoder-i8_s.gguf + + # confirm an already converted package needs no further remapping + python3 tools/community_models/convert_vibeasr_gguf.py \ + --input models/vibeasr/vae_encoder-i8_s.gguf --check +""" +import argparse +import struct +import sys +from pathlib import Path + +GGUF_MAGIC = b"GGUF" + +# VibeASR.cpp fork id -> audio.cpp id. See external/ggml/include/ggml.h for why +# audio.cpp cannot reuse 36/37. +TYPE_REMAP = {36: 43, 37: 42} + +TYPE_NAMES = {0: "f32", 1: "f16", 8: "q8_0", 14: "q6_k", 42: "i8_s", 43: "i2_s"} + +# GGUF metadata value type ids. +( + KV_UINT8, + KV_INT8, + KV_UINT16, + KV_INT16, + KV_UINT32, + KV_INT32, + KV_FLOAT32, + KV_BOOL, + KV_STRING, + KV_ARRAY, + KV_UINT64, + KV_INT64, + KV_FLOAT64, +) = range(13) + +KV_FIXED_SIZE = { + KV_UINT8: 1, + KV_INT8: 1, + KV_UINT16: 2, + KV_INT16: 2, + KV_UINT32: 4, + KV_INT32: 4, + KV_FLOAT32: 4, + KV_BOOL: 1, + KV_UINT64: 8, + KV_INT64: 8, + KV_FLOAT64: 8, +} + + +class Reader: + """Minimal forward-only GGUF header reader that tracks field offsets.""" + + def __init__(self, data: bytes): + self.data = data + self.pos = 0 + + def take(self, n: int) -> bytes: + if self.pos + n > len(self.data): + raise ValueError("GGUF header is truncated") + chunk = self.data[self.pos : self.pos + n] + self.pos += n + return chunk + + def u32(self) -> int: + return struct.unpack(" int: + return struct.unpack(" str: + return self.take(self.u64()).decode("utf-8", errors="replace") + + def skip_kv_value(self, kv_type: int) -> None: + if kv_type in KV_FIXED_SIZE: + self.take(KV_FIXED_SIZE[kv_type]) + elif kv_type == KV_STRING: + self.string() + elif kv_type == KV_ARRAY: + item_type = self.u32() + count = self.u64() + if item_type in KV_FIXED_SIZE: + self.take(KV_FIXED_SIZE[item_type] * count) + elif item_type == KV_STRING: + for _ in range(count): + self.string() + else: + raise ValueError(f"unsupported GGUF array element type {item_type}") + else: + raise ValueError(f"unsupported GGUF metadata type {kv_type}") + + +def parse_tensor_infos(data: bytes): + """Return (tensor_infos, alignment). Each info records where its type field lives.""" + reader = Reader(data) + if reader.take(4) != GGUF_MAGIC: + raise ValueError("not a GGUF file") + version = reader.u32() + if version != 3: + raise ValueError(f"unsupported GGUF version {version}") + n_tensors = reader.u64() + n_kv = reader.u64() + + alignment = 32 + for _ in range(n_kv): + key = reader.string() + kv_type = reader.u32() + if key == "general.alignment" and kv_type == KV_UINT32: + alignment = reader.u32() + else: + reader.skip_kv_value(kv_type) + + infos = [] + for _ in range(n_tensors): + name = reader.string() + n_dims = reader.u32() + dims = [reader.u64() for _ in range(n_dims)] + type_offset = reader.pos + type_id = reader.u32() + data_offset = reader.u64() + infos.append( + { + "name": name, + "dims": dims, + "type": type_id, + "type_offset": type_offset, + "data_offset": data_offset, + } + ) + return infos, alignment + + +def type_name(type_id: int) -> str: + return TYPE_NAMES.get(type_id, f"type#{type_id}") + + +def list_tensors(infos) -> None: + histogram = {} + for info in infos: + histogram[info["type"]] = histogram.get(info["type"], 0) + 1 + print(f" {info['name']:<64} {type_name(info['type']):>6} {info['dims']}") + print(f"{len(infos)} tensors") + for type_id in sorted(histogram): + print(f" {type_name(type_id):>6}: {histogram[type_id]}") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--input", type=Path, required=True, help="VibeASR.cpp GGUF (VAE encoder or LM)") + parser.add_argument("--output", type=Path, help="audio.cpp GGUF to write") + parser.add_argument("--in-place", action="store_true", help="rewrite --input itself instead of writing a copy") + parser.add_argument("--list", action="store_true", help="print the tensor table and exit") + parser.add_argument("--check", action="store_true", help="exit non-zero if any tensor still needs remapping") + args = parser.parse_args() + + if args.in_place: + if args.output is not None: + parser.error("--in-place and --output are mutually exclusive") + args.output = args.input + + data = bytearray(args.input.read_bytes()) + infos, alignment = parse_tensor_infos(bytes(data)) + + if args.list: + list_tensors(infos) + return 0 + + stale = [info for info in infos if info["type"] in TYPE_REMAP] + if args.check: + if stale: + print(f"{args.input}: {len(stale)} tensors still use VibeASR fork type ids", file=sys.stderr) + return 1 + print(f"{args.input}: type ids are already audio.cpp native") + return 0 + + if args.output is None: + parser.error("--output or --in-place is required unless --list or --check is given") + + # Guard against a double conversion: the fork ids and the audio.cpp ids are + # both valid ggml types, so a second pass would silently corrupt nothing but + # would also hide a mistake in the source package. + already = [info for info in infos if info["type"] in set(TYPE_REMAP.values())] + if already and stale: + raise SystemExit("input mixes VibeASR fork type ids with audio.cpp ids") + if not stale: + print(f"{args.input}: nothing to remap, copying through") + + for info in stale: + struct.pack_into("